//! Prio3 (`draft-irtf-cfrg-vdaf-22` §7.2): sharding, verification, aggregation and
//! unsharding over the FLP of §7.3, the message encodings of §7.2.7, and the ping-pong
//! topology of §5.7.1 for two aggregators.
//!
//! [`Prio3`] is generic over the field and the circuit; [`Vdaf`] holds the four variants this
//! module implements (§7.4.1–§7.4.4) behind one byte-level interface, the one DAP carries.

use std::marker::PhantomData;

use super::field::{decode_vec, encode_vec, vec_add, vec_sub, Field, F128, F64};
use super::flp::{decide, encode_range_checked_int, prove, query, Count, Histogram, Sum, SumVec, Valid};
use super::xof::{format_dst, Xof, SEED_SIZE};

pub const USAGE_MEAS_SHARE: u16 = 1;
pub const USAGE_PROOF_SHARE: u16 = 2;
pub const USAGE_JOINT_RANDOMNESS: u16 = 3;
pub const USAGE_PROVE_RANDOMNESS: u16 = 4;
pub const USAGE_QUERY_RANDOMNESS: u16 = 5;
pub const USAGE_JOINT_RAND_SEED: u16 = 6;
pub const USAGE_JOINT_RAND_PART: u16 = 7;

pub const NONCE_SIZE: usize = 16;
pub const VERIFY_KEY_SIZE: usize = SEED_SIZE;

/// The VDAF identifiers of the draft's Table 19.
pub const ID_COUNT: u32 = 1;
pub const ID_SUM: u32 = 2;
pub const ID_SUMVEC: u32 = 3;
pub const ID_HISTOGRAM: u32 = 4;

