vid/avidm_gf2/

namespaced.rs

//! This file implements the namespaced AvidmGf2 scheme.

use std::ops::Range;

use jf_merkle_tree::MerkleTreeScheme;
use p3_maybe_rayon::prelude::*;
use serde::{Deserialize, Serialize};

use super::{AvidmGf2Commit, AvidmGf2Share};
use crate::{
    VidError, VidResult, VidScheme,
    avidm_gf2::{AvidmGf2Scheme, MerkleTree},
};

/// Dummy struct for namespaced AvidmGf2 scheme
pub struct NsAvidmGf2Scheme;

/// Namespaced commitment type
pub type NsAvidmGf2Commit = super::AvidmGf2Commit;
/// Namespaced parameter type
pub type NsAvidmGf2Param = super::AvidmGf2Param;

/// VID Common data that needs to be broadcasted to all storage nodes
#[derive(Clone, Debug, Hash, Serialize, Deserialize, Eq, PartialEq)]
pub struct NsAvidmGf2Common {
    /// The AvidmGf2 parameters
    pub param: NsAvidmGf2Param,
    /// The list of all namespace commitments
    pub ns_commits: Vec<AvidmGf2Commit>,
    /// The size of each namespace
    pub ns_lens: Vec<usize>,
}

impl NsAvidmGf2Common {
    /// Return the total payload byte length
    pub fn payload_byte_len(&self) -> usize {
        self.ns_lens.iter().sum()
    }
}

/// Namespaced share for each storage node, contains one [`AvidmGf2Share`] for each namespace.
#[derive(Clone, Debug, Hash, Serialize, Deserialize, Eq, PartialEq, Default)]
pub struct NsAvidmGf2Share(pub(crate) Vec<AvidmGf2Share>);

impl NsAvidmGf2Share {
    /// Return the number of namespaces in this share
    pub fn num_nss(&self) -> usize {
        self.0.len()
    }

    /// Return the weight of this share
    pub fn weight(&self) -> usize {
        self.0.first().map_or(0, |share| share.weight())
    }

    /// Validate the share structure
    pub fn validate(&self) -> bool {
        let weight = self.weight();
        self.0
            .iter()
            .all(|share| share.validate() && share.weight() == weight)
    }

    /// Check whether this share contains a given namespace
    pub fn contains_ns(&self, ns_index: usize) -> bool {
        ns_index < self.num_nss()
    }

    /// The per-namespace shares, in namespace order.
    pub fn ns_shares(&self) -> &[AvidmGf2Share] {
        &self.0
    }

    /// Return the inner share for a given namespace if there exists one.
    pub fn inner_ns_share(&self, ns_index: usize) -> Option<AvidmGf2Share> {
        self.0.get(ns_index).cloned()
    }

    /// The shard range this share covers, identical across all namespaces.
    /// `None` if the share is empty or its namespaces disagree on the range.
    pub fn range(&self) -> Option<&Range<usize>> {
        let first = self.0.first()?.range();
        self.0
            .iter()
            .all(|share| share.range() == first)
            .then_some(first)
    }
}

impl From<Vec<AvidmGf2Share>> for NsAvidmGf2Share {
    fn from(ns_shares: Vec<AvidmGf2Share>) -> Self {
        Self(ns_shares)
    }
}

impl NsAvidmGf2Scheme {
    /// Setup an instance for AVID-M scheme
    pub fn setup(recovery_threshold: usize, total_weights: usize) -> VidResult<NsAvidmGf2Param> {
        NsAvidmGf2Param::new(recovery_threshold, total_weights)
    }

