Misc fixes

Cargo.toml

@@ -17,6 +17,7 @@ ark-serialize = { version = "0.5", default-features = false, features = ["derive
 ark-scale = { version = "0.0.13", default-features = false }
 ark-vrf = { version = "0.5.0", default-features = false, features = ["bandersnatch", "ring"] }
 spin = { version = "0.9", default-features = false, features = ["once"] }
+smallvec = { version = "1", default-features = false }
 sha2 = { version = "0.10", default-features = false, optional = true }
 
 [dev-dependencies]

src/lib.rs

@@ -2,13 +2,13 @@
 #![allow(clippy::result_unit_err)]
 
 extern crate alloc;
-extern crate core;
 
 use alloc::vec::Vec;
 
 use core::{fmt::Debug, ops::Range};
 use parity_scale_codec::{Decode, DecodeWithMemTracking, Encode, FullCodec, MaxEncodedLen};
 use scale_info::*;
+use smallvec::SmallVec;
 
 #[cfg(feature = "mock")]
 pub mod mock;
@@ -25,6 +25,28 @@ pub type Alias = [u8; 32];
 /// Entropy supplied for the creation of a secret key.
 pub type Entropy = [u8; 32];
 
+/// Maximum number of contexts the library supports anywhere.
+///
+/// Wire-format hard limit, enforced by deserializers. Must fit in a `u8`
+/// because wire formats encode the count with a single-byte length prefix.
+pub const MAX_CONTEXTS: usize = 16;
+
+const _: () = assert!(MAX_CONTEXTS <= u8::MAX as usize);
+
+/// SmallVec sized for one element per context.
+///
+/// Inline storage covers the common case (up to 3 contexts per proof);
+/// counts in 4..=[`MAX_CONTEXTS`] spill to the heap. Used wherever an
+/// internal collection scales with the number of contexts (aliases, VRF
+/// inputs, VRF outputs).
+pub type ContextVec<T> = SmallVec<[T; 3]>;
+
+/// Collection of aliases returned from the multi-context proof methods.
+///
+/// The single-context default wrappers (`validate`, `create`) rely on this
+/// being inline for `N=1` to avoid heap allocation.
+pub type AliasVec = ContextVec<Alias>;
+
 /// A single item in a batch proof validation request.
 ///
 /// Groups together a proof with the context and message it was created for,
@@ -46,11 +68,10 @@ pub struct BatchProofItem<Proof> {
 /// contexts without exposing the underlying member who is proving it and giving an unlinkable
 /// deterministic pseudonymic "alias" under each context.
 ///
-/// A value of this type represents a proof. It can be created using the `Self::create` function
-/// from the `Self::Secret` value associated with a `Self::Member` value who exists within a set of
-/// members identified with a `Self::Members` value. It can later be validated with the
-/// `Self::is_valid` function using `self` together with the same information used to create it
-/// (except the secret, of course!).
+/// A `Self::Proof` is created using the `Self::create` function from the `Self::Secret` value
+/// associated with a `Self::Member` value who exists within a set of members identified with a
+/// `Self::Members` value. The proof can later be validated with the `Self::is_valid` function
+/// using the same information used to create it (except the secret, of course!).
 ///
 /// A convenience [`Receipt`] type is provided for typical use cases which bundles the proof along
 /// with needed witness information describing the message and alias.
@@ -157,7 +178,7 @@ pub trait GenerateVerifiable {
 	/// of the `commitment`.
 	///
 	/// The proof will be specific to a given `context` (which determines the resultant `Alias` of
-	/// the member in a way unlinkable to the member's original identifiaction and aliases in any
+	/// the member in a way unlinkable to the member's original identification and aliases in any
 	/// other contexts) together with a provided `message` which entirely at the choice of the
 	/// individual.
 	///
@@ -191,9 +212,9 @@ pub trait GenerateVerifiable {
 		secret: &Self::Secret,
 		contexts: &[&[u8]],
 		message: &[u8],
-	) -> Result<(Self::Proof, Vec<Alias>), ()>;
+	) -> Result<(Self::Proof, AliasVec), ()>;
 
-	/// Check whether `self` is a valid proof of membership in `members` in the given `context`;
+	/// Check whether `proof` is a valid proof of membership in `members` in the given `context`;
 	/// if so, ensure that the member is necessarily associated with `alias` in this `context` and
 	/// that they elected to opine `message`.
 	fn is_valid(
@@ -222,7 +243,7 @@ pub trait GenerateVerifiable {
 		message: &[u8],
 	) -> bool {
 		match Self::validate_multi_context(config, proof, members, contexts, message) {
-			Ok(a) => a == aliases,
+			Ok(a) => a.as_slice() == aliases,
 			Err(()) => false,
 		}
 	}
@@ -249,7 +270,7 @@ pub trait GenerateVerifiable {
 		members: &Self::Members,
 		contexts: &[&[u8]],
 		message: &[u8],
-	) -> Result<Vec<Alias>, ()>;
+	) -> Result<AliasVec, ()>;
 
 	/// Check whether all of the proofs in this batch are valid, returning the `Alias` for each one,
 	/// in order of input.

src/mock.rs

@@ -20,7 +20,6 @@ use bounded_collections::{BoundedVec, ConstU32};
 use sha2::{Digest, Sha256};
 
 pub const MAX_MEMBERS: u32 = 1024;
-pub const MAX_CONTEXTS: u32 = 3;
 
 const TAG_ALIAS: &[u8] = b"verifiable-mock:v1:alias";
 const TAG_SIG: &[u8] = b"verifiable-mock:v1:sig";
@@ -69,7 +68,7 @@ fn make_proof_tag(
 pub struct MockProof {
 	pub tag: [u8; 32],
 	pub member: [u8; 32],
-	pub aliases: BoundedVec<Alias, ConstU32<MAX_CONTEXTS>>,
+	pub aliases: BoundedVec<Alias, ConstU32<{ MAX_CONTEXTS as u32 }>>,
 }
 
 /// Mock [`GenerateVerifiable`] implementation.
@@ -136,15 +135,15 @@ impl GenerateVerifiable for Mock {
 		secret: &Self::Secret,
 		contexts: &[&[u8]],
 		message: &[u8],
-	) -> Result<(Self::Proof, Vec<Alias>), ()> {
+	) -> Result<(Self::Proof, AliasVec), ()> {
 		if &member != secret {
 			return Err(());
 		}
-		if contexts.len() > MAX_CONTEXTS as usize {
+		if contexts.len() > MAX_CONTEXTS {
 			return Err(());
 		}
-		let aliases: Vec<Alias> = contexts.iter().map(|ctx| make_alias(secret, ctx)).collect();
-		let bounded = BoundedVec::try_from(aliases.clone()).map_err(|_| ())?;
+		let aliases: AliasVec = contexts.iter().map(|ctx| make_alias(secret, ctx)).collect();
+		let bounded = BoundedVec::try_from(aliases.to_vec()).map_err(|_| ())?;
 		let tag = make_proof_tag(secret, contexts, &aliases, message);
 		let proof = MockProof {
 			tag,
@@ -160,7 +159,7 @@ impl GenerateVerifiable for Mock {
 		members: &Self::Members,
 		contexts: &[&[u8]],
 		message: &[u8],
-	) -> Result<Vec<Alias>, ()> {
+	) -> Result<AliasVec, ()> {
 		let MockProof {
 			tag,
 			member,
@@ -181,7 +180,7 @@ impl GenerateVerifiable for Mock {
 		if tag != &expected_tag {
 			return Err(());
 		}
-		Ok(aliases.to_vec())
+		Ok(aliases.iter().copied().collect())
 	}
 
 	fn alias_in_context(secret: &Self::Secret, context: &[u8]) -> Result<Alias, ()> {

src/ring/bandersnatch.rs

@@ -65,7 +65,6 @@ impl RingSuiteExt for ark_vrf::suites::bandersnatch::BandersnatchSha512Ell2 {
 	const SIGNATURE_SIZE: usize = 64;
 	const RING_PROOF_SIZE: usize = 752;
 	const VRF_OUTPUT_SIZE: usize = 32;
-	const MAX_VRF_CONTEXTS: u8 = 16;
 
 	type CurveParams = Bls12_381Params;
 

src/ring/mod.rs

@@ -239,8 +239,6 @@ pub trait RingSuiteExt: RingSuite + Debug + 'static {
 	const RING_PROOF_SIZE: usize;
 	/// Compressed size of a single `ark_vrf::Output<S>`.
 	const VRF_OUTPUT_SIZE: usize;
-	/// Maximum number of VRF contexts in a multi-context proof.
-	const MAX_VRF_CONTEXTS: u8;
 
 	/// Byte array type for encoded public keys.
 	type PublicKeyBytes: FixedBytes;
@@ -270,7 +268,7 @@ pub struct MaxRingVrfSignatureLen<S: RingSuiteExt>(PhantomData<S>);
 
 impl<S: RingSuiteExt> Get<u32> for MaxRingVrfSignatureLen<S> {
 	fn get() -> u32 {
-		ring_signature_size::<S>(S::MAX_VRF_CONTEXTS) as u32
+		ring_signature_size::<S>(MAX_CONTEXTS as u8) as u32
 	}
 }
 
@@ -424,29 +422,15 @@ impl<S: RingSuiteExt> core::fmt::Debug for ProverState<S> {
 
 /// VRF outputs within a ring signature.
 ///
-/// For the common single-context case, the output is stored inline on the stack.
-/// For multi-context proofs, outputs are heap-allocated.
+/// Holds one output per context. Backed by a [`ContextVec`] — inline for small
+/// counts, heap-allocated only beyond the inline capacity.
 ///
-/// The wire format uses a `u8` length prefix followed by the outputs (no enum
-/// discriminant). On deserialization, length == 1 produces `Single`, length > 1
-/// produces `Multi`.
-enum RingVrfOutputs<S: RingSuiteExt> {
-	Single(ark_vrf::Output<S>),
-	Multi(Vec<ark_vrf::Output<S>>),
-}
-
-impl<S: RingSuiteExt> RingVrfOutputs<S> {
-	fn as_slice(&self) -> &[ark_vrf::Output<S>] {
-		match self {
-			Self::Single(o) => core::slice::from_ref(o),
-			Self::Multi(v) => v.as_slice(),
-		}
-	}
-}
+/// The wire format is a `u8` length prefix followed by the outputs.
+struct RingVrfOutputs<S: RingSuiteExt>(ContextVec<ark_vrf::Output<S>>);
 
 impl<S: RingSuiteExt> ark_serialize::Valid for RingVrfOutputs<S> {
 	fn check(&self) -> Result<(), ark_serialize::SerializationError> {
-		ark_vrf::Output::<S>::batch_check(self.as_slice().iter())
+		ark_vrf::Output::<S>::batch_check(self.0.iter())
 	}
 }
 
@@ -456,7 +440,7 @@ impl<S: RingSuiteExt> CanonicalSerialize for RingVrfOutputs<S> {
 		mut writer: W,
 		compress: ark_serialize::Compress,
 	) -> Result<(), ark_serialize::SerializationError> {
-		let slice = self.as_slice();
+		let slice = self.0.as_slice();
 		(slice.len() as u8).serialize_with_mode(&mut writer, compress)?;
 		for output in slice {
 			output.serialize_with_mode(&mut writer, compress)?;
@@ -465,7 +449,7 @@ impl<S: RingSuiteExt> CanonicalSerialize for RingVrfOutputs<S> {
 	}
 
 	fn serialized_size(&self, compress: ark_serialize::Compress) -> usize {
-		let slice = self.as_slice();
+		let slice = self.0.as_slice();
 		let item_size = slice
 			.iter()
 			.next()
@@ -481,27 +465,21 @@ impl<S: RingSuiteExt> CanonicalDeserialize for RingVrfOutputs<S> {
 		validate: ark_serialize::Validate,
 	) -> Result<Self, ark_serialize::SerializationError> {
 		let len = u8::deserialize_with_mode(&mut reader, compress, validate)?;
-		if len > S::MAX_VRF_CONTEXTS {
+		if len as usize > MAX_CONTEXTS {
 			return Err(ark_serialize::SerializationError::InvalidData);
 		}
-		if len == 1 {
-			let output =
-				ark_vrf::Output::<S>::deserialize_with_mode(&mut reader, compress, validate)?;
-			Ok(Self::Single(output))
-		} else {
-			let mut outputs = Vec::with_capacity(len as usize);
-			for _ in 0..len {
-				outputs.push(ark_vrf::Output::<S>::deserialize_with_mode(
-					&mut reader,
-					compress,
-					ark_serialize::Validate::No,
-				)?);
-			}
-			if let ark_serialize::Validate::Yes = validate {
-				ark_vrf::Output::<S>::batch_check(outputs.iter())?;
-			}
-			Ok(Self::Multi(outputs))
+		let mut outputs = ContextVec::with_capacity(len as usize);
+		for _ in 0..len {
+			outputs.push(ark_vrf::Output::<S>::deserialize_with_mode(
+				&mut reader,
+				compress,
+				ark_serialize::Validate::No,
+			)?);
+		}
+		if matches!(validate, ark_serialize::Validate::Yes) {
+			ark_vrf::Output::<S>::batch_check(outputs.iter())?;
 		}
+		Ok(Self(outputs))
 	}
 }
 
@@ -607,18 +585,18 @@ impl<S: RingSuiteExt> GenerateVerifiable for RingVrfVerifiable<S> {
 		members: &Self::Members,
 		contexts: &[&[u8]],
 		message: &[u8],
-	) -> Result<Vec<Alias>, ()> {
+	) -> Result<AliasVec, ()> {
 		let verifier_params = S::VerifierCache::get(config);
 		let ring_verifier = verifier_params.ring_verifier(members.0.clone());
 
 		let signature = RingVrfSignature::<S>::deserialize_canonical(proof.as_slice())?;
 
-		let outputs = signature.outputs.as_slice();
+		let outputs = signature.outputs.0.as_slice();
 		if contexts.len() != outputs.len() {
 			return Err(());
 		}
 
-		let (ios, aliases): (Vec<_>, Vec<_>) = contexts
+		let (ios, aliases): (ContextVec<_>, AliasVec) = contexts
 			.iter()
 			.zip(outputs.iter().copied())
 			.map(|(ctx, output)| {
@@ -663,10 +641,11 @@ impl<S: RingSuiteExt> GenerateVerifiable for RingVrfVerifiable<S> {
 			let input = ark_vrf::Input::<S>::new(&input_msg[..]).expect("H2C can't fail here");
 			let signature = RingVrfSignature::<S>::deserialize_canonical(proof.as_slice())?;
 
-			let output = match signature.outputs {
-				RingVrfOutputs::Single(o) => o,
-				_ => return Err(()),
-			};
+			let outputs = signature.outputs.0.as_slice();
+			if outputs.len() != 1 {
+				return Err(());
+			}
+			let output = outputs[0];
 
 			aliases.push(make_alias(&output));
 
@@ -708,8 +687,11 @@ impl<S: RingSuiteExt> GenerateVerifiable for RingVrfVerifiable<S> {
 		secret: &Self::Secret,
 		contexts: &[&[u8]],
 		message: &[u8],
-	) -> Result<(Self::Proof, Vec<Alias>), ()> {
+	) -> Result<(Self::Proof, AliasVec), ()> {
 		use ark_vrf::ring::Prover;
+		if contexts.len() > MAX_CONTEXTS {
+			return Err(());
+		}
 		let domain_size = RingDomainSize::try_from(commitment.domain_size).map_err(|_| ())?;
 		let prover_params = S::ProverCache::get(domain_size);
 		if commitment.prover_idx >= prover_params.max_ring_size() as u32 {
@@ -719,7 +701,7 @@ impl<S: RingSuiteExt> GenerateVerifiable for RingVrfVerifiable<S> {
 		let ring_prover =
 			prover_params.ring_prover(commitment.prover_key, commitment.prover_idx as usize);
 
-		let (ios, aliases, outputs): (Vec<_>, Vec<_>, Vec<_>) = contexts
+		let (ios, aliases, outputs): (ContextVec<_>, AliasVec, ContextVec<_>) = contexts
 			.iter()
 			.map(|ctx| {
 				let input_msg = [S::VRF_INPUT_DOMAIN, ctx].concat();
@@ -738,12 +720,10 @@ impl<S: RingSuiteExt> GenerateVerifiable for RingVrfVerifiable<S> {
 
 		let proof = secret.prove(ios, message, &ring_prover);
 
-		let outputs = if outputs.len() == 1 {
-			RingVrfOutputs::Single(outputs.into_iter().next().unwrap())
-		} else {
-			RingVrfOutputs::Multi(outputs)
+		let signature = RingVrfSignature::<S> {
+			outputs: RingVrfOutputs(outputs),
+			proof,
 		};
-		let signature = RingVrfSignature::<S> { outputs, proof };
 
 		let mut buf = vec![];
 		signature.serialize_compressed(&mut buf).map_err(|_| ())?;