/// Why a VDAF step fails; each makes the report invalid (§5.2).
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VdafError(pub &'static str);

type Seed = [u8; SEED_SIZE];

/// A verifier share, decoded: the verifiers' share and the joint randomness part.
pub type VerifierShare<F> = (Vec<F>, Option<Seed>);
/// The output of sharding: the joint randomness parts, where used, and each input share.
pub type Sharded<F> = (Option<Vec<Seed>>, Vec<InputShare<F>>);
/// An input share expanded: measurement share, proofs share and blind.
type Expanded<F> = (Vec<F>, Vec<F>, Option<Seed>);

fn seed(b: &[u8]) -> Seed {
    let mut s = [0u8; SEED_SIZE];
    s.copy_from_slice(b);
    s
}

/// One aggregator's input share, decoded (§7.2.7.2).
#[derive(Clone, Debug)]
pub enum InputShare<F> {
    Leader { meas: Vec<F>, proofs: Vec<F>, blind: Option<Seed> },
    Helper { share: Seed, blind: Option<Seed> },
}

/// Prio3 over a field and a validity circuit.
pub struct Prio3<F: Field, V: Valid<F>> {
    pub id: u32,
    pub shares: usize,
    pub proofs: usize,
    pub valid: V,
    _f: PhantomData<F>,
}

/// An aggregator's verification state (§7.2.2): its output share and the joint randomness
/// seed it computed, where the circuit uses joint randomness.
#[derive(Clone, Debug)]
pub struct VerifyState<F> {
    pub out_share: Vec<F>,
    pub corrected_seed: Option<Seed>,
}

impl<F: Field, V: Valid<F>> Prio3<F, V> {
    pub fn new(id: u32, shares: usize, proofs: usize, valid: V) -> Self {
        assert!((2..256).contains(&shares) && (1..256).contains(&proofs));
        Prio3 { id, shares, proofs, valid, _f: PhantomData }
    }

    fn jr(&self) -> bool {
        self.valid.joint_rand_len() > 0
    }

    /// The length of an output share and of an aggregate share.
    pub fn out_len(&self) -> usize {
        self.valid.output_len()
    }

    pub fn rand_size(&self) -> usize {
        SEED_SIZE * self.shares * if self.jr() { 2 } else { 1 }
    }

    fn dst(&self, usage: u16, ctx: &[u8]) -> Vec<u8> {
        let mut d = format_dst(0, self.id, usage);
        d.extend_from_slice(ctx);
        d
    }

    fn helper_meas_share(&self, ctx: &[u8], agg_id: usize, share: &[u8]) -> Vec<F> {
        Xof::expand_into_vec(share, &self.dst(USAGE_MEAS_SHARE, ctx), &[agg_id as u8], self.valid.meas_len())
    }

    fn helper_proofs_share(&self, ctx: &[u8], agg_id: usize, share: &[u8]) -> Vec<F> {
        Xof::expand_into_vec(share, &self.dst(USAGE_PROOF_SHARE, ctx), &[self.proofs as u8, agg_id as u8], self.valid.proof_len() * self.proofs)
    }

    fn prove_rands(&self, ctx: &[u8], s: &[u8]) -> Vec<F> {
        Xof::expand_into_vec(s, &self.dst(USAGE_PROVE_RANDOMNESS, ctx), &[self.proofs as u8], self.valid.prove_rand_len() * self.proofs)
    }

    fn query_rands(&self, vk: &[u8], ctx: &[u8], nonce: &[u8]) -> Vec<F> {
        let mut binder = vec![self.proofs as u8];
        binder.extend_from_slice(nonce);
        Xof::expand_into_vec(vk, &self.dst(USAGE_QUERY_RANDOMNESS, ctx), &binder, self.valid.query_rand_len() * self.proofs)
    }

    fn joint_rand_part(&self, ctx: &[u8], agg_id: usize, blind: &[u8], meas_share: &[F], nonce: &[u8]) -> Seed {
        let mut binder = vec![agg_id as u8];
        binder.extend_from_slice(nonce);
        binder.extend_from_slice(&encode_vec(meas_share));
        Xof::derive_seed(blind, &self.dst(USAGE_JOINT_RAND_PART, ctx), &binder)
    }

    fn joint_rand_seed(&self, ctx: &[u8], parts: &[Seed]) -> Seed {
        let binder: Vec<u8> = parts.concat();
        Xof::derive_seed(&[0u8; SEED_SIZE], &self.dst(USAGE_JOINT_RAND_SEED, ctx), &binder)
    }

    fn joint_rands(&self, ctx: &[u8], s: &[u8]) -> Vec<F> {
        Xof::expand_into_vec(s, &self.dst(USAGE_JOINT_RANDOMNESS, ctx), &[self.proofs as u8], self.valid.joint_rand_len() * self.proofs)
    }

    /// shard (§7.2.1) of an encoded measurement. The encoding is the caller's: an honest
    /// client encodes with the circuit's encoder; a test of the proof shards one that is not.
    pub fn shard(&self, ctx: &[u8], meas: &[F], nonce: &[u8], rand: &[u8]) -> Result<Sharded<F>, VdafError> {
        if nonce.len() != NONCE_SIZE {
            return Err(VdafError("incorrect nonce size"));
        }
        if rand.len() != self.rand_size() {
            return Err(VdafError("incorrect size of random bytes argument"));
        }
        if meas.len() != self.valid.meas_len() {
            return Err(VdafError("incorrect measurement length"));
        }
        let seeds: Vec<Seed> = rand.chunks(SEED_SIZE).map(seed).collect();
        let n = self.shares;
        let (helper_shares, helper_blinds, leader_blind, prove_seed): (Vec<Seed>, Vec<Option<Seed>>, Option<Seed>, Seed) = if self.jr() {
            let hs: Vec<Seed> = (0..n - 1).map(|i| seeds[2 * i]).collect();
            let hb: Vec<Option<Seed>> = (0..n - 1).map(|i| Some(seeds[2 * i + 1])).collect();
            (hs, hb, Some(seeds[2 * (n - 1)]), seeds[2 * (n - 1) + 1])
        } else {
            (seeds[..n - 1].to_vec(), vec![None; n - 1], None, seeds[n - 1])
        };
        let mut leader_meas = meas.to_vec();
        let mut parts: Vec<Seed> = Vec::new();
        for j in 0..n - 1 {
            let hm = self.helper_meas_share(ctx, j + 1, &helper_shares[j]);
            leader_meas = vec_sub(&leader_meas, &hm);
            if let Some(b) = helper_blinds[j] {
                parts.push(self.joint_rand_part(ctx, j + 1, &b, &hm, nonce));
            }
        }
        let mut joint_rands = Vec::new();
        if let Some(lb) = leader_blind {
            parts.insert(0, self.joint_rand_part(ctx, 0, &lb, &leader_meas, nonce));
            joint_rands = self.joint_rands(ctx, &self.joint_rand_seed(ctx, &parts));
        }
        let prove_rands = self.prove_rands(ctx, &prove_seed);
        let (prl, jrl) = (self.valid.prove_rand_len(), self.valid.joint_rand_len());
        let mut proofs = Vec::new();
        for i in 0..self.proofs {
            let jr = if jrl > 0 { &joint_rands[i * jrl..(i + 1) * jrl] } else { &[][..] };
            proofs.extend(prove(&self.valid, meas, &prove_rands[i * prl..(i + 1) * prl], jr));
        }
        for (j, hs) in helper_shares.iter().enumerate() {
            proofs = vec_sub(&proofs, &self.helper_proofs_share(ctx, j + 1, hs));
        }
        let mut shares = vec![InputShare::Leader { meas: leader_meas, proofs, blind: leader_blind }];
        for j in 0..n - 1 {
            shares.push(InputShare::Helper { share: helper_shares[j], blind: helper_blinds[j] });
        }
        Ok((self.jr().then_some(parts), shares))
    }

    fn expand(&self, ctx: &[u8], agg_id: usize, s: &InputShare<F>) -> Result<Expanded<F>, VdafError> {
        match (agg_id, s) {
            (0, InputShare::Leader { meas, proofs, blind }) => Ok((meas.clone(), proofs.clone(), *blind)),
            (i, InputShare::Helper { share, blind }) if i > 0 => {
                Ok((self.helper_meas_share(ctx, i, share), self.helper_proofs_share(ctx, i, share), *blind))
            }
            _ => Err(VdafError("input share does not match the aggregator")),
        }
    }

    /// verify_init (§7.2.2).
    pub fn verify_init(
        &self,
        vk: &[u8],
        ctx: &[u8],
        agg_id: usize,
        nonce: &[u8],
        public: &Option<Vec<Seed>>,
        input: &InputShare<F>,
    ) -> Result<(VerifyState<F>, VerifierShare<F>), VdafError> {
        if vk.len() != VERIFY_KEY_SIZE || nonce.len() != NONCE_SIZE || agg_id >= self.shares {
            return Err(VdafError("bad verify_init parameters"));
        }
        let (meas, proofs, blind) = self.expand(ctx, agg_id, input)?;
        let out_share = self.valid.truncate(&meas);
        let (mut joint_rands, mut corrected, mut part) = (Vec::new(), None, None);
        if self.jr() {
            let blind = blind.ok_or(VdafError("no blind"))?;
            let mut parts = public.clone().ok_or(VdafError("no joint randomness parts"))?;
            if parts.len() != self.shares {
                return Err(VdafError("joint randomness parts"));
            }
            let p = self.joint_rand_part(ctx, agg_id, &blind, &meas, nonce);
            parts[agg_id] = p;
            let s = self.joint_rand_seed(ctx, &parts);
            joint_rands = self.joint_rands(ctx, &s);
            corrected = Some(s);
            part = Some(p);
        }
        let query_rands = self.query_rands(vk, ctx, nonce);
        let (pl, ql, jl) = (self.valid.proof_len(), self.valid.query_rand_len(), self.valid.joint_rand_len());
        let mut verifiers = Vec::new();
        for i in 0..self.proofs {
            let jr = if jl > 0 { &joint_rands[i * jl..(i + 1) * jl] } else { &[][..] };
            let v = query(&self.valid, &meas, &proofs[i * pl..(i + 1) * pl], &query_rands[i * ql..(i + 1) * ql], jr, self.shares)
                .map_err(|_| VdafError("test point is a root of unity"))?;
            verifiers.extend(v);
        }
        Ok((VerifyState { out_share, corrected_seed: corrected }, (verifiers, part)))
    }

    /// verifier_shares_to_message (§7.2.2).
    pub fn verifier_shares_to_message(&self, ctx: &[u8], shares: &[VerifierShare<F>]) -> Result<Option<Seed>, VdafError> {
        let vl = self.valid.verifier_len();
        let mut verifiers = vec![F::zero(); vl * self.proofs];
        let mut parts = Vec::new();
        for (v, p) in shares {
            if v.len() != verifiers.len() {
                return Err(VdafError("verifier share length"));
            }
            verifiers = vec_add(&verifiers, v);
            if self.jr() {
                parts.push(p.ok_or(VdafError("no joint randomness part"))?);
            }
        }
        for i in 0..self.proofs {
            if !decide(&self.valid, &verifiers[i * vl..(i + 1) * vl]) {
                return Err(VdafError("proof verifier check failed"));
            }
        }
        Ok(self.jr().then(|| self.joint_rand_seed(ctx, &parts)))
    }

    /// verify_next (§7.2.2): the output share, once the joint randomness checks.
    pub fn verify_next(&self, state: &VerifyState<F>, msg: &Option<Seed>) -> Result<Vec<F>, VdafError> {
        if *msg != state.corrected_seed {
            return Err(VdafError("joint randomness check failed"));
        }
        Ok(state.out_share.clone())
    }

    // Encodings (§7.2.7).

    pub fn encode_public_share(&self, p: &Option<Vec<Seed>>) -> Vec<u8> {
        p.as_ref().map(|v| v.concat()).unwrap_or_default()
    }

    pub fn decode_public_share(&self, b: &[u8]) -> Result<Option<Vec<Seed>>, VdafError> {
        if !self.jr() {
            return if b.is_empty() { Ok(None) } else { Err(VdafError("public share")) };
        }
        if b.len() != SEED_SIZE * self.shares {
            return Err(VdafError("public share"));
        }
        Ok(Some(b.chunks(SEED_SIZE).map(seed).collect()))
    }

    pub fn encode_input_share(&self, s: &InputShare<F>) -> Vec<u8> {
        let mut out = Vec::new();
        match s {
            InputShare::Leader { meas, proofs, blind } => {
                out.extend(encode_vec(meas));
                out.extend(encode_vec(proofs));
                if let Some(b) = blind {
                    out.extend_from_slice(b);
                }
            }
            InputShare::Helper { share, blind } => {
                out.extend_from_slice(share);
                if let Some(b) = blind {
                    out.extend_from_slice(b);
                }
            }
        }
        out
    }

    pub fn decode_input_share(&self, agg_id: usize, b: &[u8]) -> Result<InputShare<F>, VdafError> {
        let blind_len = if self.jr() { SEED_SIZE } else { 0 };
        let bad = VdafError("input share");
        if agg_id == 0 {
            let ml = self.valid.meas_len() * F::ENCODED_SIZE;
            let pl = self.valid.proof_len() * self.proofs * F::ENCODED_SIZE;
            if b.len() != ml + pl + blind_len {
                return Err(bad);
            }
            let meas = decode_vec(&b[..ml]).ok_or(bad.clone())?;
            let proofs = decode_vec(&b[ml..ml + pl]).ok_or(bad)?;
            let blind = self.jr().then(|| seed(&b[ml + pl..]));
            Ok(InputShare::Leader { meas, proofs, blind })
        } else {
            if b.len() != SEED_SIZE + blind_len {
                return Err(bad);
            }
            Ok(InputShare::Helper { share: seed(&b[..SEED_SIZE]), blind: self.jr().then(|| seed(&b[SEED_SIZE..])) })
        }
    }

    pub fn encode_verifier_share(&self, s: &VerifierShare<F>) -> Vec<u8> {
        let mut out = encode_vec(&s.0);
        if let Some(p) = s.1 {
            out.extend_from_slice(&p);
        }
        out
    }

    pub fn decode_verifier_share(&self, b: &[u8]) -> Result<VerifierShare<F>, VdafError> {
        let vl = self.valid.verifier_len() * self.proofs * F::ENCODED_SIZE;
        let jl = if self.jr() { SEED_SIZE } else { 0 };
        if b.len() != vl + jl {
            return Err(VdafError("verifier share"));
        }
        let v = decode_vec(&b[..vl]).ok_or(VdafError("verifier share"))?;
        Ok((v, self.jr().then(|| seed(&b[vl..]))))
    }

    pub fn decode_verifier_message(&self, b: &[u8]) -> Result<Option<Seed>, VdafError> {
        match (self.jr(), b.len()) {
            (false, 0) => Ok(None),
            (true, SEED_SIZE) => Ok(Some(seed(b))),
            _ => Err(VdafError("verifier message")),
        }
    }
}

/// The variants of Prio3 this module implements, behind one byte-level interface: every
/// share, message and aggregate is its §7.2.7 encoding.
pub enum Vdaf {
    Count(Prio3<F64, Count>),
    Sum(Prio3<F64, Sum>),
    SumVec(Prio3<F128, SumVec>),
    Histogram(Prio3<F128, Histogram>),
}

macro_rules! each {
    ($self:expr, $p:ident => $body:expr) => {
        match $self {
            Vdaf::Count($p) => $body,
            Vdaf::Sum($p) => $body,
            Vdaf::SumVec($p) => $body,
            Vdaf::Histogram($p) => $body,
        }
    };
}

/// An aggregator's state after verify_init: its encoded output share and the joint
/// randomness seed it computed.
#[derive(Clone, Debug)]
pub struct State {
    out_share: Vec<u8>,
    seed: Option<Seed>,
}

fn zeros<F: Field, V: Valid<F>>(p: &Prio3<F, V>) -> Vec<u8> {
    encode_vec(&vec![F::zero(); p.valid.output_len()])
}

fn to_ints<F: Field>(v: &[F]) -> Vec<u128> {
    v.iter().map(|x| x.int()).collect()
}

fn from_ints<F: Field>(v: &[u128]) -> Result<Vec<F>, VdafError> {
    v.iter().map(|&x| if x < F::MODULUS { Ok(F::from_u128(x)) } else { Err(VdafError("element past the modulus")) }).collect()
}

impl Vdaf {
    pub fn count(shares: usize) -> Vdaf {
        Vdaf::Count(Prio3::new(ID_COUNT, shares, 1, Count::default()))
    }
    pub fn sum(shares: usize, max: u64) -> Vdaf {
        Vdaf::Sum(Prio3::new(ID_SUM, shares, 1, Sum::new(max as u128)))
    }
    pub fn sumvec(shares: usize, length: usize, max: u64, chunk: usize) -> Vdaf {
        Vdaf::SumVec(Prio3::new(ID_SUMVEC, shares, 1, SumVec::new(length, max as u128, chunk)))
    }
    pub fn histogram(shares: usize, length: usize, chunk: usize) -> Vdaf {
        Vdaf::Histogram(Prio3::new(ID_HISTOGRAM, shares, 1, Histogram::new(length, chunk)))
    }

    /// The VDAF of a DAP task (`draft-ietf-ppm-dap-20` Appendix B), two aggregators; `None`
    /// for a type this module does not implement or a configuration that does not read.
    pub fn from_dap(vdaf_type: u32, cfg: &[u8]) -> Option<Vdaf> {
        let u32_at = |i: usize| cfg.get(i..i + 4).map(|b| u32::from_be_bytes(b.try_into().unwrap()) as usize);
        let u64_at = |i: usize| cfg.get(i..i + 8).map(|b| u64::from_be_bytes(b.try_into().unwrap()));
        match vdaf_type {
            ID_COUNT if cfg.is_empty() => Some(Vdaf::count(2)),
            ID_SUM if cfg.len() == 8 => {
                let max = u64_at(0)?;
                (max > 0).then(|| Vdaf::sum(2, max))
            }
            ID_SUMVEC if cfg.len() == 16 => {
                let (length, max, chunk) = (u32_at(0)?, u64_at(4)?, u32_at(12)?);
                (length > 0 && max > 0 && chunk > 0).then(|| Vdaf::sumvec(2, length, max, chunk))
            }
            ID_HISTOGRAM if cfg.len() == 8 => {
                let (length, chunk) = (u32_at(0)?, u32_at(4)?);
                (length > 0 && chunk > 0).then(|| Vdaf::histogram(2, length, chunk))
            }
            _ => None,
        }
    }

    pub fn id(&self) -> u32 {
        each!(self, p => p.id)
    }

    /// The DAP configuration encoding (Appendix B).
    pub fn dap_config(&self) -> Vec<u8> {
        match self {
            Vdaf::Count(_) => vec![],
            Vdaf::Sum(p) => (p.valid.max as u64).to_be_bytes().to_vec(),
            Vdaf::SumVec(p) => {
                let mut v = (p.valid.length as u32).to_be_bytes().to_vec();
                v.extend_from_slice(&(p.valid.max as u64).to_be_bytes());
                v.extend_from_slice(&(p.valid.chunk as u32).to_be_bytes());
                v
            }
            Vdaf::Histogram(p) => {
                let mut v = (p.valid.length as u32).to_be_bytes().to_vec();
                v.extend_from_slice(&(p.valid.chunk as u32).to_be_bytes());
                v
            }
        }
    }

    pub fn name(&self) -> &'static str {
        match self {
            Vdaf::Count(_) => "Prio3Count",
            Vdaf::Sum(_) => "Prio3Sum",
            Vdaf::SumVec(_) => "Prio3SumVec",
            Vdaf::Histogram(_) => "Prio3Histogram",
        }
    }

    pub fn shares(&self) -> usize {
        each!(self, p => p.shares)
    }

    pub fn rand_size(&self) -> usize {
        each!(self, p => p.rand_size())
    }

    pub fn output_len(&self) -> usize {
        each!(self, p => p.out_len())
    }

    /// The circuit's encoding of a measurement (§7.4): Count `[b]`, Sum `[v]`, SumVec the
    /// vector, Histogram `[bucket]`. `None` for a measurement out of range.
    pub fn encode(&self, m: &[u128]) -> Option<Vec<u128>> {
        match self {
            Vdaf::Count(_) => (m.len() == 1 && m[0] <= 1).then(|| m.to_vec()),
            Vdaf::Sum(p) => (m.len() == 1).then(|| encode_range_checked_int::<F64>(m[0], p.valid.max)).flatten().map(|v| to_ints(&v)),
            Vdaf::SumVec(p) => {
                if m.len() != p.valid.length {
                    return None;
                }
                let mut out = Vec::new();
                for x in m {
                    out.extend(to_ints(&encode_range_checked_int::<F128>(*x, p.valid.max)?));
                }
                Some(out)
            }
            Vdaf::Histogram(p) => (m.len() == 1 && m[0] < p.valid.length as u128).then(|| {
                let mut v = vec![0u128; p.valid.length];
                v[m[0] as usize] = 1;
                v
            }),
        }
    }

    /// shard (§7.2.1) of an encoded measurement: the public share and each input share.
    pub fn shard(&self, ctx: &[u8], encoded: &[u128], nonce: &[u8], rand: &[u8]) -> Result<(Vec<u8>, Vec<Vec<u8>>), VdafError> {
        each!(self, p => {
            let (ps, is) = p.shard(ctx, &from_ints(encoded)?, nonce, rand)?;
            Ok((p.encode_public_share(&ps), is.iter().map(|s| p.encode_input_share(s)).collect()))
        })
    }

    /// verify_init (§7.2.2) over encoded shares.
    pub fn verify_init(&self, vk: &[u8], ctx: &[u8], agg_id: usize, nonce: &[u8], public: &[u8], input: &[u8]) -> Result<(State, Vec<u8>), VdafError> {
        each!(self, p => {
            let ps = p.decode_public_share(public)?;
            let is = p.decode_input_share(agg_id, input)?;
            let (st, vs) = p.verify_init(vk, ctx, agg_id, nonce, &ps, &is)?;
            Ok((State { out_share: encode_vec(&st.out_share), seed: st.corrected_seed }, p.encode_verifier_share(&vs)))
        })
    }

    pub fn verifier_shares_to_message(&self, ctx: &[u8], shares: &[Vec<u8>]) -> Result<Vec<u8>, VdafError> {
        each!(self, p => {
            let vs = shares.iter().map(|s| p.decode_verifier_share(s)).collect::<Result<Vec<_>, _>>()?;
            Ok(p.verifier_shares_to_message(ctx, &vs)?.map(|s| s.to_vec()).unwrap_or_default())
        })
    }

    /// verify_next (§7.2.2): the encoded output share.
    pub fn verify_next(&self, state: &State, msg: &[u8]) -> Result<Vec<u8>, VdafError> {
        each!(self, p => {
            let m = p.decode_verifier_message(msg)?;
            if m != state.seed {
                return Err(VdafError("joint randomness check failed"));
            }
            Ok(state.out_share.clone())
        })
    }

    /// agg_init (§7.2.4): the encoded zero aggregate share.
    pub fn agg_init(&self) -> Vec<u8> {
        each!(self, p => zeros(p))
    }

    /// agg_update and merge (§7.2.4): two encoded shares added element-wise.
    pub fn merge(&self, a: &[u8], b: &[u8]) -> Result<Vec<u8>, VdafError> {
        each!(self, p => {
            let n = p.out_len();
            let x = self.decode_agg(a, n, p)?;
            let y = self.decode_agg(b, n, p)?;
            Ok(encode_vec(&vec_add(&x, &y)))
        })
    }

    fn decode_agg<F: Field, V: Valid<F>>(&self, b: &[u8], n: usize, _p: &Prio3<F, V>) -> Result<Vec<F>, VdafError> {
        let v: Vec<F> = decode_vec(b).ok_or(VdafError("aggregate share"))?;
        if v.len() != n {
            return Err(VdafError("aggregate share"));
        }
        Ok(v)
    }

    /// Whether `b` is an aggregate share of this VDAF (§7.2.7.5).
    pub fn is_agg_share(&self, b: &[u8]) -> bool {
        each!(self, p => self.decode_agg(b, p.out_len(), p).is_ok())
    }

    /// unshard (§7.2.5): the aggregate result, each element as an integer.
    pub fn unshard(&self, shares: &[Vec<u8>]) -> Result<Vec<u128>, VdafError> {
        let mut agg = self.agg_init();
        for s in shares {
            agg = self.merge(&agg, s)?;
        }
        each!(self, p => {
            let v = self.decode_agg(&agg, p.out_len(), p)?;
            Ok(to_ints(&v))
        })
    }
}