    /// Commit to a payload given namespace table.
    /// WARN: it assumes that the namespace table is well formed, i.e. ranges
    /// are non-overlapping and cover the whole payload.
    pub fn commit(
        param: &NsAvidmGf2Param,
        payload: &[u8],
        ns_table: impl IntoIterator<Item = Range<usize>>,
    ) -> VidResult<(NsAvidmGf2Commit, NsAvidmGf2Common)> {
        let ns_table = ns_table.into_iter().collect::<Vec<_>>();
        let ns_lens = ns_table.iter().map(|r| r.len()).collect::<Vec<_>>();
        let ns_commits = ns_table
            .into_par_iter()
            .map(|ns_range| AvidmGf2Scheme::commit(param, &payload[ns_range]))
            .collect::<Result<Vec<_>, _>>()?;
        let common = NsAvidmGf2Common {
            param: param.clone(),
            ns_commits,
            ns_lens,
        };
        let commit = Self::aggregate_commit(&common.ns_commits)?;
        Ok((commit, common))
    }

    /// Aggregate per-namespace commitments into the top-level namespaced
    /// commitment (the Merkle root over the namespace commits).
    ///
    /// Lets a caller that has already computed each namespace's commit (e.g. via
    /// [`Self::ns_disperse_one`]) form the block commitment without re-encoding.
    pub fn aggregate_commit(ns_commits: &[AvidmGf2Commit]) -> VidResult<NsAvidmGf2Commit> {
        let commit = MerkleTree::from_elems(None, ns_commits.iter().map(|c| c.commit))
            .map_err(|err| VidError::Internal(err.into()))?
            .commitment();
        Ok(NsAvidmGf2Commit { commit })
    }

    /// Check whether the namespaced commitment is consistent with the common data
    pub fn is_consistent(commit: &NsAvidmGf2Commit, common: &NsAvidmGf2Common) -> bool {
        Self::aggregate_commit(&common.ns_commits).is_ok_and(|c| c == *commit)
    }

    /// Disperse a payload according to a distribution table and a namespace
    /// table.
    /// WARN: it assumes that the namespace table is well formed, i.e. ranges
    /// are non-overlapping and cover the whole payload.
    pub fn ns_disperse(
        param: &NsAvidmGf2Param,
        distribution: &[u32],
        payload: &[u8],
        ns_table: impl IntoIterator<Item = Range<usize>>,
    ) -> VidResult<(NsAvidmGf2Commit, NsAvidmGf2Common, Vec<NsAvidmGf2Share>)> {
        let num_storage_nodes = distribution.len();
        let ns_ranges: Vec<Range<usize>> = ns_table.into_iter().collect();
        let ns_lens: Vec<usize> = ns_ranges.iter().map(|r| r.len()).collect();

        // Per-namespace dispersals are independent; run them on rayon.
        // Inner `AvidmGf2Scheme::disperse` also uses `par_iter` internally for
        // leaf hashing and share assembly — rayon's work-stealing pool
        // coordinates the nested parallelism.
        let per_ns: Vec<(AvidmGf2Commit, Vec<AvidmGf2Share>)> = ns_ranges
            .par_iter()
            .map(|ns_range| {
                AvidmGf2Scheme::disperse(param, distribution, &payload[ns_range.clone()])
            })
            .collect::<VidResult<Vec<_>>>()?;

        let (ns_commits, disperses): (Vec<_>, Vec<_>) = per_ns.into_iter().unzip();

        let common = NsAvidmGf2Common {
            param: param.clone(),
            ns_commits,
            ns_lens,
        };
        let commit = Self::aggregate_commit(&common.ns_commits)?;
        let mut shares = vec![NsAvidmGf2Share::default(); num_storage_nodes];
        disperses.into_iter().for_each(|ns_disperse| {
            shares
                .iter_mut()
                .zip(ns_disperse)
                .for_each(|(share, ns_share)| share.0.push(ns_share))
        });
        Ok((commit, common, shares))
    }

    /// Disperse a single namespace's slice of the payload.
    ///
    /// `ns_payload` is the namespace's bytes and `ns_index` its position in the
    /// namespace table.
    pub fn ns_disperse_one(
        param: &NsAvidmGf2Param,
        distribution: &[u32],
        ns_payload: &[u8],
        ns_index: usize,
    ) -> VidResult<NsDispersal> {
        let payload_byte_len = ns_payload.len();
        let (commit, shares) = AvidmGf2Scheme::disperse(param, distribution, ns_payload)?;
        Ok(NsDispersal {
            ns_index,
            payload_byte_len,
            commit,
            shares,
        })
    }

