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
9use std::{collections::BTreeMap, fmt::Debug, hash::Hash, marker::PhantomData, time::Duration};
10
11use alloy::primitives::U256;
12use hotshot_utils::anytrace::*;
13use jf_vid::{VidDisperse as JfVidDisperse, VidScheme};
14use serde::{Deserialize, Serialize};
15use tokio::{task::spawn_blocking, time::Instant};
16
17use super::ns_table::parse_ns_table;
18use crate::{
19    epoch_membership::{EpochMembership, EpochMembershipCoordinator},
20    impl_has_epoch,
21    message::Proposal,
22    simple_vote::HasEpoch,
23    stake_table::HSStakeTable,
24    traits::{
25        block_contents::EncodeBytes,
26        node_implementation::NodeType,
27        signature_key::{SignatureKey, StakeTableEntryType},
28        BlockPayload,
29    },
30    vid::{
31        advz::{advz_scheme, ADVZCommitment, ADVZCommon, ADVZScheme, ADVZShare},
32        avidm::{init_avidm_param, AvidMCommitment, AvidMCommon, AvidMScheme, AvidMShare},
33    },
34    vote::HasViewNumber,
35};
36
37impl_has_epoch!(
38    ADVZDisperse<TYPES>,
39    AvidMDisperse<TYPES>,
40    VidDisperseShare2<TYPES>
41);
42
43/// ADVZ dispersal data
44#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
45pub struct ADVZDisperse<TYPES: NodeType> {
46    /// The view number for which this VID data is intended
47    pub view_number: TYPES::View,
48    /// Epoch the data of this proposal belongs to
49    pub epoch: Option<TYPES::Epoch>,
50    /// Epoch to which the recipients of this VID belong to
51    pub target_epoch: Option<TYPES::Epoch>,
52    /// VidCommitment calculated based on the number of nodes in `target_epoch`.
53    pub payload_commitment: ADVZCommitment,
54    /// A storage node's key and its corresponding VID share
55    pub shares: BTreeMap<TYPES::SignatureKey, ADVZShare>,
56    /// VID common data sent to all storage nodes
57    pub common: ADVZCommon,
58}
59
60impl<TYPES: NodeType> HasViewNumber<TYPES> for ADVZDisperse<TYPES> {
61    fn view_number(&self) -> TYPES::View {
62        self.view_number
63    }
64}
65
66impl<TYPES: NodeType> ADVZDisperse<TYPES> {
67    /// Create VID dispersal from a specified membership for the target epoch.
68    /// Uses the specified function to calculate share dispersal
69    /// Allows for more complex stake table functionality
70    async fn from_membership(
71        view_number: TYPES::View,
72        mut vid_disperse: JfVidDisperse<ADVZScheme>,
73        membership: &EpochMembershipCoordinator<TYPES>,
74        target_epoch: Option<TYPES::Epoch>,
75        data_epoch: Option<TYPES::Epoch>,
76    ) -> Self {
77        let shares = membership
78            .stake_table_for_epoch(target_epoch)
79            .await
80            .unwrap()
81            .stake_table()
82            .await
83            .iter()
84            .map(|entry| entry.stake_table_entry.public_key())
85            .map(|node| (node.clone(), vid_disperse.shares.remove(0)))
86            .collect();
87
88        Self {
89            view_number,
90            shares,
91            common: vid_disperse.common,
92            payload_commitment: vid_disperse.commit,
93            epoch: data_epoch,
94            target_epoch,
95        }
96    }
97
98    /// Calculate the vid disperse information from the payload given a view, epoch and membership,
99    /// If the sender epoch is missing, it means it's the same as the target epoch.
100    ///
101    /// # Errors
102    /// Returns an error if the disperse or commitment calculation fails
103    #[allow(clippy::panic)]
104    pub async fn calculate_vid_disperse(
105        payload: &TYPES::BlockPayload,
106        membership: &EpochMembershipCoordinator<TYPES>,
107        view: TYPES::View,
108        target_epoch: Option<TYPES::Epoch>,
109        data_epoch: Option<TYPES::Epoch>,
110    ) -> Result<(Self, Duration)> {
111        let num_nodes = membership
112            .stake_table_for_epoch(target_epoch)
113            .await?
114            .total_nodes()
115            .await;
116
117        let txns = payload.encode();
118
119        let now = Instant::now();
120        let vid_disperse = spawn_blocking(move || advz_scheme(num_nodes).disperse(&txns))
121            .await
122            .wrap()
123            .context(error!("Join error"))?
124            .wrap()
125            .context(|err| error!("Failed to calculate VID disperse. Error: {err}"))?;
126        let advz_scheme_duration = now.elapsed();
127
128        Ok((
129            Self::from_membership(view, vid_disperse, membership, target_epoch, data_epoch).await,
130            advz_scheme_duration,
131        ))
132    }
133
134    /// Returns the payload length in bytes.
135    pub fn payload_byte_len(&self) -> u32 {
136        ADVZScheme::get_payload_byte_len(&self.common)
137    }
138}
139
140#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
141/// ADVZ share and associated metadata for a single node
142pub struct ADVZDisperseShare<TYPES: NodeType> {
143    /// The view number for which this VID data is intended
144    pub view_number: TYPES::View,
145    /// Block payload commitment
146    pub payload_commitment: ADVZCommitment,
147    /// A storage node's key and its corresponding VID share
148    pub share: ADVZShare,
149    /// VID common data sent to all storage nodes
150    pub common: ADVZCommon,
151    /// a public key of the share recipient
152    pub recipient_key: TYPES::SignatureKey,
153}
154
155impl<TYPES: NodeType> HasViewNumber<TYPES> for ADVZDisperseShare<TYPES> {
156    fn view_number(&self) -> TYPES::View {
157        self.view_number
158    }
159}
160
161impl<TYPES: NodeType> ADVZDisperseShare<TYPES> {
162    /// Create a vector of `VidDisperseShare` from `VidDisperse`
163    pub fn from_advz_disperse(vid_disperse: ADVZDisperse<TYPES>) -> Vec<Self> {
164        vid_disperse
165            .shares
166            .into_iter()
167            .map(|(recipient_key, share)| Self {
168                share,
169                recipient_key,
170                view_number: vid_disperse.view_number,
171                common: vid_disperse.common.clone(),
172                payload_commitment: vid_disperse.payload_commitment,
173            })
174            .collect()
175    }
176
177    /// Consume `self` and return a `Proposal`
178    pub fn to_proposal(
179        self,
180        private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
181    ) -> Option<Proposal<TYPES, Self>> {
182        let Ok(signature) =
183            TYPES::SignatureKey::sign(private_key, self.payload_commitment.as_ref())
184        else {
185            tracing::error!("VID: failed to sign dispersal share payload");
186            return None;
187        };
188        Some(Proposal {
189            signature,
190            _pd: PhantomData,
191            data: self,
192        })
193    }
194
195    /// Create `VidDisperse` out of an iterator to `VidDisperseShare`s
196    pub fn to_advz_disperse<'a, I>(mut it: I) -> Option<ADVZDisperse<TYPES>>
197    where
198        I: Iterator<Item = &'a Self>,
199    {
200        let first_vid_disperse_share = it.next()?.clone();
201        let mut share_map = BTreeMap::new();
202        share_map.insert(
203            first_vid_disperse_share.recipient_key,
204            first_vid_disperse_share.share,
205        );
206        let mut vid_disperse = ADVZDisperse {
207            view_number: first_vid_disperse_share.view_number,
208            epoch: None,
209            target_epoch: None,
210            payload_commitment: first_vid_disperse_share.payload_commitment,
211            common: first_vid_disperse_share.common,
212            shares: share_map,
213        };
214        let _ = it.map(|vid_disperse_share| {
215            vid_disperse.shares.insert(
216                vid_disperse_share.recipient_key.clone(),
217                vid_disperse_share.share.clone(),
218            )
219        });
220        Some(vid_disperse)
221    }
222
223    /// Split a VID share proposal into a proposal for each recipient.
224    pub fn to_vid_share_proposals(
225        vid_disperse: ADVZDisperse<TYPES>,
226        signature: &<TYPES::SignatureKey as SignatureKey>::PureAssembledSignatureType,
227    ) -> Vec<Proposal<TYPES, Self>> {
228        vid_disperse
229            .shares
230            .into_iter()
231            .map(|(recipient_key, share)| Proposal {
232                data: Self {
233                    share,
234                    recipient_key,
235                    view_number: vid_disperse.view_number,
236                    common: vid_disperse.common.clone(),
237                    payload_commitment: vid_disperse.payload_commitment,
238                },
239                signature: signature.clone(),
240                _pd: PhantomData,
241            })
242            .collect()
243    }
244
245    /// Internally verify the share given necessary information
246    ///
247    /// # Errors
248    /// Verification fail
249    #[allow(clippy::result_unit_err)]
250    pub fn verify_share(&self, total_weight: usize) -> std::result::Result<(), ()> {
251        advz_scheme(total_weight)
252            .verify_share(&self.share, &self.common, &self.payload_commitment)
253            .unwrap_or(Err(()))
254    }
255
256    /// Returns the payload length in bytes.
257    pub fn payload_byte_len(&self) -> u32 {
258        ADVZScheme::get_payload_byte_len(&self.common)
259    }
260}
261
262/// ADVZ dispersal data
263#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
264pub struct AvidMDisperse<TYPES: NodeType> {
265    /// The view number for which this VID data is intended
266    pub view_number: TYPES::View,
267    /// Epoch the data of this proposal belongs to
268    pub epoch: Option<TYPES::Epoch>,
269    /// Epoch to which the recipients of this VID belong to
270    pub target_epoch: Option<TYPES::Epoch>,
271    /// VidCommitment calculated based on the number of nodes in `target_epoch`.
272    pub payload_commitment: AvidMCommitment,
273    /// A storage node's key and its corresponding VID share
274    pub shares: BTreeMap<TYPES::SignatureKey, AvidMShare>,
275    /// Length of payload in bytes
276    pub payload_byte_len: usize,
277    /// VID common data sent to all storage nodes
278    pub common: AvidMCommon,
279}
280
281impl<TYPES: NodeType> HasViewNumber<TYPES> for AvidMDisperse<TYPES> {
282    fn view_number(&self) -> TYPES::View {
283        self.view_number
284    }
285}
286
287/// The target total stake to scale to for VID.
288pub const VID_TARGET_TOTAL_STAKE: u32 = 1000;
289
290/// The weights and total weight used in VID calculations
291struct Weights {
292    // weights, in stake table order
293    weights: Vec<u32>,
294
295    // total weight
296    total_weight: usize,
297}
298
299pub fn vid_total_weight<TYPES: NodeType>(
300    stake_table: &HSStakeTable<TYPES>,
301    epoch: Option<TYPES::Epoch>,
302) -> usize {
303    if epoch.is_none() {
304        stake_table
305            .iter()
306            .fold(U256::ZERO, |acc, entry| {
307                acc + entry.stake_table_entry.stake()
308            })
309            .to::<usize>()
310    } else {
311        approximate_weights(stake_table).total_weight
312    }
313}
314
315fn approximate_weights<TYPES: NodeType>(stake_table: &HSStakeTable<TYPES>) -> Weights {
316    let total_stake = stake_table.iter().fold(U256::ZERO, |acc, entry| {
317        acc + entry.stake_table_entry.stake()
318    });
319
320    let mut total_weight: usize = 0;
321
322    // don't attempt to scale if the total stake is small enough
323    if total_stake <= U256::from(VID_TARGET_TOTAL_STAKE) {
324        let weights = stake_table
325            .iter()
326            .map(|entry| entry.stake_table_entry.stake().to::<u32>())
327            .collect();
328
329        // Note: this panics if `total_stake` exceeds `usize::MAX`, but this shouldn't happen.
330        total_weight = total_stake.to::<usize>();
331
332        Weights {
333            weights,
334            total_weight,
335        }
336    } else {
337        let weights = stake_table
338            .iter()
339            .map(|entry| {
340                let weight: U256 = ((entry.stake_table_entry.stake()
341                    * U256::from(VID_TARGET_TOTAL_STAKE))
342                    / total_stake)
343                    + U256::ONE;
344
345                // Note: this panics if `weight` exceeds `usize::MAX`, but this shouldn't happen.
346                total_weight += weight.to::<usize>();
347
348                // Note: this panics if `weight` exceeds `u32::MAX`, but this shouldn't happen
349                // and would likely cause a stack overflow in the VID calculation anyway
350                weight.to::<u32>()
351            })
352            .collect();
353
354        Weights {
355            weights,
356            total_weight,
357        }
358    }
359}
360
361impl<TYPES: NodeType> AvidMDisperse<TYPES> {
362    /// Create VID dispersal from a specified membership for the target epoch.
363    /// Uses the specified function to calculate share dispersal
364    /// Allows for more complex stake table functionality
365    async fn from_membership(
366        view_number: TYPES::View,
367        commit: AvidMCommitment,
368        shares: &[AvidMShare],
369        common: AvidMCommon,
370        membership: &EpochMembership<TYPES>,
371        target_epoch: Option<TYPES::Epoch>,
372        data_epoch: Option<TYPES::Epoch>,
373    ) -> Self {
374        let payload_byte_len = shares[0].payload_byte_len();
375        let shares = membership
376            .coordinator
377            .stake_table_for_epoch(target_epoch)
378            .await
379            .unwrap()
380            .stake_table()
381            .await
382            .iter()
383            .map(|entry| entry.stake_table_entry.public_key())
384            .zip(shares)
385            .map(|(node, share)| (node.clone(), share.clone()))
386            .collect();
387
388        Self {
389            view_number,
390            shares,
391            payload_commitment: commit,
392            epoch: data_epoch,
393            target_epoch,
394            payload_byte_len,
395            common,
396        }
397    }
398
399    /// Calculate the vid disperse information from the payload given a view, epoch and membership,
400    /// If the sender epoch is missing, it means it's the same as the target epoch.
401    ///
402    /// # Errors
403    /// Returns an error if the disperse or commitment calculation fails
404    #[allow(clippy::panic)]
405    #[allow(clippy::single_range_in_vec_init)]
406    pub async fn calculate_vid_disperse(
407        payload: &TYPES::BlockPayload,
408        membership: &EpochMembershipCoordinator<TYPES>,
409        view: TYPES::View,
410        target_epoch: Option<TYPES::Epoch>,
411        data_epoch: Option<TYPES::Epoch>,
412        metadata: &<TYPES::BlockPayload as BlockPayload<TYPES>>::Metadata,
413    ) -> Result<(Self, Duration)> {
414        let target_mem = membership.stake_table_for_epoch(target_epoch).await?;
415        let stake_table = target_mem.stake_table().await;
416        let approximate_weights = approximate_weights(&stake_table);
417
418        let txns = payload.encode();
419        let num_txns = txns.len();
420
421        let avidm_param = init_avidm_param(approximate_weights.total_weight)?;
422        let common = avidm_param.clone();
423
424        let ns_table = parse_ns_table(num_txns, &metadata.encode());
425        let ns_table_clone = ns_table.clone();
426
427        let now = Instant::now();
428        let (commit, shares) = spawn_blocking(move || {
429            AvidMScheme::ns_disperse(
430                &avidm_param,
431                &approximate_weights.weights,
432                &txns,
433                ns_table_clone,
434            )
435        })
436        .await
437        .wrap()
438        .context(error!("Join error"))?
439        .wrap()
440        .context(|err| error!("Failed to calculate VID disperse. Error: {err}"))?;
441        let ns_disperse_duration = now.elapsed();
442
443        Ok((
444            Self::from_membership(
445                view,
446                commit,
447                &shares,
448                common,
449                &target_mem,
450                target_epoch,
451                data_epoch,
452            )
453            .await,
454            ns_disperse_duration,
455        ))
456    }
457
458    /// Returns the payload length in bytes.
459    pub fn payload_byte_len(&self) -> u32 {
460        self.payload_byte_len as u32
461    }
462}
463
464#[derive(Debug, Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
465/// VID share and associated metadata for a single node
466pub struct VidDisperseShare2<TYPES: NodeType> {
467    /// The view number for which this VID data is intended
468    pub view_number: TYPES::View,
469    /// The epoch number for which this VID data belongs to
470    pub epoch: Option<TYPES::Epoch>,
471    /// The epoch number to which the recipient of this VID belongs to
472    pub target_epoch: Option<TYPES::Epoch>,
473    /// Block payload commitment
474    pub payload_commitment: AvidMCommitment,
475    /// A storage node's key and its corresponding VID share
476    pub share: AvidMShare,
477    /// a public key of the share recipient
478    pub recipient_key: TYPES::SignatureKey,
479    /// VID common data sent to all storage nodes
480    pub common: AvidMCommon,
481}
482
483impl<TYPES: NodeType> HasViewNumber<TYPES> for VidDisperseShare2<TYPES> {
484    fn view_number(&self) -> TYPES::View {
485        self.view_number
486    }
487}
488
489impl<TYPES: NodeType> VidDisperseShare2<TYPES> {
490    /// Create a vector of `VidDisperseShare` from `VidDisperse`
491    pub fn from_vid_disperse(vid_disperse: AvidMDisperse<TYPES>) -> Vec<Self> {
492        vid_disperse
493            .shares
494            .into_iter()
495            .map(|(recipient_key, share)| Self {
496                share,
497                recipient_key,
498                view_number: vid_disperse.view_number,
499                payload_commitment: vid_disperse.payload_commitment,
500                epoch: vid_disperse.epoch,
501                target_epoch: vid_disperse.target_epoch,
502                common: vid_disperse.common.clone(),
503            })
504            .collect()
505    }
506
507    /// Consume `self` and return a `Proposal`
508    pub fn to_proposal(
509        self,
510        private_key: &<TYPES::SignatureKey as SignatureKey>::PrivateKey,
511    ) -> Option<Proposal<TYPES, Self>> {
512        let Ok(signature) =
513            TYPES::SignatureKey::sign(private_key, self.payload_commitment.as_ref())
514        else {
515            tracing::error!("VID: failed to sign dispersal share payload");
516            return None;
517        };
518        Some(Proposal {
519            signature,
520            _pd: PhantomData,
521            data: self,
522        })
523    }
524
525    /// Create `VidDisperse` out of an iterator to `VidDisperseShare`s
526    pub fn to_vid_disperse<'a, I>(mut it: I) -> Option<AvidMDisperse<TYPES>>
527    where
528        I: Iterator<Item = &'a Self>,
529    {
530        let first_vid_disperse_share = it.next()?.clone();
531        let payload_byte_len = first_vid_disperse_share.share.payload_byte_len();
532        let mut share_map = BTreeMap::new();
533        share_map.insert(
534            first_vid_disperse_share.recipient_key,
535            first_vid_disperse_share.share,
536        );
537        let mut vid_disperse = AvidMDisperse {
538            view_number: first_vid_disperse_share.view_number,
539            epoch: first_vid_disperse_share.epoch,
540            target_epoch: first_vid_disperse_share.target_epoch,
541            payload_commitment: first_vid_disperse_share.payload_commitment,
542            shares: share_map,
543            payload_byte_len,
544            common: first_vid_disperse_share.common,
545        };
546        let _ = it.map(|vid_disperse_share| {
547            vid_disperse.shares.insert(
548                vid_disperse_share.recipient_key.clone(),
549                vid_disperse_share.share.clone(),
550            )
551        });
552        Some(vid_disperse)
553    }
554
555    /// Returns the payload length in bytes.
556    pub fn payload_byte_len(&self) -> u32 {
557        self.share.payload_byte_len() as u32
558    }
559
560    /// Split a VID share proposal into a proposal for each recipient.
561    pub fn to_vid_share_proposals(
562        vid_disperse: AvidMDisperse<TYPES>,
563        signature: &<TYPES::SignatureKey as SignatureKey>::PureAssembledSignatureType,
564    ) -> Vec<Proposal<TYPES, Self>> {
565        vid_disperse
566            .shares
567            .into_iter()
568            .map(|(recipient_key, share)| Proposal {
569                data: Self {
570                    share,
571                    recipient_key,
572                    view_number: vid_disperse.view_number,
573                    payload_commitment: vid_disperse.payload_commitment,
574                    epoch: vid_disperse.epoch,
575                    target_epoch: vid_disperse.target_epoch,
576                    common: vid_disperse.common.clone(),
577                },
578                signature: signature.clone(),
579                _pd: PhantomData,
580            })
581            .collect()
582    }
583
584    /// Internally verify the share given necessary information
585    ///
586    /// # Errors
587    #[allow(clippy::result_unit_err)]
588    pub fn verify_share(&self, total_weight: usize) -> std::result::Result<(), ()> {
589        let avidm_param = init_avidm_param(total_weight).map_err(|_| ())?;
590        AvidMScheme::verify_share(&avidm_param, &self.payload_commitment, &self.share)
591            .unwrap_or(Err(()))
592    }
593}
hotshot_types/data/vid_disperse.rs

hotshot_types/data/
vid_disperse.rs
