Use Iterator instead of Expander

carloskiki · carloskiki · commit 9ed6117d3c55 · 2025-07-20T14:00:03.000-04:00
diff --git a/hash2curve/src/group_digest.rs b/hash2curve/src/group_digest.rs
@@ -32,9 +32,9 @@ pub trait GroupDigest: MapToCurve {
     ///
     /// [`ExpandMsgXmd`]: crate::ExpandMsgXmd
     /// [`ExpandMsgXof`]: crate::ExpandMsgXof
-    fn hash_from_bytes<X>(msg: &[&[u8]], dst: &[&[u8]]) -> Result<ProjectivePoint<Self>>
+    fn hash_from_bytes<'dst, X>(msg: &[&[u8]], dst: &'dst [&[u8]]) -> Result<ProjectivePoint<Self>>
     where
-        X: ExpandMsg<Self::K>,
+        X: ExpandMsg<'dst, Self::K>,
     {
         let [u0, u1] = hash_to_field::<2, X, _, Self::FieldElement>(msg, dst)?;
         let q0 = Self::map_to_curve(u0);
@@ -62,9 +62,12 @@ pub trait GroupDigest: MapToCurve {
     ///
     /// [`ExpandMsgXmd`]: crate::ExpandMsgXmd
     /// [`ExpandMsgXof`]: crate::ExpandMsgXof
-    fn encode_from_bytes<X>(msg: &[&[u8]], dst: &[&[u8]]) -> Result<ProjectivePoint<Self>>
+    fn encode_from_bytes<'dst, X>(
+        msg: &[&[u8]],
+        dst: &'dst [&[u8]],
+    ) -> Result<ProjectivePoint<Self>>
     where
-        X: ExpandMsg<Self::K>,
+        X: ExpandMsg<'dst, Self::K>,
     {
         let [u] = hash_to_field::<1, X, _, Self::FieldElement>(msg, dst)?;
         let q0 = Self::map_to_curve(u);
@@ -85,9 +88,9 @@ pub trait GroupDigest: MapToCurve {
     ///
     /// [`ExpandMsgXmd`]: crate::ExpandMsgXmd
     /// [`ExpandMsgXof`]: crate::ExpandMsgXof
-    fn hash_to_scalar<X>(msg: &[&[u8]], dst: &[&[u8]]) -> Result<Self::Scalar>
+    fn hash_to_scalar<'dst, X>(msg: &[&[u8]], dst: &'dst [&[u8]]) -> Result<Self::Scalar>
     where
-        X: ExpandMsg<Self::K>,
+        X: ExpandMsg<'dst, Self::K>,
     {
         let [u] = hash_to_field::<1, X, _, Self::Scalar>(msg, dst)?;
         Ok(u)
diff --git a/hash2curve/src/hash2field.rs b/hash2curve/src/hash2field.rs
@@ -38,20 +38,24 @@ pub trait FromOkm {
 /// [`ExpandMsgXmd`]: crate::hash2field::ExpandMsgXmd
 /// [`ExpandMsgXof`]: crate::hash2field::ExpandMsgXof
 #[doc(hidden)]
-pub fn hash_to_field<const N: usize, E, K, T>(data: &[&[u8]], domain: &[&[u8]]) -> Result<[T; N]>
+pub fn hash_to_field<'dst, const N: usize, E, K, T>(
+    data: &[&[u8]],
+    domain: &'dst [&[u8]],
+) -> Result<[T; N]>
 where
-    E: ExpandMsg<K>,
+    E: ExpandMsg<'dst, K>,
     T: FromOkm + Default,
 {
     let len_in_bytes = T::Length::USIZE
         .checked_mul(N)
         .and_then(|len| len.try_into().ok())
         .and_then(NonZeroU16::new)
         .ok_or(Error)?;
-    let mut tmp = Array::<u8, <T as FromOkm>::Length>::default();
     let mut expander = E::expand_message(data, domain, len_in_bytes)?;
     Ok(core::array::from_fn(|_| {
-        expander.fill_bytes(&mut tmp);
+        let tmp = Array::<u8, <T as FromOkm>::Length>::from_iter(
+            expander.by_ref().take(T::Length::USIZE),
+        );
         T::from_okm(&tmp)
     }))
 }
diff --git a/hash2curve/src/hash2field/expand_msg.rs b/hash2curve/src/hash2field/expand_msg.rs
@@ -22,25 +22,16 @@ const MAX_DST_LEN: usize = 255;
 ///
 /// # Errors
 /// See implementors of [`ExpandMsg`] for errors.
-pub trait ExpandMsg<K> {
-    /// Type holding data for the [`Expander`].
-    type Expander<'dst>: Expander + Sized;
-
+pub trait ExpandMsg<'dst, K>: Iterator<Item = u8> + Sized {
     /// Expands `msg` to the required number of bytes.
     ///
     /// Returns an expander that can be used to call `read` until enough
     /// bytes have been consumed
-    fn expand_message<'dst>(
+    fn expand_message(
         msg: &[&[u8]],
         dst: &'dst [&[u8]],
         len_in_bytes: NonZero<u16>,
-    ) -> Result<Self::Expander<'dst>>;
-}
-
-/// Expander that, call `read` until enough bytes have been consumed.
-pub trait Expander {
-    /// Fill the array with the expanded bytes
-    fn fill_bytes(&mut self, okm: &mut [u8]);
+    ) -> Result<Self>;
 }
 
 /// The domain separation tag
diff --git a/hash2curve/src/hash2field/expand_msg/xmd.rs b/hash2curve/src/hash2field/expand_msg/xmd.rs
@@ -1,8 +1,8 @@
 //! `expand_message_xmd` based on a hash function.
 
-use core::{marker::PhantomData, num::NonZero, ops::Mul};
+use core::{num::NonZero, ops::Mul};
 
-use super::{Domain, ExpandMsg, Expander};
+use super::{Domain, ExpandMsg};
 use digest::{
     FixedOutput, HashMarker,
     array::{
@@ -21,12 +21,20 @@ use elliptic_curve::{Error, Result};
 /// - `dst > 255 && HashT::OutputSize > 255`
 /// - `len_in_bytes > 255 * HashT::OutputSize`
 #[derive(Debug)]
-pub struct ExpandMsgXmd<HashT>(PhantomData<HashT>)
+pub struct ExpandMsgXmd<'a, HashT>
 where
     HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
-    HashT::OutputSize: IsLessOrEqual<HashT::BlockSize, Output = True>;
+    HashT::OutputSize: IsLessOrEqual<HashT::BlockSize, Output = True>,
+{
+    b_0: Array<u8, HashT::OutputSize>,
+    b_vals: Array<u8, HashT::OutputSize>,
+    domain: Domain<'a, HashT::OutputSize>,
+    index: u8,
+    offset: usize,
+    length: u16,
+}
 
-impl<HashT, K> ExpandMsg<K> for ExpandMsgXmd<HashT>
+impl<'dst, HashT, K> ExpandMsg<'dst, K> for ExpandMsgXmd<'dst, HashT>
 where
     HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
     // The number of bits output by `HashT` MUST be at most `HashT::BlockSize`.
@@ -37,23 +45,18 @@ where
     K: Mul<U2>,
     HashT::OutputSize: IsGreaterOrEqual<Prod<K, U2>, Output = True>,
 {
-    type Expander<'dst> = ExpanderXmd<'dst, HashT>;
-
-    fn expand_message<'dst>(
+    fn expand_message(
         msg: &[&[u8]],
         dst: &'dst [&[u8]],
         len_in_bytes: NonZero<u16>,
-    ) -> Result<Self::Expander<'dst>> {
+    ) -> Result<Self> {
         let b_in_bytes = HashT::OutputSize::USIZE;
 
         // `255 * <b_in_bytes>` can not exceed `u16::MAX`
         if usize::from(len_in_bytes.get()) > 255 * b_in_bytes {
             return Err(Error);
         }
 
-        let ell = u8::try_from(usize::from(len_in_bytes.get()).div_ceil(b_in_bytes))
-            .expect("should never pass the previous check");
-
         let domain = Domain::xmd::<HashT>(dst)?;
         let mut b_0 = HashT::default();
         b_0.update(&Array::<u8, HashT::BlockSize>::default());
@@ -75,74 +78,51 @@ where
         b_vals.update(&[domain.len()]);
         let b_vals = b_vals.finalize_fixed();
 
-        Ok(ExpanderXmd {
+        Ok(Self {
             b_0,
             b_vals,
             domain,
             index: 1,
             offset: 0,
-            ell,
+            length: len_in_bytes.get(),
         })
     }
 }
 
-/// [`Expander`] type for [`ExpandMsgXmd`].
-#[derive(Debug)]
-pub struct ExpanderXmd<'a, HashT>
+impl<HashT> Iterator for ExpandMsgXmd<'_, HashT>
 where
     HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
     HashT::OutputSize: IsLessOrEqual<HashT::BlockSize, Output = True>,
 {
-    b_0: Array<u8, HashT::OutputSize>,
-    b_vals: Array<u8, HashT::OutputSize>,
-    domain: Domain<'a, HashT::OutputSize>,
-    index: u8,
-    offset: usize,
-    ell: u8,
-}
+    type Item = u8;
 
-impl<HashT> ExpanderXmd<'_, HashT>
-where
-    HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
-    HashT::OutputSize: IsLessOrEqual<HashT::BlockSize, Output = True>,
-{
-    fn next(&mut self) -> bool {
-        if self.index < self.ell {
-            self.index += 1;
-            self.offset = 0;
-            // b_0 XOR b_(idx - 1)
-            let mut tmp = Array::<u8, HashT::OutputSize>::default();
-            self.b_0
-                .iter()
-                .zip(&self.b_vals[..])
-                .enumerate()
-                .for_each(|(j, (b0val, bi1val))| tmp[j] = b0val ^ bi1val);
-            let mut b_vals = HashT::default();
-            b_vals.update(&tmp);
-            b_vals.update(&[self.index]);
-            self.domain.update_hash(&mut b_vals);
-            b_vals.update(&[self.domain.len()]);
-            self.b_vals = b_vals.finalize_fixed();
-            true
-        } else {
-            false
-        }
-    }
-}
-
-impl<HashT> Expander for ExpanderXmd<'_, HashT>
-where
-    HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
-    HashT::OutputSize: IsLessOrEqual<HashT::BlockSize, Output = True>,
-{
-    fn fill_bytes(&mut self, okm: &mut [u8]) {
-        for b in okm {
-            if self.offset == self.b_vals.len() && !self.next() {
-                return;
-            }
-            *b = self.b_vals[self.offset];
+    fn next(&mut self) -> Option<u8> {
+        if (self.index as u16 - 1) * HashT::OutputSize::U16 + self.offset as u16
+            == self.length
+        {
+            return None;
+        } else if self.offset != self.b_vals.len() {
+            let byte = self.b_vals[self.offset];
             self.offset += 1;
+            return Some(byte);
         }
+
+        self.index += 1;
+        self.offset = 1;
+        // b_0 XOR b_(idx - 1)
+        let mut tmp = Array::<u8, HashT::OutputSize>::default();
+        self.b_0
+            .iter()
+            .zip(&self.b_vals[..])
+            .enumerate()
+            .for_each(|(j, (b0val, bi1val))| tmp[j] = b0val ^ bi1val);
+        let mut b_vals = HashT::default();
+        b_vals.update(&tmp);
+        b_vals.update(&[self.index]);
+        self.domain.update_hash(&mut b_vals);
+        b_vals.update(&[self.domain.len()]);
+        self.b_vals = b_vals.finalize_fixed();
+        Some(self.b_vals[0])
     }
 }
 
@@ -210,15 +190,13 @@ mod test {
             assert_message::<HashT>(self.msg, domain, L::U16, self.msg_prime);
 
             let dst = [dst];
-            let mut expander = <ExpandMsgXmd<HashT> as ExpandMsg<U4>>::expand_message(
+            let expander = <ExpandMsgXmd<HashT> as ExpandMsg<U4>>::expand_message(
                 &[self.msg],
                 &dst,
                 NonZero::new(L::U16).ok_or(Error)?,
             )?;
 
-            let mut uniform_bytes = Array::<u8, L>::default();
-            expander.fill_bytes(&mut uniform_bytes);
-
+            let uniform_bytes = Array::<u8, L>::from_iter(expander);
             assert_eq!(uniform_bytes.as_slice(), self.uniform_bytes);
             Ok(())
         }
diff --git a/hash2curve/src/hash2field/expand_msg/xof.rs b/hash2curve/src/hash2field/expand_msg/xof.rs
@@ -1,9 +1,10 @@
 //! `expand_message_xof` for the `ExpandMsg` trait
 
-use super::{Domain, ExpandMsg, Expander};
+use super::{Domain, ExpandMsg};
 use core::{fmt, num::NonZero, ops::Mul};
+use digest::XofReader;
 use digest::{
-    CollisionResistance, ExtendableOutput, HashMarker, Update, XofReader, typenum::IsGreaterOrEqual,
+    CollisionResistance, ExtendableOutput, HashMarker, Update, typenum::IsGreaterOrEqual,
 };
 use elliptic_curve::Result;
 use elliptic_curve::array::{
@@ -21,7 +22,8 @@ pub struct ExpandMsgXof<HashT>
 where
     HashT: Default + ExtendableOutput + Update + HashMarker,
 {
-    reader: <HashT as ExtendableOutput>::Reader,
+    reader: HashT::Reader,
+    length: u16,
 }
 
 impl<HashT> fmt::Debug for ExpandMsgXof<HashT>
@@ -36,7 +38,7 @@ where
     }
 }
 
-impl<HashT, K> ExpandMsg<K> for ExpandMsgXof<HashT>
+impl<HashT, K> ExpandMsg<'_, K> for ExpandMsgXof<HashT>
 where
     HashT: Default + ExtendableOutput + Update + HashMarker,
     // If DST is larger than 255 bytes, the length of the computed DST is calculated by `K * 2`.
@@ -46,13 +48,7 @@ where
     // https://www.rfc-editor.org/rfc/rfc9380.html#section-5.3.2-2.1
     HashT: CollisionResistance<CollisionResistance: IsGreaterOrEqual<K, Output = True>>,
 {
-    type Expander<'dst> = Self;
-
-    fn expand_message<'dst>(
-        msg: &[&[u8]],
-        dst: &'dst [&[u8]],
-        len_in_bytes: NonZero<u16>,
-    ) -> Result<Self::Expander<'dst>> {
+    fn expand_message(msg: &[&[u8]], dst: &[&[u8]], len_in_bytes: NonZero<u16>) -> Result<Self> {
         let len_in_bytes = len_in_bytes.get();
 
         let domain = Domain::<Prod<K, U2>>::xof::<HashT>(dst)?;
@@ -66,16 +62,27 @@ where
         domain.update_hash(&mut reader);
         reader.update(&[domain.len()]);
         let reader = reader.finalize_xof();
-        Ok(Self { reader })
+        Ok(Self {
+            reader,
+            length: len_in_bytes,
+        })
     }
 }
 
-impl<HashT> Expander for ExpandMsgXof<HashT>
+impl<HashT> Iterator for ExpandMsgXof<HashT>
 where
     HashT: Default + ExtendableOutput + Update + HashMarker,
 {
-    fn fill_bytes(&mut self, okm: &mut [u8]) {
-        self.reader.read(okm);
+    type Item = u8;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.length == 0 {
+            return None;
+        }
+        self.length -= 1;
+        let mut buf = [0u8; 1];
+        self.reader.read(&mut buf);
+        Some(buf[0])
     }
 }
 
@@ -130,15 +137,13 @@ mod test {
         {
             assert_message(self.msg, domain, L::to_u16(), self.msg_prime);
 
-            let mut expander = <ExpandMsgXof<HashT> as ExpandMsg<U16>>::expand_message(
+            let expander = <ExpandMsgXof<HashT> as ExpandMsg<U16>>::expand_message(
                 &[self.msg],
                 &[dst],
                 NonZero::new(L::U16).ok_or(Error)?,
             )?;
 
-            let mut uniform_bytes = Array::<u8, L>::default();
-            expander.fill_bytes(&mut uniform_bytes);
-
+            let uniform_bytes = Array::<u8, L>::from_iter(expander);
             assert_eq!(uniform_bytes.as_slice(), self.uniform_bytes);
             Ok(())
         }
diff --git a/hash2curve/src/oprf.rs b/hash2curve/src/oprf.rs
@@ -25,5 +25,5 @@ pub trait OprfParameters: GroupDigest + PrimeCurve {
     /// and `HashToScalar` as defined in [section 4 of RFC9497][oprf].
     ///
     /// [oprf]: https://www.rfc-editor.org/rfc/rfc9497.html#name-ciphersuites
-    type ExpandMsg: ExpandMsg<<Self as GroupDigest>::K>;
+    type ExpandMsg<'a>: ExpandMsg<'a, <Self as GroupDigest>::K>;
 }

Original file line number	Diff line number	Diff line change
`@@ -25,5 +25,5 @@ pub trait OprfParameters: GroupDigest + PrimeCurve {`
`25`	`25`	/// and `HashToScalar` as defined in [section 4 of RFC9497][oprf].
`26`	`26`	`///`
`27`	`27`	`/// [oprf]: https://www.rfc-editor.org/rfc/rfc9497.html#name-ciphersuites`
`28`		`- type ExpandMsg: ExpandMsg<<Self as GroupDigest>::K>;`
	`28`	`+ type ExpandMsg<'a>: ExpandMsg<'a, <Self as GroupDigest>::K>;`
`29`	`29`	`}`