Merge pull request #424 from Evalir/evalir/copy-to-buf

feat: add `copy_*e_to_slice` family

CHANGELOG.md

@@ -10,6 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 - Support for borsh @ 1.5 ([#416])
+- `copy_le_to_slice` family to allow easier writing to pre-allocated buffers ([#424])
+
+[#416]: https://github.com/recmo/uint/pull/416
+[#424]: https://github.com/recmo/uint/pull/424
 
 ## [1.12.4] - 2024-12-16
 

Cargo.toml

@@ -156,7 +156,7 @@ arbitrary = ["dep:arbitrary", "std"]
 ark-ff = ["dep:ark-ff-03"]
 ark-ff-04 = ["dep:ark-ff-04"]
 bn-rs = ["dep:bn-rs", "std"]
-borsh = ["alloc", "dep:borsh"]
+borsh = ["dep:borsh"]
 bytemuck = ["dep:bytemuck"]
 der = ["dep:der", "alloc"]                                             # TODO: also have alloc free der impls.
 diesel = ["dep:diesel", "std", "dep:thiserror"]

src/bits.rs

@@ -375,7 +375,7 @@ impl<const BITS: usize, const LIMBS: usize> Uint<BITS, LIMBS> {
             return Self::ZERO;
         }
         let rhs = rhs % BITS;
-        self << rhs | self >> (BITS - rhs)
+        (self << rhs) | (self >> (BITS - rhs))
     }
 
     /// Shifts the bits to the right by a specified amount, `rhs`, wrapping the

src/bytes.rs

@@ -345,6 +345,112 @@ impl<const BITS: usize, const LIMBS: usize> Uint<BITS, LIMBS> {
         }
         Some(Self::from_limbs(limbs))
     }
+
+    /// Writes the little-endian representation of the [`Uint`] to the given
+    /// buffer. The buffer must be large enough to hold [`Self::BYTES`] bytes.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the buffer is not large enough to hold [`Self::BYTES`] bytes.
+    ///
+    /// # Returns
+    ///
+    /// The number of bytes written to the buffer (always equal to
+    /// [`Self::BYTES`], but often useful to make explicit for encoders).
+    #[inline]
+    pub fn copy_le_bytes_to(&self, buf: &mut [u8]) -> usize {
+        // This is debug only. Release panics occur later in copy_from_slice
+        debug_assert!(
+            buf.len() >= Self::BYTES,
+            "Buffer is too small to hold the bytes of the Uint"
+        );
+
+        #[cfg(target_endian = "little")]
+        buf[..Self::BYTES].copy_from_slice(self.as_le_slice());
+
+        #[cfg(target_endian = "big")]
+        {
+            let chunks = buf[..Self::BYTES].chunks_mut(8);
+
+            self.limbs.iter().zip(chunks).for_each(|(&limb, chunk)| {
+                let le = limb.to_le_bytes();
+                chunk.copy_from_slice(&le[..chunk.len()]);
+            });
+        }
+
+        Self::BYTES
+    }
+
+    /// Writes the little-endian representation of the [`Uint`] to the given
+    /// buffer. The buffer must be large enough to hold [`Self::BYTES`] bytes.
+    ///
+    /// # Returns
+    ///
+    /// [`None`], if the buffer is not large enough to hold [`Self::BYTES`]
+    /// bytes, and does not modify the buffer.
+    ///
+    /// [`Some`] with the number of bytes written to the buffer (always
+    /// equal to [`Self::BYTES`], but often useful to make explicit for
+    /// encoders).
+    #[inline]
+    pub fn checked_copy_le_bytes_to(&self, buf: &mut [u8]) -> Option<usize> {
+        if buf.len() < Self::BYTES {
+            return None;
+        }
+
+        Some(self.copy_le_bytes_to(buf))
+    }
+
+    /// Writes the big-endian representation of the [`Uint`] to the given
+    /// buffer. The buffer must be large enough to hold [`Self::BYTES`] bytes.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the buffer is not large enough to hold [`Self::BYTES`] bytes.
+    ///
+    /// # Returns
+    ///
+    /// The number of bytes written to the buffer (always equal to
+    /// [`Self::BYTES`], but often useful to make explicit for encoders).
+    #[inline]
+    pub fn copy_be_bytes_to(&self, buf: &mut [u8]) -> usize {
+        // This is debug only. Release panics occur later in copy_from_slice
+        debug_assert!(
+            buf.len() >= Self::BYTES,
+            "Buffer is too small to hold the bytes of the Uint"
+        );
+
+        // start from the end of the slice
+        let chunks = buf[..Self::BYTES].rchunks_mut(8);
+
+        self.limbs.iter().zip(chunks).for_each(|(&limb, chunk)| {
+            let be = limb.to_be_bytes();
+            let copy_from = 8 - chunk.len();
+            chunk.copy_from_slice(&be[copy_from..]);
+        });
+
+        Self::BYTES
+    }
+
+    /// Writes the big-endian representation of the [`Uint`] to the given
+    /// buffer. The buffer must be large enough to hold [`Self::BYTES`] bytes.
+    ///
+    /// # Returns
+    ///
+    /// [`None`], if the buffer is not large enough to hold [`Self::BYTES`]
+    /// bytes, and does not modify the buffer.
+    ///
+    /// [`Some`] with the number of bytes written to the buffer (always
+    /// equal to [`Self::BYTES`], but often useful to make explicit for
+    /// encoders).
+    #[inline]
+    pub fn checked_copy_be_bytes_to(&self, buf: &mut [u8]) -> Option<usize> {
+        if buf.len() < Self::BYTES {
+            return None;
+        }
+
+        Some(self.copy_be_bytes_to(buf))
+    }
 }
 
 /// Number of bytes required to represent the given number of bits.
@@ -488,4 +594,43 @@ mod tests {
             });
         });
     }
