Fix(engine): Simple cache to stop consecutive duplicate reads (#446)

engine-sdk/src/dup_cache.rs

@@ -0,0 +1,41 @@
+/// The intention of this struct is to prevent repeating duplicate computations/IO with the
+/// same input (key). However, unlike memoization or typical caching, this only remembers the
+/// most recent key-value pair. This means it is optimized for consecutive duplicate lookups,
+/// as opposed to general duplicated lookups. The benefit is that its memory footprint and
+/// internal logic are both minimal, and the drawback is that its use case is very narrow.
+#[derive(Default)]
+pub struct DupCache<K, V> {
+    key: K,
+    value: V,
+}
+
+impl<K: Copy + Eq, V> DupCache<K, V> {
+    pub fn get_or_insert_with<F: FnOnce() -> V>(&mut self, k: &K, f: F) -> &mut V {
+        if &self.key != k {
+            let new_value = f();
+            self.value = new_value;
+            self.key = *k;
+        }
+
+        &mut self.value
+    }
+}
+
+/// Same as `DupCache` but optimized for the case that `K = (K1, K2)`.
+#[derive(Default)]
+pub struct PairDupCache<K1, K2, V> {
+    key: (K1, K2),
+    value: V,
+}
+
+impl<K1: Copy + Eq, K2: Copy + Eq, V> PairDupCache<K1, K2, V> {
+    pub fn get_or_insert_with<F: FnOnce() -> V>(&mut self, k: (&K1, &K2), f: F) -> &mut V {
+        if (&self.key.0 != k.0) || (&self.key.1 != k.1) {
+            let new_value = f();
+            self.value = new_value;
+            self.key = (*k.0, *k.1);
+        }
+
+        &mut self.value
+    }
+}

engine-sdk/src/lib.rs

@@ -7,6 +7,7 @@ use crate::prelude::Address;
 use crate::prelude::{H256, STORAGE_PRICE_PER_BYTE};
 pub use types::keccak;
 
+pub mod dup_cache;
 pub mod env;
 pub mod error;
 pub mod io;

engine-tests/src/tests/one_inch.rs

@@ -23,11 +23,11 @@ fn test_1inch_liquidity_protocol() {
 
     let (result, profile, deployer_address) = helper.create_mooniswap_deployer();
     assert!(result.gas_used >= 5_100_000); // more than 5.1M EVM gas used
-    assert_gas_bound(profile.all_gas(), 14); // less than 14 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 12); // less than 12 NEAR Tgas used
 
     let (result, profile, pool_factory) = helper.create_pool_factory(&deployer_address);
     assert!(result.gas_used >= 2_800_000); // more than 2.8M EVM gas used
-    assert_gas_bound(profile.all_gas(), 13); // less than 13 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 12); // less than 12 NEAR Tgas used
 
     // create some ERC-20 tokens to have a liquidity pool for
     let signer_address = test_utils::address_from_secret_key(&helper.signer.secret_key);
@@ -39,7 +39,7 @@ fn test_1inch_liquidity_protocol() {
     let (result, profile, pool) =
         helper.create_pool(&pool_factory, token_a.0.address, token_b.0.address);
     assert!(result.gas_used >= 4_500_000); // more than 4.5M EVM gas used
-    assert_gas_bound(profile.all_gas(), 40); // less than 40 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 34); // less than 34 NEAR Tgas used
 
     // Approve giving ERC-20 tokens to the pool
     helper.approve_erc20_tokens(&token_a, pool.address());
@@ -58,7 +58,7 @@ fn test_1inch_liquidity_protocol() {
         },
     );
     assert!(result.gas_used >= 302_000); // more than 302k EVM gas used
-    assert_gas_bound(profile.all_gas(), 55); // less than 55 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 45); // less than 45 NEAR Tgas used
 
     // Same here
     helper.runner.context.block_timestamp += 10_000_000 * 1_000_000_000;
@@ -73,7 +73,7 @@ fn test_1inch_liquidity_protocol() {
         },
     );
     assert!(result.gas_used >= 210_000); // more than 210k EVM gas used
-    assert_gas_bound(profile.all_gas(), 57); // less than 57 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 48); // less than 48 NEAR Tgas used
 
     let (result, profile) = helper.pool_withdraw(
         &pool,
@@ -84,7 +84,7 @@ fn test_1inch_liquidity_protocol() {
         },
     );
     assert!(result.gas_used >= 150_000); // more than 150k EVM gas used
-    assert_gas_bound(profile.all_gas(), 49); // less than 49 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 39); // less than 39 NEAR Tgas used
 }
 
 #[test]
@@ -100,13 +100,13 @@ fn test_1_inch_limit_order_deploy() {
 
     // more than 3.5 million Ethereum gas used
     assert!(result.gas_used > 3_500_000);
-    // less than 14 NEAR Tgas used
-    assert_gas_bound(profile.all_gas(), 14);
+    // less than 12 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 12);
     // at least 45% of which is from wasm execution
     let wasm_fraction = 100 * profile.wasm_gas() / profile.all_gas();
     assert!(
-        30 <= wasm_fraction && wasm_fraction <= 40,
-        "{}% is not between 30% and 40%",
+        40 <= wasm_fraction && wasm_fraction <= 50,
+        "{}% is not between 40% and 50%",
         wasm_fraction
     );
 }