    /// Test-only: like [`Self::ns_disperse`] but commits to a non-codeword in
    /// every namespace (see [`AvidmGf2Scheme::disperse_non_codeword`]). Every
    /// returned share verifies against the returned common, yet no
    /// threshold-covering subset recovers a payload that re-commits to the
    /// returned commitment.
    #[cfg(any(test, feature = "testing"))]
    pub fn ns_disperse_non_codeword(
        param: &NsAvidmGf2Param,
        distribution: &[u32],
        payload: &[u8],
        ns_table: impl IntoIterator<Item = Range<usize>>,
    ) -> VidResult<(NsAvidmGf2Commit, NsAvidmGf2Common, Vec<NsAvidmGf2Share>)> {
        let num_storage_nodes = distribution.len();
        let ns_ranges: Vec<Range<usize>> = ns_table.into_iter().collect();
        let ns_lens: Vec<usize> = ns_ranges.iter().map(|r| r.len()).collect();
        let per_ns: Vec<(AvidmGf2Commit, Vec<AvidmGf2Share>)> = ns_ranges
            .iter()
            .map(|ns_range| {
                AvidmGf2Scheme::disperse_non_codeword(
                    param,
                    distribution,
                    &payload[ns_range.clone()],
                )
            })
            .collect::<VidResult<Vec<_>>>()?;
        let (ns_commits, disperses): (Vec<_>, Vec<_>) = per_ns.into_iter().unzip();
        let common = NsAvidmGf2Common {
            param: param.clone(),
            ns_commits,
            ns_lens,
        };
        let commit = Self::aggregate_commit(&common.ns_commits)?;
        let mut shares = vec![NsAvidmGf2Share::default(); num_storage_nodes];
        disperses.into_iter().for_each(|ns_disperse| {
            shares
                .iter_mut()
                .zip(ns_disperse)
                .for_each(|(share, ns_share)| share.0.push(ns_share))
        });
        Ok((commit, common, shares))
    }

    /// Verify a namespaced share given already-verified common data.
    ///
    /// Besides each namespace's Merkle proofs, this checks that each
    /// namespace's shards have the size its `ns_lens` entry disperses to. The
    /// commitment binds `ns_commits` but not `ns_lens`, so without this a
    /// share could declare any payload length and have it believed. With it,
    /// each declared length is pinned to within `2 * recovery_threshold`
    /// bytes of the one the shards were built from.
    ///
    /// That window is per namespace, and the disperser chooses how many
    /// namespaces there are, so the misdeclared total adds up: a payload split
    /// into namespaces shorter than a window can declare each as a single
    /// byte and still verify. This rejects gross forgeries early, such as a
    /// shrunken one-namespace block or an absurd length; it does not bound
    /// the total. Exact binding needs a verifier that knows the payload's
    /// namespace table, or a commitment that covers `ns_lens`.
    ///
    /// # Safety Contract
    /// Caller MUST ensure `is_consistent(commit, common)` returned `true`
    /// before calling this. Without that check, a malicious common could
    /// make an invalid share appear valid.
    pub fn verify_share_with_verified_common(
        common: &NsAvidmGf2Common,
        share: &NsAvidmGf2Share,
    ) -> VidResult<crate::VerificationResult> {
        common.param.validate()?;
        if !(common.ns_commits.len() == common.ns_lens.len()
            && common.ns_commits.len() == share.num_nss()
            && share.validate())
        {
            return Err(VidError::InvalidShare);
        }
        // Per-namespace verifications are independent. `find_any` short-
        // circuits on the first failing namespace and avoids allocating an
        // intermediate `Vec<VerificationResult>`.
        match common
            .ns_commits
            .par_iter()
            .zip(common.ns_lens.par_iter())
            .zip(share.0.par_iter())
            .map(|((commit, &ns_len), content)| {
                if !Self::ns_len_matches_shards(&common.param, ns_len, content) {
                    return Ok(Err(()));
                }
                AvidmGf2Scheme::verify_share(&common.param, commit, content)
            })
            .find_any(|r| !matches!(r, Ok(Ok(()))))
        {
            None => Ok(Ok(())),
            Some(Ok(v)) => Ok(v),
            Some(Err(e)) => Err(e),
        }
    }

