Improve codegen for enums with many cases (#9122)

* Improve codegen for enums with many cases

This commit improves the compile time of generating bindings for enums
with many cases in them (e.g. 1000+). This is done by optimizing for
enums specifically rather than handling them generically like other
variants which can reduce the amount of code going into rustc to O(1)
instead of O(N) with the number of cases. This in turn can greatly
reduce compile time.

The tradeoff made in this commit is that enums are now required to have
`#[repr(...)]` annotations along with no Rust-level discriminants
specified. This enables the use of a `transmute` to lift a discriminant
into Rust with a simple bounds check. Previously this was one large
`match` statement.

Closes #9081

* Fix some tests

* Add repr tag in fuzzing

* Fix syntax for Rust 1.78

crates/component-macro/tests/expanded/simple-wasi.rs

@@ -114,6 +114,7 @@ pub mod foo {
             #[derive(wasmtime::component::Lower)]
             #[component(enum)]
             #[derive(Clone, Copy, Eq, PartialEq)]
+            #[repr(u8)]
             pub enum Errno {
                 #[component(name = "e")]
                 E,

crates/component-macro/tests/expanded/simple-wasi_async.rs

@@ -121,6 +121,7 @@ pub mod foo {
             #[derive(wasmtime::component::Lower)]
             #[component(enum)]
             #[derive(Clone, Copy, Eq, PartialEq)]
+            #[repr(u8)]
             pub enum Errno {
                 #[component(name = "e")]
                 E,

crates/component-macro/tests/expanded/small-anonymous.rs

@@ -105,6 +105,7 @@ pub mod foo {
             #[derive(wasmtime::component::Lower)]
             #[component(enum)]
             #[derive(Clone, Copy, Eq, PartialEq)]
+            #[repr(u8)]
             pub enum Error {
                 #[component(name = "success")]
                 Success,
@@ -207,6 +208,7 @@ pub mod exports {
                 #[derive(wasmtime::component::Lower)]
                 #[component(enum)]
                 #[derive(Clone, Copy, Eq, PartialEq)]
+                #[repr(u8)]
                 pub enum Error {
                     #[component(name = "success")]
                     Success,

crates/component-macro/tests/expanded/small-anonymous_async.rs

@@ -112,6 +112,7 @@ pub mod foo {
             #[derive(wasmtime::component::Lower)]
             #[component(enum)]
             #[derive(Clone, Copy, Eq, PartialEq)]
+            #[repr(u8)]
             pub enum Error {
                 #[component(name = "success")]
                 Success,
@@ -220,6 +221,7 @@ pub mod exports {
                 #[derive(wasmtime::component::Lower)]
                 #[component(enum)]
                 #[derive(Clone, Copy, Eq, PartialEq)]
+                #[repr(u8)]
                 pub enum Error {
                     #[component(name = "success")]
                     Success,

crates/component-macro/tests/expanded/variants.rs

@@ -105,6 +105,7 @@ pub mod foo {
             #[derive(wasmtime::component::Lower)]
             #[component(enum)]
             #[derive(Clone, Copy, Eq, PartialEq)]
+            #[repr(u8)]
             pub enum E1 {
                 #[component(name = "a")]
                 A,
@@ -313,6 +314,7 @@ pub mod foo {
             #[derive(wasmtime::component::Lower)]
             #[component(enum)]
             #[derive(Clone, Copy, Eq, PartialEq)]
+            #[repr(u8)]
             pub enum MyErrno {
                 #[component(name = "bad1")]
                 Bad1,
@@ -846,6 +848,7 @@ pub mod exports {
                 #[derive(wasmtime::component::Lower)]
                 #[component(enum)]
                 #[derive(Clone, Copy, Eq, PartialEq)]
+                #[repr(u8)]
                 pub enum E1 {
                     #[component(name = "a")]
                     A,
@@ -1109,6 +1112,7 @@ pub mod exports {
                 #[derive(wasmtime::component::Lower)]
                 #[component(enum)]
                 #[derive(Clone, Copy, Eq, PartialEq)]
+                #[repr(u8)]
                 pub enum MyErrno {
                     #[component(name = "bad1")]
                     Bad1,

crates/component-macro/tests/expanded/variants_async.rs

@@ -112,6 +112,7 @@ pub mod foo {
             #[derive(wasmtime::component::Lower)]
             #[component(enum)]
             #[derive(Clone, Copy, Eq, PartialEq)]
+            #[repr(u8)]
             pub enum E1 {
                 #[component(name = "a")]
                 A,
@@ -320,6 +321,7 @@ pub mod foo {
             #[derive(wasmtime::component::Lower)]
             #[component(enum)]
             #[derive(Clone, Copy, Eq, PartialEq)]
+            #[repr(u8)]
             pub enum MyErrno {
                 #[component(name = "bad1")]
                 Bad1,
@@ -862,6 +864,7 @@ pub mod exports {
                 #[derive(wasmtime::component::Lower)]
                 #[component(enum)]
                 #[derive(Clone, Copy, Eq, PartialEq)]
+                #[repr(u8)]
                 pub enum E1 {
                     #[component(name = "a")]
                     A,
@@ -1125,6 +1128,7 @@ pub mod exports {
                 #[derive(wasmtime::component::Lower)]
                 #[component(enum)]
                 #[derive(Clone, Copy, Eq, PartialEq)]
+                #[repr(u8)]
                 pub enum MyErrno {
                     #[component(name = "bad1")]
                     Bad1,

crates/environ/src/component/types.rs

@@ -735,6 +735,24 @@ impl CanonicalAbiInfo {
         }
     }
 
+    /// Calculates ABI information for an enum with `cases` cases.
+    pub const fn enum_(cases: usize) -> CanonicalAbiInfo {
+        // NB: this is basically a duplicate definition of
+        // `CanonicalAbiInfo::variant`, these should be kept in sync.
+
+        let discrim_size = match DiscriminantSize::from_count(cases) {
+            Some(size) => size.byte_size(),
+            None => unreachable!(),
+        };
+        CanonicalAbiInfo {
+            size32: discrim_size,
+            align32: discrim_size,
+            size64: discrim_size,
+            align64: discrim_size,
+            flat_count: Some(1),
+        }
+    }
+
     /// Returns the flat count of this ABI information so long as the count
     /// doesn't exceed the `max` specified.
     pub fn flat_count(&self, max: usize) -> Option<usize> {

crates/misc/component-fuzz-util/src/lib.rs

@@ -569,10 +569,16 @@ pub fn rust_type(ty: &Type, name_counter: &mut u32, declarations: &mut TokenStre
                 .collect::<TokenStream>();
 
             let name = make_rust_name(name_counter);
+            let repr = match count.ilog2() {
+                0..=7 => quote!(u8),
+                8..=15 => quote!(u16),
+                _ => quote!(u32),
+            };
 
             declarations.extend(quote! {
                 #[derive(ComponentType, Lift, Lower, PartialEq, Debug, Copy, Clone, Arbitrary)]
                 #[component(enum)]
+                #[repr(#repr)]
                 enum #name {
                     #cases
                 }

crates/wasmtime/src/runtime/component/mod.rs

@@ -143,6 +143,7 @@ pub mod __internal {
     pub use anyhow;
     #[cfg(feature = "async")]
     pub use async_trait::async_trait;
+    pub use core::mem::transmute;
     pub use wasmtime_environ;
     pub use wasmtime_environ::component::{CanonicalAbiInfo, ComponentTypes, InterfaceType};
 }
@@ -512,6 +513,7 @@ pub use wasmtime_component_macro::bindgen;
 ///
 /// #[derive(ComponentType)]
 /// #[component(enum)]
+/// #[repr(u8)]
 /// enum Setting {
 ///     #[component(name = "yes")]
 ///     Yes,

crates/wit-bindgen/src/lib.rs

@@ -1930,6 +1930,13 @@ impl<'a> InterfaceGenerator<'a> {
         self.push_str(&derives.into_iter().collect::<Vec<_>>().join(", "));
         self.push_str(")]\n");
 
+        let repr = match enum_.cases.len().ilog2() {
+            0..=7 => "u8",
+            8..=15 => "u16",
+            _ => "u32",
+        };
+        uwriteln!(self.src, "#[repr({repr})]");
+
         self.push_str(&format!("pub enum {name} {{\n"));
         for case in enum_.cases.iter() {
             self.rustdoc(&case.docs);

tests/all/component_model/macros.rs

@@ -236,6 +236,7 @@ fn variant_derive() -> Result<()> {
 fn enum_derive() -> Result<()> {
     #[derive(ComponentType, Lift, Lower, PartialEq, Eq, Debug, Copy, Clone)]
     #[component(enum)]
+    #[repr(u8)]
     enum Foo {
         #[component(name = "foo-bar-baz")]
         A,
@@ -299,6 +300,7 @@ fn enum_derive() -> Result<()> {
     #[add_variants(257)]
     #[derive(ComponentType, Lift, Lower, PartialEq, Eq, Debug, Copy, Clone)]
     #[component(enum)]
+    #[repr(u16)]
     enum Many {}
 
     let component = Component::new(
@@ -330,6 +332,7 @@ fn enum_derive() -> Result<()> {
     // #[add_variants(65537)]
     // #[derive(ComponentType, Lift, Lower, PartialEq, Eq, Debug, Copy, Clone)]
     // #[component(enum)]
+    // #[repr(u32)]
     // enum ManyMore {}
 
     Ok(())