Add --enable-method-chaining for the Rust generator. (#1586)

QuantumSegfault · web-flow · commit be1402f49e61 · 2026-04-14T22:22:34.000Z
* Implement optional `--enable-method-chaining` for the Rust generator.

* Add `enable_method_chaining` to the `generate` macro

* Factor out `should_return_self`
diff --git a/crates/guest-rust/macro/src/lib.rs b/crates/guest-rust/macro/src/lib.rs
@@ -171,6 +171,9 @@ impl Parse for Config {
                         }
                         opts.async_ = val;
                     }
+                    Opt::EnableMethodChaining(enable) => {
+                        opts.enable_method_chaining = enable.value();
+                    }
                 }
             }
         } else {
@@ -324,6 +327,7 @@ mod kw {
     syn::custom_keyword!(disable_custom_section_link_helpers);
     syn::custom_keyword!(imports);
     syn::custom_keyword!(debug);
+    syn::custom_keyword!(enable_method_chaining);
 }
 
 #[derive(Clone)]
@@ -404,6 +408,7 @@ enum Opt {
     DisableCustomSectionLinkHelpers(syn::LitBool),
     Async(AsyncFilterSet, Span),
     Debug(syn::LitBool),
+    EnableMethodChaining(syn::LitBool),
 }
 
 impl Parse for Opt {
@@ -568,6 +573,10 @@ impl Parse for Opt {
             input.parse::<kw::debug>()?;
             input.parse::<Token![:]>()?;
             Ok(Opt::Debug(input.parse()?))
+        } else if l.peek(kw::enable_method_chaining) {
+            input.parse::<kw::enable_method_chaining>()?;
+            input.parse::<Token![:]>()?;
+            Ok(Opt::EnableMethodChaining(input.parse()?))
         } else if l.peek(Token![async]) {
             let span = input.parse::<Token![async]>()?.span;
             input.parse::<Token![:]>()?;
diff --git a/crates/guest-rust/src/lib.rs b/crates/guest-rust/src/lib.rs
@@ -860,6 +860,11 @@ extern crate std;
 ///         "-export:wasi:http/handler@0.3.0-draft#handle",
 ///         "all",
 ///     ],
+///
+///     // All resource methods with empty returns are instead generated as
+///     // returning `-> &Self`, to permit method chaining (e.g. for builder).
+///     // This expectation is also imposed on exports.
+///     enable_method_chaining: true,
 /// });
 /// ```
 ///
diff --git a/crates/rust/src/bindgen.rs b/crates/rust/src/bindgen.rs
@@ -21,6 +21,7 @@ pub(super) struct FunctionBindgen<'a, 'b> {
     pub import_return_pointer_area_align: Alignment,
     pub handle_decls: Vec<String>,
     always_owned: bool,
+    return_self: bool,
 }
 
 pub const POINTER_SIZE_EXPRESSION: &str = "::core::mem::size_of::<*const u8>()";
@@ -31,6 +32,7 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
         params: Vec<String>,
         wasm_import_module: &'b str,
         always_owned: bool,
+        return_self: bool,
     ) -> FunctionBindgen<'a, 'b> {
         FunctionBindgen {
             r#gen,
@@ -45,6 +47,7 @@ impl<'a, 'b> FunctionBindgen<'a, 'b> {
             import_return_pointer_area_align: Default::default(),
             handle_decls: Vec::new(),
             always_owned,
+            return_self,
         }
     }
 
@@ -1048,18 +1051,26 @@ impl Bindgen for FunctionBindgen<'_, '_> {
                 }
             }
 
-            Instruction::Return { amt, .. } => match amt {
-                0 => {}
-                1 => {
-                    self.push_str(&operands[0]);
-                    self.push_str("\n");
-                }
-                _ => {
-                    self.push_str("(");
-                    self.push_str(&operands.join(", "));
-                    self.push_str(")\n");
+            Instruction::Return { amt, .. } => {
+                assert!(!self.return_self || *amt == 0);
+
+                match amt {
+                    0 => {
+                        if self.return_self {
+                            self.push_str("self\n");
+                        }
+                    }
+                    1 => {
+                        self.push_str(&operands[0]);
+                        self.push_str("\n");
+                    }
+                    _ => {
+                        self.push_str("(");
+                        self.push_str(&operands.join(", "));
+                        self.push_str(")\n");
+                    }
                 }
-            },
+            }
 
             Instruction::I32Load { offset } => {
                 let tmp = self.tmp();
diff --git a/crates/rust/src/interface.rs b/crates/rust/src/interface.rs
@@ -768,7 +768,7 @@ pub mod vtable{ordinal} {{
     }
 
     fn lower_to_memory(&mut self, address: &str, value: &str, ty: &Type, module: &str) -> String {
-        let mut f = FunctionBindgen::new(self, Vec::new(), module, true);
+        let mut f = FunctionBindgen::new(self, Vec::new(), module, true, false);
         abi::lower_to_memory(f.r#gen.resolve, &mut f, address.into(), value.into(), ty);
         format!("unsafe {{ {} }}", String::from(f.src))
     }
@@ -780,7 +780,7 @@ pub mod vtable{ordinal} {{
         indirect: bool,
         module: &str,
     ) -> String {
-        let mut f = FunctionBindgen::new(self, Vec::new(), module, true);
+        let mut f = FunctionBindgen::new(self, Vec::new(), module, true, false);
         abi::deallocate_lists_in_types(f.r#gen.resolve, types, operands, indirect, &mut f);
         format!("unsafe {{ {} }}", String::from(f.src))
     }
@@ -792,13 +792,13 @@ pub mod vtable{ordinal} {{
         indirect: bool,
         module: &str,
     ) -> String {
-        let mut f = FunctionBindgen::new(self, Vec::new(), module, true);
+        let mut f = FunctionBindgen::new(self, Vec::new(), module, true, false);
         abi::deallocate_lists_and_own_in_types(f.r#gen.resolve, types, operands, indirect, &mut f);
         format!("unsafe {{ {} }}", String::from(f.src))
     }
 
     fn lift_from_memory(&mut self, address: &str, ty: &Type, module: &str) -> String {
-        let mut f = FunctionBindgen::new(self, Vec::new(), module, true);
+        let mut f = FunctionBindgen::new(self, Vec::new(), module, true, false);
         let result = abi::lift_from_memory(f.r#gen.resolve, &mut f, address.into(), ty);
         format!("unsafe {{ {}\n{result} }}", String::from(f.src))
     }
@@ -809,7 +809,13 @@ pub mod vtable{ordinal} {{
         func: &Function,
         params: Vec<String>,
     ) {
-        let mut f = FunctionBindgen::new(self, params, module, false);
+        let mut f = FunctionBindgen::new(
+            self,
+            params,
+            module,
+            false,
+            self.r#gen.should_return_self(func),
+        );
         abi::call(
             f.r#gen.resolve,
             AbiVariant::GuestImport,
@@ -1032,7 +1038,7 @@ unsafe fn call_import(&mut self, _params: Self::ParamsLower, _results: *mut u8)
             }
             lowers.push("ParamsLower(_ptr,)".to_string());
         } else {
-            let mut f = FunctionBindgen::new(self, Vec::new(), module, true);
+            let mut f = FunctionBindgen::new(self, Vec::new(), module, true, false);
             let mut results = Vec::new();
             for (i, Param { ty, .. }) in func.params.iter().enumerate() {
                 let name = format!("_lower{i}");
@@ -1078,8 +1084,13 @@ unsafe fn call_import(&mut self, _params: Self::ParamsLower, _results: *mut u8)
         }
         uwriteln!(
             self.src,
-            "_MySubtask {{ _unused: core::marker::PhantomData }}.call(({})).await",
-            params.join(" ")
+            "_MySubtask {{ _unused: core::marker::PhantomData }}.call(({})).await{}",
+            params.join(" "),
+            if self.r#gen.should_return_self(func) {
+                ";\nself"
+            } else {
+                ""
+            }
         );
     }
 
@@ -1121,7 +1132,7 @@ unsafe fn call_import(&mut self, _params: Self::ParamsLower, _results: *mut u8)
             );
         }
 
-        let mut f = FunctionBindgen::new(self, params, self.wasm_import_module, false);
+        let mut f = FunctionBindgen::new(self, params, self.wasm_import_module, false, false);
         let variant = if async_ {
             AbiVariant::GuestExportAsync
         } else {
@@ -1191,7 +1202,7 @@ unsafe fn call_import(&mut self, _params: Self::ParamsLower, _results: *mut u8)
             let params = self.print_post_return_sig(func);
             self.src.push_str("{ unsafe {\n");
 
-            let mut f = FunctionBindgen::new(self, params, self.wasm_import_module, false);
+            let mut f = FunctionBindgen::new(self, params, self.wasm_import_module, false, false);
             abi::post_return(f.r#gen.resolve, func, &mut f);
             let FunctionBindgen {
                 needs_cleanup_list,
@@ -1451,7 +1462,11 @@ unsafe fn call_import(&mut self, _params: Self::ParamsLower, _results: *mut u8)
                 }
             }
         } else {
-            self.print_result_type(&func.result);
+            if self.r#gen.should_return_self(func) {
+                self.push_str("&Self");
+            } else {
+                self.print_result_type(&func.result);
+            }
         }
         params
     }
diff --git a/crates/rust/src/lib.rs b/crates/rust/src/lib.rs
@@ -300,6 +300,10 @@ pub struct Opts {
         arg(long, require_equals = true, value_name = "true|false")
     )]
     pub merge_structurally_equal_types: Option<Option<bool>>,
+
+    /// If true, methods normally returning `()` instead return `&Self`. This applies to both imported and exported methods.
+    #[cfg_attr(feature = "clap", arg(long))]
+    pub enable_method_chaining: bool,
 }
 
 impl Opts {
@@ -1056,6 +1060,12 @@ macro_rules! __export_{world_name}_impl {{
             .async_
             .is_async(resolve, interface, func, is_import)
     }
+
+    fn should_return_self(&self, func: &Function) -> bool {
+        self.opts.enable_method_chaining
+            && func.result.is_none()
+            && matches!(&func.kind, FunctionKind::Method(_))
+    }
 }
 
 impl WorldGenerator for RustWasm {
diff --git a/crates/rust/tests/codegen.rs b/crates/rust/tests/codegen.rs
@@ -216,3 +216,22 @@ mod retyped_list {
         }
     });
 }
+
+#[allow(unused, reason = "testing codegen, not functionality")]
+mod method_chaining {
+    wit_bindgen::generate!({
+        inline: r#"
+        package test:method-chaining;
+        world test {
+            resource a {
+                constructor();
+                set-a: func(arg: u32);
+                set-b: func(arg: bool);
+                do: func();
+            }
+        }
+        "#,
+        generate_all,
+        enable_method_chaining: true
+    });
+}
diff --git a/tests/runtime/rust/method-chaining/runner.rs b/tests/runtime/rust/method-chaining/runner.rs
@@ -0,0 +1,15 @@
+//@ args = '--enable-method-chaining'
+
+include!(env!("BINDINGS"));
+
+use crate::foo::bar::i::A;
+
+struct Component;
+export!(Component);
+
+impl Guest for Component {
+    fn run() {
+        let my_a = A::new();
+        my_a.set_a(42).set_b(true).do_();
+    }
+}
diff --git a/tests/runtime/rust/method-chaining/test.rs b/tests/runtime/rust/method-chaining/test.rs
@@ -0,0 +1,40 @@
+//@ args = '--enable-method-chaining'
+
+include!(env!("BINDINGS"));
+
+use crate::exports::foo::bar::i::{Guest, GuestA};
+use std::cell::Cell;
+
+struct Component;
+export!(Component);
+impl Guest for Component {
+    type A = MyA;
+}
+
+struct MyA {
+    prop_a: Cell<u32>,
+    prop_b: Cell<bool>,
+}
+
+impl GuestA for MyA {
+    fn new() -> MyA {
+        MyA {
+            prop_a: Cell::new(0),
+            prop_b: Cell::new(false),
+        }
+    }
+
+    fn set_a(&self, a: u32) -> &Self {
+        self.prop_a.set(a);
+        self
+    }
+
+    fn set_b(&self, b: bool) -> &Self {
+        self.prop_b.set(b);
+        self
+    }
+
+    fn do_(&self) -> &Self {
+        self
+    }
+}
diff --git a/tests/runtime/rust/method-chaining/test.wit b/tests/runtime/rust/method-chaining/test.wit
@@ -0,0 +1,18 @@
+package foo:bar;
+
+interface i {
+    resource a {
+        constructor();
+        set-a: func(arg: u32);
+        set-b: func(arg: bool);
+        do: func();
+    }
+}
+world runner {
+    import i;
+
+    export run: func();
+}
+world test {
+    export i;
+}

Original file line number	Diff line number	Diff line change
`@@ -171,6 +171,9 @@ impl Parse for Config {`
`171`	`171`	`}`
`172`	`172`	`opts.async_ = val;`
`173`	`173`	`}`
	`174`	`+ Opt::EnableMethodChaining(enable) => {`
	`175`	`+ opts.enable_method_chaining = enable.value();`
	`176`	`+ }`
`174`	`177`	`}`
`175`	`178`	`}`
`176`	`179`	`} else {`
`@@ -324,6 +327,7 @@ mod kw {`
`324`	`327`	`syn::custom_keyword!(disable_custom_section_link_helpers);`
`325`	`328`	`syn::custom_keyword!(imports);`
`326`	`329`	`syn::custom_keyword!(debug);`
	`330`	`+ syn::custom_keyword!(enable_method_chaining);`
`327`	`331`	`}`
`328`	`332`
`329`	`333`	`#[derive(Clone)]`
`@@ -404,6 +408,7 @@ enum Opt {`
`404`	`408`	`DisableCustomSectionLinkHelpers(syn::LitBool),`
`405`	`409`	`Async(AsyncFilterSet, Span),`
`406`	`410`	`Debug(syn::LitBool),`
	`411`	`+ EnableMethodChaining(syn::LitBool),`
`407`	`412`	`}`
`408`	`413`
`409`	`414`	`impl Parse for Opt {`
`@@ -568,6 +573,10 @@ impl Parse for Opt {`
`568`	`573`	`input.parse::<kw::debug>()?;`
`569`	`574`	`input.parse::<Token![:]>()?;`
`570`	`575`	`Ok(Opt::Debug(input.parse()?))`
	`576`	`+ } else if l.peek(kw::enable_method_chaining) {`
	`577`	`+ input.parse::<kw::enable_method_chaining>()?;`
	`578`	`+ input.parse::<Token![:]>()?;`
	`579`	`+ Ok(Opt::EnableMethodChaining(input.parse()?))`
`571`	`580`	`} else if l.peek(Token![async]) {`
`572`	`581`	`let span = input.parse::<Token![async]>()?.span;`
`573`	`582`	`input.parse::<Token![:]>()?;`