From dc5de4af009898a1907f5dcd08f5a395588acba3 Mon Sep 17 00:00:00 2001
From: Robin Freyler <robin.freyler@gmail.com>
Date: Mon, 27 Nov 2023 12:16:33 +0100
Subject: [PATCH] refactor and cleanup of new fuel metering API

---
 .../src/engine/regmach/translator/mod.rs      | 25 +++++++++++++------
 .../src/engine/regmach/translator/visit.rs    |  6 ++---
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/crates/wasmi/src/engine/regmach/translator/mod.rs b/crates/wasmi/src/engine/regmach/translator/mod.rs
index 3e0e0aee1d..13f22cf1e0 100644
--- a/crates/wasmi/src/engine/regmach/translator/mod.rs
+++ b/crates/wasmi/src/engine/regmach/translator/mod.rs
@@ -161,7 +161,7 @@ impl<'parser> FuncTranslator<'parser> {
         let func_type = self.res.get_type_of_func(self.func);
         let block_type = BlockType::func_type(func_type);
         let end_label = self.alloc.instr_encoder.new_label();
-        let consume_fuel = self.make_consume_fuel_instr()?;
+        let consume_fuel = self.make_fuel_instr()?;
         // Note: we use a dummy `RegisterSpan` as placeholder.
         //
         // We can do this since the branch parameters of the function enclosing block
@@ -277,7 +277,7 @@ impl<'parser> FuncTranslator<'parser> {
     /// Pushes a [`Instruction::ConsumeFuel`] with base costs if fuel metering is enabled.
     ///
     /// Returns `None` if fuel metering is disabled.
-    fn make_consume_fuel_instr(&mut self) -> Result<Option<Instr>, TranslationError> {
+    fn make_fuel_instr(&mut self) -> Result<Option<Instr>, TranslationError> {
         let Some(fuel_costs) = self.fuel_costs() else {
             // Fuel metering is disabled so there is no need to create an `Instruction::ConsumeFuel`.
             return Ok(None);
@@ -291,10 +291,24 @@ impl<'parser> FuncTranslator<'parser> {
     /// Returns the most recent [`Instruction::ConsumeFuel`] in the translation process.
     ///
     /// Returns `None` if fuel metering is disabled.
-    fn consume_fuel_instr(&self) -> Option<Instr> {
+    fn fuel_instr(&self) -> Option<Instr> {
         self.alloc.control_stack.last().consume_fuel_instr()
     }
 
+    /// Returns the [`FuelCosts`] and the most recent [`Intruction::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 {
+            // Fuel metering is disabled so we can bail out.
+            return None;
+        };
+        let fuel_instr = self
+            .fuel_instr()
+            .expect("fuel metering is enabled but there is no Instruction::ConsumeFuel");
+        Some((fuel_costs, fuel_instr))
+    }
+
     /// Adds fuel to the most recent [`Instruction::ConsumeFuel`] in the translation process.
     ///
     /// Does nothing if gas metering is disabled.
@@ -302,13 +316,10 @@ impl<'parser> FuncTranslator<'parser> {
     where
         F: FnOnce(&FuelCosts) -> u64,
     {
-        let Some(fuel_costs) = self.fuel_costs() else {
+        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_instr = self
-            .consume_fuel_instr()
-            .expect("fuel metering is enabled but there is no Instruction::ConsumeFuel");
         let fuel_consumed = f(fuel_costs);
         self.alloc
             .instr_encoder
diff --git a/crates/wasmi/src/engine/regmach/translator/visit.rs b/crates/wasmi/src/engine/regmach/translator/visit.rs
index be056b0b31..2158da6814 100644
--- a/crates/wasmi/src/engine/regmach/translator/visit.rs
+++ b/crates/wasmi/src/engine/regmach/translator/visit.rs
@@ -175,7 +175,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator<'a> {
         // Optionally create the loop's [`Instruction::ConsumeFuel`].
         //
         // This is handling the fuel required for a single iteration of the loop.
-        let consume_fuel = self.make_consume_fuel_instr()?;
+        let consume_fuel = self.make_fuel_instr()?;
         // Finally create the loop control frame.
         self.alloc.control_stack.push_frame(LoopControlFrame::new(
             block_type,
@@ -253,7 +253,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator<'a> {
                 //
                 // The [`Instruction::ConsumeFuel`] for the `else` branch is
                 // created on the fly when visiting the `else` block.
-                let consume_fuel = self.make_consume_fuel_instr()?;
+                let consume_fuel = self.make_fuel_instr()?;
                 (reachability, consume_fuel)
             }
         };
@@ -307,7 +307,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator<'a> {
             }
             self.reachable = true;
             self.alloc.instr_encoder.pin_label(else_label);
-            if let Some(fuel_instr) = self.make_consume_fuel_instr()? {
+            if let Some(fuel_instr) = self.make_fuel_instr()? {
                 frame.update_consume_fuel_instr(fuel_instr);
             }
             // At this point we can restore the `else` branch parameters