Merge remote-tracking branch 'NYBACHOK/quick-fix-caching' into joneskm/staking3

Author: joneskm

gears/src/baseapp/abci.rs

@@ -278,6 +278,7 @@ impl<DB: Database, PSK: ParamsSubspaceKey, H: ABCIHandler, AI: ApplicationInfo>
         let consensus_params = self.baseapp_params_keeper.consensus_params(&ctx);
 
         state.replace_meter(Gas::from(max_gas));
+        state.take_block_cache(&mut multi_store);
 
         let mut ctx = BlockContext::new(
             &mut multi_store,

gears/src/baseapp/state.rs

@@ -49,6 +49,21 @@ impl<DB: Database, AH: ABCIHandler> ApplicationState<DB, AH> {
             .append_block_cache(multi_store);
     }
 
+    pub fn take_block_cache(&mut self, multi_store: &mut ApplicationMultiBank<DB, AH::StoreKey>) {
+        let list = self.deliver_mode.multi_store.take_block_cache();
+
+        for (key, (insert_list, delete_list)) in list {
+            let kv_store = multi_store.kv_store_mut(&key);
+
+            delete_list.into_iter().for_each(|this| {
+                kv_store.delete(&this);
+            });
+            insert_list
+                .into_iter()
+                .for_each(|(key, value)| kv_store.set(key, value));
+        }
+    }
+
     pub fn commit(&mut self, multi_store: &mut ApplicationMultiBank<DB, AH::StoreKey>) -> [u8; 32] {
         self.check_mode.multi_store.tx_cache_clear();
         self.check_mode.multi_store.block_cache_clear();

kv_store/src/bank/kv/transaction.rs

@@ -1,6 +1,6 @@
 use std::{
     borrow::Cow,
-    collections::HashMap,
+    collections::{BTreeMap, HashMap, HashSet},
     ops::RangeBounds,
     sync::{Arc, RwLock},
 };
@@ -61,6 +61,10 @@ impl<DB: Database> TransactionKVBank<DB> {
         }
     }
 
+    pub fn take_block_cache(&mut self) -> (BTreeMap<Vec<u8>, Vec<u8>>, HashSet<Vec<u8>>) {
+        self.block.take()
+    }
+
     /// Delete value from storage
     #[inline]
     pub fn delete(&mut self, k: &[u8]) -> Option<Vec<u8>> {

kv_store/src/bank/multi/transaction.rs

@@ -1,3 +1,5 @@
+use std::collections::{BTreeMap, HashSet};
+
 use database::{prefix::PrefixDB, Database};
 
 use crate::{bank::kv::transaction::TransactionKVBank, StoreKey};
@@ -43,4 +45,15 @@ impl<DB: Database, SK: StoreKey> MultiBank<DB, SK, TransactionStore<DB, SK>> {
             store.append_block_cache(other.kv_store_mut(sk))
         }
     }
+
+    pub fn take_block_cache(
+        &mut self,
+    ) -> HashMap<SK, (BTreeMap<Vec<u8>, Vec<u8>>, HashSet<Vec<u8>>)> {
+        let mut set = HashMap::with_capacity(self.backend.0.len());
+        for (sk, store) in &mut self.backend.0 {
+            set.insert(sk.clone(), store.take_block_cache());
+        }
+
+        set
+    }
 }