ab_farmer_components/
shard_commitment.rs

1//! Utilities related to shard commitments
2
3use ab_core_primitives::hashes::Blake3Hash;
4use ab_core_primitives::segments::{HistorySize, SegmentIndex};
5use ab_core_primitives::shard::NumShards;
6use ab_core_primitives::solutions::{
7    ShardCommitmentHash, ShardMembershipEntropy, SolutionShardCommitment,
8};
9use ab_io_type::trivial_type::TrivialType;
10use ab_merkle_tree::unbalanced::UnbalancedMerkleTree;
11use blake3::Hasher;
12use parking_lot::RwLock;
13use schnellru::{ByLength, LruMap};
14use std::iter;
15use std::mem::MaybeUninit;
16use std::sync::Arc;
17
18const SHARD_COMMITMENTS_ROOT_CACHE_SIZE: ByLength = ByLength::new(32);
19
20#[derive(Debug)]
21struct Inner {
22    shard_commitments_seed: Blake3Hash,
23    lru: LruMap<HistorySize, ShardCommitmentHash, ByLength>,
24}
25
26/// Cache for shard commitments roots to avoid recomputing them repeatedly
27#[derive(Debug, Clone)]
28pub struct ShardCommitmentsRootsCache {
29    inner: Arc<RwLock<Inner>>,
30}
31
32impl ShardCommitmentsRootsCache {
33    /// Create a new instance
34    pub fn new(shard_commitments_seed: Blake3Hash) -> Self {
35        Self {
36            inner: Arc::new(RwLock::new(Inner {
37                shard_commitments_seed,
38                lru: LruMap::new(SHARD_COMMITMENTS_ROOT_CACHE_SIZE),
39            })),
40        }
41    }
42
43    /// Seed used during instantiation
44    pub fn shard_commitments_seed(&self) -> Blake3Hash {
45        self.inner.read().shard_commitments_seed
46    }
47
48    /// Get root for a specified history size.
49    ///
50    /// Root will be recomputed unless already known.
51    pub fn get(&self, history_size: HistorySize) -> ShardCommitmentHash {
52        if let Some(root) = self.inner.read().lru.peek(&history_size).copied() {
53            return root;
54        }
55
56        let inner = &mut *self.inner.write();
57        // NOTE: See https://github.com/koute/schnellru/issues/7 for an explanation of the
58        // `Option` return type
59        *inner
60            .lru
61            .get_or_insert(history_size, || {
62                derive_shard_commitments_root(&inner.shard_commitments_seed, history_size)
63            })
64            .expect("Not limited by memory; qed")
65    }
66}
67
68/// Derive shard commitments root from the seed and history size
69pub fn derive_shard_commitments_root(
70    shard_commitments_seed: &Blake3Hash,
71    history_size: HistorySize,
72) -> ShardCommitmentHash {
73    let mut stream = {
74        let mut hasher = Hasher::new_keyed(shard_commitments_seed);
75        hasher.update(SegmentIndex::from(history_size).as_bytes());
76        hasher.finalize_xof()
77    };
78
79    let mut index = 0;
80    let leaves = iter::from_fn(|| {
81        if index < SolutionShardCommitment::NUM_LEAVES {
82            let mut bytes = [0; ShardCommitmentHash::SIZE];
83            stream.fill(&mut bytes);
84
85            index += 1;
86
87            Some(bytes)
88        } else {
89            None
90        }
91    });
92
93    // NOTE: Using unbalanced implementation since balanced implementation requires allocation of
94    // leaves
95    let root = UnbalancedMerkleTree::compute_root_only::<
96        { SolutionShardCommitment::NUM_LEAVES as u64 },
97        _,
98        _,
99    >(leaves)
100    .expect("List of leaves is not empty; qed");
101
102    ShardCommitmentHash::new(root)
103}
104
105/// Derive solution shard commitment
106pub fn derive_solution_shard_commitment(
107    public_key_hash: &Blake3Hash,
108    shard_commitments_seed: &Blake3Hash,
109    shard_commitments_root: &ShardCommitmentHash,
110    history_size: HistorySize,
111    shard_membership_entropy: &ShardMembershipEntropy,
112    num_shards: NumShards,
113) -> SolutionShardCommitment {
114    let mut stream = {
115        let mut hasher = Hasher::new_keyed(shard_commitments_seed);
116        hasher.update(SegmentIndex::from(history_size).as_bytes());
117        hasher.finalize_xof()
118    };
119
120    let leaf_index = num_shards.derive_shard_commitment_index(
121        public_key_hash,
122        shard_commitments_root,
123        shard_membership_entropy,
124        history_size,
125    ) as usize;
126
127    let mut leaf = [0; _];
128    let mut index = 0;
129    let leaves = iter::from_fn(|| {
130        if index < SolutionShardCommitment::NUM_LEAVES {
131            let mut bytes = [0; ShardCommitmentHash::SIZE];
132            stream.fill(&mut bytes);
133
134            if index == leaf_index {
135                leaf = bytes;
136            }
137
138            index += 1;
139
140            Some(bytes)
141        } else {
142            None
143        }
144    });
145
146    let mut proof = [MaybeUninit::uninit(); _];
147    // NOTE: Using unbalanced implementation since balanced implementation requires an allocation
148    // and uses a lot more RAM
149    let (_root, computed_proof) = UnbalancedMerkleTree::compute_root_and_proof_in::<
150        { SolutionShardCommitment::NUM_LEAVES as u64 },
151        _,
152        _,
153    >(leaves, leaf_index, &mut proof)
154    .expect("Index is always within the list of leaves; qed");
155    debug_assert_eq!(computed_proof.len(), proof.len());
156
157    // SAFETY: Checked above that it is fully initialized
158    let proof = unsafe { MaybeUninit::array_assume_init(proof) };
159
160    SolutionShardCommitment {
161        root: *shard_commitments_root,
162        proof: ShardCommitmentHash::array_from_repr(proof),
163        leaf: ShardCommitmentHash::new(leaf),
164    }
165}
ab_farmer_components/shard_commitment.rs

ab_farmer_components/
shard_commitment.rs
