Engine optimization: cache all reads from NEAR state

engine-sdk/src/dup_cache.rs

@@ -1,3 +1,17 @@
+use aurora_engine_types::BTreeMap;
+
+/// All entries are cached (unlike `DupCache`).
+#[derive(Default)]
+pub struct FullCache<K, V> {
+    inner: BTreeMap<K, V>,
+}
+
+impl<K: Ord, V> FullCache<K, V> {
+    pub fn get_or_insert_with<F: FnOnce() -> V>(&mut self, k: K, f: F) -> &mut V {
+        self.inner.entry(k).or_insert_with(f)
+    }
+}
+
 /// 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,

engine-tests/src/test_utils/mod.rs

@@ -827,11 +827,12 @@ pub fn panic_on_fail(status: TransactionStatus) {
 }
 
 pub fn assert_gas_bound(total_gas: u64, tgas_bound: u64) {
-    let bound = tgas_bound * 1_000_000_000_000;
+    // Add 1 to round up
+    let tgas_used = (total_gas / 1_000_000_000_000) + 1;
     assert!(
-        total_gas <= bound,
-        "{} Tgas is not less than {} Tgas",
-        total_gas / 1_000_000_000_000,
+        tgas_used == tgas_bound,
+        "{} Tgas is not equal to {} Tgas",
+        tgas_used,
         tgas_bound,
     );
 }

engine-tests/src/tests/one_inch.rs

@@ -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(), 33); // less than 33 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 29);
 
     // 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(), 44); // less than 44 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 35);
 
     // 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(), 46); // less than 46 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 37);
 
     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(), 38); // less than 38 NEAR Tgas used
+    assert_gas_bound(profile.all_gas(), 32);
 }
 
 #[test]

engine-tests/src/tests/repro.rs