    /// Verify a namespaced share
    pub fn verify_share(
        commit: &NsAvidmGf2Commit,
        common: &NsAvidmGf2Common,
        share: &NsAvidmGf2Share,
    ) -> VidResult<crate::VerificationResult> {
        if !Self::is_consistent(commit, common) {
            return Ok(Err(()));
        }
        Self::verify_share_with_verified_common(common, share)
    }

    /// Recover the entire payload from enough share
    pub fn recover(common: &NsAvidmGf2Common, shares: &[NsAvidmGf2Share]) -> VidResult<Vec<u8>> {
        if shares.is_empty() {
            return Err(VidError::InsufficientShares);
        }
        // Each `ns_recover` is independent: distinct decoder state, distinct
        // output bytes. Run them on the rayon pool so multi-namespace blocks
        // actually use more than one core.
        let per_ns: Vec<Vec<u8>> = (0..common.ns_lens.len())
            .into_par_iter()
            .map(|ns_index| Self::ns_recover(common, ns_index, shares))
            .collect::<VidResult<Vec<_>>>()?;
        Ok(per_ns.concat())
    }

    /// Recover the payload for a given namespace.
    /// Given namespace ID should be valid for all shares, i.e. `ns_commits` and `content` have
    /// at least `ns_index` elements for all shares.
    pub fn ns_recover(
        common: &NsAvidmGf2Common,
        ns_index: usize,
        shares: &[NsAvidmGf2Share],
    ) -> VidResult<Vec<u8>> {
        common.param.validate()?;
        if shares.is_empty() {
            return Err(VidError::InsufficientShares);
        }
        let (Some(ns_commit), Some(&ns_len)) = (
            common.ns_commits.get(ns_index),
            common.ns_lens.get(ns_index),
        ) else {
            return Err(VidError::IndexOutOfBound);
        };
        if !shares.iter().all(|share| share.contains_ns(ns_index)) {
            return Err(VidError::IndexOutOfBound);
        }
        let shares: Vec<_> = shares
            .iter()
            .filter_map(|share| share.inner_ns_share(ns_index))
            .collect();
        // The commitment does not bind `ns_lens`, so shards sized for another
        // namespace length must be rejected before decoding.
        if !shares.iter().all(|share| {
            share.payload.is_empty() || Self::ns_len_matches_shards(&common.param, ns_len, share)
        }) {
            return Err(VidError::InvalidShare);
        }
        AvidmGf2Scheme::recover(&common.param, ns_commit, &shares)
    }

    /// Whether `content`'s shards have the size an `ns_len`-byte namespace
    /// disperses to under `param`.
    fn ns_len_matches_shards(
        param: &NsAvidmGf2Param,
        ns_len: usize,
        content: &AvidmGf2Share,
    ) -> bool {
        content.shard_bytes().is_some_and(|shard_bytes| {
            AvidmGf2Scheme::shard_bytes(ns_len, param.recovery_threshold) == Some(shard_bytes)
        })
    }
}

/// A namespaced dispersal.
#[derive(Clone, Debug)]
#[non_exhaustive]
pub struct NsDispersal {
    /// Index of this namespace in the namespace table.
    pub ns_index: usize,
    /// Byte length of this namespace's slice of the payload.
    pub payload_byte_len: usize,
    /// Commitment to this namespace's shards.
    pub commit: AvidmGf2Commit,
    /// One share per storage node, in distribution order.
    pub shares: Vec<AvidmGf2Share>,
}

/// Unit tests
#[cfg(test)]
pub mod tests {
    use rand::{RngCore, SeedableRng, rngs::StdRng, seq::SliceRandom};

    use crate::{
        VidError,
        avidm_gf2::{AvidmGf2Param, namespaced::NsAvidmGf2Scheme},
    };

