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

ab_farmer_components/
shard_commitment.rs
