Skip to main content

hotshot_types/data/
vid_disperse.rs

1// Copyright (c) 2021-2024 Espresso Systems (espressosys.com)
2// This file is part of the HotShot repository.
3
4// You should have received a copy of the MIT License
5// along with the HotShot repository. If not, see <https://mit-license.org/>.
6
7//! This module provides types for VID disperse related data structures.
8//!
9//! We have three types of VID disperse related structs:
10//!
11//! 1. `ADVZ*`: VID V0, The most original VID scheme, which has guaranteed recovery but very inefficient.
12//! 2. `AvidM*`: VID V1, the efficient VID scheme, where we use it after the epocn upgrade. It's more
13//!    efficient but doesn't guarantee recovery. A VID V1 commitment could correspond to some junk
14//!    data, there'll be a proof of incorrect encoding in this case.
15//! 3. `AvidmGf2*`: VID V2, almost the same as VID V1 but we have a much more efficient recovery
16//!    implementation.
17
18use std::{
19    collections::BTreeMap, fmt::Debug, hash::Hash, marker::PhantomData, ops::Range, sync::Arc,
20    time::Duration,
21};
22
23use alloy_primitives::U256;
24use hotshot_utils::anytrace::*;
25use jf_advz::{VidDisperse as JfVidDisperse, VidScheme};
26use serde::{Deserialize, Serialize};
27use tokio::{task::spawn_blocking, time::Instant};
28
29use super::ns_table::parse_ns_table;
30use crate::{
31    PeerConfig,
32    data::{EpochNumber, ViewNumber},
33    epoch_membership::{EpochMembership, EpochMembershipCoordinator},
34    message::Proposal,
35    simple_vote::HasEpoch,
36    traits::{
37        BlockPayload,
38        block_contents::EncodeBytes,
39        node_implementation::NodeType,
40        signature_key::{SignatureKey, StakeTableEntryType},
41    },
42    vid::{
43        advz::{ADVZCommitment, ADVZCommon, ADVZScheme, ADVZShare, advz_scheme},
44        avidm::{AvidMCommitment, AvidMCommon, AvidMScheme, AvidMShare, init_avidm_param},
45        avidm_gf2::{
46            AvidmGf2Commitment, AvidmGf2Common, AvidmGf2Param, AvidmGf2Scheme, AvidmGf2Share,
47            init_avidm_gf2_param,
48        },
49    },
50    vote::HasViewNumber,
51};
52
53impl<NODE: NodeType> HasEpoch for ADVZDisperse<NODE> {
54    fn epoch(&self) -> Option<EpochNumber> {
55        self.epoch
56    }
57}
58
59impl<NODE: NodeType> HasEpoch for AvidMDisperse<NODE> {
60    fn epoch(&self) -> Option<EpochNumber> {
61        self.epoch
62    }
63}
64
65impl<NODE: NodeType> HasEpoch for AvidMDisperseShare<NODE> {
66    fn epoch(&self) -> Option<EpochNumber> {
67        self.epoch
68    }
69}
70
71impl<NODE: NodeType> HasEpoch for AvidmGf2Disperse<NODE> {
72    fn epoch(&self) -> Option<EpochNumber> {
73        self.epoch
74    }
75}
76
77impl<NODE: NodeType> HasEpoch for AvidmGf2DisperseShare<NODE> {
78    fn epoch(&self) -> Option<EpochNumber> {
79        self.epoch
80    }
81}
82
83/// ADVZ dispersal data
84#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
85pub struct ADVZDisperse<TYPES: NodeType> {
86    /// The view number for which this VID data is intended
87    pub view_number: ViewNumber,
88    /// Epoch the data of this proposal belongs to
89    pub epoch: Option<EpochNumber>,
90    /// Epoch to which the recipients of this VID belong to
91    pub target_epoch: Option<EpochNumber>,
92    /// VidCommitment calculated based on the number of nodes in `target_epoch`.
93    pub payload_commitment: ADVZCommitment,
94    /// A storage node's key and its corresponding VID share
95    pub shares: BTreeMap<TYPES::SignatureKey, ADVZShare>,
96    /// VID common data sent to all storage nodes
97    pub common: ADVZCommon,
98}
99
100impl<TYPES: NodeType> HasViewNumber for ADVZDisperse<TYPES> {
101    fn view_number(&self) -> ViewNumber {
102        self.view_number
103    }
104}
105
106impl<TYPES: NodeType> ADVZDisperse<TYPES> {
107    /// Create VID dispersal from a specified membership for the target epoch.
108    /// Uses the specified function to calculate share dispersal
109    /// Allows for more complex stake table functionality
110    async fn from_membership(
111        view_number: ViewNumber,
112        mut vid_disperse: JfVidDisperse<ADVZScheme>,
113        membership: &EpochMembershipCoordinator<TYPES>,
114        target_epoch: Option<EpochNumber>,
115        data_epoch: Option<EpochNumber>,
116    ) -> Result<Self> {
117        let shares = membership
118            .stake_table_for_epoch(target_epoch)?
119            .stake_table()
120            .map(|entry| entry.stake_table_entry.public_key())
121            .map(|node| (node.clone(), vid_disperse.shares.remove(0)))
122            .collect();
123
124        Ok(Self {
125            view_number,
126            shares,
127            common: vid_disperse.common,
128            payload_commitment: vid_disperse.commit,
129            epoch: data_epoch,
130            target_epoch,
131        })
132    }
133
134    /// Calculate the vid disperse information from the payload given a view, epoch and membership,
135    /// If the sender epoch is missing, it means it's the same as the target epoch.
136    ///
137    /// # Errors
138    /// Returns an error if the disperse or commitment calculation fails
139    #[allow(clippy::panic)]
140    pub async fn calculate_vid_disperse(
141        payload: &TYPES::BlockPayload,
142        membership: &EpochMembershipCoordinator<TYPES>,
143        view: ViewNumber,
144        target_epoch: Option<EpochNumber>,
145        data_epoch: Option<EpochNumber>,
146    ) -> Result<(Self, Duration)> {
147        let num_nodes = membership
148            .stake_table_for_epoch(target_epoch)?
149            .total_nodes();
150
151        let txns = payload.encode();
152
153        let now = Instant::now();
154        let vid_disperse = spawn_blocking(move || advz_scheme(num_nodes).disperse(&txns))
155            .await
156            .wrap()
157            .context(error!("Join error"))?
158            .wrap()
159            .context(|err| error!("Failed to calculate VID disperse. Error: {err}"))?;
160        let advz_scheme_duration = now.elapsed();
161
162        Ok((
163            Self::from_membership(view, vid_disperse, membership, target_epoch, data_epoch).await?,
164            advz_scheme_duration,
165        ))
166    }
167
168    /// This function splits a VID disperse into individual shares.
169    pub fn to_shares(self) -> Vec<ADVZDisperseShare<TYPES>> {
170        self.shares
171            .into_iter()
172            .map(|(recipient_key, share)| ADVZDisperseShare {
173                share,
174                recipient_key,
175                view_number: self.view_number,
176                common: self.common.clone(),
177                payload_commitment: self.payload_commitment,
178            })
179            .collect()
180    }
181
182    /// Split a VID disperse into a share proposal for each recipient.
183    pub fn to_share_proposals(
184        self,
185        signature: &<<TYPES as NodeType>::SignatureKey as SignatureKey>::PureAssembledSignatureType,
186    ) -> Vec<Proposal<TYPES, ADVZDisperseShare<TYPES>>> {
187        self.shares
188            .into_iter()
189            .map(|(recipient_key, share)| Proposal {
190                data: ADVZDisperseShare {
191                    share,
192                    recipient_key,
193                    view_number: self.view_number,
194                    payload_commitment: self.payload_commitment,
195                    common: self.common.clone(),
196                },
197                signature: signature.clone(),
198                _pd: PhantomData,
199            })
200            .collect()
201    }
202
203    /// Returns the payload length in bytes.
204    pub fn payload_byte_len(&self) -> u32 {
205        ADVZScheme::get_payload_byte_len(&self.common)
206    }
207}
208
209#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
210/// ADVZ share and associated metadata for a single node
211pub struct ADVZDisperseShare<TYPES: NodeType> {
212    /// The view number for which this VID data is intended
213    pub view_number: ViewNumber,
214    /// Block payload commitment
215    pub payload_commitment: ADVZCommitment,
216    /// A storage node's key and its corresponding VID share
217    pub share: ADVZShare,
218    /// VID common data sent to all storage nodes
219    pub common: ADVZCommon,
220    /// a public key of the share recipient
221    pub recipient_key: TYPES::SignatureKey,
222}
223
224impl<NODE: NodeType> HasEpoch for ADVZDisperseShare<NODE> {
225    fn epoch(&self) -> Option<EpochNumber> {
226        None
227    }
228}
229
230impl<TYPES: NodeType> HasViewNumber for ADVZDisperseShare<TYPES> {
231    fn view_number(&self) -> ViewNumber {
232        self.view_number
233    }
234}
235
236impl<TYPES: NodeType> ADVZDisperseShare<TYPES> {
237    /// Consume `self` and return a `Proposal`
238    pub fn to_proposal(
239        self,
240        private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
241    ) -> Option<Proposal<TYPES, Self>> {
242        let Ok(signature) =
243            TYPES::SignatureKey::sign(private_key, self.payload_commitment.as_ref())
244        else {
245            tracing::error!("VID: failed to sign dispersal share payload");
246            return None;
247        };
248        Some(Proposal {
249            signature,
250            _pd: PhantomData,
251            data: self,
252        })
253    }
254
255    /// Create `VidDisperse` out of an iterator to `VidDisperseShare`s
256    pub fn to_advz_disperse<'a, I>(mut it: I) -> Option<ADVZDisperse<TYPES>>
257    where
258        I: Iterator<Item = &'a Self>,
259    {
260        let first_vid_disperse_share = it.next()?.clone();
261        let mut share_map = BTreeMap::new();
262        share_map.insert(
263            first_vid_disperse_share.recipient_key,
264            first_vid_disperse_share.share,
265        );
266        let mut vid_disperse = ADVZDisperse {
267            view_number: first_vid_disperse_share.view_number,
268            epoch: None,
269            target_epoch: None,
270            payload_commitment: first_vid_disperse_share.payload_commitment,
271            common: first_vid_disperse_share.common,
272            shares: share_map,
273        };
274        it.for_each(|vid_disperse_share| {
275            vid_disperse.shares.insert(
276                vid_disperse_share.recipient_key.clone(),
277                vid_disperse_share.share.clone(),
278            );
279        });
280        Some(vid_disperse)
281    }
282
283    /// Check if vid common is consistent with the commitment.
284    pub fn is_consistent(&self) -> bool {
285        ADVZScheme::is_consistent(&self.payload_commitment, &self.common).is_ok()
286    }
287
288    /// Verify share assuming common data is already verified consistent.
289    /// Caller MUST call `is_consistent()` first.
290    pub fn verify_with_verified_common(&self) -> bool {
291        let total_weight = ADVZScheme::get_num_storage_nodes(&self.common) as usize;
292        advz_scheme(total_weight)
293            .verify_share(&self.share, &self.common, &self.payload_commitment)
294            .is_ok_and(|r| r.is_ok())
295    }
296
297    /// Internally verify the share given necessary information
298    pub fn verify(&self, _total_weight: usize) -> bool {
299        self.is_consistent() && self.verify_with_verified_common()
300    }
301
302    /// Returns the payload length in bytes.
303    pub fn payload_byte_len(&self) -> u32 {
304        ADVZScheme::get_payload_byte_len(&self.common)
305    }
306}
307
308/// AvidM dispersal data
309#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
310pub struct AvidMDisperse<TYPES: NodeType> {
311    /// The view number for which this VID data is intended
312    pub view_number: ViewNumber,
313    /// Epoch the data of this proposal belongs to
314    pub epoch: Option<EpochNumber>,
315    /// Epoch to which the recipients of this VID belong to
316    pub target_epoch: Option<EpochNumber>,
317    /// VidCommitment calculated based on the number of nodes in `target_epoch`.
318    pub payload_commitment: AvidMCommitment,
319    /// A storage node's key and its corresponding VID share
320    pub shares: BTreeMap<TYPES::SignatureKey, AvidMShare>,
321    /// Length of payload in bytes
322    pub payload_byte_len: usize,
323    /// VID common data sent to all storage nodes
324    pub common: AvidMCommon,
325}
326
327impl<TYPES: NodeType> HasViewNumber for AvidMDisperse<TYPES> {
328    fn view_number(&self) -> ViewNumber {
329        self.view_number
330    }
331}
332
333/// The target total stake to scale to for VID.
334pub const VID_TARGET_TOTAL_STAKE: u32 = 1000;
335
336/// The weights and total weight used in VID calculations
337struct Weights {
338    // weights, in stake table order
339    weights: Vec<u32>,
340
341    // total weight
342    total_weight: usize,
343}
344
345pub fn vid_total_weight<'a, T, I>(stake_table: I, epoch: Option<EpochNumber>) -> usize
346where
347    T: NodeType,
348    I: Iterator<Item = &'a PeerConfig<T>>,
349{
350    if epoch.is_none() {
351        stake_table
352            .fold(U256::ZERO, |acc, entry| {
353                acc + entry.stake_table_entry.stake()
354            })
355            .to::<usize>()
356    } else {
357        let stake_table = stake_table.cloned().collect::<Vec<_>>();
358        approximate_weights(&stake_table[..]).total_weight
359    }
360}
361
362fn approximate_weights<TYPES: NodeType>(stake_table: &[PeerConfig<TYPES>]) -> Weights {
363    let total_stake = stake_table.iter().fold(U256::ZERO, |acc, entry| {
364        acc + entry.stake_table_entry.stake()
365    });
366
367    let mut total_weight: usize = 0;
368
369    // don't attempt to scale if the total stake is small enough
370    if total_stake <= U256::from(VID_TARGET_TOTAL_STAKE) {
371        let weights = stake_table
372            .iter()
373            .map(|entry| entry.stake_table_entry.stake().to::<u32>())
374            .collect();
375
376        // Note: this panics if `total_stake` exceeds `usize::MAX`, but this shouldn't happen.
377        total_weight = total_stake.to::<usize>();
378
379        Weights {
380            weights,
381            total_weight,
382        }
383    } else {
384        let weights = stake_table
385            .iter()
386            .map(|entry| {
387                let weight: U256 = ((entry.stake_table_entry.stake()
388                    * U256::from(VID_TARGET_TOTAL_STAKE))
389                    / total_stake)
390                    + U256::ONE;
391
392                // Note: this panics if `weight` exceeds `usize::MAX`, but this shouldn't happen.
393                total_weight += weight.to::<usize>();
394
395                // Note: this panics if `weight` exceeds `u32::MAX`, but this shouldn't happen
396                // and would likely cause a stack overflow in the VID calculation anyway
397                weight.to::<u32>()
398            })
399            .collect();
400
401        Weights {
402            weights,
403            total_weight,
404        }
405    }
406}
407
408impl<TYPES: NodeType> AvidMDisperse<TYPES> {
409    /// Create VID dispersal from a specified membership for the target epoch.
410    /// Uses the specified function to calculate share dispersal
411    /// Allows for more complex stake table functionality
412    async fn from_membership(
413        view_number: ViewNumber,
414        commit: AvidMCommitment,
415        shares: &[AvidMShare],
416        common: AvidMCommon,
417        membership: &EpochMembership<TYPES>,
418        target_epoch: Option<EpochNumber>,
419        data_epoch: Option<EpochNumber>,
420    ) -> Result<Self> {
421        let payload_byte_len = shares[0].payload_byte_len();
422        let shares = membership
423            .coordinator
424            .stake_table_for_epoch(target_epoch)?
425            .stake_table()
426            .map(|entry| entry.stake_table_entry.public_key())
427            .zip(shares)
428            .map(|(node, share)| (node.clone(), share.clone()))
429            .collect();
430
431        Ok(Self {
432            view_number,
433            shares,
434            payload_commitment: commit,
435            epoch: data_epoch,
436            target_epoch,
437            payload_byte_len,
438            common,
439        })
440    }
441
442    /// Calculate the vid disperse information from the payload given a view, epoch and membership,
443    /// If the sender epoch is missing, it means it's the same as the target epoch.
444    ///
445    /// # Errors
446    /// Returns an error if the disperse or commitment calculation fails
447    #[allow(clippy::panic)]
448    #[allow(clippy::single_range_in_vec_init)]
449    pub async fn calculate_vid_disperse(
450        payload: &TYPES::BlockPayload,
451        membership: &EpochMembershipCoordinator<TYPES>,
452        view: ViewNumber,
453        target_epoch: Option<EpochNumber>,
454        data_epoch: Option<EpochNumber>,
455        metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
456    ) -> Result<(Self, Duration)> {
457        let target_mem = membership.stake_table_for_epoch(target_epoch)?;
458        let stake_table: Vec<_> = target_mem.stake_table().cloned().collect();
459        let approximate_weights = approximate_weights(&stake_table);
460
461        let txns = payload.encode();
462        let num_txns = txns.len();
463
464        let avidm_param = init_avidm_param(approximate_weights.total_weight)?;
465        let common = avidm_param.clone();
466
467        let ns_table = parse_ns_table(num_txns, &metadata.encode());
468        let ns_table_clone = ns_table.clone();
469
470        let now = Instant::now();
471        let (commit, shares) = spawn_blocking(move || {
472            AvidMScheme::ns_disperse(
473                &avidm_param,
474                &approximate_weights.weights,
475                &txns,
476                ns_table_clone,
477            )
478        })
479        .await
480        .wrap()
481        .context(error!("Join error"))?
482        .wrap()
483        .context(|err| error!("Failed to calculate VID disperse. Error: {err}"))?;
484        let ns_disperse_duration = now.elapsed();
485
486        Ok((
487            Self::from_membership(
488                view,
489                commit,
490                &shares,
491                common,
492                &target_mem,
493                target_epoch,
494                data_epoch,
495            )
496            .await?,
497            ns_disperse_duration,
498        ))
499    }
500
501    /// This function splits a VID disperse into individual shares.
502    pub fn to_shares(self) -> Vec<AvidMDisperseShare<TYPES>> {
503        self.shares
504            .into_iter()
505            .map(|(recipient_key, share)| AvidMDisperseShare {
506                share,
507                recipient_key,
508                view_number: self.view_number,
509                payload_commitment: self.payload_commitment,
510                epoch: self.epoch,
511                target_epoch: self.target_epoch,
512                common: self.common.clone(),
513            })
514            .collect()
515    }
516
517    /// Split a VID disperse into a share proposal for each recipient.
518    pub fn to_share_proposals(
519        self,
520        signature: &<<TYPES as NodeType>::SignatureKey as SignatureKey>::PureAssembledSignatureType,
521    ) -> Vec<Proposal<TYPES, AvidMDisperseShare<TYPES>>> {
522        self.shares
523            .into_iter()
524            .map(|(recipient_key, share)| Proposal {
525                data: AvidMDisperseShare {
526                    share,
527                    recipient_key,
528                    view_number: self.view_number,
529                    payload_commitment: self.payload_commitment,
530                    epoch: self.epoch,
531                    target_epoch: self.target_epoch,
532                    common: self.common.clone(),
533                },
534                signature: signature.clone(),
535                _pd: PhantomData,
536            })
537            .collect()
538    }
539
540    /// Construct a VID disperse from an iterator of disperse shares.
541    pub fn try_from_shares<'a, I>(mut it: I) -> Option<Self>
542    where
543        I: Iterator<Item = &'a AvidMDisperseShare<TYPES>>,
544    {
545        let first_vid_disperse_share = it.next()?.clone();
546        let payload_byte_len = first_vid_disperse_share.share.payload_byte_len();
547        let mut share_map = BTreeMap::new();
548        share_map.insert(
549            first_vid_disperse_share.recipient_key,
550            first_vid_disperse_share.share,
551        );
552        let mut vid_disperse = Self {
553            view_number: first_vid_disperse_share.view_number,
554            epoch: first_vid_disperse_share.epoch,
555            target_epoch: first_vid_disperse_share.target_epoch,
556            payload_commitment: first_vid_disperse_share.payload_commitment,
557            shares: share_map,
558            payload_byte_len,
559            common: first_vid_disperse_share.common,
560        };
561        it.for_each(|vid_disperse_share| {
562            vid_disperse.shares.insert(
563                vid_disperse_share.recipient_key.clone(),
564                vid_disperse_share.share.clone(),
565            );
566        });
567        Some(vid_disperse)
568    }
569
570    /// Returns the payload length in bytes.
571    pub fn payload_byte_len(&self) -> u32 {
572        self.payload_byte_len as u32
573    }
574}
575
576#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
577/// VID share and associated metadata for a single node
578pub struct AvidMDisperseShare<TYPES: NodeType> {
579    /// The view number for which this VID data is intended
580    pub view_number: ViewNumber,
581    /// The epoch number for which this VID data belongs to
582    pub epoch: Option<EpochNumber>,
583    /// The epoch number to which the recipient of this VID belongs to
584    pub target_epoch: Option<EpochNumber>,
585    /// Block payload commitment
586    pub payload_commitment: AvidMCommitment,
587    /// A storage node's key and its corresponding VID share
588    pub share: AvidMShare,
589    /// a public key of the share recipient
590    pub recipient_key: TYPES::SignatureKey,
591    /// VID common data sent to all storage nodes
592    pub common: AvidMCommon,
593}
594
595impl<TYPES: NodeType> HasViewNumber for AvidMDisperseShare<TYPES> {
596    fn view_number(&self) -> ViewNumber {
597        self.view_number
598    }
599}
600
601impl<TYPES: NodeType> AvidMDisperseShare<TYPES> {
602    /// Consume `self` and return a `Proposal`
603    pub fn to_proposal(
604        self,
605        private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
606    ) -> Option<Proposal<TYPES, Self>> {
607        let Ok(signature) =
608            TYPES::SignatureKey::sign(private_key, self.payload_commitment.as_ref())
609        else {
610            tracing::error!("VID: failed to sign dispersal share payload");
611            return None;
612        };
613        Some(Proposal {
614            signature,
615            _pd: PhantomData,
616            data: self,
617        })
618    }
619
620    /// Returns the payload length in bytes.
621    pub fn payload_byte_len(&self) -> u32 {
622        self.share.payload_byte_len() as u32
623    }
624
625    /// Check if vid common is consistent with the commitment.
626    /// For AvidM, ns_commits is inside the share, so there's no separate consistency check.
627    pub fn is_consistent(&self) -> bool {
628        true
629    }
630
631    /// Verify share assuming common data is already verified consistent.
632    /// For AvidM, this is equivalent to the full verify since there's
633    /// no separate consistency check (ns_commits is inside the share).
634    pub fn verify_with_verified_common(&self) -> bool {
635        AvidMScheme::verify_share(&self.common, &self.payload_commitment, &self.share)
636            .is_ok_and(|r| r.is_ok())
637    }
638
639    /// Internally verify the share given necessary information
640    pub fn verify(&self, _total_weight: usize) -> bool {
641        self.is_consistent() && self.verify_with_verified_common()
642    }
643}
644
645/// AvidmGf2 dispersal data
646#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
647pub struct AvidmGf2Disperse<TYPES: NodeType> {
648    /// The view number for which this VID data is intended
649    pub view_number: ViewNumber,
650    /// Epoch the data of this proposal belongs to
651    pub epoch: Option<EpochNumber>,
652    /// Epoch to which the recipients of this VID belong to
653    pub target_epoch: Option<EpochNumber>,
654    /// VidCommitment calculated based on the number of nodes in `target_epoch`.
655    pub payload_commitment: AvidmGf2Commitment,
656    /// A storage node's key and its corresponding VID share
657    pub shares: BTreeMap<TYPES::SignatureKey, AvidmGf2Share>,
658    /// Length of payload in bytes
659    pub payload_byte_len: usize,
660    /// VID common data sent to all storage nodes
661    pub common: AvidmGf2Common,
662}
663
664impl<TYPES: NodeType> HasViewNumber for AvidmGf2Disperse<TYPES> {
665    fn view_number(&self) -> ViewNumber {
666        self.view_number
667    }
668}
669
670impl<TYPES: NodeType> AvidmGf2Disperse<TYPES> {
671    /// Create VID dispersal from a specified membership for the target epoch.
672    /// Uses the specified function to calculate share dispersal
673    /// Allows for more complex stake table functionality
674    fn from_membership(
675        view_number: ViewNumber,
676        commit: AvidmGf2Commitment,
677        shares: &[AvidmGf2Share],
678        common: AvidmGf2Common,
679        membership: &EpochMembership<TYPES>,
680        target_epoch: Option<EpochNumber>,
681        data_epoch: Option<EpochNumber>,
682    ) -> Result<Self> {
683        let payload_byte_len = common.payload_byte_len();
684        let shares = membership
685            .coordinator
686            .stake_table_for_epoch(target_epoch)?
687            .stake_table()
688            .map(|entry| entry.stake_table_entry.public_key())
689            .zip(shares)
690            .map(|(node, share)| (node.clone(), share.clone()))
691            .collect();
692
693        Ok(Self {
694            view_number,
695            shares,
696            payload_commitment: commit,
697            epoch: data_epoch,
698            target_epoch,
699            payload_byte_len,
700            common,
701        })
702    }
703
704    /// Calculate the vid disperse information from the payload given a view, epoch and membership,
705    /// If the sender epoch is missing, it means it's the same as the target epoch.
706    ///
707    /// # Errors
708    /// Returns an error if the disperse or commitment calculation fails
709    pub async fn calculate_vid_disperse(
710        payload: &TYPES::BlockPayload,
711        membership: &EpochMembershipCoordinator<TYPES>,
712        view: ViewNumber,
713        target_epoch: Option<EpochNumber>,
714        data_epoch: Option<EpochNumber>,
715        metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
716    ) -> Result<(Self, Duration)> {
717        let target_mem = membership.stake_table_for_epoch(target_epoch)?;
718        let stake_table: Vec<_> = target_mem.stake_table().cloned().collect();
719        let approximate_weights = approximate_weights(&stake_table);
720
721        let txns = payload.encode();
722        let num_txns = txns.len();
723
724        let avidm_param = init_avidm_gf2_param(approximate_weights.total_weight)?;
725
726        let ns_table = parse_ns_table(num_txns, &metadata.encode());
727        let ns_table_clone = ns_table.clone();
728
729        let now = Instant::now();
730        let (commit, common, shares) = spawn_blocking(move || {
731            AvidmGf2Scheme::ns_disperse(
732                &avidm_param,
733                &approximate_weights.weights,
734                &txns,
735                ns_table_clone,
736            )
737        })
738        .await
739        .wrap()
740        .context(error!("Join error"))?
741        .wrap()
742        .context(|err| error!("Failed to calculate VID disperse. Error: {err}"))?;
743        let ns_disperse_duration = now.elapsed();
744
745        Ok((
746            Self::from_membership(
747                view,
748                commit,
749                &shares,
750                common,
751                &target_mem,
752                target_epoch,
753                data_epoch,
754            )?,
755            ns_disperse_duration,
756        ))
757    }
758
759    /// This function splits a VID disperse into individual shares.
760    pub fn to_shares(self) -> Vec<AvidmGf2DisperseShare<TYPES>> {
761        self.shares
762            .into_iter()
763            .map(|(recipient_key, share)| AvidmGf2DisperseShare {
764                share,
765                recipient_key,
766                view_number: self.view_number,
767                payload_commitment: self.payload_commitment,
768                epoch: self.epoch,
769                target_epoch: self.target_epoch,
770                common: self.common.clone(),
771            })
772            .collect()
773    }
774
775    /// Split a VID disperse into a share proposal for each recipient.
776    pub fn to_share_proposals(
777        self,
778        signature: &<<TYPES as NodeType>::SignatureKey as SignatureKey>::PureAssembledSignatureType,
779    ) -> Vec<Proposal<TYPES, AvidmGf2DisperseShare<TYPES>>> {
780        self.shares
781            .into_iter()
782            .map(|(recipient_key, share)| Proposal {
783                data: AvidmGf2DisperseShare {
784                    share,
785                    recipient_key,
786                    view_number: self.view_number,
787                    payload_commitment: self.payload_commitment,
788                    epoch: self.epoch,
789                    target_epoch: self.target_epoch,
790                    common: self.common.clone(),
791                },
792                signature: signature.clone(),
793                _pd: PhantomData,
794            })
795            .collect()
796    }
797
798    /// Construct a VID disperse from an iterator of disperse shares.
799    pub fn try_from_shares<'a, I>(mut it: I) -> Option<Self>
800    where
801        I: Iterator<Item = &'a AvidmGf2DisperseShare<TYPES>>,
802    {
803        let first_vid_disperse_share = it.next()?.clone();
804        let payload_byte_len = first_vid_disperse_share.common.payload_byte_len();
805        let mut share_map = BTreeMap::new();
806        share_map.insert(
807            first_vid_disperse_share.recipient_key,
808            first_vid_disperse_share.share,
809        );
810        let mut vid_disperse = Self {
811            view_number: first_vid_disperse_share.view_number,
812            epoch: first_vid_disperse_share.epoch,
813            target_epoch: first_vid_disperse_share.target_epoch,
814            payload_commitment: first_vid_disperse_share.payload_commitment,
815            shares: share_map,
816            payload_byte_len,
817            common: first_vid_disperse_share.common,
818        };
819        it.for_each(|vid_disperse_share| {
820            vid_disperse.shares.insert(
821                vid_disperse_share.recipient_key.clone(),
822                vid_disperse_share.share.clone(),
823            );
824        });
825        Some(vid_disperse)
826    }
827
828    /// Returns the payload length in bytes.
829    pub fn payload_byte_len(&self) -> u32 {
830        self.payload_byte_len as u32
831    }
832
833    /// Resolve the inputs needed to disperse `payload` to `target_epoch`'s
834    /// committee one namespace at a time.
835    ///
836    /// Mirrors the setup in [`Self::calculate_vid_disperse`] but returns the
837    /// owned parameters instead of performing the dispersal, leaving the (heavy,
838    /// per-namespace) computation to the caller.
839    pub fn disperse_params(
840        payload: &TYPES::BlockPayload,
841        membership: &EpochMembershipCoordinator<TYPES>,
842        target_epoch: Option<EpochNumber>,
843        metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
844    ) -> Result<AvidmGf2DisperseParams<TYPES>> {
845        let target_mem = membership.stake_table_for_epoch(target_epoch)?;
846        let stake_table: Vec<_> = target_mem.stake_table().cloned().collect();
847        let Weights {
848            weights,
849            total_weight,
850        } = approximate_weights(&stake_table);
851        let recipients = stake_table
852            .iter()
853            .map(|entry| entry.stake_table_entry.public_key())
854            .collect();
855        let payload = payload.encode();
856        let ns_table = parse_ns_table(payload.len(), &metadata.encode());
857        let param = init_avidm_gf2_param(total_weight)?;
858        Ok(AvidmGf2DisperseParams {
859            param,
860            weights,
861            recipients,
862            ns_table,
863            payload,
864        })
865    }
866}
867
868/// VID share and associated metadata for a single node
869#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
870pub struct AvidmGf2DisperseShare<TYPES: NodeType> {
871    /// The view number for which this VID data is intended
872    pub view_number: ViewNumber,
873    /// The epoch number for which this VID data belongs to
874    pub epoch: Option<EpochNumber>,
875    /// The epoch number to which the recipient of this VID belongs to
876    pub target_epoch: Option<EpochNumber>,
877    /// Block payload commitment
878    pub payload_commitment: AvidmGf2Commitment,
879    /// A storage node's key and its corresponding VID share
880    pub share: AvidmGf2Share,
881    /// a public key of the share recipient
882    pub recipient_key: TYPES::SignatureKey,
883    /// VID common data sent to all storage nodes
884    pub common: AvidmGf2Common,
885}
886
887impl<TYPES: NodeType> HasViewNumber for AvidmGf2DisperseShare<TYPES> {
888    fn view_number(&self) -> ViewNumber {
889        self.view_number
890    }
891}
892
893impl<TYPES: NodeType> AvidmGf2DisperseShare<TYPES> {
894    /// Consume `self` and return a `Proposal`
895    pub fn to_proposal(
896        self,
897        private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
898    ) -> Option<Proposal<TYPES, Self>> {
899        let Ok(signature) =
900            TYPES::SignatureKey::sign(private_key, self.payload_commitment.as_ref())
901        else {
902            tracing::error!("VID: failed to sign dispersal share payload");
903            return None;
904        };
905        Some(Proposal {
906            signature,
907            _pd: PhantomData,
908            data: self,
909        })
910    }
911    /// Returns the payload length in bytes.
912    pub fn payload_byte_len(&self) -> u32 {
913        self.common.payload_byte_len() as u32
914    }
915    /// Check if vid common is consistent with the commitment.
916    pub fn is_consistent(&self) -> bool {
917        AvidmGf2Scheme::is_consistent(&self.payload_commitment, &self.common)
918    }
919
920    /// Verify share assuming common data is already verified consistent.
921    /// Caller MUST call `is_consistent()` first.
922    pub fn verify_with_verified_common(&self) -> bool {
923        AvidmGf2Scheme::verify_share_with_verified_common(&self.common, &self.share)
924            .is_ok_and(|r| r.is_ok())
925    }
926
927    /// Internally verify the share given necessary information
928    pub fn verify(&self, total_weight: usize) -> bool {
929        // A share's commitment hash-binds its `ns_commits` (via `is_consistent`)
930        // but not its `param`, so check `param` against the committee-derived
931        // expectation; otherwise a forged param would pass verification.
932        init_avidm_gf2_param(total_weight).is_ok_and(|expected| self.common.param == expected)
933            && self.is_consistent()
934            && self.verify_with_verified_common()
935    }
936}
937
938/// VID shares consist of fragments as the unit of transmission.
939#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
940pub struct AvidmGf2DisperseShareFragment<T: NodeType> {
941    /// The view number for which this VID data is intended.
942    pub view_number: ViewNumber,
943    /// The epoch number for which this VID data belongs to.
944    pub epoch: Option<EpochNumber>,
945    /// The epoch number to which the recipient of this VID belongs to.
946    pub target_epoch: Option<EpochNumber>,
947    /// Block payload commitment (the aggregate over all namespaces).
948    pub payload_commitment: AvidmGf2Commitment,
949    /// A public key of the share recipient.
950    pub recipient_key: T::SignatureKey,
951    /// VID erasure parameters; identical across a view's fragments.
952    pub param: AvidmGf2Param,
953    /// Total number of namespaces in the block.
954    pub num_namespaces: usize,
955    /// Namespace share fragment pieces.
956    pub namespaces: Vec<AvidmGf2NamespacePiece>,
957}
958
959impl<T: NodeType> HasViewNumber for AvidmGf2DisperseShareFragment<T> {
960    fn view_number(&self) -> ViewNumber {
961        self.view_number
962    }
963}
964
965impl<T: NodeType> HasEpoch for AvidmGf2DisperseShareFragment<T> {
966    fn epoch(&self) -> Option<EpochNumber> {
967        self.epoch
968    }
969}
970
971/// VID share fragments hold pieces.
972///
973/// A piece is the smallest unit. Fragments group multiple pieces together,
974/// depending on the cumulative length of their payload lengths.
975#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
976pub struct AvidmGf2NamespacePiece {
977    /// This namespace's index in the namespace table, in `0..num_namespaces`.
978    pub ns_index: usize,
979    /// Byte length of this namespace's slice of the payload.
980    pub ns_payload_byte_len: usize,
981    /// Commitment to this namespace's shards.
982    pub ns_commit: vid::avidm_gf2::AvidmGf2Commit,
983    /// This recipient's share of this namespace.
984    pub ns_share: vid::avidm_gf2::AvidmGf2Share,
985}
986
987/// Parameters for a per-namespace AvidmGf2 dispersal.
988///
989/// [`AvidmGf2Disperse::disperse_params`] resolves the VID parameters and
990/// recipient ordering up front so a caller can send each namespace's
991/// fragments as they are produced.
992///
993/// `recipients[i]` is the storage node whose weight is `weights[i]` and
994/// whose share is the `i`th entry of each namespace's dispersal.
995#[non_exhaustive]
996pub struct AvidmGf2DisperseParams<T: NodeType> {
997    /// VID erasure parameters.
998    pub param: AvidmGf2Param,
999    /// Per-node weights, in stake-table order.
1000    pub weights: Vec<u32>,
1001    /// Recipients, in the same order as `weights`.
1002    pub recipients: Vec<T::SignatureKey>,
1003    /// Namespace byte ranges over the encoded payload.
1004    pub ns_table: Vec<Range<usize>>,
1005    /// The encoded payload.
1006    pub payload: Arc<[u8]>,
1007}