engine-tests/src/tests/sanity.rs

@@ -92,8 +92,8 @@ fn test_deploy_largest_contract() {
         result.gas_used,
     );
 
-    // Less than 14 NEAR Tgas
-    test_utils::assert_gas_bound(profile.all_gas(), 14);
+    // Less than 12 NEAR Tgas
+    test_utils::assert_gas_bound(profile.all_gas(), 12);
 }
 
 #[test]

engine-tests/src/tests/uniswap.rs

@@ -28,7 +28,7 @@ fn test_uniswap_input_multihop() {
     let mut context = UniswapTestContext::new("uniswap");
 
     // evm_gas = 970k
-    // near total gas = 258 Tgas
+    // near total gas = 214 Tgas
     // Wish: optimize so that this transaction costs less than 200 Tgas.
     // For now we just have to increase the burnt gas limit to make it run to completion.
     context.runner.wasm_config.limit_config.max_gas_burnt = 500_000_000_000_000;
@@ -54,7 +54,7 @@ fn test_uniswap_exact_output() {
 
     let (_result, profile) =
         context.add_equal_liquidity(LIQUIDITY_AMOUNT.into(), &token_a, &token_b);
-    test_utils::assert_gas_bound(profile.all_gas(), 76);
+    test_utils::assert_gas_bound(profile.all_gas(), 58);
     let wasm_fraction = 100 * profile.wasm_gas() / profile.all_gas();
     assert!(
         20 <= wasm_fraction && wasm_fraction <= 30,
@@ -64,7 +64,7 @@ fn test_uniswap_exact_output() {
 
     let (_amount_in, profile) =
         context.exact_output_single(&token_a, &token_b, OUTPUT_AMOUNT.into());
-    test_utils::assert_gas_bound(profile.all_gas(), 38);
+    test_utils::assert_gas_bound(profile.all_gas(), 32);
     let wasm_fraction = 100 * profile.wasm_gas() / profile.all_gas();
     assert!(
         25 <= wasm_fraction && wasm_fraction <= 35,

engine/src/engine.rs

@@ -6,6 +6,7 @@ use evm::{Config, CreateScheme, ExitError, ExitFatal, ExitReason};
 
 use crate::connector::EthConnectorContract;
 use crate::map::BijectionMap;
+use aurora_engine_sdk::dup_cache::{DupCache, PairDupCache};
 use aurora_engine_sdk::env::Env;
 use aurora_engine_sdk::io::{StorageIntermediate, IO};
 use aurora_engine_sdk::promise::{PromiseHandler, PromiseId};
@@ -406,6 +407,8 @@ pub struct Engine<'env, I: IO, E: Env> {
     io: I,
     env: &'env E,
     generation_cache: RefCell<BTreeMap<Address, u32>>,
+    account_info_cache: RefCell<DupCache<Address, Basic>>,
+    contract_storage_cache: RefCell<PairDupCache<Address, H256, H256>>,
 }
 
 pub(crate) const CONFIG: &Config = &Config::london();
@@ -438,6 +441,8 @@ impl<'env, I: IO + Copy, E: Env> Engine<'env, I, E> {
             io,
             env,
             generation_cache: RefCell::new(BTreeMap::new()),
+            account_info_cache: RefCell::new(DupCache::default()),
+            contract_storage_cache: RefCell::new(PairDupCache::default()),
         }
     }
 
@@ -1421,10 +1426,15 @@ impl<'env, I: IO + Copy, E: Env> evm::backend::Backend for Engine<'env, I, E> {
     /// Returns basic account information.
     fn basic(&self, address: H160) -> Basic {
         let address = Address::new(address);
-        Basic {
-            nonce: get_nonce(&self.io, &address),
-            balance: get_balance(&self.io, &address).raw(),
-        }
+        let result = self
+            .account_info_cache
+            .borrow_mut()
+            .get_or_insert_with(&address, || Basic {
+                nonce: get_nonce(&self.io, &address),
+                balance: get_balance(&self.io, &address).raw(),
+            })
+            .clone();
+        result
     }
 
     /// Returns the code of the contract from an address.
@@ -1440,7 +1450,13 @@ impl<'env, I: IO + Copy, E: Env> evm::backend::Backend for Engine<'env, I, E> {
             .borrow_mut()
             .entry(address)
             .or_insert_with(|| get_generation(&self.io, &address));
-        get_storage(&self.io, &address, &index, generation)
+        let result = *self
+            .contract_storage_cache
+            .borrow_mut()
+            .get_or_insert_with((&address, &index), || {
+                get_storage(&self.io, &address, &index, generation)
+            });
+        result
     }
 
     /// Get original storage value of address at index, if available.