add fuel metering for return instructions

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

@@ -21,6 +21,7 @@ use crate::{
             },
             translator::{stack::RegisterSpace, ValueStack},
         },
+        FuelCosts,
         TranslationError,
     },
     module::ModuleResources,
@@ -253,20 +254,6 @@ impl InstrEncoder {
         Ok(())
     }
 
-    /// Bumps consumed fuel for [`Instruction::ConsumeFuel`] of `instr` by `delta`.
-    ///
-    /// # Errors
-    ///
-    /// If consumed fuel is out of bounds after this operation.
-    #[allow(dead_code)] // TODO: remove
-    pub fn bump_fuel_consumption(
-        &mut self,
-        instr: Instr,
-        delta: u64,
-    ) -> Result<(), TranslationError> {
-        self.instrs.get_mut(instr).bump_fuel_consumption(delta)
-    }
-
     /// Push the [`Instruction`] to the [`InstrEncoder`].
     pub fn push_instr(&mut self, instr: Instruction) -> Result<Instr, TranslationError> {
         let last_instr = self.instrs.push(instr)?;
@@ -480,11 +467,36 @@ impl InstrEncoder {
         false
     }
 
+    /// Bumps consumed fuel for [`Instruction::ConsumeFuel`] of `instr` by `delta`.
+    ///
+    /// # Errors
+    ///
+    /// If consumed fuel is out of bounds after this operation.
+    pub fn bump_fuel_consumption<F>(
+        &mut self,
+        fuel_info: Option<(FuelCosts, Instr)>,
+        f: F,
+    ) -> Result<(), TranslationError>
+    where
+        F: FnOnce(&FuelCosts) -> u64,
+    {
+        let Some((fuel_costs, fuel_instr)) = fuel_info else {
+            // Fuel metering is disabled so we can bail out.
+            return Ok(());
+        };
+        let fuel_consumed = f(&fuel_costs);
+        self.instrs
+            .get_mut(fuel_instr)
+            .bump_fuel_consumption(fuel_consumed)?;
+        Ok(())
+    }
+
     /// Encodes an unconditional `return` instruction.
     pub fn encode_return(
         &mut self,
         stack: &mut ValueStack,
         values: &[TypedProvider],
+        fuel_info: Option<(FuelCosts, Instr)>,
     ) -> Result<(), TranslationError> {
         let instr = match values {
             [] => Instruction::Return,
@@ -517,6 +529,9 @@ impl InstrEncoder {
             }
             [v0, v1, v2, rest @ ..] => {
                 debug_assert!(!rest.is_empty());
+                self.bump_fuel_consumption(fuel_info, |costs| {
+                    costs.fuel_for_copies(rest.len() as u64 + 3)
+                })?;
                 if let Some(span) = RegisterSpanIter::from_providers(values) {
                     self.push_instr(Instruction::return_span(span))?;
                     return Ok(());
@@ -529,6 +544,7 @@ impl InstrEncoder {
                 return Ok(());
             }
         };
+        self.bump_fuel_consumption(fuel_info, FuelCosts::base)?;
         self.push_instr(instr)?;
         Ok(())
     }
@@ -539,7 +555,13 @@ impl InstrEncoder {
         stack: &mut ValueStack,
         condition: Register,
         values: &[TypedProvider],
+        fuel_info: Option<(FuelCosts, Instr)>,
     ) -> Result<(), TranslationError> {
+        // 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),
@@ -565,6 +587,9 @@ impl InstrEncoder {
             }
             [v0, v1, rest @ ..] => {
                 debug_assert!(!rest.is_empty());
+                self.bump_fuel_consumption(fuel_info, |costs| {
+                    costs.fuel_for_copies(rest.len() as u64 + 3)
+                })?;
                 if let Some(span) = RegisterSpanIter::from_providers(values) {
                     self.push_instr(Instruction::return_nez_span(condition, span))?;
                     return Ok(());
@@ -576,6 +601,7 @@ impl InstrEncoder {
                 return Ok(());
             }
         };
+        self.bump_fuel_consumption(fuel_info, FuelCosts::base)?;
         self.push_instr(instr)?;
         Ok(())
     }

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

@@ -286,8 +286,8 @@ impl<'parser> FuncTranslator<'parser> {
     /// Returns the [`FuelCosts`] and the most recent [`Instruction::ConsumeFuel`] in the translation process.
     ///
     /// Returns `None` if fuel metering is disabled.
-    fn fuel_costs_and_instr(&self) -> Option<(&FuelCosts, Instr)> {
-        let Some(fuel_costs) = self.fuel_costs() else {
+    fn fuel_costs_and_instr(&self) -> Option<(FuelCosts, Instr)> {
+        let Some(&fuel_costs) = self.fuel_costs() else {
             // Fuel metering is disabled so we can bail out.
             return None;
         };
@@ -318,14 +318,10 @@ impl<'parser> FuncTranslator<'parser> {
     where
         F: FnOnce(&FuelCosts) -> u64,
     {
-        let Some((fuel_costs, fuel_instr)) = self.fuel_costs_and_instr() else {
-            // Fuel metering is disabled so we can bail out.
-            return Ok(());
-        };
-        let fuel_consumed = f(fuel_costs);
+        let fuel_info = self.fuel_costs_and_instr();
         self.alloc
             .instr_encoder
-            .bump_fuel_consumption(fuel_instr, fuel_consumed)?;
+            .bump_fuel_consumption(fuel_info, f)?;
         Ok(())
     }
 
@@ -2094,11 +2090,12 @@ impl<'parser> FuncTranslator<'parser> {
     pub fn translate_return(&mut self) -> Result<(), TranslationError> {
         let func_type = self.func_type();
         let results = func_type.results();
+        let fuel_info = self.fuel_costs_and_instr();
         let values = &mut self.alloc.buffer;
         self.alloc.stack.pop_n(results.len(), values);
         self.alloc
             .instr_encoder
-            .encode_return(&mut self.alloc.stack, values)?;
+            .encode_return(&mut self.alloc.stack, values, fuel_info)?;
         self.reachable = false;
         Ok(())
     }
@@ -2107,10 +2104,14 @@ impl<'parser> FuncTranslator<'parser> {
     pub fn translate_return_if(&mut self, condition: Register) -> Result<(), TranslationError> {
         bail_unreachable!(self);
         let len_results = self.func_type().results().len();
+        let fuel_info = self.fuel_costs_and_instr();
         let values = &mut self.alloc.buffer;
         self.alloc.stack.peek_n(len_results, values);
-        self.alloc
-            .instr_encoder
-            .encode_return_nez(&mut self.alloc.stack, condition, values)
+        self.alloc.instr_encoder.encode_return_nez(
+            &mut self.alloc.stack,
+            condition,
+            values,
+            fuel_info,
+        )
     }
 }

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

@@ -573,9 +573,12 @@ impl<'a> VisitOperator<'a> for FuncTranslator<'a> {
             self.alloc.instr_encoder.pin_label(label);
             match self.alloc.control_stack.acquire_target(depth) {
                 AcquiredTarget::Return(_frame) => {
+                    // Note: We do not use fuel metering for the below
+                    //       `encode_return` instruction since it is part
+                    //       of the `br_table` instruction above.
                     self.alloc
                         .instr_encoder
-                        .encode_return(&mut self.alloc.stack, values)?;
+                        .encode_return(&mut self.alloc.stack, values, None)?;
                 }
                 AcquiredTarget::Branch(frame) => {
                     frame.bump_branches();