From 1f30ca8a52c329b1af7ca94fd85ddfd278fd7f2c Mon Sep 17 00:00:00 2001 From: evalir Date: Mon, 30 Dec 2024 17:10:10 +0100 Subject: [PATCH] feat: add `copy_*e_to_slice` family --- CHANGELOG.md | 4 ++ src/bits.rs | 2 +- src/bytes.rs | 145 ++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 3 + src/support/borsh.rs | 23 ++++++- src/support/postgres.rs | 2 +- 6 files changed, 174 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index afcc35e..c25bacf 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/bits.rs b/src/bits.rs index 94d10c7..339ad15 100644 --- a/src/bits.rs +++ b/src/bits.rs @@ -375,7 +375,7 @@ impl Uint { 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 diff --git a/src/bytes.rs b/src/bytes.rs index 57f94c1..bc0d04d 100644 --- a/src/bytes.rs +++ b/src/bytes.rs @@ -345,6 +345,112 @@ impl Uint { } 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 { + 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 { + 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)|{ + let mut buf = [0; BYTES]; + value.copy_le_bytes_to(&mut buf); + assert_eq!(buf, value.to_le_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::()); + 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)|{ + 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"); + } + }); + }); + } } diff --git a/src/lib.rs b/src/lib.rs index f40111a..5cc19c9 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/support/borsh.rs b/src/support/borsh.rs index 8cd5d76..7636c83 100644 --- a/src/support/borsh.rs +++ b/src/support/borsh.rs @@ -41,9 +41,25 @@ impl BorshDeserialize for Uint BorshSerialize for Uint { #[inline] fn serialize(&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::(), Self::BYTES) + }; + self.copy_le_bytes_to(&mut buf); + writer.write_all(&buf) + } } } @@ -60,6 +76,7 @@ impl BorshSerialize for Bits self.as_uint().serialize(writer) } } + #[cfg(test)] mod test { use super::*; diff --git a/src/support/postgres.rs b/src/support/postgres.rs index 58a5107..06dd758 100644 --- a/src/support/postgres.rs +++ b/src/support/postgres.rs @@ -221,7 +221,7 @@ impl<'a, const BITS: usize, const LIMBS: usize> FromSql<'a> for Uint 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; }