    fn disperse_with_payload(
        payload: &[u8],
    ) -> (
        crate::avidm_gf2::namespaced::NsAvidmGf2Commit,
        crate::avidm_gf2::namespaced::NsAvidmGf2Common,
        Vec<crate::avidm_gf2::namespaced::NsAvidmGf2Share>,
    ) {
        let num_storage_nodes = 9;
        let ns_table = [(0usize..15), (15..48)];

        let mut rng = jf_utils::test_rng();
        let weights: Vec<u32> = (0..num_storage_nodes)
            .map(|_| rng.next_u32() % 5 + 1)
            .collect();
        let total_weights: u32 = weights.iter().sum();
        let recovery_threshold = total_weights.div_ceil(3) as usize;
        let params = NsAvidmGf2Scheme::setup(recovery_threshold, total_weights as usize).unwrap();

        NsAvidmGf2Scheme::ns_disperse(&params, &weights, payload, ns_table.iter().cloned()).unwrap()
    }

    fn setup_test_data() -> (
        crate::avidm_gf2::namespaced::NsAvidmGf2Commit,
        crate::avidm_gf2::namespaced::NsAvidmGf2Common,
        Vec<crate::avidm_gf2::namespaced::NsAvidmGf2Share>,
    ) {
        let payload: Vec<u8> = (0u8..48).collect();
        disperse_with_payload(&payload)
    }

    #[test]
    fn verify_share_with_verified_common_accepts_valid() {
        let (commit, common, shares) = setup_test_data();
        assert!(NsAvidmGf2Scheme::is_consistent(&commit, &common));
        for share in &shares {
            assert!(
                NsAvidmGf2Scheme::verify_share_with_verified_common(&common, share)
                    .is_ok_and(|r| r.is_ok())
            );
        }
    }

    #[test]
    fn verify_share_with_verified_common_rejects_tampered_share() {
        let (_commit, common, shares) = setup_test_data();
        // Create a tampered share by removing one namespace entry
        let mut tampered = shares[0].clone();
        tampered.0.pop();
        assert!(NsAvidmGf2Scheme::verify_share_with_verified_common(&common, &tampered).is_err());

        // Create a tampered share by dispersing a different payload and swapping
        let (_commit2, _common2, shares2) = disperse_with_payload(&[0xAB; 48]);
        let mut mixed = shares[0].clone();
        mixed.0[0] = shares2[0].0[0].clone();
        assert!(
            NsAvidmGf2Scheme::verify_share_with_verified_common(&common, &mixed)
                .is_ok_and(|r| r.is_err())
        );
    }

    /// Disperse `payload` under a known `recovery_threshold`, so a test knows
    /// the shard-size window the `ns_lens` check tolerates.
    fn disperse_with_threshold(
        recovery_threshold: usize,
        payload: &[u8],
        ns_table: &[std::ops::Range<usize>],
    ) -> (
        crate::avidm_gf2::namespaced::NsAvidmGf2Commit,
        crate::avidm_gf2::namespaced::NsAvidmGf2Common,
        Vec<crate::avidm_gf2::namespaced::NsAvidmGf2Share>,
    ) {
        let total_weights = 3 * recovery_threshold;
        let weights = vec![1u32; total_weights];
        let params = NsAvidmGf2Scheme::setup(recovery_threshold, total_weights).unwrap();
        NsAvidmGf2Scheme::ns_disperse(&params, &weights, payload, ns_table.iter().cloned()).unwrap()
    }

