ab_core_primitives/
shard.rs

1//! Shard-related primitives
2
3use crate::hashes::Blake3Hash;
4use crate::nano_u256::NanoU256;
5use crate::segments::HistorySize;
6use crate::solutions::{ShardCommitmentHash, ShardMembershipEntropy, SolutionShardCommitment};
7use ab_blake3::single_block_keyed_hash;
8use ab_io_type::trivial_type::TrivialType;
9use core::num::{NonZeroU16, NonZeroU32, NonZeroU128};
10use core::ops::RangeInclusive;
11use derive_more::Display;
12#[cfg(feature = "scale-codec")]
13use parity_scale_codec::{Decode, Encode, Input, MaxEncodedLen};
14#[cfg(feature = "serde")]
15use serde::{Deserialize, Deserializer, Serialize};
16
17const INTERMEDIATE_SHARDS_RANGE: RangeInclusive<u32> = 1..=1023;
18const INTERMEDIATE_SHARD_BITS: u32 = 10;
19const INTERMEDIATE_SHARD_MASK: u32 = u32::MAX >> (u32::BITS - INTERMEDIATE_SHARD_BITS);
20
21/// A kind of shard.
22///
23/// Schematically, the hierarchy of shards is as follows:
24/// ```text
25///                          Beacon chain
26///                          /          \
27///      Intermediate shard 1            Intermediate shard 2
28///              /  \                            /  \
29/// Leaf shard 11   Leaf shard 12   Leaf shard 22   Leaf shard 22
30/// ```
31///
32/// Beacon chain has index 0, intermediate shards indices 1..=1023. Leaf shards share the same least
33/// significant 10 bits as their respective intermediate shards, so leaf shards of intermediate
34/// shard 1 have indices like 1025,2049,3097,etc.
35///
36/// If represented as least significant bits first (as it will be in the formatted address):
37/// ```text
38/// Intermediate shard bits
39///     \            /
40///      1000_0000_0001_0000_0000
41///                 /            \
42///                Leaf shard bits
43/// ```
44///
45/// Note that shards that have 10 least significant bits set to 0 (corresponds to the beacon chain)
46/// are not leaf shards, in fact, they are not even physical shards that have nodes in general. The
47/// meaning of these shards is TBD, currently they are called "phantom" shards and may end up
48/// containing some system contracts with special meaning, but no blocks will ever exist for such
49/// shards.
50#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
51pub enum ShardKind {
52    /// Beacon chain shard
53    BeaconChain,
54    /// Intermediate shard directly below the beacon chain that has child shards
55    IntermediateShard,
56    /// Leaf shard, which doesn't have child shards
57    LeafShard,
58    /// TODO
59    Phantom,
60}
61
62impl ShardKind {
63    /// Try to convert to real shard kind.
64    ///
65    /// Returns `None` for phantom shard.
66    #[inline(always)]
67    pub fn to_real(self) -> Option<RealShardKind> {
68        match self {
69            ShardKind::BeaconChain => Some(RealShardKind::BeaconChain),
70            ShardKind::IntermediateShard => Some(RealShardKind::IntermediateShard),
71            ShardKind::LeafShard => Some(RealShardKind::LeafShard),
72            ShardKind::Phantom => None,
73        }
74    }
75}
76
77/// Real shard kind for which a block may exist, see [`ShardKind`] for more details
78#[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)]
79pub enum RealShardKind {
80    /// Beacon chain shard
81    BeaconChain,
82    /// Intermediate shard directly below the beacon chain that has child shards
83    IntermediateShard,
84    /// Leaf shard, which doesn't have child shards
85    LeafShard,
86}
87
88impl From<RealShardKind> for ShardKind {
89    #[inline(always)]
90    fn from(shard_kind: RealShardKind) -> Self {
91        match shard_kind {
92            RealShardKind::BeaconChain => ShardKind::BeaconChain,
93            RealShardKind::IntermediateShard => ShardKind::IntermediateShard,
94            RealShardKind::LeafShard => ShardKind::LeafShard,
95        }
96    }
97}
98
99/// Shard index
100#[derive(Debug, Display, Copy, Clone, Hash, Ord, PartialOrd, Eq, PartialEq, TrivialType)]
101#[cfg_attr(feature = "scale-codec", derive(Encode, Decode, MaxEncodedLen))]
102#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
103#[repr(C)]
104pub struct ShardIndex(u32);
105
106const impl Default for ShardIndex {
107    #[inline(always)]
108    fn default() -> Self {
109        Self::BEACON_CHAIN
110    }
111}
112
113const impl From<ShardIndex> for u32 {
114    #[inline(always)]
115    fn from(shard_index: ShardIndex) -> Self {
116        shard_index.0
117    }
118}
119
120impl ShardIndex {
121    /// Beacon chain
122    pub const BEACON_CHAIN: Self = Self(0);
123    /// Max possible shard index
124    pub const MAX_SHARD_INDEX: u32 = Self::MAX_SHARDS.get() - 1;
125    /// Max possible number of shards
126    pub const MAX_SHARDS: NonZeroU32 = NonZeroU32::new(2u32.pow(20)).expect("Not zero; qed");
127    /// Max possible number of addresses per shard
128    pub const MAX_ADDRESSES_PER_SHARD: NonZeroU128 =
129        NonZeroU128::new(2u128.pow(108)).expect("Not zero; qed");
130
131    /// Create shard index from `u32`.
132    ///
133    /// Returns `None` if `shard_index > ShardIndex::MAX_SHARD_INDEX`
134    ///
135    /// This is typically only necessary for low-level code.
136    #[inline(always)]
137    pub const fn new(shard_index: u32) -> Option<Self> {
138        if shard_index > Self::MAX_SHARD_INDEX {
139            return None;
140        }
141
142        Some(Self(shard_index))
143    }
144
145    /// Whether the shard index corresponds to the beacon chain
146    #[inline(always)]
147    pub const fn is_beacon_chain(&self) -> bool {
148        self.0 == Self::BEACON_CHAIN.0
149    }
150
151    /// Whether the shard index corresponds to an intermediate shard
152    #[inline(always)]
153    pub const fn is_intermediate_shard(&self) -> bool {
154        self.0 >= *INTERMEDIATE_SHARDS_RANGE.start() && self.0 <= *INTERMEDIATE_SHARDS_RANGE.end()
155    }
156
157    /// Whether the shard index corresponds to an intermediate shard
158    #[inline(always)]
159    pub const fn is_leaf_shard(&self) -> bool {
160        if self.0 <= *INTERMEDIATE_SHARDS_RANGE.end() || self.0 > Self::MAX_SHARD_INDEX {
161            return false;
162        }
163
164        self.0 & INTERMEDIATE_SHARD_MASK != 0
165    }
166
167    /// Whether the shard index corresponds to a real shard
168    #[inline(always)]
169    pub const fn is_real(&self) -> bool {
170        !self.is_phantom_shard()
171    }
172
173    /// Whether the shard index corresponds to a phantom shard
174    #[inline(always)]
175    pub const fn is_phantom_shard(&self) -> bool {
176        if self.0 <= *INTERMEDIATE_SHARDS_RANGE.end() || self.0 > Self::MAX_SHARD_INDEX {
177            return false;
178        }
179
180        self.0 & INTERMEDIATE_SHARD_MASK == 0
181    }
182
183    /// Check if this shard is a child shard of `parent`
184    #[inline(always)]
185    pub const fn is_child_of(self, parent: Self) -> bool {
186        match self.shard_kind() {
187            Some(ShardKind::BeaconChain) => false,
188            Some(ShardKind::IntermediateShard | ShardKind::Phantom) => parent.is_beacon_chain(),
189            Some(ShardKind::LeafShard) => {
190                // Check that the least significant bits match
191                self.0 & INTERMEDIATE_SHARD_MASK == parent.0
192            }
193            None => false,
194        }
195    }
196
197    /// Get index of the parent shard (for leaf and intermediate shards)
198    #[inline(always)]
199    pub const fn parent_shard(self) -> Option<ShardIndex> {
200        match self.shard_kind()? {
201            ShardKind::BeaconChain => None,
202            ShardKind::IntermediateShard | ShardKind::Phantom => Some(ShardIndex::BEACON_CHAIN),
203            ShardKind::LeafShard => Some(Self(self.0 & INTERMEDIATE_SHARD_MASK)),
204        }
205    }
206
207    /// Get shard kind
208    #[inline(always)]
209    pub const fn shard_kind(&self) -> Option<ShardKind> {
210        if self.0 == Self::BEACON_CHAIN.0 {
211            Some(ShardKind::BeaconChain)
212        } else if self.0 >= *INTERMEDIATE_SHARDS_RANGE.start()
213            && self.0 <= *INTERMEDIATE_SHARDS_RANGE.end()
214        {
215            Some(ShardKind::IntermediateShard)
216        } else if self.0 > Self::MAX_SHARD_INDEX {
217            None
218        } else if self.0 & INTERMEDIATE_SHARD_MASK == 0 {
219            // Check if the least significant bits correspond to the beacon chain
220            Some(ShardKind::Phantom)
221        } else {
222            Some(ShardKind::LeafShard)
223        }
224    }
225}
226
227/// Unchecked number of shards in the network.
228///
229/// Should be converted into [`NumShards`] before use.
230#[derive(Debug, Copy, Clone, Eq, PartialEq, TrivialType)]
231#[cfg_attr(feature = "scale-codec", derive(Encode, Decode, MaxEncodedLen))]
232#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
233#[repr(C)]
234pub struct NumShardsUnchecked {
235    /// The number of intermediate shards
236    pub intermediate_shards: u16,
237    /// The number of leaf shards per intermediate shard
238    pub leaf_shards_per_intermediate_shard: u16,
239}
240
241impl From<NumShards> for NumShardsUnchecked {
242    fn from(value: NumShards) -> Self {
243        Self {
244            intermediate_shards: value.intermediate_shards.get(),
245            leaf_shards_per_intermediate_shard: value.leaf_shards_per_intermediate_shard.get(),
246        }
247    }
248}
249
250/// Number of shards in the network
251#[derive(Debug, Copy, Clone, Eq, PartialEq)]
252#[cfg_attr(feature = "scale-codec", derive(Encode, MaxEncodedLen))]
253#[cfg_attr(feature = "serde", derive(Serialize))]
254pub struct NumShards {
255    /// The number of intermediate shards
256    intermediate_shards: NonZeroU16,
257    /// The number of leaf shards per intermediate shard
258    leaf_shards_per_intermediate_shard: NonZeroU16,
259}
260
261#[cfg(feature = "scale-codec")]
262impl Decode for NumShards {
263    fn decode<I>(input: &mut I) -> Result<Self, parity_scale_codec::Error>
264    where
265        I: Input,
266    {
267        let intermediate_shards = Decode::decode(input)
268            .map_err(|error| error.chain("Could not decode `NumShards::intermediate_shards`"))?;
269        let leaf_shards_per_intermediate_shard = Decode::decode(input).map_err(|error| {
270            error.chain("Could not decode `NumShards::leaf_shards_per_intermediate_shard`")
271        })?;
272
273        Self::new(intermediate_shards, leaf_shards_per_intermediate_shard)
274            .ok_or_else(|| "Invalid `NumShards`".into())
275    }
276}
277
278#[cfg(feature = "serde")]
279impl<'de> Deserialize<'de> for NumShards {
280    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
281    where
282        D: Deserializer<'de>,
283    {
284        #[derive(Deserialize)]
285        struct NumShards {
286            intermediate_shards: NonZeroU16,
287            leaf_shards_per_intermediate_shard: NonZeroU16,
288        }
289
290        let num_shards_inner = NumShards::deserialize(deserializer)?;
291
292        Self::new(
293            num_shards_inner.intermediate_shards,
294            num_shards_inner.leaf_shards_per_intermediate_shard,
295        )
296        .ok_or_else(|| serde::de::Error::custom("Invalid `NumShards`"))
297    }
298}
299
300impl TryFrom<NumShardsUnchecked> for NumShards {
301    type Error = ();
302
303    fn try_from(value: NumShardsUnchecked) -> Result<Self, Self::Error> {
304        Self::new(
305            NonZeroU16::new(value.intermediate_shards).ok_or(())?,
306            NonZeroU16::new(value.leaf_shards_per_intermediate_shard).ok_or(())?,
307        )
308        .ok_or(())
309    }
310}
311
312impl NumShards {
313    /// Create a new instance from a number of intermediate shards and leaf shards per
314    /// intermediate shard.
315    ///
316    /// Returns `None` if inputs are invalid.
317    ///
318    /// This is typically only necessary for low-level code.
319    #[inline(always)]
320    pub const fn new(
321        intermediate_shards: NonZeroU16,
322        leaf_shards_per_intermediate_shard: NonZeroU16,
323    ) -> Option<Self> {
324        if intermediate_shards.get()
325            > (*INTERMEDIATE_SHARDS_RANGE.end() - *INTERMEDIATE_SHARDS_RANGE.start() + 1) as u16
326        {
327            return None;
328        }
329
330        let num_shards = Self {
331            intermediate_shards,
332            leaf_shards_per_intermediate_shard,
333        };
334
335        if num_shards.leaf_shards() > ShardIndex::MAX_SHARDS {
336            return None;
337        }
338
339        Some(num_shards)
340    }
341
342    /// The number of intermediate shards
343    #[inline(always)]
344    pub const fn intermediate_shards(self) -> NonZeroU16 {
345        self.intermediate_shards
346    }
347    /// The number of leaf shards per intermediate shard
348    #[inline(always)]
349    pub const fn leaf_shards_per_intermediate_shard(self) -> NonZeroU16 {
350        self.leaf_shards_per_intermediate_shard
351    }
352
353    /// Total number of leaf shards in the network
354    #[inline(always)]
355    pub const fn leaf_shards(&self) -> NonZeroU32 {
356        NonZeroU32::new(
357            self.intermediate_shards.get() as u32
358                * self.leaf_shards_per_intermediate_shard.get() as u32,
359        )
360        .expect("Not zero; qed")
361    }
362
363    /// Iterator over all intermediate shards
364    #[inline(always)]
365    pub fn iter_intermediate_shards(&self) -> impl Iterator<Item = ShardIndex> {
366        INTERMEDIATE_SHARDS_RANGE
367            .take(usize::from(self.intermediate_shards.get()))
368            .map(ShardIndex)
369    }
370
371    /// Iterator over all intermediate shards
372    #[inline(always)]
373    pub fn iter_leaf_shards(&self) -> impl Iterator<Item = ShardIndex> {
374        self.iter_intermediate_shards()
375            .flat_map(|intermediate_shard| {
376                (0..u32::from(self.leaf_shards_per_intermediate_shard.get())).map(
377                    move |leaf_shard_index| {
378                        ShardIndex(
379                            (leaf_shard_index << INTERMEDIATE_SHARD_BITS) | intermediate_shard.0,
380                        )
381                    },
382                )
383            })
384    }
385
386    /// Derive shard index that should be used in a solution
387    #[inline]
388    pub fn derive_shard_index(
389        &self,
390        public_key_hash: &Blake3Hash,
391        shard_commitments_root: &ShardCommitmentHash,
392        shard_membership_entropy: &ShardMembershipEntropy,
393        history_size: HistorySize,
394    ) -> ShardIndex {
395        let hash = single_block_keyed_hash(public_key_hash, &{
396            let mut bytes_to_hash = [0u8; {
397                ShardCommitmentHash::SIZE
398                    + ShardMembershipEntropy::SIZE
399                    + HistorySize::SIZE as usize
400            }];
401            bytes_to_hash[..ShardCommitmentHash::SIZE]
402                .copy_from_slice(shard_commitments_root.as_bytes());
403            bytes_to_hash[ShardCommitmentHash::SIZE..][..ShardMembershipEntropy::SIZE]
404                .copy_from_slice(shard_membership_entropy.as_bytes());
405            bytes_to_hash[ShardCommitmentHash::SIZE + ShardMembershipEntropy::SIZE..]
406                .copy_from_slice(history_size.as_bytes());
407            bytes_to_hash
408        })
409        .expect("Input is smaller than block size; qed");
410        // Going through `NanoU256` because the total number of shards is not guaranteed to be a
411        // power of two
412        let shard_index_offset =
413            NanoU256::from_le_bytes(hash) % u64::from(self.leaf_shards().get());
414
415        self.iter_leaf_shards()
416            .nth(shard_index_offset as usize)
417            .unwrap_or(ShardIndex::BEACON_CHAIN)
418    }
419
420    /// Derive shard commitment index that should be used in a solution.
421    ///
422    /// Returned index is always within `0`..[`SolutionShardCommitment::NUM_LEAVES`] range.
423    #[inline]
424    pub fn derive_shard_commitment_index(
425        &self,
426        public_key_hash: &Blake3Hash,
427        shard_commitments_root: &ShardCommitmentHash,
428        shard_membership_entropy: &ShardMembershipEntropy,
429        history_size: HistorySize,
430    ) -> u32 {
431        let hash = single_block_keyed_hash(public_key_hash, &{
432            let mut bytes_to_hash = [0u8; {
433                ShardCommitmentHash::SIZE
434                    + ShardMembershipEntropy::SIZE
435                    + HistorySize::SIZE as usize
436            }];
437            bytes_to_hash[..ShardCommitmentHash::SIZE]
438                .copy_from_slice(shard_commitments_root.as_bytes());
439            bytes_to_hash[ShardCommitmentHash::SIZE..][..ShardMembershipEntropy::SIZE]
440                .copy_from_slice(shard_membership_entropy.as_bytes());
441            bytes_to_hash[ShardCommitmentHash::SIZE + ShardMembershipEntropy::SIZE..]
442                .copy_from_slice(history_size.as_bytes());
443            bytes_to_hash
444        })
445        .expect("Input is smaller than block size; qed");
446        const {
447            assert!(SolutionShardCommitment::NUM_LEAVES.is_power_of_two());
448        }
449        u32::from_le_bytes([hash[0], hash[1], hash[2], hash[3]])
450            % SolutionShardCommitment::NUM_LEAVES as u32
451    }
452
453    /// More efficient version of [`Self::derive_shard_index()`] and
454    /// [`Self::derive_shard_commitment_index()`] in a single call, see those functions for details
455    #[inline]
456    pub fn derive_shard_index_and_shard_commitment_index(
457        &self,
458        public_key_hash: &Blake3Hash,
459        shard_commitments_root: &ShardCommitmentHash,
460        shard_membership_entropy: &ShardMembershipEntropy,
461        history_size: HistorySize,
462    ) -> (ShardIndex, u32) {
463        let hash = single_block_keyed_hash(public_key_hash, &{
464            let mut bytes_to_hash = [0u8; {
465                ShardCommitmentHash::SIZE
466                    + ShardMembershipEntropy::SIZE
467                    + HistorySize::SIZE as usize
468            }];
469            bytes_to_hash[..ShardCommitmentHash::SIZE]
470                .copy_from_slice(shard_commitments_root.as_bytes());
471            bytes_to_hash[ShardCommitmentHash::SIZE..][..ShardMembershipEntropy::SIZE]
472                .copy_from_slice(shard_membership_entropy.as_bytes());
473            bytes_to_hash[ShardCommitmentHash::SIZE + ShardMembershipEntropy::SIZE..]
474                .copy_from_slice(history_size.as_bytes());
475            bytes_to_hash
476        })
477        .expect("Input is smaller than block size; qed");
478
479        // Going through `NanoU256` because the total number of shards is not guaranteed to be a
480        // power of two
481        let shard_index_offset =
482            NanoU256::from_le_bytes(hash) % u64::from(self.leaf_shards().get());
483
484        let shard_index = self
485            .iter_leaf_shards()
486            .nth(shard_index_offset as usize)
487            .unwrap_or(ShardIndex::BEACON_CHAIN);
488
489        const {
490            assert!(SolutionShardCommitment::NUM_LEAVES.is_power_of_two());
491        }
492        let shard_commitment_index = u32::from_le_bytes([hash[0], hash[1], hash[2], hash[3]])
493            % SolutionShardCommitment::NUM_LEAVES as u32;
494
495        (shard_index, shard_commitment_index)
496    }
497}
ab_core_primitives/shard.rs

ab_core_primitives/
shard.rs
