Skip to main content

hotshot_new_protocol/vid/
fragments.rs

1use std::collections::{BTreeMap, BTreeSet};
2
3use hotshot_types::{
4    data::{
5        EpochNumber, VidCommitment2, VidDisperseShare2, ViewNumber,
6        vid_disperse::{AvidmGf2DisperseShareFragment, AvidmGf2NamespacePiece},
7    },
8    traits::node_implementation::NodeType,
9    vid::avidm_gf2::{AvidmGf2Common, AvidmGf2Param},
10};
11
12pub struct VidFragmentAccumulator<T: NodeType> {
13    pending: BTreeMap<ViewNumber, PendingShare<T>>,
14    completed: BTreeSet<ViewNumber>,
15    lower_bound: ViewNumber,
16}
17
18#[derive(Debug, thiserror::Error)]
19pub enum VidFragmentError {
20    #[error("fragment disagrees with the view's pinned metadata")]
21    Inconsistent,
22
23    #[error("namespace index {index} out of range for {num_namespaces} namespaces")]
24    IndexOutOfRange { index: usize, num_namespaces: usize },
25
26    #[error("duplicate fragment for namespace index {0}")]
27    DuplicateIndex(usize),
28
29    #[error("fragment contains no namespaces")]
30    Empty,
31}
32
33/// A view's partially-collected namespace pieces, keyed by namespace index.
34struct PendingShare<T: NodeType> {
35    epoch: Option<EpochNumber>,
36    target_epoch: Option<EpochNumber>,
37    payload_commitment: VidCommitment2,
38    recipient_key: T::SignatureKey,
39    param: AvidmGf2Param,
40    num_namespaces: usize,
41    pieces: BTreeMap<usize, AvidmGf2NamespacePiece>,
42}
43
44impl<T: NodeType> Default for VidFragmentAccumulator<T> {
45    fn default() -> Self {
46        Self::new()
47    }
48}
49
50impl<T: NodeType> VidFragmentAccumulator<T> {
51    pub fn new() -> Self {
52        Self {
53            pending: BTreeMap::new(),
54            completed: BTreeSet::new(),
55            lower_bound: ViewNumber::genesis(),
56        }
57    }
58
59    /// Buffer a `fragment` addressed to this node.
60    ///
61    /// Returns `Ok(None)` while namespaces are still outstanding,
62    /// `Ok(Some(share))` once the final namespace completes the view, and
63    /// `Err` if the fragment is malformed or inconsistent with the view's
64    /// already-pinned metadata.
65    pub fn accept(
66        &mut self,
67        fragment: AvidmGf2DisperseShareFragment<T>,
68    ) -> Result<Option<VidDisperseShare2<T>>, VidFragmentError> {
69        let view = fragment.view_number;
70        if view < self.lower_bound || self.completed.contains(&view) {
71            return Ok(None);
72        }
73        if fragment.num_namespaces == 0 {
74            return Err(VidFragmentError::Empty);
75        }
76        let pending = self.pending.entry(view).or_insert_with(|| PendingShare {
77            epoch: fragment.epoch,
78            target_epoch: fragment.target_epoch,
79            payload_commitment: fragment.payload_commitment,
80            recipient_key: fragment.recipient_key.clone(),
81            param: fragment.param.clone(),
82            num_namespaces: fragment.num_namespaces,
83            pieces: BTreeMap::new(),
84        });
85        if pending.num_namespaces != fragment.num_namespaces
86            || pending.epoch != fragment.epoch
87            || pending.target_epoch != fragment.target_epoch
88            || pending.payload_commitment != fragment.payload_commitment
89            || pending.recipient_key != fragment.recipient_key
90            || pending.param != fragment.param
91        {
92            return Err(VidFragmentError::Inconsistent);
93        }
94        for piece in fragment.namespaces {
95            let ns_index = piece.ns_index;
96            if ns_index >= pending.num_namespaces {
97                return Err(VidFragmentError::IndexOutOfRange {
98                    index: ns_index,
99                    num_namespaces: pending.num_namespaces,
100                });
101            }
102            if pending.pieces.contains_key(&ns_index) {
103                return Err(VidFragmentError::DuplicateIndex(ns_index));
104            }
105            pending.pieces.insert(ns_index, piece);
106        }
107        if pending.pieces.len() != pending.num_namespaces {
108            return Ok(None);
109        }
110        // Every namespace is present and indices are distinct and in range, so
111        // they cover `0..num_namespaces` exactly; the `BTreeMap` yields them in
112        // that order.
113        let pending = self.pending.remove(&view).expect("just inserted above");
114        self.completed.insert(view);
115        let mut ns_commits = Vec::with_capacity(pending.num_namespaces);
116        let mut ns_lens = Vec::with_capacity(pending.num_namespaces);
117        let mut ns_shares = Vec::with_capacity(pending.num_namespaces);
118        for piece in pending.pieces.into_values() {
119            ns_commits.push(piece.ns_commit);
120            ns_lens.push(piece.ns_payload_byte_len);
121            ns_shares.push(piece.ns_share);
122        }
123        Ok(Some(VidDisperseShare2 {
124            view_number: view,
125            epoch: pending.epoch,
126            target_epoch: pending.target_epoch,
127            payload_commitment: pending.payload_commitment,
128            share: ns_shares.into(),
129            recipient_key: pending.recipient_key,
130            common: AvidmGf2Common {
131                param: pending.param,
132                ns_commits,
133                ns_lens,
134            },
135        }))
136    }
137
138    pub fn gc(&mut self, view_number: ViewNumber) {
139        self.pending = self.pending.split_off(&view_number);
140        self.completed = self.completed.split_off(&view_number);
141        self.lower_bound = view_number;
142    }
143}