Refactor some ABI bits in wit-bindgen-core (#1241)

Pass some variables around instead of storing them on `Generator`,
removing the need to synthesize dummy values in a few contexts.

Author: alexcrichton

crates/c/src/lib.rs

@@ -1877,7 +1877,7 @@ impl InterfaceGenerator<'_> {
 
             let mut f = FunctionBindgen::new(self, c_sig, &import_name);
             f.params = params;
-            abi::post_return(f.gen.resolve, func, &mut f, false);
+            abi::post_return(f.gen.resolve, func, &mut f);
             let FunctionBindgen { src, .. } = f;
             self.src.c_fns(&src);
             self.src.c_fns("}\n");

crates/core/src/abi.rs

@@ -738,7 +738,7 @@ pub fn call(
     bindgen: &mut impl Bindgen,
     async_: bool,
 ) {
-    Generator::new(resolve, variant, lift_lower, bindgen, async_).call(func);
+    Generator::new(resolve, bindgen).call(func, variant, lift_lower, async_);
 }
 
 pub fn lower_to_memory<B: Bindgen>(
@@ -748,14 +748,12 @@ pub fn lower_to_memory<B: Bindgen>(
     value: B::Operand,
     ty: &Type,
 ) {
-    // TODO: refactor so we don't need to pass in a bunch of unused dummy parameters:
-    let mut generator = Generator::new(
-        resolve,
-        AbiVariant::GuestImport,
-        LiftLower::LowerArgsLiftResults,
-        bindgen,
-        true,
-    );
+    let mut generator = Generator::new(resolve, bindgen);
+    // TODO: make this configurable? Right this this function is only called for
+    // future/stream callbacks so it's appropriate to skip realloc here as it's
+    // all "lower for wasm import", but this might get reused for something else
+    // in the future.
+    generator.realloc = Some(Realloc::None);
     generator.stack.push(value);
     generator.write_to_memory(ty, address, Default::default());
 }
@@ -766,14 +764,7 @@ pub fn lift_from_memory<B: Bindgen>(
     address: B::Operand,
     ty: &Type,
 ) -> B::Operand {
-    // TODO: refactor so we don't need to pass in a bunch of unused dummy parameters:
-    let mut generator = Generator::new(
-        resolve,
-        AbiVariant::GuestImport,
-        LiftLower::LowerArgsLiftResults,
-        bindgen,
-        true,
-    );
+    let mut generator = Generator::new(resolve, bindgen);
     generator.read_from_memory(ty, address, Default::default());
     generator.stack.pop().unwrap()
 }
@@ -784,16 +775,10 @@ pub fn lift_from_memory<B: Bindgen>(
 /// This is only intended to be used in guest generators for exported
 /// functions and will primarily generate `GuestDeallocate*` instructions,
 /// plus others used as input to those instructions.
-pub fn post_return(resolve: &Resolve, func: &Function, bindgen: &mut impl Bindgen, async_: bool) {
-    Generator::new(
-        resolve,
-        AbiVariant::GuestExport,
-        LiftLower::LiftArgsLowerResults,
-        bindgen,
-        async_,
-    )
-    .post_return(func);
+pub fn post_return(resolve: &Resolve, func: &Function, bindgen: &mut impl Bindgen) {
+    Generator::new(resolve, bindgen).post_return(func);
 }
+
 /// Returns whether the `Function` specified needs a post-return function to
 /// be generated in guest code.
 ///
@@ -846,47 +831,54 @@ fn needs_post_return(resolve: &Resolve, ty: &Type) -> bool {
     }
 }
 
+#[derive(Copy, Clone)]
+pub enum Realloc {
+    None,
+    Export(&'static str),
+}
+
 struct Generator<'a, B: Bindgen> {
-    variant: AbiVariant,
-    lift_lower: LiftLower,
     bindgen: &'a mut B,
-    async_: bool,
     resolve: &'a Resolve,
     operands: Vec<B::Operand>,
     results: Vec<B::Operand>,
     stack: Vec<B::Operand>,
     return_pointer: Option<B::Operand>,
+    realloc: Option<Realloc>,
 }
 
 impl<'a, B: Bindgen> Generator<'a, B> {
-    fn new(
-        resolve: &'a Resolve,
-        variant: AbiVariant,
-        lift_lower: LiftLower,
-        bindgen: &'a mut B,
-        async_: bool,
-    ) -> Generator<'a, B> {
+    fn new(resolve: &'a Resolve, bindgen: &'a mut B) -> Generator<'a, B> {
         Generator {
             resolve,
-            variant,
-            lift_lower,
             bindgen,
-            async_,
             operands: Vec::new(),
             results: Vec::new(),
             stack: Vec::new(),
             return_pointer: None,
+            realloc: None,
         }
     }
 
-    fn call(&mut self, func: &Function) {
+    fn call(&mut self, func: &Function, variant: AbiVariant, lift_lower: LiftLower, async_: bool) {
         const MAX_FLAT_PARAMS: usize = 16;
 
-        let sig = self.resolve.wasm_signature(self.variant, func);
+        let sig = self.resolve.wasm_signature(variant, func);
+
+        // Lowering parameters calling a wasm import _or_ returning a result
+        // from an async-lifted wasm export means we don't need to pass
+        // ownership, but we pass ownership in all other cases.
+        let realloc = match (variant, lift_lower, async_) {
+            (AbiVariant::GuestImport, LiftLower::LowerArgsLiftResults, _)
+            | (AbiVariant::GuestExport, LiftLower::LiftArgsLowerResults, true) => Realloc::None,
+            _ => Realloc::Export("cabi_realloc"),
+        };
+        assert!(self.realloc.is_none());
 
-        match self.lift_lower {
+        match lift_lower {
             LiftLower::LowerArgsLiftResults => {
-                if let (AbiVariant::GuestExport, true) = (self.variant, self.async_) {
+                self.realloc = Some(realloc);
+                if let (AbiVariant::GuestExport, true) = (variant, async_) {
                     unimplemented!("host-side code generation for async lift/lower not supported");
                 }
 
@@ -902,7 +894,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                     self_.stack.push(ptr);
                 };
 
-                if self.async_ {
+                if async_ {
                     let ElementInfo { size, align } = self
                         .bindgen
                         .sizes()
@@ -926,7 +918,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                             .bindgen
                             .sizes()
                             .record(func.params.iter().map(|t| &t.1));
-                        let ptr = match self.variant {
+                        let ptr = match variant {
                             // When a wasm module calls an import it will provide
                             // space that isn't explicitly deallocated.
                             AbiVariant::GuestImport => self.bindgen.return_pointer(size, align),
@@ -949,8 +941,9 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                         lower_to_memory(self, ptr);
                     }
                 }
+                self.realloc = None;
 
-                if self.async_ {
+                if async_ {
                     let ElementInfo { size, align } =
                         self.bindgen.sizes().record(func.result.iter());
                     let ptr = self.bindgen.return_pointer(size, align);
@@ -964,7 +957,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                 } else {
                     // If necessary we may need to prepare a return pointer for
                     // this ABI.
-                    if self.variant == AbiVariant::GuestImport && sig.retptr {
+                    if variant == AbiVariant::GuestImport && sig.retptr {
                         let info = self.bindgen.sizes().params(&func.result);
                         let ptr = self.bindgen.return_pointer(info.size, info.align);
                         self.return_pointer = Some(ptr.clone());
@@ -978,22 +971,22 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                     });
                 }
 
-                if !(sig.retptr || self.async_) {
+                if !(sig.retptr || async_) {
                     // With no return pointer in use we can simply lift the
                     // result(s) of the function from the result of the core
                     // wasm function.
                     if let Some(ty) = &func.result {
                         self.lift(ty)
                     }
                 } else {
-                    let ptr = match self.variant {
+                    let ptr = match variant {
                         // imports into guests means it's a wasm module
                         // calling an imported function. We supplied the
                         // return pointer as the last argument (saved in
                         // `self.return_pointer`) so we use that to read
                         // the result of the function from memory.
                         AbiVariant::GuestImport => {
-                            assert!(sig.results.is_empty() || self.async_);
+                            assert!(sig.results.is_empty() || async_);
                             self.return_pointer.take().unwrap()
                         }
 
@@ -1025,7 +1018,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                 });
             }
             LiftLower::LiftArgsLowerResults => {
-                if let (AbiVariant::GuestImport, true) = (self.variant, self.async_) {
+                if let (AbiVariant::GuestImport, true) = (variant, async_) {
                     todo!("implement host-side support for async lift/lower");
                 }
 
@@ -1065,10 +1058,10 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                 // ... and that allows us to call the interface types function
                 self.emit(&Instruction::CallInterface {
                     func,
-                    async_: self.async_,
+                    async_: async_,
                 });
 
-                let (lower_to_memory, async_results) = if self.async_ {
+                let (lower_to_memory, async_results) = if async_ {
                     self.emit(&Instruction::AsyncPostCallInterface { func });
 
                     let mut results = Vec::new();
@@ -1083,8 +1076,8 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                 // This was dynamically allocated by the caller (or async start
                 // function) so after it's been read by the guest we need to
                 // deallocate it.
-                if let AbiVariant::GuestExport = self.variant {
-                    if sig.indirect_params && !self.async_ {
+                if let AbiVariant::GuestExport = variant {
+                    if sig.indirect_params && !async_ {
                         let ElementInfo { size, align } = self
                             .bindgen
                             .sizes()
@@ -1094,14 +1087,16 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                     }
                 }
 
+                self.realloc = Some(realloc);
+
                 if !lower_to_memory {
                     // With no return pointer in use we simply lower the
                     // result(s) and return that directly from the function.
                     if let Some(ty) = &func.result {
                         self.lower(ty);
                     }
                 } else {
-                    match self.variant {
+                    match variant {
                         // When a function is imported to a guest this means
                         // it's a host providing the implementation of the
                         // import. The result is stored in the pointer
@@ -1159,9 +1154,12 @@ impl<'a, B: Bindgen> Generator<'a, B> {
                         amt: sig.results.len(),
                     });
                 }
+                self.realloc = None;
             }
         }
 
+        assert!(self.realloc.is_none());
+
         assert!(
             self.stack.is_empty(),
             "stack has {} items remaining",
@@ -1170,7 +1168,7 @@ impl<'a, B: Bindgen> Generator<'a, B> {
     }
 
     fn post_return(&mut self, func: &Function) {
-        let sig = self.resolve.wasm_signature(self.variant, func);
+        let sig = self.resolve.wasm_signature(AbiVariant::GuestExport, func);
 
         // Currently post-return is only used for lists and lists are always
         // returned indirectly through memory due to their flat representation
@@ -1424,13 +1422,9 @@ impl<'a, B: Bindgen> Generator<'a, B> {
     }
 
     fn list_realloc(&self) -> Option<&'static str> {
-        // Lowering parameters calling a wasm import _or_ returning a result
-        // from an async-lifted wasm export means we don't need to pass
-        // ownership, but we pass ownership in all other cases.
-        match (self.variant, self.lift_lower, self.async_) {
-            (AbiVariant::GuestImport, LiftLower::LowerArgsLiftResults, _)
-            | (AbiVariant::GuestExport, LiftLower::LiftArgsLowerResults, true) => None,
-            _ => Some("cabi_realloc"),
+        match self.realloc.expect("realloc should be configured") {
+            Realloc::None => None,
+            Realloc::Export(s) => Some(s),
         }
     }
 

crates/csharp/src/interface.rs

@@ -495,7 +495,7 @@ impl InterfaceGenerator<'_> {
                 ParameterType::ABI,
             );
 
-            abi::post_return(bindgen.interface_gen.resolve, func, &mut bindgen, false);
+            abi::post_return(bindgen.interface_gen.resolve, func, &mut bindgen);
 
             let src = bindgen.src;
 

crates/moonbit/src/lib.rs

@@ -969,7 +969,7 @@ impl InterfaceGenerator<'_> {
                 (0..sig.results.len()).map(|i| format!("p{i}")).collect(),
             );
 
-            abi::post_return(bindgen.gen.resolve, func, &mut bindgen, false);
+            abi::post_return(bindgen.gen.resolve, func, &mut bindgen);
 
             let src = bindgen.src;
 

crates/rust/src/interface.rs

@@ -1065,7 +1065,7 @@ pub mod vtable{ordinal} {{
             self.src.push_str("{ unsafe {\n");
 
             let mut f = FunctionBindgen::new(self, params, async_, self.wasm_import_module, false);
-            abi::post_return(f.r#gen.resolve, func, &mut f, async_);
+            abi::post_return(f.r#gen.resolve, func, &mut f);
             let FunctionBindgen {
                 needs_cleanup_list,
                 src,