Add support for C-style enums with implicit discriminants (#4152)

RunDevelopment · web-flow · commit 02679e1b95cb · 2024-10-10T16:53:42.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,15 @@
 # `wasm-bindgen` Change Log
 --------------------------------------------------------------------------------
 
+## Unreleased
+
+### Added
+
+* Added support for implicit discriminants in enums.
+  [#4152](https://github.com/rustwasm/wasm-bindgen/pull/4152)
+
+--------------------------------------------------------------------------------
+
 ## [0.2.94](https://github.com/rustwasm/wasm-bindgen/compare/0.2.93...0.2.94)
 
 Released 2024-10-09
diff --git a/crates/cli/tests/reference/enums.d.ts b/crates/cli/tests/reference/enums.d.ts
@@ -37,6 +37,12 @@ export enum Color {
    */
   Red = 2,
 }
+export enum ImplicitDiscriminant {
+  A = 0,
+  B = 1,
+  C = 42,
+  D = 43,
+}
 /**
  * The name of a color.
  */
diff --git a/crates/cli/tests/reference/enums.js b/crates/cli/tests/reference/enums.js
@@ -79,6 +79,8 @@ Yellow:1,"1":"Yellow",
  */
 Red:2,"2":"Red", });
 
+export const ImplicitDiscriminant = Object.freeze({ A:0,"0":"A",B:1,"1":"B",C:42,"42":"C",D:43,"43":"D", });
+
 const __wbindgen_enum_ColorName = ["green", "yellow", "red"];
 
 const __wbindgen_enum_FooBar = ["foo", "bar"];
diff --git a/crates/cli/tests/reference/enums.rs b/crates/cli/tests/reference/enums.rs
@@ -47,15 +47,21 @@ pub fn option_string_enum_echo(color: Option<ColorName>) -> Option<ColorName> {
 
 /// An unused string enum.
 #[wasm_bindgen(js_name = "FooBar")]
-#[derive(PartialEq, Debug)]
 pub enum UnusedStringEnum {
     Foo = "foo",
     Bar = "bar",
 }
 
 #[wasm_bindgen]
-#[derive(PartialEq, Debug)]
 enum PrivateStringEnum {
     Foo = "foo",
     Bar = "bar",
 }
+
+#[wasm_bindgen]
+pub enum ImplicitDiscriminant {
+    A,
+    B,
+    C = 42,
+    D,
+}
diff --git a/crates/macro-support/src/parser.rs b/crates/macro-support/src/parser.rs
@@ -1,5 +1,6 @@
 use std::cell::{Cell, RefCell};
 use std::char;
+use std::collections::HashMap;
 use std::str::Chars;
 
 use ast::OperationKind;
@@ -1399,28 +1400,18 @@ impl<'a> MacroParse<(&'a mut TokenStream, BindgenAttrs)> for syn::ItemEnum {
             return string_enum(self, program, js_name, generate_typescript, comments);
         }
 
-        let has_discriminant = self.variants[0].discriminant.is_some();
-
         match self.vis {
             syn::Visibility::Public(_) => {}
             _ => bail_span!(self, "only public enums are allowed with #[wasm_bindgen]"),
         }
 
+        let mut last_discriminant: Option<u32> = None;
+        let mut discriminate_map: HashMap<u32, &syn::Variant> = HashMap::new();
+
         let variants = self
             .variants
             .iter()
