support variant, enum, and union derives (#4359)

* support variant, enum, and union derives

This is the second stage of implementing #4308.  It adds support for deriving
variant, enum, and union impls for `ComponentType`, `Lift`, and `Lower`.  It
also fixes derived record impls for generic `struct`s, which I had intended to
support in my previous commit, but forgot to test.

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

* deduplicate component-macro code

Thanks to @jameysharp for the suggestion!

Signed-off-by: Joel Dice <joel.dice@fermyon.com>

crates/component-macro/src/lib.rs

@@ -1,16 +1,32 @@
 use proc_macro2::{Literal, TokenStream, TokenTree};
 use quote::{format_ident, quote};
 use std::collections::HashSet;
+use std::fmt;
 use syn::{parse_macro_input, parse_quote, Data, DeriveInput, Error, Result};
 
-#[derive(Debug)]
-enum Style {
-    Record,
+#[derive(Debug, Copy, Clone)]
+enum VariantStyle {
     Variant,
     Enum,
     Union,
 }
 
+impl fmt::Display for VariantStyle {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.write_str(match self {
+            Self::Variant => "variant",
+            Self::Enum => "enum",
+            Self::Union => "union",
+        })
+    }
+}
+
+#[derive(Debug, Copy, Clone)]
+enum Style {
+    Record,
+    Variant(VariantStyle),
+}
+
 fn find_style(input: &DeriveInput) -> Result<Style> {
     let mut style = None;
 
@@ -50,9 +66,9 @@ fn find_style(input: &DeriveInput) -> Result<Style> {
 
         style = Some(match style_string.as_ref() {
             "record" => Style::Record,
-            "variant" => Style::Variant,
-            "enum" => Style::Enum,
-            "union" => Style::Union,
+            "variant" => Style::Variant(VariantStyle::Variant),
+            "enum" => Style::Variant(VariantStyle::Enum),
+            "union" => Style::Variant(VariantStyle::Union),
             "flags" => {
                 return Err(Error::new_spanned(
                     &attribute.tokens,
@@ -73,10 +89,10 @@ fn find_style(input: &DeriveInput) -> Result<Style> {
     style.ok_or_else(|| Error::new_spanned(input, "missing `component` attribute"))
 }
 
-fn find_rename(field: &syn::Field) -> Result<Option<Literal>> {
+fn find_rename(attributes: &[syn::Attribute]) -> Result<Option<Literal>> {
     let mut name = None;
 
-    for attribute in &field.attrs {
+    for attribute in attributes {
         if attribute.path.leading_colon.is_some() || attribute.path.segments.len() != 1 {
             continue;
         }
@@ -119,33 +135,41 @@ fn find_rename(field: &syn::Field) -> Result<Option<Literal>> {
     Ok(name)
 }
 
-fn add_trait_bounds(generics: &syn::Generics) -> syn::Generics {
+fn add_trait_bounds(generics: &syn::Generics, bound: syn::TypeParamBound) -> syn::Generics {
     let mut generics = generics.clone();
     for param in &mut generics.params {
         if let syn::GenericParam::Type(ref mut type_param) = *param {
-            type_param
-                .bounds
-                .push(parse_quote!(wasmtime::component::ComponentType));
+            type_param.bounds.push(bound.clone());
         }
     }
     generics
 }
 
+struct VariantCase<'a> {
+    attrs: &'a [syn::Attribute],
+    ident: &'a syn::Ident,
+    ty: Option<&'a syn::Type>,
+}
+
 trait Expander {
     fn expand_record(&self, input: &DeriveInput, fields: &syn::FieldsNamed) -> Result<TokenStream>;
 
-    // TODO: expand_variant, expand_enum, etc.
+    fn expand_variant(
+        &self,
+        input: &DeriveInput,
+        cases: &[VariantCase],
+        style: VariantStyle,
+    ) -> Result<TokenStream>;
 }
 
-fn expand(expander: &dyn Expander, input: DeriveInput) -> Result<TokenStream> {
-    match find_style(&input)? {
+fn expand(expander: &dyn Expander, input: &DeriveInput) -> Result<TokenStream> {
+    match find_style(input)? {
         Style::Record => expand_record(expander, input),
-
-        style => Err(Error::new_spanned(input, format!("todo: expand {style:?}"))),
+        Style::Variant(style) => expand_variant(expander, input, style),
     }
 }
 
-fn expand_record(expander: &dyn Expander, input: DeriveInput) -> Result<TokenStream> {
+fn expand_record(expander: &dyn Expander, input: &DeriveInput) -> Result<TokenStream> {
     let name = &input.ident;
 
     let body = if let Data::Struct(body) = &input.data {
@@ -153,12 +177,12 @@ fn expand_record(expander: &dyn Expander, input: DeriveInput) -> Result<TokenStr
     } else {
         return Err(Error::new(
             name.span(),
-            "`record` component types can only be derived for `struct`s",
+            "`record` component types can only be derived for Rust `struct`s",
         ));
     };
 
     match &body.fields {
-        syn::Fields::Named(fields) => expander.expand_record(&input, fields),
+        syn::Fields::Named(fields) => expander.expand_record(input, fields),
 
         syn::Fields::Unnamed(_) | syn::Fields::Unit => Err(Error::new(
             name.span(),
@@ -167,9 +191,72 @@ fn expand_record(expander: &dyn Expander, input: DeriveInput) -> Result<TokenStr
     }
 }
 
+fn expand_variant(
+    expander: &dyn Expander,
+    input: &DeriveInput,
+    style: VariantStyle,
+) -> Result<TokenStream> {
+    let name = &input.ident;
+
+    let body = if let Data::Enum(body) = &input.data {
+        body
+    } else {
+        return Err(Error::new(
+            name.span(),
+            format!(
+                "`{}` component types can only be derived for Rust `enum`s",
+                style
+            ),
+        ));
+    };
+
+    if body.variants.is_empty() {
+        return Err(Error::new(
+            name.span(),
+            format!("`{}` component types can only be derived for Rust `enum`s with at least one variant", style),
+        ));
+    }
+
+    let cases = body
+        .variants
+        .iter()
+        .map(
+            |syn::Variant {
+                 attrs,
+                 ident,
+                 fields,
+                 ..
+             }| {
+                Ok(VariantCase {
+                    attrs,
+                    ident,
+                    ty: match fields {
+                        syn::Fields::Unnamed(fields) if fields.unnamed.len() == 1 => {
+                            Some(&fields.unnamed[0].ty)
+                        }
+                        syn::Fields::Unit => None,
+                        _ => {
+                            return Err(Error::new(
+                                name.span(),
+                                format!(
+                                    "`{}` component types can only be derived for Rust `enum`s \
+                                     containing variants with at most one unnamed field each",
+                                    style
+                                ),
+                            ))
+                        }
+                    },
+                })
+            },
+        )
+        .collect::<Result<Vec<_>>>()?;
+
+    expander.expand_variant(input, &cases, style)
+}
+
 #[proc_macro_derive(Lift, attributes(component))]
 pub fn lift(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
-    expand(&LiftExpander, parse_macro_input!(input as DeriveInput))
+    expand(&LiftExpander, &parse_macro_input!(input as DeriveInput))
         .unwrap_or_else(Error::into_compile_error)
         .into()
 }
@@ -183,21 +270,21 @@ impl Expander for LiftExpander {
         let mut lifts = TokenStream::new();
         let mut loads = TokenStream::new();
 
-        for syn::Field { ty, ident, .. } in &fields.named {
+        for syn::Field { ident, ty, .. } in &fields.named {
             lifts.extend(quote!(#ident: <#ty as wasmtime::component::Lift>::lift(
-                    store, options, &src.#ident
-                )?,));
+                store, options, &src.#ident
+            )?,));
 
             loads.extend(quote!(#ident: <#ty as wasmtime::component::Lift>::load(
-                    memory,
-                    &bytes
-                        [#internal::next_field::<#ty>(&mut offset)..]
-                        [..<#ty as wasmtime::component::ComponentType>::SIZE32]
-                )?,));
+                memory,
+                &bytes
+                    [#internal::next_field::<#ty>(&mut offset)..]
+                    [..<#ty as wasmtime::component::ComponentType>::SIZE32]
+            )?,));
         }
 
         let name = &input.ident;
-        let generics = add_trait_bounds(&input.generics);
+        let generics = add_trait_bounds(&input.generics, parse_quote!(wasmtime::component::Lift));
         let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
 
         let expanded = quote! {
@@ -230,11 +317,86 @@ impl Expander for LiftExpander {
 
         Ok(expanded)
     }
+
+    fn expand_variant(
+        &self,
+        input: &DeriveInput,
+        cases: &[VariantCase],
+        _style: VariantStyle,
+    ) -> Result<TokenStream> {
+        let internal = quote!(wasmtime::component::__internal);
+
+        let mut lifts = TokenStream::new();
+        let mut loads = TokenStream::new();
+
+        for (index, VariantCase { ident, ty, .. }) in cases.iter().enumerate() {
+            let index_u8 = u8::try_from(index).map_err(|_| {
+                Error::new(
+                    input.ident.span(),
+                    "`enum`s with more than 256 variants not yet supported",
+                )
+            })?;
+
+            let index_i32 = index_u8 as i32;
+
+            if let Some(ty) = ty {
+                lifts.extend(
+                    quote!(#index_i32 => Self::#ident(<#ty as wasmtime::component::Lift>::lift(
+                        store, options, unsafe { &src.payload.#ident }
+                    )?),),
+                );
+
+                loads.extend(
+                    quote!(#index_u8 => Self::#ident(<#ty as wasmtime::component::Lift>::load(
+                        memory, &payload[..<#ty as wasmtime::component::ComponentType>::SIZE32]
+                    )?),),
+                );
+            } else {
+                lifts.extend(quote!(#index_i32 => Self::#ident,));
+
+                loads.extend(quote!(#index_u8 => Self::#ident,));
+            }
+        }
+
+        let name = &input.ident;
+        let generics = add_trait_bounds(&input.generics, parse_quote!(wasmtime::component::Lift));
+        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+
+        let expanded = quote! {
+            unsafe impl #impl_generics wasmtime::component::Lift for #name #ty_generics #where_clause {
+                #[inline]
+                fn lift(
+                    store: &#internal::StoreOpaque,
+                    options: &#internal::Options,
+                    src: &Self::Lower,
+                ) -> #internal::anyhow::Result<Self> {
+                    Ok(match src.tag.get_i32() {
+                        #lifts
+                        discrim => #internal::anyhow::bail!("unexpected discriminant: {}", discrim),
+                    })
+                }
+
+                #[inline]
+                fn load(memory: &#internal::Memory, bytes: &[u8]) -> #internal::anyhow::Result<Self> {
+                    let align = <Self as wasmtime::component::ComponentType>::ALIGN32;
+                    debug_assert!((bytes.as_ptr() as usize) % (align as usize) == 0);
+                    let discrim = bytes[0];
+                    let payload = &bytes[#internal::align_to(1, align)..];
+                    Ok(match discrim {
+                        #loads
+                        discrim => #internal::anyhow::bail!("unexpected discriminant: {}", discrim),
+                    })
+                }
+            }
+        };
+
+        Ok(expanded)
+    }
 }
 
 #[proc_macro_derive(Lower, attributes(component))]
 pub fn lower(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
-    expand(&LowerExpander, parse_macro_input!(input as DeriveInput))
+    expand(&LowerExpander, &parse_macro_input!(input as DeriveInput))
         .unwrap_or_else(Error::into_compile_error)
         .into()
 }
@@ -248,18 +410,18 @@ impl Expander for LowerExpander {
         let mut lowers = TokenStream::new();
         let mut stores = TokenStream::new();
 
-        for syn::Field { ty, ident, .. } in &fields.named {
+        for syn::Field { ident, ty, .. } in &fields.named {
             lowers.extend(quote!(wasmtime::component::Lower::lower(
-                    &self.#ident, store, options, #internal::map_maybe_uninit!(dst.#ident)
-                )?;));
+                &self.#ident, store, options, #internal::map_maybe_uninit!(dst.#ident)
+            )?;));
 
             stores.extend(quote!(wasmtime::component::Lower::store(
-                    &self.#ident, memory, #internal::next_field::<#ty>(&mut offset)
-                )?;));
+                &self.#ident, memory, #internal::next_field::<#ty>(&mut offset)
+            )?;));
         }
 
         let name = &input.ident;
-        let generics = add_trait_bounds(&input.generics);
+        let generics = add_trait_bounds(&input.generics, parse_quote!(wasmtime::component::Lower));
         let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
 
         let expanded = quote! {
@@ -281,6 +443,7 @@ impl Expander for LowerExpander {
                     memory: &mut #internal::MemoryMut<'_, T>,
                     mut offset: usize
                 ) -> #internal::anyhow::Result<()> {
+                    debug_assert!(offset % (<Self as wasmtime::component::ComponentType>::ALIGN32 as usize) == 0);
                     #stores
                     Ok(())
                 }
@@ -289,13 +452,104 @@ impl Expander for LowerExpander {
 
         Ok(expanded)
     }
+
+    fn expand_variant(
+        &self,
+        input: &DeriveInput,
+        cases: &[VariantCase],
+        _style: VariantStyle,
+    ) -> Result<TokenStream> {
+        let internal = quote!(wasmtime::component::__internal);
+
+        let mut lowers = TokenStream::new();
+        let mut stores = TokenStream::new();
+
+        for (index, VariantCase { ident, ty, .. }) in cases.iter().enumerate() {
+            let index_u8 = u8::try_from(index).map_err(|_| {
+                Error::new(
+                    input.ident.span(),
+                    "`enum`s with more than 256 variants not yet supported",
+                )
+            })?;
+
+            let index_i32 = index_u8 as i32;
+
+            let pattern;
+            let lower;
+            let store;
+
+            if ty.is_some() {
+                pattern = quote!(Self::#ident(value));
+                lower = quote!(value.lower(store, options, #internal::map_maybe_uninit!(dst.payload.#ident)));
+                store = quote!(value.store(
+                    memory,
+                    offset + #internal::align_to(1, <Self as wasmtime::component::ComponentType>::ALIGN32)
+                ));
+            } else {
+                pattern = quote!(Self::#ident);
+                lower = quote!(Ok(()));
+                store = quote!(Ok(()));
+            }
+
+            lowers.extend(quote!(#pattern => {
+                #internal::map_maybe_uninit!(dst.tag).write(wasmtime::ValRaw::i32(#index_i32));
+                #lower
+            }));
+
+            stores.extend(quote!(#pattern => {
+                memory.get::<1>(offset)[0] = #index_u8;
+                #store
+            }));
+        }
+
+        let name = &input.ident;
+        let generics = add_trait_bounds(&input.generics, parse_quote!(wasmtime::component::Lower));
+        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+
+        let expanded = quote! {
+            unsafe impl #impl_generics wasmtime::component::Lower for #name #ty_generics #where_clause {
+                #[inline]
+                fn lower<T>(
+                    &self,
+                    store: &mut wasmtime::StoreContextMut<T>,
+                    options: &#internal::Options,
+                    dst: &mut std::mem::MaybeUninit<Self::Lower>,
+                ) -> #internal::anyhow::Result<()> {
+                    // See comment in <Result<T, E> as Lower>::lower for why we zero out the payload here
+                    unsafe {
+                        #internal::map_maybe_uninit!(dst.payload)
+                            .as_mut_ptr()
+                            .write_bytes(0u8, 1);
+                    }
+
+                    match self {
+                        #lowers
+                    }
+                }
+
+                #[inline]
+                fn store<T>(
+                    &self,
+                    memory: &mut #internal::MemoryMut<'_, T>,
+                    mut offset: usize
+                ) -> #internal::anyhow::Result<()> {
+                    debug_assert!(offset % (<Self as wasmtime::component::ComponentType>::ALIGN32 as usize) == 0);
+                    match self {
+                        #stores
+                    }
+                }
+            }
+        };
+
+        Ok(expanded)
+    }
 }
 
 #[proc_macro_derive(ComponentType, attributes(component))]
 pub fn component_type(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     expand(
         &ComponentTypeExpander,
-        parse_macro_input!(input as DeriveInput),
+        &parse_macro_input!(input as DeriveInput),
     )
     .unwrap_or_else(Error::into_compile_error)
     .into()
@@ -308,20 +562,31 @@ impl Expander for ComponentTypeExpander {
         let internal = quote!(wasmtime::component::__internal);
 
         let mut field_names_and_checks = TokenStream::new();
+        let mut lower_generic_params = TokenStream::new();
+        let mut lower_generic_args = TokenStream::new();
         let mut lower_field_declarations = TokenStream::new();
         let mut sizes = TokenStream::new();
         let mut unique_types = HashSet::new();
 
-        for field @ syn::Field { ty, ident, .. } in &fields.named {
-            lower_field_declarations
-                .extend(quote!(#ident: <#ty as wasmtime::component::ComponentType>::Lower,));
-
-            let literal = find_rename(field)?
+        for (
+            index,
+            syn::Field {
+                attrs, ident, ty, ..
+            },
+        ) in fields.named.iter().enumerate()
+        {
+            let name = find_rename(attrs)?
                 .unwrap_or_else(|| Literal::string(&ident.as_ref().unwrap().to_string()));
 
-            field_names_and_checks.extend(
-                quote!((#literal, <#ty as wasmtime::component::ComponentType>::typecheck),),
-            );
+            let generic = format_ident!("T{}", index);
+
+            lower_generic_params.extend(quote!(#generic: Copy,));
+            lower_generic_args.extend(quote!(<#ty as wasmtime::component::ComponentType>::Lower,));
+
+            lower_field_declarations.extend(quote!(#ident: #generic,));
+
+            field_names_and_checks
+                .extend(quote!((#name, <#ty as wasmtime::component::ComponentType>::typecheck),));
 
             sizes.extend(quote!(
                 size = #internal::align_to(size, <#ty as wasmtime::component::ComponentType>::ALIGN32);
@@ -342,20 +607,38 @@ impl Expander for ComponentTypeExpander {
             .collect::<TokenStream>();
 
         let name = &input.ident;
-        let generics = add_trait_bounds(&input.generics);
+        let generics = add_trait_bounds(
+            &input.generics,
+            parse_quote!(wasmtime::component::ComponentType),
+        );
         let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
-        let lower = format_ident!("_Lower{}", name);
+        let lower = format_ident!("Lower{}", name);
+
+        // You may wonder why we make the types of all the fields of the #lower struct generic.  This is to work
+        // around the lack of [perfect derive support in
+        // rustc](https://smallcultfollowing.com/babysteps//blog/2022/04/12/implied-bounds-and-perfect-derive/#what-is-perfect-derive)
+        // as of this writing.
+        //
+        // If the struct we're deriving a `ComponentType` impl for has any generic parameters, then #lower needs
+        // generic parameters too.  And if we just copy the parameters and bounds from the impl to #lower, then the
+        // `#[derive(Clone, Copy)]` will fail unless the original generics were declared with those bounds, which
+        // we don't want to require.
+        //
+        // Alternatively, we could just pass the `Lower` associated type of each generic type as arguments to
+        // #lower, but that would require distinguishing between generic and concrete types when generating
+        // #lower_field_declarations, which would require some form of symbol resolution.  That doesn't seem worth
+        // the trouble.
 
         let expanded = quote! {
             #[doc(hidden)]
             #[derive(Clone, Copy)]
             #[repr(C)]
-            pub struct #lower {
+            pub struct #lower <#lower_generic_params> {
                 #lower_field_declarations
             }
 
             unsafe impl #impl_generics wasmtime::component::ComponentType for #name #ty_generics #where_clause {
-                type Lower = #lower;
+                type Lower = #lower <#lower_generic_args>;
 
                 const SIZE32: usize = {
                     let mut size = 0;
@@ -381,4 +664,159 @@ impl Expander for ComponentTypeExpander {
 
         Ok(quote!(const _: () = { #expanded };))
     }
+
+    fn expand_variant(
+        &self,
+        input: &DeriveInput,
+        cases: &[VariantCase],
+        style: VariantStyle,
+    ) -> Result<TokenStream> {
+        let internal = quote!(wasmtime::component::__internal);
+
+        let mut case_names_and_checks = TokenStream::new();
+        let mut lower_payload_generic_params = TokenStream::new();
+        let mut lower_payload_generic_args = TokenStream::new();
+        let mut lower_payload_case_declarations = TokenStream::new();
+        let mut lower_generic_args = TokenStream::new();
+        let mut sizes = TokenStream::new();
+        let mut unique_types = HashSet::new();
+
+        for (index, VariantCase { attrs, ident, ty }) in cases.iter().enumerate() {
+            let rename = find_rename(attrs)?;
+
+            if let (Some(_), VariantStyle::Union) = (&rename, style) {
+                return Err(Error::new(
+                    ident.span(),
+                    "renaming `union` cases is not permitted; only the type is used",
+                ));
+            }
+
+            let name = rename.unwrap_or_else(|| Literal::string(&ident.to_string()));
+
+            if let Some(ty) = ty {
+                sizes.extend({
+                    let size = quote!(<#ty as wasmtime::component::ComponentType>::SIZE32);
+                    quote!(if #size > size {
+                        size = #size;
+                    })
+                });
+
+                case_names_and_checks.extend(match style {
+                    VariantStyle::Variant => {
+                        quote!((#name, <#ty as wasmtime::component::ComponentType>::typecheck),)
+                    }
+                    VariantStyle::Union => {
+                        quote!(<#ty as wasmtime::component::ComponentType>::typecheck,)
+                    }
+                    VariantStyle::Enum => {
+                        return Err(Error::new(
+                            ident.span(),
+                            "payloads are not permitted for `enum` cases",
+                        ))
+                    }
+                });
+
+                let generic = format_ident!("T{}", index);
+
+                lower_payload_generic_params.extend(quote!(#generic: Copy,));
+                lower_payload_generic_args.extend(quote!(#generic,));
+                lower_payload_case_declarations.extend(quote!(#ident: #generic,));
+                lower_generic_args
+                    .extend(quote!(<#ty as wasmtime::component::ComponentType>::Lower,));
+
+                unique_types.insert(ty);
+            } else {
+                case_names_and_checks.extend(match style {
+                    VariantStyle::Variant => {
+                        quote!((#name, <() as wasmtime::component::ComponentType>::typecheck),)
+                    }
+                    VariantStyle::Union => {
+                        quote!(<() as wasmtime::component::ComponentType>::typecheck,)
+                    }
+                    VariantStyle::Enum => quote!(#name,),
+                });
+            }
+        }
+
+        if lower_payload_case_declarations.is_empty() {
+            lower_payload_case_declarations.extend(quote!(_dummy: ()));
+        }
+
+        let alignments = unique_types
+            .into_iter()
+            .map(|ty| {
+                let align = quote!(<#ty as wasmtime::component::ComponentType>::ALIGN32);
+                quote!(if #align > align {
+                    align = #align;
+                })
+            })
+            .collect::<TokenStream>();
+
+        let typecheck = match style {
+            VariantStyle::Variant => quote!(typecheck_variant),
+            VariantStyle::Union => quote!(typecheck_union),
+            VariantStyle::Enum => quote!(typecheck_enum),
+        };
+
+        let name = &input.ident;
+        let generics = add_trait_bounds(
+            &input.generics,
+            parse_quote!(wasmtime::component::ComponentType),
+        );
+        let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
+        let lower = format_ident!("Lower{}", name);
+        let lower_payload = format_ident!("LowerPayload{}", name);
+
+        // You may wonder why we make the types of all the fields of the #lower struct and #lower_payload union
+        // generic.  This is to work around a [normalization bug in
+        // rustc](https://github.com/rust-lang/rust/issues/90903) such that the compiler does not understand that
+        // e.g. `<i32 as ComponentType>::Lower` is `Copy` despite the bound specified in `ComponentType`'s
+        // definition.
+        //
+        // See also the comment in `Self::expand_record` above for another reason why we do this.
+
+        let expanded = quote! {
+            #[doc(hidden)]
+            #[derive(Clone, Copy)]
+            #[repr(C)]
+            pub struct #lower<#lower_payload_generic_params> {
+                tag: wasmtime::ValRaw,
+                payload: #lower_payload<#lower_payload_generic_args>
+            }
+
+            #[doc(hidden)]
+            #[allow(non_snake_case)]
+            #[derive(Clone, Copy)]
+            #[repr(C)]
+            union #lower_payload<#lower_payload_generic_params> {
+                #lower_payload_case_declarations
+            }
+
+            unsafe impl #impl_generics wasmtime::component::ComponentType for #name #ty_generics #where_clause {
+                type Lower = #lower<#lower_generic_args>;
+
+                #[inline]
+                fn typecheck(
+                    ty: &#internal::InterfaceType,
+                    types: &#internal::ComponentTypes,
+                ) -> #internal::anyhow::Result<()> {
+                    #internal::#typecheck(ty, types, &[#case_names_and_checks])
+                }
+
+                const SIZE32: usize = {
+                    let mut size = 0;
+                    #sizes
+                    #internal::align_to(1, Self::ALIGN32) + size
+                };
+
+                const ALIGN32: u32 = {
+                    let mut align = 1;
+                    #alignments
+                    align
+                };
+            }
+        };
+
+        Ok(quote!(const _: () = { #expanded };))
+    }
 }

crates/wasmtime/src/component/func/typed.rs

@@ -1492,6 +1492,95 @@ pub fn typecheck_record(
     }
 }
 
+/// Verify that the given wasm type is a variant with the expected cases in the right order and with the right
+/// names.
+pub fn typecheck_variant(
+    ty: &InterfaceType,
+    types: &ComponentTypes,
+    expected: &[(&str, fn(&InterfaceType, &ComponentTypes) -> Result<()>)],
+) -> Result<()> {
+    match ty {
+        InterfaceType::Variant(index) => {
+            let cases = &types[*index].cases;
+
+            if cases.len() != expected.len() {
+                bail!(
+                    "expected variant of {} cases, found {} cases",
+                    expected.len(),
+                    cases.len()
+                );
+            }
+
+            for (case, &(name, check)) in cases.iter().zip(expected) {
+                check(&case.ty, types)
+                    .with_context(|| format!("type mismatch for case {}", name))?;
+
+                if case.name != name {
+                    bail!("expected variant case named {}, found {}", name, case.name);
+                }
+            }
+
+            Ok(())
+        }
+        other => bail!("expected `variant` found `{}`", desc(other)),
+    }
+}
+
+/// Verify that the given wasm type is a enum with the expected cases in the right order and with the right
+/// names.
+pub fn typecheck_enum(ty: &InterfaceType, types: &ComponentTypes, expected: &[&str]) -> Result<()> {
+    match ty {
+        InterfaceType::Enum(index) => {
+            let names = &types[*index].names;
+
+            if names.len() != expected.len() {
+                bail!(
+                    "expected enum of {} names, found {} names",
+                    expected.len(),
+                    names.len()
+                );
+            }
+
+            for (name, expected) in names.iter().zip(expected) {
+                if name != expected {
+                    bail!("expected enum case named {}, found {}", expected, name);
+                }
+            }
+
+            Ok(())
+        }
+        other => bail!("expected `enum` found `{}`", desc(other)),
+    }
+}
+
+/// Verify that the given wasm type is a union with the expected cases in the right order.
+pub fn typecheck_union(
+    ty: &InterfaceType,
+    types: &ComponentTypes,
+    expected: &[fn(&InterfaceType, &ComponentTypes) -> Result<()>],
+) -> Result<()> {
+    match ty {
+        InterfaceType::Union(index) => {
+            let union_types = &types[*index].types;
+
+            if union_types.len() != expected.len() {
+                bail!(
+                    "expected union of {} types, found {} types",
+                    expected.len(),
+                    union_types.len()
+                );
+            }
+
+            for (index, (ty, check)) in union_types.iter().zip(expected).enumerate() {
+                check(ty, types).with_context(|| format!("type mismatch for case {}", index))?;
+            }
+
+            Ok(())
+        }
+        other => bail!("expected `union` found `{}`", desc(other)),
+    }
+}
+
 unsafe impl<T> ComponentType for Option<T>
 where
     T: ComponentType,

crates/wasmtime/src/component/mod.rs

@@ -24,7 +24,8 @@ pub use wasmtime_component_macro::{ComponentType, Lift, Lower};
 #[doc(hidden)]
 pub mod __internal {
     pub use super::func::{
-        align_to, next_field, typecheck_record, MaybeUninitExt, Memory, MemoryMut, Options,
+        align_to, next_field, typecheck_enum, typecheck_record, typecheck_union, typecheck_variant,
+        MaybeUninitExt, Memory, MemoryMut, Options,
     };
     pub use crate::map_maybe_uninit;
     pub use crate::store::StoreOpaque;

tests/all/component_model/macros.rs

@@ -81,7 +81,6 @@ fn record_derive() -> Result<()> {
 
     let engine = super::engine();
     let mut store = Store::new(&engine, ());
-    let input = Foo { a: -42, b: 73 };
 
     // Happy path: component type matches field count, names, and types
 
@@ -94,6 +93,7 @@ fn record_derive() -> Result<()> {
     )?;
     let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
 
+    let input = Foo { a: -42, b: 73 };
     let output = instance
         .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")?
         .call_and_post_return(&mut store, (input,))?;
@@ -150,6 +150,327 @@ fn record_derive() -> Result<()> {
     )?;
     let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
 
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Happy path redux, with generics this time
+
+    #[derive(ComponentType, Lift, Lower, PartialEq, Eq, Debug, Copy, Clone)]
+    #[component(record)]
+    struct Generic<A, B> {
+        #[component(name = "foo-bar-baz")]
+        a: A,
+        b: B,
+    }
+
+    let input = Generic {
+        a: -43_i32,
+        b: 74_u32,
+    };
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(
+            r#"(type $Foo (record (field "foo-bar-baz" s32) (field "b" u32)))"#,
+            8,
+        ),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    let output = instance
+        .get_typed_func::<(Generic<i32, u32>,), Generic<i32, u32>, _>(&mut store, "echo")?
+        .call_and_post_return(&mut store, (input,))?;
+
+    assert_eq!(input, output);
+
+    Ok(())
+}
+
+#[test]
+fn union_derive() -> Result<()> {
+    #[derive(ComponentType, Lift, Lower, PartialEq, Debug, Copy, Clone)]
+    #[component(union)]
+    enum Foo {
+        A(i32),
+        B(u32),
+        C(i32),
+    }
+
+    let engine = super::engine();
+    let mut store = Store::new(&engine, ());
+
+    // Happy path: component type matches case count and types
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (union s32 u32 s32))"#, 8),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let func = instance.get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")?;
+
+    for &input in &[Foo::A(-42), Foo::B(73), Foo::C(314159265)] {
+        let output = func.call_and_post_return(&mut store, (input,))?;
+
+        assert_eq!(input, output);
+    }
+
+    // Sad path: case count mismatch (too few)
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (union s32 u32))"#, 8),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Sad path: case count mismatch (too many)
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (union s32 u32 s32 s32))"#, 8),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Sad path: case type mismatch
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (union s32 s32 s32))"#, 8),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Happy path redux, with generics this time
+
+    #[derive(ComponentType, Lift, Lower, PartialEq, Debug, Copy, Clone)]
+    #[component(union)]
+    enum Generic<A, B, C> {
+        A(A),
+        B(B),
+        C(C),
+    }
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (union s32 u32 s32))"#, 8),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let func = instance.get_typed_func::<(Generic<i32, u32, i32>,), Generic<i32, u32, i32>, _>(
+        &mut store, "echo",
+    )?;
+
+    for &input in &[
+        Generic::<i32, u32, i32>::A(-42),
+        Generic::B(73),
+        Generic::C(314159265),
+    ] {
+        let output = func.call_and_post_return(&mut store, (input,))?;
+
+        assert_eq!(input, output);
+    }
+
+    Ok(())
+}
+
+#[test]
+fn variant_derive() -> Result<()> {
+    #[derive(ComponentType, Lift, Lower, PartialEq, Eq, Debug, Copy, Clone)]
+    #[component(variant)]
+    enum Foo {
+        #[component(name = "foo-bar-baz")]
+        A(i32),
+        B(u32),
+        C,
+    }
+
+    let engine = super::engine();
+    let mut store = Store::new(&engine, ());
+
+    // Happy path: component type matches case count, names, and types
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(
+            r#"(type $Foo (variant (case "foo-bar-baz" s32) (case "B" u32) (case "C" unit)))"#,
+            8,
+        ),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let func = instance.get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")?;
+
+    for &input in &[Foo::A(-42), Foo::B(73), Foo::C] {
+        let output = func.call_and_post_return(&mut store, (input,))?;
+
+        assert_eq!(input, output);
+    }
+
+    // Sad path: case count mismatch (too few)
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(
+            r#"(type $Foo (variant (case "foo-bar-baz" s32) (case "B" u32)))"#,
+            8,
+        ),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Sad path: case count mismatch (too many)
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(
+            r#"(type $Foo (variant (case "foo-bar-baz" s32) (case "B" u32) (case "C" unit) (case "D" u32)))"#,
+            8,
+        ),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Sad path: case name mismatch
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(
+            r#"(type $Foo (variant (case "A" s32) (case "B" u32) (case "C" unit)))"#,
+            8,
+        ),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Sad path: case type mismatch
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(
+            r#"(type $Foo (variant (case "foo-bar-baz" s32) (case "B" s32) (case "C" unit)))"#,
+            8,
+        ),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Happy path redux, with generics this time
+
+    #[derive(ComponentType, Lift, Lower, PartialEq, Eq, Debug, Copy, Clone)]
+    #[component(variant)]
+    enum Generic<A, B> {
+        #[component(name = "foo-bar-baz")]
+        A(A),
+        B(B),
+        C,
+    }
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(
+            r#"(type $Foo (variant (case "foo-bar-baz" s32) (case "B" u32) (case "C" unit)))"#,
+            8,
+        ),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let func = instance
+        .get_typed_func::<(Generic<i32, u32>,), Generic<i32, u32>, _>(&mut store, "echo")?;
+
+    for &input in &[Generic::<i32, u32>::A(-42), Generic::B(73), Generic::C] {
+        let output = func.call_and_post_return(&mut store, (input,))?;
+
+        assert_eq!(input, output);
+    }
+
+    Ok(())
+}
+
+#[test]
+fn enum_derive() -> Result<()> {
+    #[derive(ComponentType, Lift, Lower, PartialEq, Eq, Debug, Copy, Clone)]
+    #[component(enum)]
+    enum Foo {
+        #[component(name = "foo-bar-baz")]
+        A,
+        B,
+        C,
+    }
+
+    let engine = super::engine();
+    let mut store = Store::new(&engine, ());
+
+    // Happy path: component type matches case count and names
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (enum "foo-bar-baz" "B" "C"))"#, 4),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+    let func = instance.get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")?;
+
+    for &input in &[Foo::A, Foo::B, Foo::C] {
+        let output = func.call_and_post_return(&mut store, (input,))?;
+
+        assert_eq!(input, output);
+    }
+
+    // Sad path: case count mismatch (too few)
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (enum "foo-bar-baz" "B"))"#, 4),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Sad path: case count mismatch (too many)
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (enum "foo-bar-baz" "B" "C" "D"))"#, 4),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
+    assert!(instance
+        .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
+        .is_err());
+
+    // Sad path: case name mismatch
+
+    let component = Component::new(
+        &engine,
+        make_echo_component(r#"(type $Foo (enum "A" "B" "C"))"#, 4),
+    )?;
+    let instance = Linker::new(&engine).instantiate(&mut store, &component)?;
+
     assert!(instance
         .get_typed_func::<(Foo,), Foo, _>(&mut store, "echo")
         .is_err());