From 649515a08afdd4d80121b6789326540a3e27e3e5 Mon Sep 17 00:00:00 2001
From: Alex Su <alexanderytsu@gmail.com>
Date: Mon, 30 Oct 2023 17:08:01 +1100
Subject: [PATCH] simulate machine context fields in the TestVM

---
 test_vm/src/lib.rs | 55 ++++++++++++++++++++++++++++++++--------------
 vm_api/src/lib.rs  | 36 ++++++++++++++++++++----------
 2 files changed, 62 insertions(+), 29 deletions(-)

diff --git a/test_vm/src/lib.rs b/test_vm/src/lib.rs
index 8497ed5bc..0f05bf2c2 100644
--- a/test_vm/src/lib.rs
+++ b/test_vm/src/lib.rs
@@ -54,12 +54,15 @@ pub struct TestVM {
     pub primitives: FakePrimitives,
     pub store: Rc<MemoryBlockstore>,
     pub state_root: RefCell<Cid>,
-    circulating_supply: RefCell<TokenAmount>,
     actors_dirty: RefCell<bool>,
     actors_cache: RefCell<HashMap<Address, ActorState>>,
+    invocations: RefCell<Vec<InvocationTrace>>,
+    // MachineContext equivalents
     network_version: NetworkVersion,
     curr_epoch: RefCell<ChainEpoch>,
-    invocations: RefCell<Vec<InvocationTrace>>,
+    circulating_supply: RefCell<TokenAmount>,
+    base_fee: RefCell<TokenAmount>,
+    timestamp: RefCell<u64>,
 }
 
 impl TestVM {
@@ -81,6 +84,8 @@ impl TestVM {
             network_version: NetworkVersion::V16,
             curr_epoch: RefCell::new(ChainEpoch::zero()),
             invocations: RefCell::new(vec![]),
+            base_fee: RefCell::new(TokenAmount::zero()),
+            timestamp: RefCell::new(0),
         }
     }
 
@@ -274,10 +279,6 @@ impl VM for TestVM {
         self.store.as_ref()
     }
 
-    fn epoch(&self) -> ChainEpoch {
-        *self.curr_epoch.borrow()
-    }
-
     fn execute_message(
         &self,
         from: &Address,
@@ -361,10 +362,6 @@ impl VM for TestVM {
         st.resolve_address(&self.store, address).unwrap()
     }
 
-    fn set_epoch(&self, epoch: ChainEpoch) {
-        self.curr_epoch.replace(epoch);
-    }
-
     fn balance(&self, address: &Address) -> TokenAmount {
         let a = self.actor(address);
         a.map_or(TokenAmount::zero(), |a| a.balance)
@@ -400,13 +397,6 @@ impl VM for TestVM {
     fn actor_manifest(&self) -> BTreeMap<Cid, Type> {
         ACTOR_TYPES.clone()
     }
-    fn circulating_supply(&self) -> TokenAmount {
-        self.circulating_supply.borrow().clone()
-    }
-
-    fn set_circulating_supply(&self, supply: TokenAmount) {
-        self.circulating_supply.replace(supply);
-    }
 
     fn actor_states(&self) -> BTreeMap<Address, ActorState> {
         let map = self.actor_map();
@@ -419,4 +409,35 @@ impl VM for TestVM {
 
         tree
     }
+
+    fn epoch(&self) -> ChainEpoch {
+        *self.curr_epoch.borrow()
+    }
+
+    fn set_epoch(&self, epoch: ChainEpoch) {
+        self.curr_epoch.replace(epoch);
+    }
+    fn circulating_supply(&self) -> TokenAmount {
+        self.circulating_supply.borrow().clone()
+    }
+
+    fn set_circulating_supply(&self, supply: TokenAmount) {
+        self.circulating_supply.replace(supply);
+    }
+
+    fn base_fee(&self) -> TokenAmount {
+        self.base_fee.borrow().clone()
+    }
+
+    fn set_base_fee(&self, amount: TokenAmount) {
+        self.base_fee.replace(amount);
+    }
+
+    fn timestamp(&self) -> u64 {
+        *self.timestamp.borrow()
+    }
+
+    fn set_timestamp(&self, timestamp: u64) {
+        self.timestamp.replace(timestamp);
+    }
 }
diff --git a/vm_api/src/lib.rs b/vm_api/src/lib.rs
index e6ee50bfc..a59ff427a 100644
--- a/vm_api/src/lib.rs
+++ b/vm_api/src/lib.rs
@@ -35,12 +35,6 @@ pub trait VM {
     /// Returns the underlying blockstore of the VM
     fn blockstore(&self) -> &dyn Blockstore;
 
-    /// Get the current chain epoch
-    fn epoch(&self) -> ChainEpoch;
-
-    /// Sets the epoch to the specified value
-    fn set_epoch(&self, epoch: ChainEpoch);
-
     /// Get information about an actor
     fn actor(&self, address: &Address) -> Option<ActorState>;
 
@@ -76,12 +70,6 @@ pub trait VM {
     /// Take all the invocations that have been made since the last call to this method
     fn take_invocations(&self) -> Vec<InvocationTrace>;
 
-    /// Set the circulating supply constant for the network
-    fn set_circulating_supply(&self, supply: TokenAmount);
-
-    /// Get the circulating supply constant for the network
-    fn circulating_supply(&self) -> TokenAmount;
-
     /// Provides access to VM primitives
     fn primitives(&self) -> &dyn Primitives;
 
@@ -90,6 +78,30 @@ pub trait VM {
 
     /// Returns a map of all actor addresses to their corresponding states
     fn actor_states(&self) -> BTreeMap<Address, ActorState>;
+
+    /// Get the current chain epoch
+    fn epoch(&self) -> ChainEpoch;
+
+    /// Sets the epoch to the specified value
+    fn set_epoch(&self, epoch: ChainEpoch);
+
+    /// Get the circulating supply constant for the network
+    fn circulating_supply(&self) -> TokenAmount;
+
+    /// Set the circulating supply constant for the network
+    fn set_circulating_supply(&self, supply: TokenAmount);
+
+    /// Get the current base fee
+    fn base_fee(&self) -> TokenAmount;
+
+    /// Set the current base fee
+    fn set_base_fee(&self, amount: TokenAmount);
+
+    /// Get the current timestamp
+    fn timestamp(&self) -> u64;
+
+    /// Set the current timestamp
+    fn set_timestamp(&self, timestamp: u64);
 }
 
 #[derive(Clone, PartialEq, Eq, Debug)]