-            .enumerate()
-            .map(|(i, v)| {
-                // Require that everything either has a discriminant or doesn't.
-                // We don't really want to get in the business of emulating how
-                // rustc assigns values to enums.
-                if v.discriminant.is_some() != has_discriminant {
-                    bail_span!(
-                        v,
-                        "must either annotate discriminant of all variants or none"
-                    );
-                }
-
+            .map(|v| {
                 let value = match &v.discriminant {
                     Some((_, expr)) => match get_expr(expr) {
                         syn::Expr::Lit(syn::ExprLit {
@@ -1442,8 +1433,33 @@ impl<'a> MacroParse<(&'a mut TokenStream, BindgenAttrs)> for syn::ItemEnum {
                              number literal values",
                         ),
                     },
-                    None => i as u32,
+                    None => {
+                        // Use the same algorithm as rustc to determine the next discriminant
+                        // https://doc.rust-lang.org/reference/items/enumerations.html#implicit-discriminants
+                        if let Some(last) = last_discriminant {
+                            if let Some(value) = last.checked_add(1) {
+                                value
+                            } else {
+                                bail_span!(
+                                    v,
+                                    "the discriminants of C-style enums with #[wasm_bindgen] must be representable as u32"
+                                );
+                            }
+                        } else {
+                            0
+                        }
+                    }
                 };
+                last_discriminant = Some(value);
+
+                if let Some(old) = discriminate_map.insert(value, v) {
+                    bail_span!(
+                        v,
+                        "discriminant value `{}` is already used by {} in this enum",
+                        value,
+                        old.ident
+                    );
+                }
 
                 let comments = extract_doc_comments(&v.attrs);
                 Ok(ast::Variant {
@@ -1454,21 +1470,9 @@ impl<'a> MacroParse<(&'a mut TokenStream, BindgenAttrs)> for syn::ItemEnum {
             })
             .collect::<Result<Vec<_>, Diagnostic>>()?;
 
-        let mut values = variants.iter().map(|v| v.value).collect::<Vec<_>>();
-        values.sort();
-        let hole = values
-            .windows(2)
-            .find_map(|window| {
-                if window[0] + 1 != window[1] {
-                    Some(window[0] + 1)
-                } else {
-                    None
-                }
-            })
-            .unwrap_or(*values.last().unwrap() + 1);
-        for value in values {
-            assert!(hole != value);
-        }
+        let hole = (0..=u32::MAX)
+            .find(|v| !discriminate_map.contains_key(v))
+            .unwrap();
 
         self.to_tokens(tokens);
 
diff --git a/crates/macro/ui-tests/invalid-enums.rs b/crates/macro/ui-tests/invalid-enums.rs
@@ -37,4 +37,30 @@ enum G {
     C,
 }
 
+#[wasm_bindgen]
+pub enum H {
+    A = 1,
+    B = 1, // collision
+}
+
+#[wasm_bindgen]
+pub enum I {
+    A = 4294967294, // = u32::MAX - 1
+    B, // would be u32::MAX
+    C, // would be u32::MAX + 1
+}
+
+#[wasm_bindgen]
+pub enum J {
+    A, // = 0
+    B = 0, // collision
+}
+
+#[wasm_bindgen]
+pub enum K {
+    A = 3,
+    B = 2,
+    C, // = 3 -> collision
+}
+
 fn main() {}
diff --git a/crates/macro/ui-tests/invalid-enums.stderr b/crates/macro/ui-tests/invalid-enums.stderr
@@ -39,3 +39,27 @@ error: all variants of a string enum must have a string value
    |
 37 |     C,
    |     ^
+
+error: discriminant value `1` is already used by A in this enum
+  --> ui-tests/invalid-enums.rs:43:5
+   |
+43 |     B = 1, // collision
+   |     ^^^^^
+
+error: the discriminants of C-style enums with #[wasm_bindgen] must be representable as u32
+  --> ui-tests/invalid-enums.rs:50:5
+   |
+50 |     C, // would be u32::MAX + 1
+   |     ^
+
+error: discriminant value `0` is already used by A in this enum
+  --> ui-tests/invalid-enums.rs:56:5
+   |
+56 |     B = 0, // collision
+   |     ^^^^^
+
+error: discriminant value `3` is already used by A in this enum
+  --> ui-tests/invalid-enums.rs:63:5
+   |
+63 |     C, // = 3 -> collision
+   |     ^
