Send to all shards & add trie key col

chain/chain/src/chain.rs

@@ -4427,9 +4427,15 @@ impl Chain {
     ) -> HashMap<ShardId, Vec<Receipt>> {
         let mut result = HashMap::new();
         for receipt in receipts {
-            let shard_id = shard_layout.account_id_to_shard_id(receipt.receiver_id());
-            let entry = result.entry(shard_id).or_insert_with(Vec::new);
-            entry.push(receipt)
+            if receipt.send_to_all_shards() {
+                for shard_id in shard_layout.shard_ids() {
+                    result.entry(shard_id).or_insert_with(Vec::new).push(receipt.clone());
+                }
+            } else {
+                let shard_id = shard_layout.account_id_to_shard_id(receipt.receiver_id());
+                let entry = result.entry(shard_id).or_insert_with(Vec::new);
+                entry.push(receipt);
+            }
         }
         result
     }
@@ -4450,13 +4456,22 @@ impl Chain {
         }
         let mut cache = HashMap::new();
         for receipt in receipts {
-            let &mut shard_id = cache
-                .entry(receipt.receiver_id())
-                .or_insert_with(|| shard_layout.account_id_to_shard_id(receipt.receiver_id()));
-            // This unwrap should be safe as we pre-populated the map with all
-            // valid shard ids.
-            let shard_index = shard_layout.get_shard_index(shard_id).unwrap();
-            result_map.get_mut(&shard_index).unwrap().1.push(receipt);
+            if receipt.send_to_all_shards() {
+                for shard_id in shard_layout.shard_ids() {
+                    // This unwrap should be safe as we pre-populated the map with all
+                    // valid shard ids.
+                    let shard_index = shard_layout.get_shard_index(shard_id).unwrap();
+                    result_map.get_mut(&shard_index).unwrap().1.push(receipt);
+                }
+            } else {
+                let &mut shard_id = cache
+                    .entry(receipt.receiver_id())
+                    .or_insert_with(|| shard_layout.account_id_to_shard_id(receipt.receiver_id()));
+                // This unwrap should be safe as we pre-populated the map with all
+                // valid shard ids.
+                let shard_index = shard_layout.get_shard_index(shard_id).unwrap();
+                result_map.get_mut(&shard_index).unwrap().1.push(receipt);
+            };
         }
 
         let mut result_vec = vec![];

chain/chain/src/runtime/tests.rs

@@ -372,8 +372,14 @@ impl TestEnv {
         let shard_layout = self.epoch_manager.get_shard_layout_from_prev_block(&new_hash).unwrap();
         let mut new_receipts = HashMap::<_, Vec<Receipt>>::new();
         for receipt in all_receipts {
-            let shard_id = shard_layout.account_id_to_shard_id(receipt.receiver_id());
-            new_receipts.entry(shard_id).or_default().push(receipt);
+            if receipt.send_to_all_shards() {
+                for shard_id in shard_layout.shard_ids() {
+                    new_receipts.entry(shard_id).or_default().push(receipt.clone());
+                }
+            } else {
+                let shard_id = shard_layout.account_id_to_shard_id(receipt.receiver_id());
+                new_receipts.entry(shard_id).or_default().push(receipt);
+            }
         }
         self.last_receipts = new_receipts;
         self.last_proposals = all_proposals;

chain/chain/src/store/mod.rs

@@ -387,9 +387,10 @@ pub fn filter_incoming_receipts_for_shard(
         let mut filtered_receipts = vec![];
         let ReceiptProof(receipts, shard_proof) = receipt_proof.clone();
         for receipt in receipts {
-            let receiver_shard_id =
-                target_shard_layout.account_id_to_shard_id(receipt.receiver_id());
-            if receiver_shard_id == target_shard_id {
+            if receipt.send_to_all_shards()
+                || target_shard_layout.account_id_to_shard_id(receipt.receiver_id())
+                    == target_shard_id
+            {
                 tracing::trace!(target: "chain", receipt_id=?receipt.receipt_id(), "including receipt");
                 filtered_receipts.push(receipt);
             } else {

chain/chain/src/test_utils.rs

@@ -297,7 +297,8 @@ mod test {
             let shard_receipts: Vec<Receipt> = receipts
                 .iter()
                 .filter(|&receipt| {
-                    shard_layout.account_id_to_shard_id(receipt.receiver_id()) == shard_id
+                    receipt.send_to_all_shards()
+                        || shard_layout.account_id_to_shard_id(receipt.receiver_id()) == shard_id
                 })
                 .cloned()
                 .collect();

core/primitives/src/receipt.rs

@@ -485,6 +485,10 @@ impl Receipt {
         *self.receipt_id()
     }
 
+    pub fn send_to_all_shards(&self) -> bool {
+        matches!(self.receipt(), ReceiptEnum::GlobalContractDitribution(..))
+    }
+
     /// Generates a receipt with a transfer from system for a given balance without a receipt_id.
     /// This should be used for token refunds instead of gas refunds. It inherits priority from the parent receipt.
     /// It doesn't refund the allowance of the access key. For gas refunds use `new_gas_refund`.

core/primitives/src/trie_key.rs

@@ -62,6 +62,7 @@ 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;
+    pub const GLOBAL_CONTRACT_DATA: 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 +78,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,6 +96,7 @@ pub mod col {
         (BANDWIDTH_SCHEDULER_STATE, "BandwidthSchedulerState"),
         (BUFFERED_RECEIPT_GROUPS_QUEUE_DATA, "BufferedReceiptGroupsQueueData"),
         (BUFFERED_RECEIPT_GROUPS_QUEUE_ITEM, "BufferedReceiptGroupsQueueItem"),
+        (GLOBAL_CONTRACT_DATA, "GlobalContractData"),
     ];
 }
 
@@ -193,6 +195,9 @@ pub enum TrieKey {
         receiving_shard: ShardId,
         index: u64,
     },
+    GlobalContractCode {
+        code_hash: CryptoHash,
+    },
 }
 
 /// Provides `len` function.
@@ -277,6 +282,9 @@ impl TrieKey {
                     + std::mem::size_of::<u64>()
                     + std::mem::size_of_val(index)
             }
+            TrieKey::GlobalContractCode { code_hash } => {
+                col::GLOBAL_CONTRACT_DATA.len() + code_hash.as_ref().len()
+            }
         }
     }
 
@@ -370,6 +378,10 @@ impl TrieKey {
                 buf.extend(&receiving_shard.to_le_bytes());
                 buf.extend(&index.to_le_bytes());
             }
+            TrieKey::GlobalContractCode { code_hash } => {
+                buf.push(col::GLOBAL_CONTRACT_DATA);
+                buf.extend(code_hash.as_ref());
+            },
         };
         debug_assert_eq!(expected_len, buf.len() - start_len);
     }
@@ -401,6 +413,7 @@ impl TrieKey {
             TrieKey::BandwidthSchedulerState => None,
             TrieKey::BufferedReceiptGroupsQueueData { .. } => None,
             TrieKey::BufferedReceiptGroupsQueueItem { .. } => None,
+            TrieKey::GlobalContractCode { .. } => None,
         }
     }
 }

core/primitives/src/types.rs

@@ -360,6 +360,7 @@ impl StateChanges {
                 TrieKey::BandwidthSchedulerState => {}
                 TrieKey::BufferedReceiptGroupsQueueData { .. } => {}
                 TrieKey::BufferedReceiptGroupsQueueItem { .. } => {}
+                TrieKey::GlobalContractCode { .. } => {}
             }
         }