@@ -46,7 +46,7 @@ fn repro_GdASJ3KESs() {
         SubmitResult::try_from_slice(&outcome.return_data.as_value().unwrap()).unwrap();
 
     assert_eq!(submit_result.gas_used, 706713);
-    assert_eq!(239, profile.all_gas() / 1_000_000_000_000);
+    assert_eq!(173, profile.all_gas() / 1_000_000_000_000);
 
     // Also validate the SubmitResult in the standalone engine
     let mut standalone = standalone::StandaloneRunner::default();
@@ -99,7 +99,7 @@ fn repro_8ru7VEA() {
         SubmitResult::try_from_slice(&outcome.return_data.as_value().unwrap()).unwrap();
 
     assert_eq!(submit_result.gas_used, 1732181);
-    assert_eq!(411, profile.all_gas() / 1_000_000_000_000);
+    assert_eq!(309, profile.all_gas() / 1_000_000_000_000);
 
     // Also validate the SubmitResult in the standalone engine
     let mut standalone = standalone::StandaloneRunner::default();

engine-tests/src/tests/uniswap.rs

@@ -28,10 +28,7 @@ fn test_uniswap_input_multihop() {
     let mut context = UniswapTestContext::new("uniswap");
 
     // evm_gas = 970k
-    // near total gas = 204 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;
+    // near total gas = 163 Tgas
 
     let tokens = context.create_tokens(10, MINT_AMOUNT.into());
     for (token_a, token_b) in tokens.iter().zip(tokens.iter().skip(1)) {
@@ -41,9 +38,7 @@ fn test_uniswap_input_multihop() {
 
     let (_amount_out, _evm_gas, profile) = context.exact_input(&tokens, INPUT_AMOUNT.into());
 
-    println!("{:?}", profile.host_breakdown);
-    println!("NEAR_GAS_WASM {:?}", profile.wasm_gas());
-    println!("NEAR_GAS_TOTAL {:?}", profile.all_gas());
+    assert_eq!(163, profile.all_gas() / 1_000_000_000_000);
 }
 
 #[test]
@@ -54,7 +49,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(), 57);
+    test_utils::assert_gas_bound(profile.all_gas(), 47);
     let wasm_fraction = 100 * profile.wasm_gas() / profile.all_gas();
     assert!(
         25 <= wasm_fraction && wasm_fraction <= 35,
@@ -64,7 +59,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(), 31);
+    test_utils::assert_gas_bound(profile.all_gas(), 26);
     let wasm_fraction = 100 * profile.wasm_gas() / profile.all_gas();
     assert!(
         25 <= wasm_fraction && wasm_fraction <= 35,

engine/src/engine.rs

@@ -6,7 +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::dup_cache::FullCache;
 use aurora_engine_sdk::env::Env;
 use aurora_engine_sdk::io::{StorageIntermediate, IO};
 use aurora_engine_sdk::promise::{PromiseHandler, PromiseId};
@@ -422,9 +422,9 @@ 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_code_cache: RefCell<DupCache<Address, Vec<u8>>>,
-    contract_storage_cache: RefCell<PairDupCache<Address, H256, H256>>,
+    account_info_cache: RefCell<FullCache<Address, Basic>>,
+    contract_code_cache: RefCell<FullCache<Address, Vec<u8>>>,
+    contract_storage_cache: RefCell<FullCache<(Address, H256), H256>>,
 }
 
 pub(crate) const CONFIG: &Config = &Config::london();
@@ -457,9 +457,9 @@ 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_code_cache: RefCell::new(DupCache::default()),
-            contract_storage_cache: RefCell::new(PairDupCache::default()),
+            account_info_cache: RefCell::new(FullCache::default()),
+            contract_code_cache: RefCell::new(FullCache::default()),
+            contract_storage_cache: RefCell::new(FullCache::default()),
         }
     }
 
@@ -1519,15 +1519,15 @@ impl<'env, I: IO + Copy, E: Env> evm::backend::Backend for Engine<'env, I, E> {
     fn exists(&self, address: H160) -> bool {
         let address = Address::new(address);
         let mut cache = self.account_info_cache.borrow_mut();
-        let basic_info = cache.get_or_insert_with(&address, || Basic {
+        let basic_info = cache.get_or_insert_with(address, || Basic {
             nonce: get_nonce(&self.io, &address),
             balance: get_balance(&self.io, &address).raw(),
         });
         if !basic_info.balance.is_zero() || !basic_info.nonce.is_zero() {
             return true;
         }
         let mut cache = self.contract_code_cache.borrow_mut();
-        let code = cache.get_or_insert_with(&address, || get_code(&self.io, &address));
+        let code = cache.get_or_insert_with(address, || get_code(&self.io, &address));
         !code.is_empty()
     }
 
@@ -1537,7 +1537,7 @@ impl<'env, I: IO + Copy, E: Env> evm::backend::Backend for Engine<'env, I, E> {
         let result = self
             .account_info_cache
             .borrow_mut()
-            .get_or_insert_with(&address, || Basic {
+            .get_or_insert_with(address, || Basic {
                 nonce: get_nonce(&self.io, &address),
                 balance: get_balance(&self.io, &address).raw(),
             })
@@ -1550,7 +1550,7 @@ impl<'env, I: IO + Copy, E: Env> evm::backend::Backend for Engine<'env, I, E> {
         let address = Address::new(address);
         self.contract_code_cache
             .borrow_mut()
-            .get_or_insert_with(&address, || get_code(&self.io, &address))
+            .get_or_insert_with(address, || get_code(&self.io, &address))
             .clone()
     }
 
@@ -1565,7 +1565,7 @@ impl<'env, I: IO + Copy, E: Env> evm::backend::Backend for Engine<'env, I, E> {
         let result = *self
             .contract_storage_cache
             .borrow_mut()
-            .get_or_insert_with((&address, &index), || {
+            .get_or_insert_with((address, index), || {
                 get_storage(&self.io, &address, &index, generation)
             });
         result