Implement optional --enable-method-chaining for the Rust generator.

Author: QuantumSegfault

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,
         }
     }
 
@@ -1049,7 +1052,11 @@ impl Bindgen for FunctionBindgen<'_, '_> {
             }
 
             Instruction::Return { amt, .. } => match amt {
-                0 => {}
+                0 => {
+                    if self.return_self {
+                        uwriteln!(self.src, "self");
+                    }
+                }
                 1 => {
                     self.push_str(&operands[0]);
                     self.push_str("\n");

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,
+            matches!(&func.kind, FunctionKind::Method(_)) && self.r#gen.opts.enable_method_chaining,
+        );
         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,17 @@ 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 let (FunctionKind::Method(_), None, true) = (
+                &func.kind,
+                func.result,
+                self.r#gen.opts.enable_method_chaining
+            ) {
+                ";\nself"
+            } else {
+                ""
+            }
         );
     }
 
@@ -1121,7 +1136,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 +1206,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,
@@ -1439,19 +1454,29 @@ unsafe fn call_import(&mut self, _params: Self::ParamsLower, _results: *mut u8)
     fn print_signature(&mut self, func: &Function, params_owned: bool, sig: &FnSig) -> Vec<String> {
         let params = self.print_docs_and_params(func, params_owned, sig);
         self.push_str(" -> ");
-        if let FunctionKind::Constructor(resource_id) = &func.kind {
-            match classify_constructor_return_type(&self.resolve, *resource_id, &func.result) {
-                ConstructorReturnType::Self_ => {
-                    self.push_str("Self");
+        match &func.kind {
+            FunctionKind::Constructor(resource_id) => {
+                match classify_constructor_return_type(&self.resolve, *resource_id, &func.result) {
+                    ConstructorReturnType::Self_ => {
+                        self.push_str("Self");
+                    }
+                    ConstructorReturnType::Result { err } => {
+                        self.push_str("Result<Self, ");
+                        self.print_result_type(&err);
+                        self.push_str("> where Self: Sized");
+                    }
                 }
-                ConstructorReturnType::Result { err } => {
-                    self.push_str("Result<Self, ");
-                    self.print_result_type(&err);
-                    self.push_str("> where Self: Sized");
+            }
+            FunctionKind::Method(_) => {
+                if self.r#gen.opts.enable_method_chaining && func.result.is_none() {
+                    self.push_str("&Self");
+                } else {
+                    self.print_result_type(&func.result);
                 }
             }
-        } else {
-            self.print_result_type(&func.result);
+            _ => {
+                self.print_result_type(&func.result);
+            }
         }
         params
     }

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 {

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_();
+    }
+}

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
+    }
+}

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;
+}