adjust Wasmi translation for conditional return removal

Author: Robbepop

crates/wasmi/src/engine/translator/instr_encoder.rs

@@ -641,67 +641,6 @@ impl InstrEncoder {
         Ok(())
     }
 
-    /// Encodes an conditional `return` instruction.
-    pub fn encode_return_nez(
-        &mut self,
-        stack: &mut ValueStack,
-        condition: Reg,
-        values: &[TypedProvider],
-        fuel_info: FuelInfo,
-    ) -> Result<(), Error> {
-        // Note: We bump fuel unconditionally even if the conditional return is not taken.
-        //       This is very conservative and may lead to more fuel costs than
-        //       actually needed for the computation. We might revisit this decision
-        //       later. An alternative solution would consume fuel during execution
-        //       time only when the return is taken.
-        let instr = match values {
-            [] => Instruction::return_nez(condition),
-            [TypedProvider::Register(reg)] => Instruction::return_nez_reg(condition, *reg),
-            [TypedProvider::Const(value)] => match value.ty() {
-                ValType::I32 => Instruction::return_nez_imm32(condition, i32::from(*value)),
-                ValType::I64 => match <Const32<i64>>::try_from(i64::from(*value)).ok() {
-                    Some(value) => Instruction::return_nez_i64imm32(condition, value),
-                    None => Instruction::return_nez_reg(condition, stack.alloc_const(*value)?),
-                },
-                ValType::F32 => Instruction::return_nez_imm32(condition, f32::from(*value)),
-                ValType::F64 => match <Const32<f64>>::try_from(f64::from(*value)).ok() {
-                    Some(value) => Instruction::return_nez_f64imm32(condition, value),
-                    None => Instruction::return_nez_reg(condition, stack.alloc_const(*value)?),
-                },
-                ValType::V128 | ValType::FuncRef | ValType::ExternRef => {
-                    Instruction::return_nez_reg(condition, stack.alloc_const(*value)?)
-                }
-            },
-            [v0, v1] => {
-                let reg0 = stack.provider2reg(v0)?;
-                let reg1 = stack.provider2reg(v1)?;
-                Instruction::return_nez_reg2_ext(condition, reg0, reg1)
-            }
-            [v0, v1, rest @ ..] => {
-                debug_assert!(!rest.is_empty());
-                // Note: The fuel for return values might result in 0 charges if there aren't
-                //       enough return values to account for at least 1 fuel. Therefore we need
-                //       to also bump by `FuelCosts::base` to charge at least 1 fuel.
-                self.bump_fuel_consumption(&fuel_info, FuelCostsProvider::base)?;
-                self.bump_fuel_consumption(&fuel_info, |costs| {
-                    costs.fuel_for_copying_values(rest.len() as u64 + 3)
-                })?;
-                if let Some(span) = BoundedRegSpan::from_providers(values) {
-                    self.push_instr(Instruction::return_nez_span(condition, span))?;
-                    return Ok(());
-                }
-                let reg0 = stack.provider2reg(v0)?;
-                let reg1 = stack.provider2reg(v1)?;
-                self.push_instr(Instruction::return_nez_many_ext(condition, reg0, reg1))?;
-                self.encode_register_list(stack, rest)?;
-                return Ok(());
-            }
-        };
-        self.bump_fuel_consumption(&fuel_info, FuelCostsProvider::base)?;
-        self.push_instr(instr)?;
-        Ok(())
-    }
-
     /// Encode the given slice of [`TypedProvider`] as a list of [`Reg`].
     ///
     /// # Note

crates/wasmi/src/engine/translator/mod.rs

@@ -796,10 +796,8 @@ impl FuncTranslator {
         let block_type = BlockType::func_type(func_type);
         let end_label = self.alloc.instr_encoder.new_label();
         let consume_fuel = self.make_fuel_instr()?;
-        // Note: we use a dummy `RegSpan` as placeholder.
-        //
-        // We can do this since the branch parameters of the function enclosing block
-        // are never used due to optimizations to directly return to the caller instead.
+        // Note: we use a dummy `RegSpan` as placeholder since the function enclosing
+        //       control block never has branch parameters.
         let branch_params = RegSpan::new(Reg::from(0));
         let block_frame = BlockControlFrame::new(
             block_type,
@@ -1043,6 +1041,42 @@ impl FuncTranslator {
                     FuelInfo::some(fuel_costs, fuel_instr)
                 }
             };
+            if self.reachable && frame.is_branched_to() {
+                // If the end of the `block` is reachable AND
+                // there are branches to the end of the `block`
+                // prior, we need to copy the results to the
+                // block result registers.
+                //
+                // # Note
+                //
+                // We can skip this step if the above condition is
+                // not met since the code at this point is either
+                // unreachable OR there is only one source of results
+                // and thus there is no need to copy the results around.
+                // self.translate_copy_branch_params(frame.branch_params(self.engine()))?;
+                let branch_params = frame.branch_params(self.engine());
+                let params = &mut self.alloc.buffer.providers;
+                self.alloc
+                    .stack
+                    .pop_n(usize::from(branch_params.len()), params);
+                self.alloc.instr_encoder.encode_copies(
+                    &mut self.alloc.stack,
+                    branch_params,
+                    &self.alloc.buffer.providers[..],
+                    &fuel_info,
+                )?;
+                // Case: branches to this block exist so we cannot treat the
+                //       basic block as a no-op and instead have to put its
+                //       block results on top of the stack.
+                self.alloc
+                    .stack
+                    .trunc(frame.block_height().into_u16() as usize);
+                for result in frame.branch_params(self.engine()) {
+                    self.alloc.stack.push_register(result)?;
+                }
+            }
+            // Since the `block` is now sealed we can pin its end label.
+            self.alloc.instr_encoder.pin_label(frame.end_label());
             // We dropped the Wasm `block` that encloses the function itself so we can return.
             return self.translate_return_with(&fuel_info);
         }
@@ -2732,21 +2766,6 @@ impl FuncTranslator {
         Ok(())
     }
 
-    /// Translates a conditional `br_if` that targets the function enclosing `block`.
-    fn translate_return_if(&mut self, condition: Reg) -> Result<(), Error> {
-        bail_unreachable!(self);
-        let len_results = self.func_type().results().len();
-        let fuel_info = self.fuel_info();
-        let values = &mut self.alloc.buffer.providers;
-        self.alloc.stack.peek_n(len_results, values);
-        self.alloc.instr_encoder.encode_return_nez(
-            &mut self.alloc.stack,
-            condition,
-            values,
-            fuel_info,
-        )
-    }
-
     /// Create either [`Instruction::CallIndirectParams`] or [`Instruction::CallIndirectParamsImm16`] depending on the inputs.
     fn call_indirect_params(
         &mut self,

crates/wasmi/src/engine/translator/visit.rs

@@ -403,7 +403,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
         };
         let fuel_info = self.fuel_info();
         let frame = match self.alloc.control_stack.acquire_target(relative_depth) {
-            AcquiredTarget::Return(_frame) => return self.translate_return_if(condition),
+            AcquiredTarget::Return(frame) => frame,
             AcquiredTarget::Branch(frame) => frame,
         };
         frame.bump_branches();