refactor: Combine contract accesses and deployments into ContractUpdates

chain/chain/src/chain_update.rs

@@ -157,8 +157,7 @@ impl<'a> ChainUpdate<'a> {
                         shard_id,
                         apply_result.proof,
                         apply_result.applied_receipts_hash,
-                        apply_result.contract_accesses,
-                        apply_result.contract_deploys,
+                        apply_result.contract_updates,
                     );
                 }
             }
@@ -188,8 +187,7 @@ impl<'a> ChainUpdate<'a> {
                         shard_uid.shard_id(),
                         apply_result.proof,
                         apply_result.applied_receipts_hash,
-                        apply_result.contract_accesses,
-                        apply_result.contract_deploys,
+                        apply_result.contract_updates,
                     );
                 }
             }

chain/chain/src/resharding/manager.rs

@@ -160,10 +160,7 @@ impl ReshardingManager {
                 new_shard_uid.shard_id(),
                 Some(partial_storage),
                 CryptoHash::default(),
-                // No contract code is accessed during resharding.
-                // TODO(#11099): Confirm if sending no contracts is ok here.
-                Default::default(),
-                // No contract code is deployed during resharding.
+                // No contract code is accessed or deployed during resharding.
                 // TODO(#11099): Confirm if sending no contracts is ok here.
                 Default::default(),
             );

chain/chain/src/runtime/mod.rs

@@ -468,10 +468,9 @@ impl NightshadeRuntime {
             processed_yield_timeouts: apply_result.processed_yield_timeouts,
             applied_receipts_hash: hash(&borsh::to_vec(receipts).unwrap()),
             congestion_info: apply_result.congestion_info,
-            contract_accesses: apply_result.contract_accesses,
             bandwidth_requests: apply_result.bandwidth_requests,
             bandwidth_scheduler_state_hash: apply_result.bandwidth_scheduler_state_hash,
-            contract_deploys: apply_result.contract_deploys,
+            contract_updates: apply_result.contract_updates,
         };
 
         Ok(result)

chain/chain/src/store/mod.rs

@@ -1,5 +1,5 @@
 use std::collections::hash_map::Entry;
-use std::collections::{BTreeSet, HashMap, HashSet};
+use std::collections::{HashMap, HashSet};
 use std::io;
 
 use borsh::{BorshDeserialize, BorshSerialize};