// ---------------------------------------------------------------------------------------
// The ping-pong topology (§5.7.1), for Prio3's one round.

/// A ping-pong message (§5.7.1): `initialize` (0) with a verifier share, `continue` (1) with a
/// verifier message and share, `finish` (2) with a verifier message.
pub fn pp_encode(kind: u8, items: &[&[u8]]) -> Vec<u8> {
    let mut v = vec![kind];
    for i in items {
        v.extend_from_slice(&(i.len() as u32).to_be_bytes());
        v.extend_from_slice(i);
    }
    v
}

pub fn pp_decode(b: &[u8]) -> Result<(u8, Vec<Vec<u8>>), VdafError> {
    let bad = VdafError("ping-pong message");
    let (&kind, mut rest) = b.split_first().ok_or(bad.clone())?;
    let n = match kind {
        0 | 2 => 1,
        1 => 2,
        _ => return Err(bad),
    };
    let mut items = Vec::new();
    for _ in 0..n {
        let len = u32::from_be_bytes(rest.get(..4).ok_or(bad.clone())?.try_into().unwrap()) as usize;
        let item = rest.get(4..4 + len).ok_or(bad.clone())?;
        items.push(item.to_vec());
        rest = &rest[4 + len..];
    }
    if !rest.is_empty() {
        return Err(bad);
    }
    Ok((kind, items))
}

