unc_primitives/
state_sync.rs

1use crate::hash::CryptoHash;
2use crate::merkle::MerklePath;
3use crate::sharding::{
4    ReceiptProof, ShardChunk, ShardChunkHeader, ShardChunkHeaderV1, ShardChunkV1,
5};
6use crate::types::{BlockHeight, EpochId, ShardId, StateRoot, StateRootNode};
7use borsh::{BorshDeserialize, BorshSerialize};
8use std::sync::Arc;
9use unc_primitives_core::types::EpochHeight;
10
11#[derive(PartialEq, Eq, Clone, Debug, BorshSerialize, BorshDeserialize)]
12pub struct ReceiptProofResponse(pub CryptoHash, pub Arc<Vec<ReceiptProof>>);
13
14#[derive(PartialEq, Eq, Clone, Debug, BorshSerialize, BorshDeserialize)]
15pub struct RootProof(pub CryptoHash, pub MerklePath);
16
17#[derive(PartialEq, Eq, Clone, Debug, BorshSerialize, BorshDeserialize)]
18pub struct StateHeaderKey(pub ShardId, pub CryptoHash);
19
20#[derive(PartialEq, Eq, Clone, Debug, BorshSerialize, BorshDeserialize)]
21pub struct StatePartKey(pub CryptoHash, pub ShardId, pub u64 /* PartId */);
22
23#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
24pub struct ShardStateSyncResponseHeaderV1 {
25    pub chunk: ShardChunkV1,
26    pub chunk_proof: MerklePath,
27    pub prev_chunk_header: Option<ShardChunkHeaderV1>,
28    pub prev_chunk_proof: Option<MerklePath>,
29    pub incoming_receipts_proofs: Vec<ReceiptProofResponse>,
30    pub root_proofs: Vec<Vec<RootProof>>,
31    pub state_root_node: StateRootNode,
32}
33
34#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
35pub struct ShardStateSyncResponseHeaderV2 {
36    pub chunk: ShardChunk,
37    pub chunk_proof: MerklePath,
38    pub prev_chunk_header: Option<ShardChunkHeader>,
39    pub prev_chunk_proof: Option<MerklePath>,
40    pub incoming_receipts_proofs: Vec<ReceiptProofResponse>,
41    pub root_proofs: Vec<Vec<RootProof>>,
42    pub state_root_node: StateRootNode,
43}
44
45#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
46pub enum CachedParts {
47    AllParts,
48    NoParts,
49    /// Represents a subset of parts cached.
50    /// Can represent both NoParts and AllParts, but in those cases use the
51    /// corresponding enum values for efficiency.
52    BitArray(BitArray),
53}
54
55/// Represents an array of boolean values in a compact form.
56#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
57pub struct BitArray {
58    data: Vec<u8>,
59    capacity: u64,
60}
61
62impl BitArray {
63    pub fn new(capacity: u64) -> Self {
64        let num_bytes = (capacity + 7) / 8;
65        Self { data: vec![0; num_bytes as usize], capacity }
66    }
67
68    pub fn set_bit(&mut self, bit: u64) {
69        assert!(bit < self.capacity);
70        self.data[(bit / 8) as usize] |= 1 << (bit % 8);
71    }
72}
73
74#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
75pub enum ShardStateSyncResponseHeader {
76    V1(ShardStateSyncResponseHeaderV1),
77    V2(ShardStateSyncResponseHeaderV2),
78}
79
80impl ShardStateSyncResponseHeader {
81    #[inline]
82    pub fn take_chunk(self) -> ShardChunk {
83        match self {
84            Self::V1(header) => ShardChunk::V1(header.chunk),
85            Self::V2(header) => header.chunk,
86        }
87    }
88
89    #[inline]
90    pub fn cloned_chunk(&self) -> ShardChunk {
91        match self {
92            Self::V1(header) => ShardChunk::V1(header.chunk.clone()),
93            Self::V2(header) => header.chunk.clone(),
94        }
95    }
96
97    #[inline]
98    pub fn cloned_prev_chunk_header(&self) -> Option<ShardChunkHeader> {
99        match self {
100            Self::V1(header) => header.prev_chunk_header.clone().map(ShardChunkHeader::V1),
101            Self::V2(header) => header.prev_chunk_header.clone(),
102        }
103    }
104
105    #[inline]
106    pub fn chunk_height_included(&self) -> BlockHeight {
107        match self {
108            Self::V1(header) => header.chunk.header.height_included,
109            Self::V2(header) => header.chunk.height_included(),
110        }
111    }
112
113    #[inline]
114    pub fn chunk_prev_state_root(&self) -> StateRoot {
115        match self {
116            Self::V1(header) => header.chunk.header.inner.prev_state_root,
117            Self::V2(header) => header.chunk.prev_state_root(),
118        }
119    }
120
121    #[inline]
122    pub fn chunk_proof(&self) -> &MerklePath {
123        match self {
124            Self::V1(header) => &header.chunk_proof,
125            Self::V2(header) => &header.chunk_proof,
126        }
127    }
128
129    #[inline]
130    pub fn prev_chunk_proof(&self) -> &Option<MerklePath> {
131        match self {
132            Self::V1(header) => &header.prev_chunk_proof,
133            Self::V2(header) => &header.prev_chunk_proof,
134        }
135    }
136
137    #[inline]
138    pub fn incoming_receipts_proofs(&self) -> &[ReceiptProofResponse] {
139        match self {
140            Self::V1(header) => &header.incoming_receipts_proofs,
141            Self::V2(header) => &header.incoming_receipts_proofs,
142        }
143    }
144
145    #[inline]
146    pub fn root_proofs(&self) -> &[Vec<RootProof>] {
147        match self {
148            Self::V1(header) => &header.root_proofs,
149            Self::V2(header) => &header.root_proofs,
150        }
151    }
152
153    #[inline]
154    pub fn state_root_node(&self) -> &StateRootNode {
155        match self {
156            Self::V1(header) => &header.state_root_node,
157            Self::V2(header) => &header.state_root_node,
158        }
159    }
160
161    pub fn num_state_parts(&self) -> u64 {
162        get_num_state_parts(self.state_root_node().memory_usage)
163    }
164}
165
166#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
167pub struct ShardStateSyncResponseV1 {
168    pub header: Option<ShardStateSyncResponseHeaderV1>,
169    pub part: Option<(u64, Vec<u8>)>,
170}
171
172#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
173pub struct ShardStateSyncResponseV2 {
174    pub header: Option<ShardStateSyncResponseHeaderV2>,
175    pub part: Option<(u64, Vec<u8>)>,
176}
177
178#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
179pub struct ShardStateSyncResponseV3 {
180    pub header: Option<ShardStateSyncResponseHeaderV2>,
181    pub part: Option<(u64, Vec<u8>)>,
182    /// Parts that can be provided **cheaply**.
183    // Can be `None` only if both `header` and `part` are `None`.
184    pub cached_parts: Option<CachedParts>,
185    /// Whether the node can provide parts for this epoch of this shard.
186    /// Assumes that a node can either provide all state parts or no state parts.
187    pub can_generate: bool,
188}
189
190#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)]
191pub enum ShardStateSyncResponse {
192    V1(ShardStateSyncResponseV1),
193    V2(ShardStateSyncResponseV2),
194    V3(ShardStateSyncResponseV3),
195}
196
197impl ShardStateSyncResponse {
198    pub fn part_id(&self) -> Option<u64> {
199        match self {
200            Self::V1(response) => response.part_id(),
201            Self::V2(response) => response.part.as_ref().map(|(part_id, _)| *part_id),
202            Self::V3(response) => response.part.as_ref().map(|(part_id, _)| *part_id),
203        }
204    }
205
206    pub fn take_header(self) -> Option<ShardStateSyncResponseHeader> {
207        match self {
208            Self::V1(response) => response.header.map(ShardStateSyncResponseHeader::V1),
209            Self::V2(response) => response.header.map(ShardStateSyncResponseHeader::V2),
210            Self::V3(response) => response.header.map(ShardStateSyncResponseHeader::V2),
211        }
212    }
213
214    pub fn part(&self) -> &Option<(u64, Vec<u8>)> {
215        match self {
216            Self::V1(response) => &response.part,
217            Self::V2(response) => &response.part,
218            Self::V3(response) => &response.part,
219        }
220    }
221
222    pub fn take_part(self) -> Option<(u64, Vec<u8>)> {
223        match self {
224            Self::V1(response) => response.part,
225            Self::V2(response) => response.part,
226            Self::V3(response) => response.part,
227        }
228    }
229
230    pub fn can_generate(&self) -> bool {
231        match self {
232            Self::V1(_response) => false,
233            Self::V2(_response) => false,
234            Self::V3(response) => response.can_generate,
235        }
236    }
237
238    pub fn cached_parts(&self) -> &Option<CachedParts> {
239        match self {
240            Self::V1(_response) => &None,
241            Self::V2(_response) => &None,
242            Self::V3(response) => &response.cached_parts,
243        }
244    }
245}
246
247impl ShardStateSyncResponseV1 {
248    pub fn part_id(&self) -> Option<u64> {
249        self.part.as_ref().map(|(part_id, _)| *part_id)
250    }
251}
252
253pub const STATE_PART_MEMORY_LIMIT: bytesize::ByteSize = bytesize::ByteSize(30 * bytesize::MIB);
254
255pub fn get_num_state_parts(memory_usage: u64) -> u64 {
256    (memory_usage + STATE_PART_MEMORY_LIMIT.as_u64() - 1) / STATE_PART_MEMORY_LIMIT.as_u64()
257}
258
259#[derive(BorshSerialize, BorshDeserialize, Debug, Clone)]
260/// Represents the progress of dumps state of a shard.
261pub enum StateSyncDumpProgress {
262    /// Represents two cases:
263    /// * An epoch dump is complete
264    /// * The node is running its first epoch and there is nothing to dump.
265    AllDumped {
266        /// The dumped state corresponds to the state at the beginning of the specified epoch.
267        epoch_id: EpochId,
268        epoch_height: EpochHeight,
269    },
270    /// Represents the case of an epoch being partially dumped.
271    InProgress {
272        /// The dumped state corresponds to the state at the beginning of the specified epoch.
273        epoch_id: EpochId,
274        epoch_height: EpochHeight,
275        /// Block hash of the first block of the epoch.
276        /// The dumped state corresponds to the state before applying this block.
277        sync_hash: CryptoHash,
278    },
279}
280
281#[cfg(test)]
282mod tests {
283    use crate::state_sync::{get_num_state_parts, STATE_PART_MEMORY_LIMIT};
284
285    #[test]
286    fn test_get_num_state_parts() {
287        assert_eq!(get_num_state_parts(0), 0);
288        assert_eq!(get_num_state_parts(1), 1);
289        assert_eq!(get_num_state_parts(STATE_PART_MEMORY_LIMIT.as_u64()), 1);
290        assert_eq!(get_num_state_parts(STATE_PART_MEMORY_LIMIT.as_u64() + 1), 2);
291        assert_eq!(get_num_state_parts(STATE_PART_MEMORY_LIMIT.as_u64() * 100), 100);
292        assert_eq!(get_num_state_parts(STATE_PART_MEMORY_LIMIT.as_u64() * 100 + 1), 101);
293    }
294}
unc_primitives/state_sync.rs

unc_primitives/
state_sync.rs