@@ -23,7 +23,7 @@ use near_primitives::sharding::{
 use near_primitives::state_sync::{
     ReceiptProofResponse, ShardStateSyncResponseHeader, StateHeaderKey, StateSyncDumpProgress,
 };
-use near_primitives::stateless_validation::contract_distribution::CodeHash;
+use near_primitives::stateless_validation::contract_distribution::ContractUpdates;
 use near_primitives::stateless_validation::stored_chunk_state_transition_data::{
     StoredChunkStateTransitionData, StoredChunkStateTransitionDataV2,
 };
@@ -2010,10 +2010,10 @@ impl<'a> ChainStoreUpdate<'a> {
         shard_id: ShardId,
         partial_storage: Option<PartialStorage>,
         applied_receipts_hash: CryptoHash,
-        contract_accesses: BTreeSet<CodeHash>,
-        contract_deploys: BTreeSet<CodeHash>,
+        contract_updates: ContractUpdates,
     ) {
         if let Some(partial_storage) = partial_storage {
+            let ContractUpdates { contract_accesses, contract_deploys } = contract_updates;
             self.state_transition_data.insert(
                 (block_hash, shard_id),
                 StoredChunkStateTransitionData::V2(StoredChunkStateTransitionDataV2 {

chain/chain/src/test_utils/kv_runtime.rs

@@ -1388,8 +1388,7 @@ impl RuntimeAdapter for KeyValueRuntime {
             congestion_info: Self::get_congestion_info(PROTOCOL_VERSION),
             bandwidth_requests: BandwidthRequests::default_for_protocol_version(PROTOCOL_VERSION),
             bandwidth_scheduler_state_hash: CryptoHash::default(),
-            contract_accesses: Default::default(),
-            contract_deploys: Default::default(),
+            contract_updates: Default::default(),
         })
     }
 

chain/chain/src/types.rs

@@ -1,5 +1,3 @@
-use std::collections::BTreeSet;
-
 use borsh::{BorshDeserialize, BorshSerialize};
 use near_async::time::{Duration, Utc};
 use near_chain_configs::GenesisConfig;
@@ -26,7 +24,7 @@ use near_primitives::receipt::{PromiseYieldTimeout, Receipt};
 use near_primitives::sandbox::state_patch::SandboxStatePatch;
 use near_primitives::shard_layout::ShardUId;
 use near_primitives::state_part::PartId;
-use near_primitives::stateless_validation::contract_distribution::CodeHash;
+use near_primitives::stateless_validation::contract_distribution::ContractUpdates;
 use near_primitives::transaction::{ExecutionOutcomeWithId, SignedTransaction};
 use near_primitives::types::validator_stake::{ValidatorStake, ValidatorStakeIter};
 use near_primitives::types::{
@@ -112,10 +110,8 @@ pub struct ApplyChunkResult {
     pub bandwidth_requests: Option<BandwidthRequests>,
     /// Used only for a sanity check.
     pub bandwidth_scheduler_state_hash: CryptoHash,
-    /// Code-hashes of the contracts accessed (called) while applying the chunk.
-    pub contract_accesses: BTreeSet<CodeHash>,
-    /// Code-hashes of the contracts deployed while applying the chunk.
-    pub contract_deploys: BTreeSet<CodeHash>,
+    /// Contracts accessed and deployed while applying the chunk.
+    pub contract_updates: ContractUpdates,
 }
 
 impl ApplyChunkResult {

chain/client/src/stateless_validation/partial_witness/partial_witness_actor.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeSet;
+use std::collections::HashSet;
 use std::sync::Arc;
 
 use itertools::Itertools;
@@ -368,7 +368,7 @@ impl PartialWitnessActor {
         let runtime_config = self
             .runtime
             .get_runtime_config(self.epoch_manager.get_epoch_protocol_version(&key.epoch_id)?)?;
-        let missing_contract_hashes = BTreeSet::from_iter(
+        let missing_contract_hashes = HashSet::from_iter(
             accesses
                 .contracts()
                 .iter()

chain/client/src/stateless_validation/partial_witness/partial_witness_tracker.rs

@@ -1,4 +1,4 @@
-use std::collections::BTreeSet;
+use std::collections::HashSet;
 use std::num::NonZeroUsize;
 use std::sync::Arc;
 
@@ -43,7 +43,7 @@ enum AccessedContractsState {
     Unknown,
     /// Received `ChunkContractAccesses` and sent `ContractCodeRequest`,
     /// waiting for response from the chunk producer.
-    Requested { contract_hashes: BTreeSet<CodeHash>, requested_at: Instant },
+    Requested { contract_hashes: HashSet<CodeHash>, requested_at: Instant },
     /// Received a valid `ContractCodeResponse`.
     Received(Vec<CodeBytes>),
 }
@@ -145,7 +145,7 @@ struct CacheEntry {
 
 enum CacheUpdate {
     WitnessPart(PartialEncodedStateWitness, Arc<WitnessEncoder>),
-    AccessedContractHashes(BTreeSet<CodeHash>),
+    AccessedContractHashes(HashSet<CodeHash>),
     AccessedContractCodes(Vec<CodeBytes>),
 }
 
@@ -234,7 +234,7 @@ impl CacheEntry {
         }
     }
 
-    fn set_requested_contracts(&mut self, contract_hashes: BTreeSet<CodeHash>) {
+    fn set_requested_contracts(&mut self, contract_hashes: HashSet<CodeHash>) {
         match &self.accessed_contracts {
             AccessedContractsState::Unknown => {
                 self.accessed_contracts = AccessedContractsState::Requested {
@@ -251,7 +251,7 @@ impl CacheEntry {
     fn set_received_contracts(&mut self, contract_codes: Vec<CodeBytes>) {
         match &self.accessed_contracts {
             AccessedContractsState::Requested { contract_hashes, requested_at } => {
-                let actual = BTreeSet::from_iter(
+                let actual = HashSet::from_iter(
                     contract_codes.iter().map(|code| CodeHash(CryptoHash::hash_bytes(&code.0))),
                 );
                 let expected = contract_hashes;
@@ -380,7 +380,7 @@ impl PartialEncodedStateWitnessTracker {
     pub fn store_accessed_contract_hashes(
         &mut self,
         key: ChunkProductionKey,
-        hashes: BTreeSet<CodeHash>,
+        hashes: HashSet<CodeHash>,
     ) -> Result<(), Error> {
         tracing::debug!(target: "client", ?key, ?hashes, "store_accessed_contract_hashes");
         let update = CacheUpdate::AccessedContractHashes(hashes);

chain/client/src/stateless_validation/state_witness_producer.rs

@@ -1,4 +1,4 @@
-use std::collections::{BTreeSet, HashMap, HashSet};
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use near_async::messaging::{CanSend, IntoSender};
@@ -12,7 +12,7 @@ use near_primitives::hash::{hash, CryptoHash};
 use near_primitives::receipt::Receipt;
 use near_primitives::sharding::{ChunkHash, ReceiptProof, ShardChunk, ShardChunkHeader};
 use near_primitives::stateless_validation::contract_distribution::{
-    ChunkContractAccesses, CodeHash,
+    ChunkContractAccesses, ContractUpdates,
 };
 use near_primitives::stateless_validation::state_witness::{
     ChunkStateTransition, ChunkStateWitness,
@@ -37,8 +37,7 @@ struct StateTransitionData {
     main_transition: ChunkStateTransition,
     implicit_transitions: Vec<ChunkStateTransition>,
     applied_receipts_hash: CryptoHash,
-    contract_accesses: BTreeSet<CodeHash>,
-    contract_deploys: BTreeSet<CodeHash>,
+    contract_updates: ContractUpdates,
 }
 
 /// Result of creating witness.
@@ -48,10 +47,8 @@ struct StateTransitionData {
 pub(crate) struct CreateWitnessResult {
     /// State witness created.
     pub(crate) state_witness: ChunkStateWitness,
-    /// Code-hashes of contracts accessed while applying the previous chunk.
-    pub(crate) contract_accesses: BTreeSet<CodeHash>,
-    /// Code-hashes of contracts deployed while applying the previous chunk.
-    pub(crate) contract_deploys: BTreeSet<CodeHash>,
+    /// Contracts accessed and deployed while applying the chunk.
+    pub contract_updates: ContractUpdates,
 }
 
 impl Client {
@@ -76,14 +73,13 @@ impl Client {
 
         let my_signer =
             validator_signer.as_ref().ok_or(Error::NotAValidator(format!("send state witness")))?;
-        let CreateWitnessResult { state_witness, contract_accesses, contract_deploys } = self
-            .create_state_witness(
-                my_signer.validator_id().clone(),
-                prev_block_header,
-                prev_chunk_header,
-                chunk,
-                transactions_storage_proof,
-            )?;
+        let CreateWitnessResult { state_witness, contract_updates } = self.create_state_witness(
+            my_signer.validator_id().clone(),
+            prev_block_header,
+            prev_chunk_header,
+            chunk,
+            transactions_storage_proof,
+        )?;
 
         if self.config.save_latest_witnesses {
             self.chain.chain_store.save_latest_chunk_state_witness(&state_witness)?;
@@ -108,12 +104,12 @@ impl Client {
         if ProtocolFeature::ExcludeContractCodeFromStateWitness.enabled(protocol_version) {
             // TODO(#11099): Currently we consume contract_deploys only for the following log message. Distribute it to validators
             // that will not validate the current witness so that they can follow-up with requesting the contract code.
-            tracing::debug!(target: "client", ?contract_accesses, ?contract_deploys, "Contract accesses and deploys while sending state witness");
-            if !contract_accesses.is_empty() {
-                self.send_contract_accesses_to_chunk_validators(
+            tracing::debug!(target: "client", ?contract_updates, "Contract accesses and deploys while sending state witness");
+            if !contract_updates.contract_accesses.is_empty() {
+                self.send_contract_updates_to_validators(
                     epoch_id,
                     &chunk_header,
-                    contract_accesses,
+                    contract_updates,
                     my_signer.as_ref(),
                 );
             }
@@ -143,8 +139,7 @@ impl Client {
             main_transition,
             implicit_transitions,
             applied_receipts_hash,
-            contract_accesses,
-            contract_deploys,
+            contract_updates,
         } = self.collect_state_transition_data(&chunk_header, prev_chunk_header)?;
 
         let new_transactions = chunk.transactions().to_vec();
@@ -178,7 +173,7 @@ impl Client {
             new_transactions,
             new_transactions_validation_state,
         );
-        Ok(CreateWitnessResult { state_witness, contract_accesses, contract_deploys })
+        Ok(CreateWitnessResult { state_witness, contract_updates })
     }
 
     /// Collect state transition data necessary to produce state witness for
@@ -219,7 +214,7 @@ impl Client {
             if current_shard_id != next_shard_id {
                 // If shard id changes, we need to get implicit state
                 // transition from current shard id to the next shard id.
-                let (chunk_state_transition, _, _, _) =
+                let (chunk_state_transition, _, _) =
                     self.get_state_transition(&current_block_hash, &next_epoch_id, next_shard_id)?;
                 implicit_transitions.push(chunk_state_transition);
             }
@@ -231,7 +226,7 @@ impl Client {
             }
 
             // Add implicit state transition.
-            let (chunk_state_transition, _, _, _) = self.get_state_transition(
+            let (chunk_state_transition, _, _) = self.get_state_transition(
                 &current_block_hash,
                 &current_epoch_id,
                 current_shard_id,
@@ -247,19 +242,17 @@ impl Client {
         implicit_transitions.reverse();
 
         // Get the main state transition.
-        let (main_transition, receipts_hash, contract_accesses, contract_deploys) =
-            if prev_chunk_header.is_genesis() {
-                self.get_genesis_state_transition(&main_block, &epoch_id, shard_id)?
-            } else {
-                self.get_state_transition(&main_block, &epoch_id, shard_id)?
-            };
+        let (main_transition, receipts_hash, contract_updates) = if prev_chunk_header.is_genesis() {
+            self.get_genesis_state_transition(&main_block, &epoch_id, shard_id)?
+        } else {
+            self.get_state_transition(&main_block, &epoch_id, shard_id)?
+        };
 
         Ok(StateTransitionData {
             main_transition,
             implicit_transitions,
             applied_receipts_hash: receipts_hash,
-            contract_accesses: BTreeSet::from_iter(contract_accesses.into_iter()),
-            contract_deploys: BTreeSet::from_iter(contract_deploys.into_iter()),
+            contract_updates,
         })
     }
 
@@ -269,8 +262,7 @@ impl Client {
         block_hash: &CryptoHash,
         epoch_id: &EpochId,
         shard_id: ShardId,
-    ) -> Result<(ChunkStateTransition, CryptoHash, BTreeSet<CodeHash>, BTreeSet<CodeHash>), Error>
-    {
+    ) -> Result<(ChunkStateTransition, CryptoHash, ContractUpdates), Error> {
         let shard_uid = self.chain.epoch_manager.shard_id_to_uid(shard_id, epoch_id)?;
         let stored_chunk_state_transition_data = self
             .chain
@@ -303,15 +295,18 @@ impl Client {
                     contract_deploys,
                 }) => (base_state, receipts_hash, contract_accesses, contract_deploys),
             };
+        let contract_updates = ContractUpdates {
+            contract_accesses: HashSet::from_iter(contract_accesses.into_iter()),
+            contract_deploys: HashSet::from_iter(contract_deploys.into_iter()),
+        };
         Ok((
             ChunkStateTransition {
                 block_hash: *block_hash,
                 base_state,
                 post_state_root: *self.chain.get_chunk_extra(block_hash, &shard_uid)?.state_root(),
             },
             receipts_hash,
-            BTreeSet::from_iter(contract_accesses.into_iter()),
-            BTreeSet::from_iter(contract_deploys.into_iter()),
+            contract_updates,
         ))
     }
 
@@ -320,8 +315,7 @@ impl Client {
         block_hash: &CryptoHash,
         epoch_id: &EpochId,
         shard_id: ShardId,
-    ) -> Result<(ChunkStateTransition, CryptoHash, BTreeSet<CodeHash>, BTreeSet<CodeHash>), Error>
-    {
+    ) -> Result<(ChunkStateTransition, CryptoHash, ContractUpdates), Error> {
         let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, &epoch_id)?;
         Ok((
             ChunkStateTransition {
@@ -331,7 +325,6 @@ impl Client {
             },
             hash(&borsh::to_vec::<[Receipt]>(&[]).unwrap()),
             Default::default(),
-            Default::default(),
         ))
     }
 
@@ -422,13 +415,15 @@ impl Client {
     /// Sends the contract accesses to the same chunk validators
     /// (except for the chunk producers that track the same shard),
     /// which will receive the state witness for the new chunk.  
-    fn send_contract_accesses_to_chunk_validators(
+    fn send_contract_updates_to_validators(
         &self,
         epoch_id: &EpochId,
         chunk_header: &ShardChunkHeader,
-        contract_accesses: BTreeSet<CodeHash>,
+        contract_updates: ContractUpdates,
         my_signer: &ValidatorSigner,
     ) {
+        let ContractUpdates { contract_accesses, .. } = contract_updates;
+
         let chunk_production_key = ChunkProductionKey {
             epoch_id: *epoch_id,
             shard_id: chunk_header.shard_id(),