    /// The commitment does not bind `ns_lens`, so a share must not verify
    /// against a common whose `ns_lens` its shards were not dispersed from.
    /// The shard size pins each length to a `2 * recovery_threshold`-byte
    /// window: lengths inside it still verify, lengths outside do not, and
    /// a mismatch is a verification failure rather than an error or a panic.
    #[test]
    fn verify_share_rejects_ns_len_outside_shard_window() {
        let k = 4;
        let payload = vec![7u8; 100];
        let (commit, common, shares) = disperse_with_threshold(k, &payload, &[0..30, 30..100]);
        assert!(NsAvidmGf2Scheme::is_consistent(&commit, &common));
        // 30 bytes plus the pad byte need ceil(31 / 4) = 8 bytes per shard,
        // already even, so those shards fit any length in 24..32.
        assert_eq!(shares[0].0[0].shard_bytes(), Some(8));
        let window = (8 - 2) * k..8 * k;

        for ns_len in [window.start, window.end - 1] {
            let mut tampered = common.clone();
            tampered.ns_lens[0] = ns_len;
            for share in &shares {
                assert_eq!(
                    NsAvidmGf2Scheme::verify_share_with_verified_common(&tampered, share).unwrap(),
                    Ok(()),
                    "{ns_len}"
                );
            }
        }
        for ns_len in [0, window.start - 1, window.end, 100, usize::MAX] {
            let mut tampered = common.clone();
            tampered.ns_lens[0] = ns_len;
            for share in &shares {
                assert_eq!(
                    NsAvidmGf2Scheme::verify_share_with_verified_common(&tampered, share).unwrap(),
                    Err(()),
                    "{ns_len}"
                );
            }
        }

        // The forgery this catches: shrink the second namespace by a full
        // window, so the share reports a smaller block than was committed to.
        // The same shrink spread over several namespaces is not caught, see
        // `verify_share_accepts_ns_lens_shrunk_inside_every_window`.
        let mut shrunk = common.clone();
        shrunk.ns_lens[1] -= 2 * k;
        assert!(shrunk.payload_byte_len() < common.payload_byte_len());
        for share in &shares {
            assert_eq!(
                NsAvidmGf2Scheme::verify_share_with_verified_common(&shrunk, share).unwrap(),
                Err(())
            );
        }
    }

    /// The window is per namespace and the disperser chooses how many there
    /// are, so the misdeclared total adds up: a payload split into namespaces
    /// shorter than a window can declare every one of them as a single byte
    /// and still verify. This pins that residual gap; a check that binds
    /// `ns_lens` exactly should flip it.
    #[test]
    fn verify_share_accepts_ns_lens_shrunk_inside_every_window() {
        let k = 4;
        // Lengths 1..=7 all disperse to 2-byte shards under k = 4.
        let ns_len = 2 * k - 1;
        let num_ns = 10;
        let payload = vec![7u8; num_ns * ns_len];
        let ns_table: Vec<_> = (0..num_ns).map(|i| i * ns_len..(i + 1) * ns_len).collect();
        let (commit, common, shares) = disperse_with_threshold(k, &payload, &ns_table);
        for share in &shares {
            assert_eq!(share.0[0].shard_bytes(), Some(2));
        }

        let mut shrunk = common.clone();
        shrunk.ns_lens = vec![1; num_ns];
        assert!(NsAvidmGf2Scheme::is_consistent(&commit, &shrunk));
        assert_eq!(shrunk.payload_byte_len(), num_ns);
        for share in &shares {
            assert_eq!(
                NsAvidmGf2Scheme::verify_share_with_verified_common(&shrunk, share).unwrap(),
                Ok(())
            );
        }
    }

    #[test]
    fn composition_equivalence() {
        let (commit, common, shares) = setup_test_data();
        for share in &shares {
            let full_result = NsAvidmGf2Scheme::verify_share(&commit, &common, share)
                .unwrap()
                .is_ok();
            let composed_result = NsAvidmGf2Scheme::is_consistent(&commit, &common)
                && NsAvidmGf2Scheme::verify_share_with_verified_common(&common, share)
                    .unwrap()
                    .is_ok();
            assert_eq!(full_result, composed_result);
        }
    }

    #[test]
    fn is_consistent_rejects_tampered_commit() {
        let (commit, common, _shares) = setup_test_data();
        // Use commit from a different dispersal
        let (different_commit, ..) = disperse_with_payload(&[0xCD; 48]);
        // Verify original is consistent
        assert!(NsAvidmGf2Scheme::is_consistent(&commit, &common));
        // Verify different commit is inconsistent with original common
        assert!(!NsAvidmGf2Scheme::is_consistent(&different_commit, &common));
    }

