feat: add trie key global contract code (#12766)

Part of #12715.

For each global contract we store its code replicated on every shard.
This PR adds a trie key to support that.

chain/chain/src/flat_storage_resharder.rs

@@ -1086,7 +1086,8 @@ fn shard_split_handle_key_value(
         col::DELAYED_RECEIPT_OR_INDICES
         | col::PROMISE_YIELD_INDICES
         | col::PROMISE_YIELD_TIMEOUT
-        | col::BANDWIDTH_SCHEDULER_STATE => {
+        | col::BANDWIDTH_SCHEDULER_STATE
+        | col::GLOBAL_CONTRACT_CODE => {
             copy_kv_to_all_children(&split_params, key, value, store_update)
         }
         col::BUFFERED_RECEIPT_INDICES

core/primitives/src/trie_key.rs

@@ -1,6 +1,6 @@
 use crate::hash::CryptoHash;
 use crate::types::AccountId;
-use borsh::{BorshDeserialize, BorshSerialize};
+use borsh::{to_vec, BorshDeserialize, BorshSerialize};
 use near_crypto::PublicKey;
 use near_primitives_core::types::ShardId;
 use near_schema_checker_lib::ProtocolSchema;
@@ -62,6 +62,9 @@ pub mod col {
     pub const BUFFERED_RECEIPT_GROUPS_QUEUE_DATA: u8 = 16;
     /// A single item of `ReceiptGroupsQueue`. Values are of type `ReceiptGroup`.
     pub const BUFFERED_RECEIPT_GROUPS_QUEUE_ITEM: u8 = 17;
+    /// Global contract code instance. Values are contract blobs,
+    /// the same as for `CONTRACT_CODE`.
+    pub const GLOBAL_CONTRACT_CODE: u8 = 18;
 
     /// All columns except those used for the delayed receipts queue, the yielded promises
     /// queue, and the outgoing receipts buffer, which are global state for the shard.
@@ -77,7 +80,7 @@ pub mod col {
         (PROMISE_YIELD_RECEIPT, "PromiseYieldReceipt"),
     ];
 
-    pub const ALL_COLUMNS_WITH_NAMES: [(u8, &'static str); 17] = [
+    pub const ALL_COLUMNS_WITH_NAMES: [(u8, &'static str); 18] = [
         (ACCOUNT, "Account"),
         (CONTRACT_CODE, "ContractCode"),
         (ACCESS_KEY, "AccessKey"),
@@ -95,9 +98,32 @@ pub mod col {
         (BANDWIDTH_SCHEDULER_STATE, "BandwidthSchedulerState"),
         (BUFFERED_RECEIPT_GROUPS_QUEUE_DATA, "BufferedReceiptGroupsQueueData"),
         (BUFFERED_RECEIPT_GROUPS_QUEUE_ITEM, "BufferedReceiptGroupsQueueItem"),
+        (GLOBAL_CONTRACT_CODE, "GlobalContractCode"),
     ];
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, BorshDeserialize, BorshSerialize, ProtocolSchema)]
+pub enum GlobalContractCodeIdentifier {
+    CodeHash(CryptoHash),
+    AccountId(AccountId),
+}
+
+impl GlobalContractCodeIdentifier {
+    pub fn len(&self) -> usize {
+        1 + match self {
+            Self::CodeHash(hash) => hash.as_bytes().len(),
+            Self::AccountId(account_id) => {
+                // Corresponds to String repr in borsh spec
+                size_of::<u32>() + account_id.len()
+            }
+        }
+    }
+
+    pub fn append_into(&self, buf: &mut Vec<u8>) {
+        buf.extend(to_vec(self).unwrap());
+    }
+}
+
 /// Describes the key of a specific key-value record in a state trie.
 #[derive(Debug, Clone, PartialEq, Eq, BorshDeserialize, BorshSerialize, ProtocolSchema)]
 pub enum TrieKey {
@@ -193,6 +219,9 @@ pub enum TrieKey {
         receiving_shard: ShardId,
         index: u64,
     },
+    GlobalContractCode {
+        identifier: GlobalContractCodeIdentifier,
+    },
 }
 
 /// Provides `len` function.
@@ -277,6 +306,9 @@ impl TrieKey {
                     + std::mem::size_of::<u64>()
                     + std::mem::size_of_val(index)
             }
+            TrieKey::GlobalContractCode { identifier } => {
+                col::GLOBAL_CONTRACT_CODE.len() + identifier.len()
+            }
         }
     }
 
@@ -370,6 +402,10 @@ impl TrieKey {
                 buf.extend(&receiving_shard.to_le_bytes());
                 buf.extend(&index.to_le_bytes());
             }
+            TrieKey::GlobalContractCode { identifier } => {
+                buf.push(col::GLOBAL_CONTRACT_CODE);
+                identifier.append_into(buf);
+            }
         };
         debug_assert_eq!(expected_len, buf.len() - start_len);
     }
@@ -401,6 +437,9 @@ impl TrieKey {
             TrieKey::BandwidthSchedulerState => None,
             TrieKey::BufferedReceiptGroupsQueueData { .. } => None,
             TrieKey::BufferedReceiptGroupsQueueItem { .. } => None,
+            // Even though global contract code might be deployed under account id, it doesn't
+            // correspond to the data stored for that account id, so always returning None here.
+            TrieKey::GlobalContractCode { .. } => None,
         }
     }
 }
