fix(trie): combining state parts memory usage

chain/chain/src/chain.rs

@@ -1741,15 +1741,18 @@ impl Chain {
         let shard_state_header = self.get_state_header(shard_id, sync_hash)?;
         let mut height = shard_state_header.chunk_height_included();
         let state_root = shard_state_header.chunk_prev_state_root();
-        let mut parts = vec![];
         for part_id in 0..num_parts {
             let key = StatePartKey(sync_hash, shard_id, part_id).try_to_vec()?;
-            parts.push(self.store.owned_store().get(ColStateParts, &key)?.unwrap());
+            let part = self.store.owned_store().get(ColStateParts, &key)?.unwrap();
+            self.runtime_adapter.apply_state_part(
+                shard_id,
+                &state_root,
+                part_id,
+                num_parts,
+                &part,
+            )?;
         }
 
-        // Confirm that state matches the parts we received
-        self.runtime_adapter.confirm_state(shard_id, &state_root, &parts)?;
-
         // Applying the chunk starts here
         let mut chain_update = self.chain_update();
         chain_update.set_state_finalize(shard_id, sync_hash, shard_state_header)?;

chain/chain/src/test_utils.rs

@@ -816,15 +816,12 @@ impl RuntimeAdapter for KeyValueRuntime {
         num_parts: u64,
     ) -> Result<Vec<u8>, Error> {
         assert!(part_id < num_parts);
+        if part_id != 0 {
+            return Ok(vec![]);
+        }
         let state = self.state.read().unwrap().get(&state_root).unwrap().clone();
         let data = state.try_to_vec().expect("should never fall");
-        let state_size = data.len() as u64;
-        let begin = state_size / num_parts * part_id;
-        let mut end = state_size / num_parts * (part_id + 1);
-        if part_id + 1 == num_parts {
-            end = state_size;
-        }
-        Ok(data[begin as usize..end as usize].to_vec())
+        Ok(data)
     }
 
     fn validate_state_part(
@@ -839,18 +836,18 @@ impl RuntimeAdapter for KeyValueRuntime {
         true
     }
 
-    fn confirm_state(
+    fn apply_state_part(
         &self,
         _shard_id: ShardId,
         state_root: &StateRoot,
-        parts: &Vec<Vec<u8>>,
+        part_id: u64,
+        _num_parts: u64,
+        data: &Vec<u8>,
     ) -> Result<(), Error> {
-        let mut data = vec![];
-        for part in parts {
-            data.push(part.clone());
+        if part_id != 0 {
+            return Ok(());
         }
-        let data_flatten: Vec<u8> = data.iter().flatten().cloned().collect();
-        let state = KVState::try_from_slice(&data_flatten).unwrap();
+        let state = KVState::try_from_slice(data).unwrap();
         self.state.write().unwrap().insert(state_root.clone(), state.clone());
         let data = state.try_to_vec()?;
         let state_size = data.len() as u64;

chain/chain/src/types.rs

@@ -566,11 +566,13 @@ pub trait RuntimeAdapter: Send + Sync {
     ) -> bool;
 
     /// Should be executed after accepting all the parts to set up a new state.
-    fn confirm_state(
+    fn apply_state_part(
         &self,
         shard_id: ShardId,
         state_root: &StateRoot,
-        parts: &Vec<Vec<u8>>,
+        part_id: u64,
+        num_parts: u64,
+        part: &Vec<u8>,
     ) -> Result<(), Error>;
 
     /// Returns StateRootNode of a state.

chain/client/tests/process_blocks.rs

@@ -1401,7 +1401,7 @@ fn test_process_block_after_state_sync() {
     env.clients[0].chain.reset_data_pre_state_sync(sync_hash).unwrap();
     env.clients[0]
         .runtime_adapter
-        .confirm_state(0, &chunk_extra.state_root, &vec![state_part])
+        .apply_state_part(0, &chunk_extra.state_root, 0, 1, &state_part)
         .unwrap();
     let block = env.clients[0].produce_block(sync_height + 1).unwrap().unwrap();
     let (_, res) = env.clients[0].process_block(block, Provenance::PRODUCED);