    #[test]
    fn is_consistent_rejects_tampered_common() {
        let (commit, common, _shares) = setup_test_data();
        // Swap in ns_commits from a different dispersal
        let (_, different_common, _) = disperse_with_payload(&[0xCD; 48]);
        let mut tampered_common = common;
        tampered_common.ns_commits = different_common.ns_commits;
        assert!(!NsAvidmGf2Scheme::is_consistent(&commit, &tampered_common));
    }

    #[test]
    fn round_trip() {
        // play with these items
        let num_storage_nodes = 9;
        let ns_lens = [15, 33];
        let ns_table = [(0usize..15), (15..48)];
        let payload_byte_len = ns_lens.iter().sum();

        let mut rng = jf_utils::test_rng();

        // more items as a function of the above
        let weights: Vec<u32> = (0..num_storage_nodes)
            .map(|_| rng.next_u32() % 5 + 1)
            .collect();
        let total_weights: u32 = weights.iter().sum();
        let recovery_threshold = total_weights.div_ceil(3) as usize;
        let params = NsAvidmGf2Scheme::setup(recovery_threshold, total_weights as usize).unwrap();

        println!(
            "recovery_threshold:: {recovery_threshold} num_storage_nodes: {num_storage_nodes} \
             payload_byte_len: {payload_byte_len}"
        );
        println!("weights: {weights:?}");

        let payload = {
            let mut bytes_random = vec![0u8; payload_byte_len];
            rng.fill_bytes(&mut bytes_random);
            bytes_random
        };

        let (commit, common, mut shares) =
            NsAvidmGf2Scheme::ns_disperse(&params, &weights, &payload, ns_table.iter().cloned())
                .unwrap();

        assert_eq!(shares.len(), num_storage_nodes);

        assert_eq!(
            commit,
            NsAvidmGf2Scheme::commit(&params, &payload, ns_table.iter().cloned())
                .unwrap()
                .0
        );

        // verify shares
        shares.iter().for_each(|share| {
            assert!(
                NsAvidmGf2Scheme::verify_share(&commit, &common, share).is_ok_and(|r| r.is_ok())
            )
        });

        // test payload recovery on a random subset of shares
        shares.shuffle(&mut rng);
        let mut cumulated_weights = 0;
        let mut cut_index = 0;
        while cumulated_weights <= recovery_threshold {
            cumulated_weights += shares[cut_index].weight();
            cut_index += 1;
        }
        let ns0_payload_recovered =
            NsAvidmGf2Scheme::ns_recover(&common, 0, &shares[..cut_index]).unwrap();
        assert_eq!(ns0_payload_recovered[..], payload[ns_table[0].clone()]);
        let ns1_payload_recovered =
            NsAvidmGf2Scheme::ns_recover(&common, 1, &shares[..cut_index]).unwrap();
        assert_eq!(ns1_payload_recovered[..], payload[ns_table[1].clone()]);
        let payload_recovered = NsAvidmGf2Scheme::recover(&common, &shares[..cut_index]).unwrap();
        assert_eq!(payload_recovered, payload);
    }

