From b78dd69ae189b9d1870fe36ab887b220c02c3718 Mon Sep 17 00:00:00 2001 From: Brian Smith Date: Tue, 19 May 2026 12:26:19 -0700 Subject: [PATCH 1/2] bigint: Replace `dynarray` usage with a `slice::Buf`. --- src/arithmetic/bigint/exp.rs | 92 ++++++++++---------- src/polyfill/dynarray.rs | 157 ----------------------------------- src/polyfill/mod.rs | 1 - src/polyfill/uninit_slice.rs | 43 ++++++++++ 4 files changed, 91 insertions(+), 202 deletions(-) delete mode 100644 src/polyfill/dynarray.rs diff --git a/src/arithmetic/bigint/exp.rs b/src/arithmetic/bigint/exp.rs index 32183b90d8..96004199e7 100644 --- a/src/arithmetic/bigint/exp.rs +++ b/src/arithmetic/bigint/exp.rs @@ -99,7 +99,7 @@ fn elem_exp_consttime_inner<'out, N, M, const STORAGE_LIMBS: usize>( }; use crate::{ bssl, c, error, - polyfill::{StartMutPtr, dynarray}, + polyfill::{StartMutPtr, slice::Buf}, }; let base_rinverse = base_mod_n.reduced_mont(out, m, tmp); @@ -148,50 +148,54 @@ fn elem_exp_consttime_inner<'out, N, M, const STORAGE_LIMBS: usize>( Ok(acc) } - let mut storage: [MaybeUninit; STORAGE_LIMBS] = + let mut table: [MaybeUninit; STORAGE_LIMBS] = [const { MaybeUninit::uninit() }; STORAGE_LIMBS]; - let table = dynarray::Uninit::new(&mut storage, STORAGE_ENTRIES, num_limbs)?.init_fold( - |init, uninit| { - let r: Result<&'_ mut [Limb], LimbSliceError> = match init.len() { - // table[0] = base**0 (i.e. 1). - 0 => Ok( - One::write_mont_identity_assuming_full_upper_limb(uninit, m)? - .leak_limbs_into_mut_less_safe(), - ), - - // table[1] = base*R == (base/R * RRR)/R - 1 => limbs_mul_mont( - ( - uninit, - base_rinverse.leak_limbs_less_safe(), - oneRRR.leak_limbs_less_safe(), - ), - m.limbs(), - m.n0(), - m.cpu_features(), - ), - - // table[2*i] = (n**i)**2/R - i if i % 2 == 0 => { - let sqrt = init.mid().unwrap_or_else(|| unreachable!()); - limbs_square_mont((uninit, sqrt), m.limbs(), m.n0(), m.cpu_features()) - } - - // table[i + 1] = n**1*n**i/R - _ => { - let one = init.get(1).unwrap_or_else(|| unreachable!()); - let previous = init.last().unwrap_or_else(|| unreachable!()); - limbs_mul_mont((uninit, one, previous), m.limbs(), m.n0(), m.cpu_features()) - } - }; - r.map_err(|e| match e { - LimbSliceError::LenMismatch(e) => e, // Also unreachable. - LimbSliceError::TooLong(_) => unreachable!(), - LimbSliceError::TooShort(_) => unreachable!(), - }) - }, - )?; - let table: &[Limb] = table.as_flattened(); + let mut table = Buf::from(Uninit::from(table.as_mut())); + + // table[0] = base**0 (i.e. 1). + table.try_write_with::(num_limbs.get(), |_init, uninit| { + Ok( + One::write_mont_identity_assuming_full_upper_limb(uninit, m)? + .leak_limbs_into_mut_less_safe(), + ) + })?; + + // table[1] = base*R == (base/R * RRR)/R + table.try_write_with(num_limbs.get(), |_init, uninit| { + limbs_mul_mont( + ( + uninit, + base_rinverse.leak_limbs_less_safe(), + oneRRR.leak_limbs_less_safe(), + ), + m.limbs(), + m.n0(), + m.cpu_features(), + ) + })?; + + for _i in 1..16 { + let n = num_limbs.get(); + + // table[2*i] = (n**i)**2/R + table.try_write_with(n, |init, uninit| { + let sqrt_start = init.len() / 2; + let sqrt = init + .get(sqrt_start..(sqrt_start + n)) + .unwrap_or_else(|| unreachable!()); + limbs_square_mont((uninit, sqrt), m.limbs(), m.n0(), m.cpu_features()) + })?; + + // table[2*i + 1] = (n**1)*(n**(2*i))/R + table.try_write_with(n, |init, uninit| { + let one = init.get(n..(n + n)).unwrap_or_else(|| unreachable!()); + let previous = init + .get((init.len() - n)..) + .unwrap_or_else(|| unreachable!()); + limbs_mul_mont((uninit, one, previous), m.limbs(), m.n0(), m.cpu_features()) + })?; + } + let table: &[Limb] = table.into_filled(); let mut tmp_slice = tmp .as_uninit(..num_limbs.get()) diff --git a/src/polyfill/dynarray.rs b/src/polyfill/dynarray.rs deleted file mode 100644 index 5c4f4a718f..0000000000 --- a/src/polyfill/dynarray.rs +++ /dev/null @@ -1,157 +0,0 @@ -// Copyright 2025 Brian Smith. -// -// Permission to use, copy, modify, and/or distribute this software for any -// purpose with or without fee is hereby granted, provided that the above -// copyright notice and this permission notice appear in all copies. -// -// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES -// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF -// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY -// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES -// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION -// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN -// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. - -#![cfg(not(target_arch = "x86_64"))] - -//! Arrays of fixed-length slices, where the length is chosen at runtime, -//! and where the backing store is provided by the user. - -#[allow(unused_imports)] -use crate::polyfill::prelude::*; - -use crate::error::LenMismatchError; -use core::{mem::MaybeUninit, num::NonZero}; - -/// An uninitialized array of slices. -/// -/// E: Copy to avoid drop issues. -pub struct Uninit<'e, E> { - storage: &'e mut [MaybeUninit], - elems_per_item: NonZero, - len: usize, -} - -/// An uninitialized array. -/// -/// `E: Copy` to avoid having to implement any `Drop` logic. -impl Uninit<'_, E> { - /// Create an uninitialized array of `num_elems` slices of length - ///`elems_per_item`, backed by `storage`. - /// - pub fn new( - storage: &'_ mut [MaybeUninit], - num_elems: usize, - elems_per_item: NonZero, - ) -> Result, LenMismatchError> { - let total_elems = num_elems - .checked_mul(elems_per_item.get()) - .ok_or_else(|| LenMismatchError::new(elems_per_item.get()))?; - // TODO: Should we split `storage` from the end instead? - //let mid = overallocated_storage.len().checked_sub(total_elems)?; - //let storage = overallocated_storage.get_mut(mid..)?; - let storage = storage - .get_mut(..total_elems) - .ok_or_else(|| LenMismatchError::new(num_elems))?; - Ok(Uninit { - storage, - elems_per_item, - len: num_elems, - }) - } - - /// Initialize the array by calling the function `f` once for each entry. - /// - /// The first argument to `f`, `init`, references the already-initialized - /// contents of the array; it's length indicates how many entries have been - /// initialized previously, which is `i` for the `i`th item. - /// - /// The second argument to `f`, `uninit`, is the entry to initialize. `f` - /// must return a reference to `uninit` after fully initializing it. - pub fn init_fold<'r, Error: Into>( - self, - mut f: impl for<'i, 'u> FnMut( - &'i mut Array<'_, E>, - super::slice::Uninit<'u, E>, - ) -> Result<&'u mut [E], Error>, - ) -> Result, LenMismatchError> - where - Self: 'r, - LenMismatchError: From, - { - for (init_len, mid) in (self.elems_per_item.get()..) - .step_by(self.elems_per_item.get()) - .enumerate() - { - let Some(init_and_current) = self.storage.get_mut(..mid) else { - break; - }; - let (init, current) = init_and_current - .split_at_mut_checked(mid - self.elems_per_item.get()) - .unwrap_or_else(|| unreachable!()); - let init = super::slice::Uninit::from(init); - let mut init = Array { - storage: unsafe { init.assume_init() }, - len: init_len, - elems_per_item: self.elems_per_item, - }; - let r = f(&mut init, super::slice::Uninit::from(current))?; - if r.len() != self.elems_per_item.get() { - return Err(LenMismatchError::new(r.len())); - } - } - let init = super::slice::Uninit::from(self.storage); - Ok(Array { - storage: unsafe { init.assume_init() }, - elems_per_item: self.elems_per_item, - len: self.len, - }) - } -} - -pub struct Array<'e, E> { - storage: &'e mut [E], - elems_per_item: NonZero, - len: usize, -} - -impl Array<'_, E> { - pub fn len(&self) -> usize { - assert_eq!(self.len, self.storage.len() / self.elems_per_item.get()); - self.len - } - - pub fn last(&self) -> Option<&[E]> { - let before = self.storage.len().checked_sub(self.elems_per_item.get())?; - self.get_(before) - } - - // `self.mid()` is equivalent to `self.get(self.len() / 2)`. - // - // Potentially this is easier to optimize since it avoids multiplication. - pub fn mid(&self) -> Option<&[E]> { - let adjust = if self.len % 2 == 0 { - 0 - } else { - self.elems_per_item.get() - }; - self.get_((self.storage.len() - adjust) / 2) - } - - // Inline so that the compiler can do strength reduction on the - // multiplication. - #[inline(always)] - pub fn get(&self, i: usize) -> Option<&[E]> { - let before = i.checked_mul(self.elems_per_item.get())?; - self.get_(before) - } - - fn get_(&self, before: usize) -> Option<&[E]> { - let after = before.checked_add(self.elems_per_item.get())?; - self.storage.get(before..after) - } - - pub fn as_flattened(&self) -> &[E] { - self.storage - } -} diff --git a/src/polyfill/mod.rs b/src/polyfill/mod.rs index 3d31b26bc0..264e65c7b1 100644 --- a/src/polyfill/mod.rs +++ b/src/polyfill/mod.rs @@ -46,7 +46,6 @@ mod cold_error; mod aliasing_slices; mod array_flat_map; mod array_split_map; -pub mod dynarray; pub mod partial_buffer; pub mod sliceutil; diff --git a/src/polyfill/uninit_slice.rs b/src/polyfill/uninit_slice.rs index 6c37e65b7a..7a1ed4cfec 100644 --- a/src/polyfill/uninit_slice.rs +++ b/src/polyfill/uninit_slice.rs @@ -264,6 +264,11 @@ impl<'target, E: Copy> Buf<'target, E> { unsafe { filled.assume_init() } } + #[cfg_attr(target_arch = "x86_64", expect(dead_code))] + pub fn into_filled(self) -> &'target [E] { + self.into_filled_mut() + } + pub fn into_filled_mut(self) -> &'target mut [E] { let (filled, _unfilled) = self.split_filled_mut(); filled @@ -293,6 +298,44 @@ impl<'target, E: Copy> Buf<'target, E> { let filled = unsafe { filled.assume_init() }; (filled, unfilled) } + + fn split_at_spare_mut(&mut self) -> (&mut [E], Uninit<'_, E>) { + let (filled, unfilled) = self + .storage + .reborrow_mut() + .split_at_mut_checked(self.filled) + .unwrap_or_else(|| unreachable!()); // by invariant. + // SAFETY: by invariant. + let filled = unsafe { filled.assume_init() }; + (filled, unfilled) + } + + /// Reserves the first `len` bytes of the unfilled space as `to_fill`, then + /// calls `f(filled, to_fill)` where `filled` is the filled space. If `f` + /// returns `Ok(filled)`, then `filled` must be `to_fill`, filled in. + #[cfg_attr(target_arch = "x86_64", expect(dead_code))] + pub fn try_write_with( + &mut self, + len: usize, + f: impl for<'a> FnOnce(&mut [E], Uninit<'a, E>) -> Result<&'a mut [E], Err>, + ) -> Result<(), Err> + where + Err: From, + { + let (filled, mut unfilled) = self.split_at_spare_mut(); + let Some(to_fill) = unfilled.get_mut(..len) else { + return Err(LenMismatchError::new(unfilled.len()))?; + }; + let len_and_ptr: *const [MaybeUninit] = ptr::from_ref(to_fill.target); + let len_and_ptr: *const [E] = len_and_ptr as *const [E]; // cast_init + let res = f(filled, to_fill)?; + if !ptr::eq(res, len_and_ptr) { + // XXX: It could have been the address that was different. + Err(LenMismatchError::new(res.len()))?; + } + self.filled += res.len(); + Ok(()) + } } /// A writable cursor analogous to `core::io::BorrowedCursor`. From bfb78d23bb16e974842ea1a591bcb69618256734 Mon Sep 17 00:00:00 2001 From: Brian Smith Date: Tue, 19 May 2026 13:20:14 -0700 Subject: [PATCH 2/2] bigint: Use `try_write_with` to simplify `IntoMont` construction. --- src/arithmetic/bigint/modulus/mont.rs | 19 ++++++------------- src/polyfill/uninit_slice.rs | 1 - 2 files changed, 6 insertions(+), 14 deletions(-) diff --git a/src/arithmetic/bigint/modulus/mont.rs b/src/arithmetic/bigint/modulus/mont.rs index 18fb17bfea..f5de8c7c45 100644 --- a/src/arithmetic/bigint/modulus/mont.rs +++ b/src/arithmetic/bigint/modulus/mont.rs @@ -149,19 +149,12 @@ impl ValidatedInput<'_> { let (_num_limbs, value) = rest.split_first().unwrap_or_else(|| unreachable!()); // Since we just wrote it. N0::write_into(n0, value); - let one_pos = out.filled().len(); - - // Placeholder for the value of 1*RR. - out.unfilled().write_repeat(limb::ZERO, num_limbs)?; - let (mont, one) = out - .filled_mut() - .split_at_mut_checked(one_pos) - .unwrap_or_else(|| unreachable!()); // Since we just wrote it. - let m = &Mont::<'_, M>::from_storage_unchecked_less_safe(mont, cpu); - - let _: elem::Mut<'_, _, RR> = - One::write_mont_identity(Uninit::from(one), m, self.len_bits())?.mul_r(m)?; // in place. - Ok(()) + out.try_write_with(num_limbs, |init, uninit| { + let m = &Mont::<'_, M>::from_storage_unchecked_less_safe(init, cpu); + let r: elem::Mut<'_, _, RR> = + One::write_mont_identity(uninit, m, self.len_bits())?.mul_r(m)?; // in place. + Ok(r.leak_limbs_into_mut_less_safe()) + }) }) } } diff --git a/src/polyfill/uninit_slice.rs b/src/polyfill/uninit_slice.rs index 7a1ed4cfec..bb39b8baa3 100644 --- a/src/polyfill/uninit_slice.rs +++ b/src/polyfill/uninit_slice.rs @@ -313,7 +313,6 @@ impl<'target, E: Copy> Buf<'target, E> { /// Reserves the first `len` bytes of the unfilled space as `to_fill`, then /// calls `f(filled, to_fill)` where `filled` is the filled space. If `f` /// returns `Ok(filled)`, then `filled` must be `to_fill`, filled in. - #[cfg_attr(target_arch = "x86_64", expect(dead_code))] pub fn try_write_with( &mut self, len: usize,