@@ -932,4 +971,20 @@ mod tests {
             "buffered receipt trie key optimization broken, must fit in a u16"
         );
     }
+
+    #[test]
+    fn test_global_contract_code_identifier_len() {
+        check_global_contract_code_identifier_len(GlobalContractCodeIdentifier::CodeHash(
+            CryptoHash::hash_bytes(&[42]),
+        ));
+        check_global_contract_code_identifier_len(GlobalContractCodeIdentifier::AccountId(
+            "alice.near".parse().unwrap(),
+        ));
+    }
+
+    fn check_global_contract_code_identifier_len(identifier: GlobalContractCodeIdentifier) {
+        let mut buf = Vec::new();
+        identifier.append_into(&mut buf);
+        assert_eq!(buf.len(), identifier.len());
+    }
 }

core/primitives/src/types.rs

@@ -368,6 +368,8 @@ impl StateChanges {
                 TrieKey::BandwidthSchedulerState => {}
                 TrieKey::BufferedReceiptGroupsQueueData { .. } => {}
                 TrieKey::BufferedReceiptGroupsQueueItem { .. } => {}
+                // Global contract code is not a part of account, so ignoring it as well.
+                TrieKey::GlobalContractCode { .. } => {}
             }
         }
 

core/store/src/trie/ops/resharding.rs

@@ -53,7 +53,8 @@ fn boundary_account_to_intervals(
             col::DELAYED_RECEIPT_OR_INDICES
             | col::PROMISE_YIELD_INDICES
             | col::PROMISE_YIELD_TIMEOUT
-            | col::BANDWIDTH_SCHEDULER_STATE => {
+            | col::BANDWIDTH_SCHEDULER_STATE
+            | col::GLOBAL_CONTRACT_CODE => {
                 // This section contains the keys that we need to copy to both shards.
                 intervals.push(get_interval_for_copy_to_both_children(prefix))
             }
@@ -361,6 +362,7 @@ mod tests {
                 ..vec![col::BUFFERED_RECEIPT_GROUPS_QUEUE_DATA + 1],
             vec![col::BUFFERED_RECEIPT_GROUPS_QUEUE_ITEM]
                 ..vec![col::BUFFERED_RECEIPT_GROUPS_QUEUE_ITEM + 1],
+            vec![col::GLOBAL_CONTRACT_CODE]..vec![col::GLOBAL_CONTRACT_CODE + 1],
         ];
         assert!(left_intervals.iter().all(|range| range.start < range.end));
         for (actual, expected) in left_intervals.iter().zip_eq(expected_left_intervals.iter()) {
@@ -385,6 +387,7 @@ mod tests {
             append_key(col::PROMISE_YIELD_RECEIPT, &alice_account)
                 ..vec![col::PROMISE_YIELD_RECEIPT + 1],
             vec![col::BANDWIDTH_SCHEDULER_STATE]..vec![col::BANDWIDTH_SCHEDULER_STATE + 1],
+            vec![col::GLOBAL_CONTRACT_CODE]..vec![col::GLOBAL_CONTRACT_CODE + 1],
         ];
         assert!(right_intervals.iter().all(|range| range.start < range.end));
         for (actual, expected) in right_intervals.iter().zip_eq(expected_right_intervals.iter()) {

tools/protocol-schema-check/res/protocol_schema.toml

@@ -82,7 +82,7 @@ CompressedEpochSyncProof = 1117061636
 CongestionInfo = 2682682461
 CongestionInfoV1 = 2571332168
 ConnectionInfoRepr = 3621760869
-ConsolidatedStateChange = 1148312806
+ConsolidatedStateChange = 2962557518
 ContractCacheKey = 1745279861
 ContractCodeRequest = 1530126649
 ContractCodeRequestInner = 1643875081
@@ -144,6 +144,7 @@ FlatStorageStatus = 3964465569
 FunctionCallAction = 2405840012
 FunctionCallError = 3652274053
 FunctionCallPermission = 1517509673
+GlobalContractCodeIdentifier = 3486582651
 GlobalContractDeployMode = 1256912586
 GlobalContractIdentifier = 31664100
 Handshake = 115352275
@@ -199,7 +200,7 @@ PromiseYieldIndices = 405847541
 PromiseYieldTimeout = 3189361393
 PublicKey = 601042198
 RawStateChange = 206262877
-RawStateChangesWithTrieKey = 36781564
+RawStateChangesWithTrieKey = 4065520827
 RawTrieNode = 4239211001
 RawTrieNodeWithSize = 1474149765
 ReasonForBan = 792112981
@@ -279,7 +280,7 @@ TransactionV0 = 2318136898
 TransactionV1 = 2359321608
 TransferAction = 1078380396
 TrieChanges = 2613580820
-TrieKey = 1352104737
+TrieKey = 544589632
 TrieQueueIndices = 2601394796
 TrieRefcountAddition = 2117109883
 TrieRefcountSubtraction = 2150368599