rust: Fix cases where with provides resources (#855)

* rust: Fix cases where `with` provides resources

This commit fixes a bug where if `with` was used to provide a resource
then the internal `resources` map wasn't updated correctly. This could
lead to "index not found" panics in #832. To fix this I've opted to
remove the `resources` map entirely as the `owned` field was already not
tracked and this additionally brings the Wasmtime and Rust generators a
bit closer together.

Instead now each `InterfaceId` has an entry of whether it was last seen
as an export or an import and this is consulted when determining whether
a resource is exported or not.

This fix then required shuffling around where the stub implementations
are generated to ensure that when the stubs for top-level exported
functions are generated the right metadata of "was this interface last
seen as an import" is generated correctly.

Closes #832

* Update crates/rust/src/lib.rs

Co-authored-by: Pat Hickey <pat@moreproductive.org>

---------

Co-authored-by: Pat Hickey <pat@moreproductive.org>

Author: alexcrichton

crates/rust/src/bindgen.rs

@@ -1,4 +1,4 @@
-use crate::{int_repr, to_rust_ident, wasm_type, Direction, InterfaceGenerator, RustFlagsRepr};
+use crate::{int_repr, to_rust_ident, wasm_type, InterfaceGenerator, RustFlagsRepr};
 use heck::*;
 use std::fmt::Write as _;
 use std::mem;
@@ -411,11 +411,12 @@ impl Bindgen for FunctionBindgen<'_, '_> {
             } => {
                 let op = &operands[0];
                 let rt = self.gen.gen.runtime_path();
-                let resource = dealias(self.gen.resolve, *resource);
-                results.push(match self.gen.gen.resources[&resource].direction {
-                    Direction::Import => format!("({op}).take_handle() as i32"),
-                    Direction::Export => format!("{rt}::Resource::take_handle(&{op}) as i32"),
-                });
+                let result = if self.gen.is_exported_resource(*resource) {
+                    format!("{rt}::Resource::take_handle(&{op}) as i32")
+                } else {
+                    format!("({op}).take_handle() as i32")
+                };
+                results.push(result);
             }
 
             Instruction::HandleLower {
@@ -435,9 +436,7 @@ impl Bindgen for FunctionBindgen<'_, '_> {
 
                 let dealiased_resource = dealias(resolve, *resource);
 
-                let result = if let Direction::Export =
-                    self.gen.gen.resources[&dealiased_resource].direction
-                {
+                let result = if self.gen.is_exported_resource(*resource) {
                     match handle {
                         Handle::Borrow(_) => {
                             let name = resolve.types[*resource]
@@ -834,7 +833,6 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                         ));
                     }
                     FunctionKind::Constructor(ty) => {
-                        self.gen.mark_resource_owned(*ty);
                         self.push_str(&format!(
                             "Own{0}::new(<_{0}Impl as Guest{0}>::new",
                             resolve.types[*ty]

crates/rust/src/interface.rs

@@ -1,7 +1,7 @@
 use crate::bindgen::FunctionBindgen;
 use crate::{
-    int_repr, to_rust_ident, to_upper_camel_case, wasm_type, Direction, ExportKey, FnSig,
-    Identifier, InterfaceName, Ownership, RustFlagsRepr, RustWasm,
+    int_repr, to_rust_ident, to_upper_camel_case, wasm_type, ExportKey, FnSig, Identifier,
+    InterfaceName, Ownership, RustFlagsRepr, RustWasm,
 };
 use anyhow::Result;
 use heck::*;
@@ -1115,7 +1115,6 @@ impl InterfaceGenerator<'_> {
             }
 
             TypeDefKind::Handle(Handle::Own(ty)) => {
-                self.mark_resource_owned(*ty);
                 self.print_ty(&Type::Id(*ty), mode);
             }
 
@@ -1596,7 +1595,6 @@ impl InterfaceGenerator<'_> {
             // aliases) is used prior to generating declarations.  For example,
             // if only borrows are used, no need to generate the `Own{name}`
             // version.
-            self.mark_resource_owned(target);
             for prefix in ["Own", ""] {
                 self.rustdoc(docs);
                 self.push_str(&format!(
@@ -1700,31 +1698,25 @@ impl InterfaceGenerator<'_> {
         self.push_str(&format!("{rt}::vec::Vec", rt = self.gen.runtime_path()));
     }
 
-    fn is_exported_resource(&self, mut ty: TypeId) -> bool {
-        loop {
-            let def = &self.resolve.types[ty];
-            if let TypeOwner::World(_) = &def.owner {
-                // Worlds cannot export types of any kind as of this writing.
-                return false;
-            }
-            match &def.kind {
-                TypeDefKind::Type(Type::Id(id)) => ty = *id,
-                _ => break,
-            }
+    pub fn is_exported_resource(&self, ty: TypeId) -> bool {
+        let ty = dealias(self.resolve, ty);
+        let ty = &self.resolve.types[ty];
+        match &ty.kind {
+            TypeDefKind::Resource => {}
+            _ => return false,
         }
 
-        matches!(
-            self.gen.resources.get(&ty).map(|info| info.direction),
-            Some(Direction::Export)
-        )
-    }
+        match ty.owner {
+            // Worlds cannot export types of any kind as of this writing.
+            TypeOwner::World(_) => false,
+
+            // Interfaces are "stateful" currently where whatever we last saw
+            // them as dictates whether it's exported or not.
+            TypeOwner::Interface(i) => !self.gen.interface_last_seen_as_import[&i],
 
-    pub fn mark_resource_owned(&mut self, resource: TypeId) {
-        self.gen
-            .resources
-            .entry(dealias(self.resolve, resource))
-            .or_default()
-            .owned = true;
+            // Shouldn't be the case for resources
+            TypeOwner::None => unreachable!(),
+        }
     }
 
     fn push_string_name(&mut self) {
@@ -1765,15 +1757,7 @@ impl<'a> wit_bindgen_core::InterfaceGenerator<'a> for InterfaceGenerator<'a> {
         self.print_typedef_record(id, record, docs);
     }
 
-    fn type_resource(&mut self, id: TypeId, name: &str, docs: &Docs) {
-        let entry = self
-            .gen
-            .resources
-            .entry(dealias(self.resolve, id))
-            .or_default();
-        if !self.in_import {
-            entry.direction = Direction::Export;
-        }
+    fn type_resource(&mut self, _id: TypeId, name: &str, docs: &Docs) {
         self.rustdoc(docs);
         let camel = to_upper_camel_case(name);
         let rt = self.gen.runtime_path();

crates/rust/src/lib.rs

@@ -9,24 +9,12 @@ use std::process::{Command, Stdio};
 use std::str::FromStr;
 use wit_bindgen_core::abi::{Bitcast, WasmType};
 use wit_bindgen_core::{
-    uwriteln, wit_parser::*, Direction, Files, InterfaceGenerator as _, Source, Types,
-    WorldGenerator,
+    uwriteln, wit_parser::*, Files, InterfaceGenerator as _, Source, Types, WorldGenerator,
 };
 
 mod bindgen;
 mod interface;
 
-#[derive(Default)]
-struct ResourceInfo {
-    // Note that a resource can be both imported and exported (e.g. when
-    // importing and exporting the same interface which contains one or more
-    // resources).  In that case, this field will be `Import` while we're
-    // importing the interface and later change to `Export` while we're
-    // exporting the interface.
-    direction: Direction,
-    owned: bool,
-}
-
 struct InterfaceName {
     /// True when this interface name has been remapped through the use of `with` in the `bindgen!`
     /// macro invocation.
@@ -45,12 +33,14 @@ struct RustWasm {
     export_modules: Vec<(String, Vec<String>)>,
     skip: HashSet<String>,
     interface_names: HashMap<InterfaceId, InterfaceName>,
-    resources: HashMap<TypeId, ResourceInfo>,
+    /// Each imported and exported interface is stored in this map. Value indicates if last use was import.
+    interface_last_seen_as_import: HashMap<InterfaceId, bool>,
     import_funcs_called: bool,
     with_name_counter: usize,
     // Track the with options that were used. Remapped interfaces provided via `with`
     // are required to be used.
     used_with_opts: HashSet<String>,
+    world: Option<WorldId>,
 }
 
 #[cfg(feature = "clap")]
@@ -360,7 +350,7 @@ fn name_package_module(resolve: &Resolve, id: PackageId) -> String {
 }
 
 impl WorldGenerator for RustWasm {
-    fn preprocess(&mut self, resolve: &Resolve, _world: WorldId) {
+    fn preprocess(&mut self, resolve: &Resolve, world: WorldId) {
         wit_bindgen_core::generated_preamble(&mut self.src, env!("CARGO_PKG_VERSION"));
 
         // Render some generator options to assist with debugging and/or to help
@@ -391,6 +381,7 @@ impl WorldGenerator for RustWasm {
             uwriteln!(self.src, "//   * with {with:?}");
         }
         self.types.analyze(resolve);
+        self.world = Some(world);
     }
 
     fn import_interface(
@@ -400,6 +391,7 @@ impl WorldGenerator for RustWasm {
         id: InterfaceId,
         _files: &mut Files,
     ) {
+        self.interface_last_seen_as_import.insert(id, true);
         let wasm_import_module = resolve.name_world_key(name);
         let mut gen = self.interface(
             Identifier::Interface(id, name),
@@ -442,6 +434,7 @@ impl WorldGenerator for RustWasm {
         id: InterfaceId,
         _files: &mut Files,
     ) -> Result<()> {
+        self.interface_last_seen_as_import.insert(id, false);
         let mut gen = self.interface(Identifier::Interface(id, name), None, resolve, false);
         let (snake, module_path) = gen.start_append_submodule(name);
         if gen.gen.name_interface(resolve, id, name, true) {
@@ -450,6 +443,31 @@ impl WorldGenerator for RustWasm {
         gen.types(id);
         gen.generate_exports(resolve.interfaces[id].functions.values())?;
         gen.finish_append_submodule(&snake, module_path);
+
+        if self.opts.stubs {
+            let (pkg, name) = match name {
+                WorldKey::Name(name) => (None, name),
+                WorldKey::Interface(id) => {
+                    let interface = &resolve.interfaces[*id];
+                    (
+                        Some(interface.package.unwrap()),
+                        interface.name.as_ref().unwrap(),
+                    )
+                }
+            };
+            for (resource, funcs) in group_by_resource(resolve.interfaces[id].functions.values()) {
+                let world_id = self.world.unwrap();
+                let mut gen = self.interface(Identifier::World(world_id), None, resolve, false);
+                let pkg = pkg.map(|pid| {
+                    let namespace = resolve.packages[pid].name.namespace.clone();
+                    let package_module = name_package_module(resolve, pid);
+                    (namespace, package_module)
+                });
+                gen.generate_stub(resource, pkg, name, true, &funcs);
+                let stub = gen.finish();
+                self.src.push_str(&stub);
+            }
+        }
         Ok(())
     }
 
@@ -464,6 +482,16 @@ impl WorldGenerator for RustWasm {
         gen.generate_exports(funcs.iter().map(|f| f.1))?;
         let src = gen.finish();
         self.src.push_str(&src);
+
+        if self.opts.stubs {
+            for (resource, funcs) in group_by_resource(funcs.iter().map(|f| f.1)) {
+                let mut gen = self.interface(Identifier::World(world), None, resolve, false);
+                let world = &resolve.worlds[world];
+                gen.generate_stub(resource, None, &world.name, false, &funcs);
+                let stub = gen.finish();
+                self.src.push_str(&stub);
+            }
+        }
         Ok(())
     }
 
@@ -541,50 +569,6 @@ impl WorldGenerator for RustWasm {
 
         if self.opts.stubs {
             self.src.push_str("\n#[derive(Debug)]\npub struct Stub;\n");
-            let world_id = world;
-            let world = &resolve.worlds[world];
-            let mut funcs = Vec::new();
-            for (name, export) in world.exports.iter() {
-                let (pkg, name) = match name {
-                    WorldKey::Name(name) => (None, name),
-                    WorldKey::Interface(id) => {
-                        let interface = &resolve.interfaces[*id];
-                        (
-                            Some(interface.package.unwrap()),
-                            interface.name.as_ref().unwrap(),
-                        )
-                    }
-                };
-                match export {
-                    WorldItem::Function(func) => {
-                        funcs.push(func);
-                    }
-                    WorldItem::Interface(id) => {
-                        for (resource, funcs) in
-                            group_by_resource(resolve.interfaces[*id].functions.values())
-                        {
-                            let mut gen =
-                                self.interface(Identifier::World(world_id), None, resolve, false);
-                            let pkg = pkg.map(|pid| {
-                                let namespace = resolve.packages[pid].name.namespace.clone();
-                                let package_module = name_package_module(resolve, pid);
-                                (namespace, package_module)
-                            });
-                            gen.generate_stub(resource, pkg, name, true, &funcs);
-                            let stub = gen.finish();
-                            self.src.push_str(&stub);
-                        }
-                    }
-                    WorldItem::Type(_) => unreachable!(),
-                }
-            }
-
-            for (resource, funcs) in group_by_resource(funcs.into_iter()) {
-                let mut gen = self.interface(Identifier::World(world_id), None, resolve, false);
-                gen.generate_stub(resource, None, &world.name, false, &funcs);
-                let stub = gen.finish();
-                self.src.push_str(&stub);
-            }
         }
 
         let mut src = mem::take(&mut self.src);

crates/rust/tests/codegen.rs

@@ -415,3 +415,45 @@ mod with {
         my::inline::bar::bar(&msg);
     }
 }
+
+mod with_and_resources {
+    wit_bindgen::generate!({
+        inline: "
+            package my:inline;
+
+            interface foo {
+                resource a;
+            }
+
+            interface bar {
+                use foo.{a};
+
+                bar: func(m: a) -> list<a>;
+            }
+
+            world baz {
+                import bar;
+            }
+        ",
+        with: {
+            "my:inline/foo": other::my::inline::foo,
+        },
+    });
+
+    pub mod other {
+        wit_bindgen::generate!({
+            inline: "
+                package my:inline;
+
+                interface foo {
+                    resource a;
+                }
+
+                world dummy {
+                    use foo.{a};
+                    import bar: func(m: a);
+                }
+            ",
+        });
+    }
+}