+
+    #[test]
+    fn copy_to() {
+        const_for!(BITS in SIZES {
+            const LIMBS: usize = nlimbs(BITS);
+            const BYTES: usize = nbytes(BITS);
+            proptest!(|(value: Uint<BITS, LIMBS>)|{
+                let mut buf = [0; BYTES];
+                value.copy_le_bytes_to(&mut buf);
+                assert_eq!(buf, value.to_le_bytes::<BYTES>());
+                assert_eq!(value, Uint::try_from_le_slice(&buf).unwrap());
+
+                let mut buf = [0; BYTES];
+                value.copy_be_bytes_to(&mut buf);
+                assert_eq!(buf, value.to_be_bytes::<BYTES>());
+                assert_eq!(value, Uint::try_from_be_slice(&buf).unwrap());
+            });
+        });
+    }
+
+    #[test]
+    fn checked_copy_to() {
+        const_for!(BITS in SIZES {
+            const LIMBS: usize = nlimbs(BITS);
+            const BYTES: usize = nbytes(BITS);
+            proptest!(|(value: Uint<BITS, LIMBS>)|{
+                if BYTES != 0 {
+                    let mut buf = [0; BYTES];
+                    let too_short = buf.len() - 1;
+
+                    assert_eq!(value.checked_copy_le_bytes_to(&mut buf[..too_short]), None);
+                    assert_eq!(buf, [0; BYTES], "buffer was modified");
+
+                    assert_eq!(value.checked_copy_be_bytes_to(&mut buf[..too_short]), None);
+                    assert_eq!(buf, [0; BYTES], "buffer was modified");
+                }
+            });
+        });
+    }
 }

src/lib.rs

@@ -34,6 +34,9 @@
 )]
 
 #[cfg(feature = "alloc")]
+#[allow(unused_imports)]
+// `unused_imports` triggers on macro_use, which is required by some support
+// modules.
 #[macro_use]
 extern crate alloc;
 

src/support/borsh.rs

@@ -41,9 +41,25 @@ impl<const BITS: usize, const LIMBS: usize> BorshDeserialize for Uint<BITS, LIMB
 impl<const BITS: usize, const LIMBS: usize> BorshSerialize for Uint<BITS, LIMBS> {
     #[inline]
     fn serialize<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
-        // TODO: non-allocating and remove the `alloc` feature requirement
-        let bytes = self.as_le_bytes();
-        writer.write_all(&bytes[..Self::BYTES])
+        #[cfg(target_endian = "little")]
+        return writer.write_all(self.as_le_slice());
+
+        // TODO: Replace the unsafety with `generic_const_exprs` when
+        // available
+        #[cfg(target_endian = "big")]
+        {
+            let mut limbs = [0u64; LIMBS];
+            // SAFETY: `limbs` is known to have identical memory layout and
+            // alignment to `[u8; LIMBS * 8]`, which is guaranteed to safely
+            // contain  [u8; Self::BYTES]`, as `LIMBS * 8 >= Self::BYTES`.
+            // Reference:
+            // https://doc.rust-lang.org/reference/type-layout.html#array-layout
+            let mut buf = unsafe {
+                core::slice::from_raw_parts_mut(limbs.as_mut_ptr().cast::<u8>(), Self::BYTES)
+            };
+            self.copy_le_bytes_to(&mut buf);
+            writer.write_all(&buf)
+        }
     }
 }
 
@@ -60,6 +76,7 @@ impl<const BITS: usize, const LIMBS: usize> BorshSerialize for Bits<BITS, LIMBS>
         self.as_uint().serialize(writer)
     }
 }
+
 #[cfg(test)]
 mod test {
     use super::*;

src/support/postgres.rs

@@ -221,7 +221,7 @@ impl<'a, const BITS: usize, const LIMBS: usize> FromSql<'a> for Uint<BITS, LIMBS
                 let mut raw = raw.to_owned();
                 if padding > 0 {
                     for i in (1..raw.len()).rev() {
-                        raw[i] = raw[i] >> padding | raw[i - 1] << (8 - padding);
+                        raw[i] = (raw[i] >> padding) | (raw[i - 1] << (8 - padding));
                     }
                     raw[0] >>= padding;
                 }