Fix references to inlined types in generated rust code for stubs

wasm-rpc-stubgen/src/lib.rs

@@ -82,7 +82,7 @@ pub struct GenerateArgs {
     /// Always inline all the data types defined in the source WIT instead of copying and depending on
     /// it from the stub WIT. This is useful for example with ComponentizeJS currently where otherwise
     /// the original component's interface would be added as an import to the final WASM.
-    #[clap(long, default_value = "false")]
+    #[clap(long, default_value_t = false)]
     pub always_inline_types: bool,
 }
 
@@ -117,7 +117,7 @@ pub struct BuildArgs {
     /// Always inline all the data types defined in the source WIT instead of copying and depending on
     /// it from the stub WIT. This is useful for example with ComponentizeJS currently where otherwise
     /// the original component's interface would be added as an import to the final WASM.
-    #[clap(long, default_value = "false")]
+    #[clap(long, default_value_t = false)]
     pub always_inline_types: bool,
 }
 
@@ -185,6 +185,7 @@ pub fn generate(args: GenerateArgs) -> anyhow::Result<()> {
         &args.world,
         &args.stub_crate_version,
         &args.wasm_rpc_path_override,
+        args.always_inline_types
     )
     .context("Failed to gather information for the stub generator. Make sure source_wit_root has a valid WIT file.")?;
 
@@ -214,6 +215,7 @@ pub async fn build(args: BuildArgs) -> anyhow::Result<()> {
         &args.world,
         &args.stub_crate_version,
         &args.wasm_rpc_path_override,
+        args.always_inline_types,
     )
     .context("Failed to gather information for the stub generator")?;
 
@@ -330,7 +332,8 @@ pub fn add_stub_dependency(args: AddStubDependencyArgs) -> anyhow::Result<()> {
             stub_root,
             &Some(world_name),
             "0.0.1", // Version is unused when it comes to re-generating stub at this stage.
-            &None, // wasm-rpc path is is unused when it comes to re-generating stub during dependency addition
+            &None, // wasm-rpc path is unused when it comes to re-generating stub during dependency addition
+            true,
         )?;
 
         // We filter the dependencies of stub that's already existing in dest_wit_root

wasm-rpc-stubgen/src/rust.rs

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::stub::{FunctionResultStub, FunctionStub, StubDefinition};
+use crate::stub::{FunctionResultStub, FunctionStub, StubDefinition, StubTypeOwner};
 use anyhow::anyhow;
 use heck::{ToShoutySnakeCase, ToSnakeCase, ToUpperCamelCase};
 use proc_macro2::{Ident, Span, TokenStream};
@@ -552,8 +552,9 @@ fn type_to_rust_ident(typ: &Type, def: &StubDefinition) -> anyhow::Result<TokenS
                     let mut path = Vec::new();
                     path.push(quote! { crate });
                     path.push(quote! { bindings });
-                    match &typedef.owner {
-                        TypeOwner::World(world_id) => {
+                    let fixed_owner = def.fix_inlined_owner(typedef);
+                    match &fixed_owner {
+                        StubTypeOwner::Source(TypeOwner::World(world_id)) => {
                             let world = def
                                 .resolve
                                 .worlds
@@ -575,7 +576,7 @@ fn type_to_rust_ident(typ: &Type, def: &StubDefinition) -> anyhow::Result<TokenS
                             path.push(quote! { #ns_ident });
                             path.push(quote! { #name_ident });
                         }
-                        TypeOwner::Interface(interface_id) => {
+                        StubTypeOwner::Source(TypeOwner::Interface(interface_id)) => {
                             let interface = def
                                 .resolve
                                 .interfaces
@@ -606,7 +607,27 @@ fn type_to_rust_ident(typ: &Type, def: &StubDefinition) -> anyhow::Result<TokenS
                             path.push(quote! { #name_ident });
                             path.push(quote! { #interface_ident });
                         }
-                        TypeOwner::None => {}
+                        StubTypeOwner::Source(TypeOwner::None) => {}
+                        StubTypeOwner::StubInterface => {
+                            let root_ns = Ident::new(
+                                &def.root_package_name.namespace.to_snake_case(),
+                                Span::call_site(),
+                            );
+                            let root_name = Ident::new(
+                                &format!("{}_stub", def.root_package_name.name.to_snake_case()),
+                                Span::call_site(),
+                            );
+                            let stub_interface_name = format!("stub-{}", def.source_world_name()?);
+                            let stub_interface_name = Ident::new(
+                                &to_rust_ident(&stub_interface_name).to_snake_case(),
+                                Span::call_site(),
+                            );
+
+                            path.push(quote! { exports });
+                            path.push(quote! { #root_ns });
+                            path.push(quote! { #root_name });
+                            path.push(quote! { #stub_interface_name });
+                        }
                     }
                     Ok(quote! { #(#path)::*::#typ })
                 }

wasm-rpc-stubgen/src/stub.rs

@@ -34,6 +34,7 @@ pub struct StubDefinition {
     pub unresolved_root: UnresolvedPackage,
     pub unresolved_deps: Vec<UnresolvedPackage>,
     pub wasm_rpc_path_override: Option<String>,
+    pub always_inline_types: bool,
 }
 
 impl StubDefinition {
@@ -43,6 +44,7 @@ impl StubDefinition {
         selected_world: &Option<String>,
         stub_crate_version: &str,
         wasm_rpc_path_override: &Option<String>,
+        always_inline_types: bool,
     ) -> anyhow::Result<Self> {
         let (root, deps) = get_unresolved_packages(source_wit_root)?;
         let root_package = root.name.clone();
@@ -71,6 +73,7 @@ impl StubDefinition {
             unresolved_root: root,
             unresolved_deps: deps,
             wasm_rpc_path_override: wasm_rpc_path_override.clone(),
+            always_inline_types,
         })
     }
 
@@ -119,6 +122,40 @@ impl StubDefinition {
         final_resolve.push(final_root.clone())?;
         Ok(())
     }
+
+    pub fn fix_inlined_owner(&self, typedef: &TypeDef) -> StubTypeOwner {
+        if self.is_inlined(typedef) {
+            StubTypeOwner::StubInterface
+        } else {
+            StubTypeOwner::Source(typedef.owner)
+        }
+    }
+
+    fn is_inlined(&self, typedef: &TypeDef) -> bool {
+        match &typedef.owner {
+            TypeOwner::Interface(interface_id) => {
+                if self.always_inline_types {
+                    if let Some(resolved_owner_interface) =
+                        self.resolve.interfaces.get(*interface_id)
+                    {
+                        if let Some(name) = resolved_owner_interface.name.as_ref() {
+                            self.interfaces
+                                .iter()
+                                .any(|interface| &interface.name == name)
+                        } else {
+                            false
+                        }
+                    } else {
+                        false
+                    }
+                } else {
+                    false
+                }
+            }
+            TypeOwner::World(_) => true,
+            TypeOwner::None => false,
+        }
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -252,6 +289,12 @@ impl FunctionResultStub {
     }
 }
 
+#[derive(Debug, Clone)]
+pub enum StubTypeOwner {
+    StubInterface,
+    Source(TypeOwner),
+}
+
 fn collect_stub_imports<'a>(
     types: impl Iterator<Item = (&'a String, &'a TypeId)>,
     resolve: &Resolve,