From b220780e1468f570c26e406573f081afe544ec85 Mon Sep 17 00:00:00 2001
From: Jay Geng <jay@stellar.org>
Date: Fri, 19 May 2023 19:22:06 -0400
Subject: [PATCH] Add accessors and modifiers for fuel amount and `FuelCosts`

---
 crates/wasmi/src/engine/config.rs |  5 +++++
 crates/wasmi/src/engine/mod.rs    |  2 +-
 crates/wasmi/src/func/caller.rs   | 18 +++++++++++++++++-
 crates/wasmi/src/lib.rs           |  1 +
 crates/wasmi/src/store.rs         | 26 +++++++++++++++++++++++++-
 5 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/crates/wasmi/src/engine/config.rs b/crates/wasmi/src/engine/config.rs
index 5340b2fb57..00ec4890b4 100644
--- a/crates/wasmi/src/engine/config.rs
+++ b/crates/wasmi/src/engine/config.rs
@@ -390,6 +390,11 @@ impl Config {
         self
     }
 
+    pub fn set_fuel_costs(&mut self, costs: FuelCosts) -> &mut Self {
+        self.fuel_costs = costs;
+        self
+    }
+
     /// Returns the [`FuelConsumptionMode`] for the [`Engine`].
     ///
     /// Returns `None` if fuel metering is disabled for the [`Engine`].
diff --git a/crates/wasmi/src/engine/mod.rs b/crates/wasmi/src/engine/mod.rs
index e870f0129d..008f4d06e3 100644
--- a/crates/wasmi/src/engine/mod.rs
+++ b/crates/wasmi/src/engine/mod.rs
@@ -19,7 +19,7 @@ mod tests;
 pub use self::{
     bytecode::DropKeep,
     code_map::CompiledFunc,
-    config::{Config, FuelConsumptionMode},
+    config::{Config, FuelConsumptionMode, FuelCosts},
     func_builder::{
         FuncBuilder,
         FuncTranslatorAllocations,
diff --git a/crates/wasmi/src/func/caller.rs b/crates/wasmi/src/func/caller.rs
index e73891cbe2..acf7cef0ed 100644
--- a/crates/wasmi/src/func/caller.rs
+++ b/crates/wasmi/src/func/caller.rs
@@ -15,7 +15,7 @@ pub struct Caller<'a, T> {
 
 impl<'a, T> Caller<'a, T> {
     /// Creates a new [`Caller`] from the given store context and [`Instance`] handle.
-    pub(crate) fn new<C>(ctx: &'a mut C, instance: Option<&Instance>) -> Self
+    pub fn new<C>(ctx: &'a mut C, instance: Option<&Instance>) -> Self
     where
         C: AsContextMut<UserState = T>,
     {
@@ -69,6 +69,13 @@ impl<'a, T> Caller<'a, T> {
         self.ctx.store.fuel_consumed()
     }
 
+    /// Returns the amount of total [`Fuel`] supplied to the [`Store`](crate::Store).
+    ///
+    /// Returns `None` if fuel metering is disabled.
+    pub fn fuel_total(&self) -> Option<u64> {
+        self.ctx.store.fuel_total()
+    }
+
     /// Synthetically consumes an amount of fuel for the [`Store`](crate::Store).
     ///
     /// Returns the remaining amount of fuel after this operation.
@@ -84,6 +91,15 @@ impl<'a, T> Caller<'a, T> {
     pub fn consume_fuel(&mut self, delta: u64) -> Result<u64, FuelError> {
         self.ctx.store.consume_fuel(delta)
     }
+
+    /// Resets the total and consumed amounts of fuel to 0 for the [`Store`](crate::Store).
+    ///
+    /// # Errors
+    ///
+    /// - If fuel metering is disabled.
+    pub fn reset_fuel(&mut self) -> Result<(), FuelError> {
+        self.ctx.store.reset_fuel()
+    }
 }
 
 impl<T> AsContext for Caller<'_, T> {
diff --git a/crates/wasmi/src/lib.rs b/crates/wasmi/src/lib.rs
index 68604a18f7..b054e9d034 100644
--- a/crates/wasmi/src/lib.rs
+++ b/crates/wasmi/src/lib.rs
@@ -125,6 +125,7 @@ pub use self::{
         Config,
         Engine,
         FuelConsumptionMode,
+        FuelCosts,
         ResumableCall,
         ResumableInvocation,
         StackLimits,
diff --git a/crates/wasmi/src/store.rs b/crates/wasmi/src/store.rs
index 7612ad72fc..eb231cb657 100644
--- a/crates/wasmi/src/store.rs
+++ b/crates/wasmi/src/store.rs
@@ -211,7 +211,7 @@ pub struct Fuel {
     /// The remaining fuel.
     remaining: u64,
     /// The total amount of fuel so far.
-    total: u64,
+    pub(crate) total: u64,
 }
 
 impl Fuel {
@@ -256,6 +256,12 @@ impl Fuel {
             .ok_or(TrapCode::OutOfFuel)?;
         Ok(self.remaining)
     }
+
+    /// Reset the total and remaining fuel amounts to 0.
+    pub fn reset_fuel(&mut self) {
+        self.remaining = 0;
+        self.total = 0;
+    }
 }
 
 impl StoreInner {
@@ -849,6 +855,14 @@ impl<T> Store<T> {
         Some(self.inner.fuel.fuel_consumed())
     }
 
+    /// Returns the amount of total [`Fuel`] supplied to the [`Store`].
+    ///
+    /// Returns `None` if fuel metering is disabled.
+    pub fn fuel_total(&self) -> Option<u64> {
+        self.check_fuel_metering_enabled().ok()?;
+        Some(self.inner.fuel.total)
+    }
+
     /// Synthetically consumes an amount of fuel for the [`Store`].
     ///
     /// Returns the remaining amount of fuel after this operation.
@@ -869,6 +883,16 @@ impl<T> Store<T> {
             .map_err(|_error| FuelError::out_of_fuel())
     }
 
+    /// Resets the total and consumed amounts of fuel to 0 for the [`Store`].
+    ///
+    /// # Errors
+    ///
+    /// - If fuel metering is disabled.
+    pub fn reset_fuel(&mut self) -> Result<(), FuelError> {
+        self.check_fuel_metering_enabled()?;
+        Ok(self.inner.fuel.reset_fuel())
+    }
+
     /// Allocates a new [`TrampolineEntity`] and returns a [`Trampoline`] reference to it.
     pub(super) fn alloc_trampoline(&mut self, func: TrampolineEntity<T>) -> Trampoline {
         let idx = self.trampolines.alloc(func);