    proptest::proptest! {
        #[test]
        fn any_threshold_covering_subset_recovers_the_payload(
            weights in proptest::collection::vec(1u32..=5, 1..=9),
            payload in proptest::collection::vec(proptest::num::u8::ANY, 0..2048),
            ns_ends in proptest::collection::vec(proptest::num::usize::ANY, 0..4),
            shuffle_seed: u64,
        ) {
            let total_weights: u32 = weights.iter().sum();
            let recovery_threshold = total_weights.div_ceil(3) as usize;
            let params =
                NsAvidmGf2Scheme::setup(recovery_threshold, total_weights as usize).unwrap();
            let mut ends: Vec<usize> = ns_ends
                .iter()
                .map(|end| end % (payload.len() + 1))
                .chain([payload.len()])
                .collect();
            ends.sort_unstable();
            ends.dedup();
            let ns_table: Vec<_> = std::iter::once(0)
                .chain(ends.iter().copied())
                .zip(ends.iter().copied())
                .map(|(start, end)| start..end)
                .filter(|range| !range.is_empty() || payload.is_empty())
                .collect();

            let (_commit, common, mut shares) =
                NsAvidmGf2Scheme::ns_disperse(&params, &weights, &payload, ns_table.clone())
                    .unwrap();
            shares.shuffle(&mut StdRng::seed_from_u64(shuffle_seed));
            let mut weight = 0;
            let subset: Vec<_> = shares
                .into_iter()
                .take_while(|share| {
                    let needed = weight < recovery_threshold;
                    weight += share.weight();
                    needed
                })
                .collect();

            proptest::prop_assert_eq!(NsAvidmGf2Scheme::recover(&common, &subset).unwrap(), payload.clone());
            for (ns_index, range) in ns_table.into_iter().enumerate() {
                proptest::prop_assert_eq!(
                    NsAvidmGf2Scheme::ns_recover(&common, ns_index, &subset).unwrap(),
                    payload[range].to_vec()
                );
            }
        }
    }

    #[test]
    fn ns_recover_rejects_ns_len_the_shards_were_not_dispersed_from() {
        let k = 4;
        let (_commit, common, shares) = disperse_with_threshold(k, &[7u8; 100], &[0..30, 30..100]);
        let mut shrunk = common.clone();
        shrunk.ns_lens[1] -= 2 * k;
        assert!(matches!(
            NsAvidmGf2Scheme::ns_recover(&shrunk, 1, &shares),
            Err(VidError::InvalidShare)
        ));
    }

    #[test]
    fn ns_recover_rejects_common_with_fewer_commits_than_lens() {
        let (_commit, mut common, shares) = setup_test_data();
        common.ns_commits.pop();
        assert!(matches!(
            NsAvidmGf2Scheme::ns_recover(&common, 1, &shares),
            Err(VidError::IndexOutOfBound)
        ));
    }

    #[test]
    fn forged_common_param_is_rejected() {
        let (_commit, common, shares) = setup_test_data();
        for (recovery_threshold, total_weights) in [(0, 4), (1, 1 << 16)] {
            let mut forged = common.clone();
            forged.param = AvidmGf2Param {
                total_weights,
                recovery_threshold,
            };
            assert!(
                matches!(
                    NsAvidmGf2Scheme::verify_share_with_verified_common(&forged, &shares[0]),
                    Err(VidError::InvalidParam)
                ),
                "verify_share_with_verified_common {:?}",
                forged.param
            );
            assert!(
                matches!(
                    NsAvidmGf2Scheme::ns_recover(&forged, 0, &shares),
                    Err(VidError::InvalidParam)
                ),
                "ns_recover {:?}",
                forged.param
            );
        }
    }

    /// A caller that encodes each namespace with `ns_disperse_one` and folds the
    /// per-namespace commits with `aggregate_commit` must get the same
    /// commitment as the monolithic `commit`.
    #[test]
    fn aggregate_commit_matches_commit() {
        let num_storage_nodes = 9;
        let ns_table = [(0usize..15), (15..48), (48..49)];
        let weights = vec![1u32; num_storage_nodes];
        let params = NsAvidmGf2Scheme::setup(3, num_storage_nodes).unwrap();
        let payload: Vec<u8> = (0..49).map(|i| i as u8).collect();

        let ns_commits: Vec<_> = ns_table
            .iter()
            .enumerate()
            .map(|(ns_index, range)| {
                NsAvidmGf2Scheme::ns_disperse_one(
                    &params,
                    &weights,
                    &payload[range.clone()],
                    ns_index,
                )
                .unwrap()
                .commit
            })
            .collect();
        let aggregated = NsAvidmGf2Scheme::aggregate_commit(&ns_commits).unwrap();

        let (expected, _) =
            NsAvidmGf2Scheme::commit(&params, &payload, ns_table.iter().cloned()).unwrap();
        assert_eq!(aggregated, expected);
    }
}