diff --git a/ml-dsa/src/lib.rs b/ml-dsa/src/lib.rs index b15bd09..16a11c8 100644 --- a/ml-dsa/src/lib.rs +++ b/ml-dsa/src/lib.rs @@ -36,6 +36,9 @@ //! # } //! ``` +#[cfg(feature = "alloc")] +extern crate alloc; + mod algebra; mod crypto; mod encode; @@ -56,8 +59,11 @@ use crate::hint::Hint; use crate::ntt::{Ntt, NttInverse}; use crate::param::{ParameterSet, QMinus1, SamplingSize, SpecQ}; use crate::sampling::{expand_a, expand_mask, expand_s, sample_in_ball}; -use core::convert::{TryFrom, TryInto}; -use core::fmt; +use core::{ + convert::{TryFrom, TryInto}, + fmt, + ops::{Deref, DerefMut}, +}; use ctutils::{Choice, CtEq}; use hybrid_array::{ Array, @@ -90,10 +96,10 @@ pub(crate) type B64 = Array; pub type Seed = B32; /// An ML-DSA signature -#[derive(Clone, PartialEq, Debug)] +#[derive(Clone, Debug, PartialEq)] pub struct Signature { c_tilde: Array, - z: Vector, + z: MaybeBox>, h: Hint

, } @@ -113,7 +119,7 @@ impl Signature

{ let (c_tilde, z, h) = P::split_sig(enc); let c_tilde = c_tilde.clone(); - let z = P::decode_z(z); + let z = MaybeBox::new(P::decode_z(z)); let h = Hint::bit_unpack(h)?; if z.infinity_norm() >= P::GAMMA1_MINUS_BETA { @@ -421,7 +427,7 @@ impl ExpandedSigningKey

{ continue; } - let z = z.mod_plus_minus::(); + let z = MaybeBox::new(z.mod_plus_minus::()); return Signature { c_tilde, z, h }; } @@ -980,6 +986,45 @@ where } } +/// Type which opportunistically uses `Box` when the `alloc` feature is available but falls back to +/// a stack-allocated type when it's unavailable. +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct MaybeBox { + #[cfg(not(feature = "alloc"))] + inner: T, + #[cfg(feature = "alloc")] + inner: alloc::boxed::Box, +} + +impl MaybeBox { + /// Create a new `MaybeBox`, using `Box` if `alloc` is available. + #[inline] + pub(crate) fn new(inner: T) -> Self { + #[cfg(not(feature = "alloc"))] + { + Self { inner } + } + #[cfg(feature = "alloc")] + Self { + inner: alloc::boxed::Box::new(inner), + } + } +} + +impl Deref for MaybeBox { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.inner + } +} + +impl DerefMut for MaybeBox { + fn deref_mut(&mut self) -> &mut Self::Target { + &mut self.inner + } +} + #[cfg(test)] mod test { use super::*;