impl Vdaf {
    /// ping_pong_leader_init: the leader's state and its `initialize` message.
    pub fn leader_init(&self, vk: &[u8], ctx: &[u8], nonce: &[u8], public: &[u8], input: &[u8]) -> Result<(State, Vec<u8>), VdafError> {
        let (st, vs) = self.verify_init(vk, ctx, 0, nonce, public, input)?;
        Ok((st, pp_encode(0, &[&vs])))
    }

    /// ping_pong_helper_init for a one-round VDAF: the helper's output share and its
    /// `finish` message, which carries the verifier message.
    pub fn helper_init(&self, vk: &[u8], ctx: &[u8], nonce: &[u8], public: &[u8], input: &[u8], inbound: &[u8]) -> Result<(Vec<u8>, Vec<u8>), VdafError> {
        let (st, own) = self.verify_init(vk, ctx, 1, nonce, public, input)?;
        let (kind, items) = pp_decode(inbound)?;
        if kind != 0 {
            return Err(VdafError("not an initialize message"));
        }
        let msg = self.verifier_shares_to_message(ctx, &[items[0].clone(), own])?;
        let out = self.verify_next(&st, &msg)?;
        Ok((out, pp_encode(2, &[&msg])))
    }

    /// ping_pong_leader_continued for a one-round VDAF: the leader's output share.
    pub fn leader_continued(&self, state: &State, inbound: &[u8]) -> Result<Vec<u8>, VdafError> {
        let (kind, items) = pp_decode(inbound)?;
        if kind != 2 {
            return Err(VdafError("not a finish message"));
        }
        self.verify_next(state, &items[0])
    }
}
