arithmetic: Make Montgomery API more flexible.

Change the computation of `R` and `RR` to support writing the value to
a larger output buffer, to support future refactorings.

Author: briansmith

src/arithmetic/bigint/exp.rs

@@ -150,7 +150,7 @@ fn elem_exp_consttime_inner<N, M, const STORAGE_LIMBS: usize>(
         |init, uninit| {
             let r: Result<&mut [Limb], LimbSliceError> = match init.len() {
                 // table[0] = base**0 (i.e. 1).
-                0 => Ok(One::fillR(uninit, m)?),
+                0 => Ok(One::write_mont_identity(&mut uninit.into_cursor(), m)?),
 
                 // table[1] = base*R == (base/R * RRR)/R
                 1 => limbs_mul_mont(
@@ -282,6 +282,8 @@ fn elem_exp_consttime_inner<N, M, const STORAGE_LIMBS: usize>(
     let (acc, rest) = state.split_at_mut(m_len);
     let (base_cached, m_cached) = rest.split_at_mut(m_len);
 
+    let mut acc = polyfill::slice::Uninit::from(acc);
+
     // "To improve cache locality" according to upstream.
     let (m_cached, _) = polyfill::slice::Uninit::from(m_cached)
         .write_copy_of_slice(m.limbs())?
@@ -327,7 +329,7 @@ fn elem_exp_consttime_inner<N, M, const STORAGE_LIMBS: usize>(
     // All entries in `table` will be Montgomery encoded.
 
     // t0 = table[0] = base**0 (i.e. 1).
-    let t0 = One::fillR(acc.as_mut().into(), m)?;
+    let t0 = One::write_mont_identity(&mut acc.reborrow_mut().into_cursor(), m)?;
     scatter5(t0, table, LeakyWindow5::_0)?;
 
     // acc = base**1 (i.e. base).

src/arithmetic/bigint/modulus/one.rs

@@ -24,7 +24,7 @@ use crate::{
     cpu,
     error::LenMismatchError,
     limb::{self, LIMB_BITS},
-    polyfill,
+    polyfill::slice::Cursor,
 };
 use core::mem::size_of;
 
@@ -45,14 +45,16 @@ impl<M, E> One<M, E> {
 }
 
 impl<M> One<M, R> {
-    pub(in super::super) fn fillR<'r>(
-        out: polyfill::slice::Uninit<'r, Limb>,
+    /// Writes the value of the Montgomery multiplication identity `R` for `m` to
+    /// `out`.
+    pub(in super::super) fn write_mont_identity<'r>(
+        out: &mut Cursor<'r, Limb>,
         m: &Mont<'_, M>,
     ) -> Result<&'r mut [Limb], LenMismatchError> {
         let r = m.limbs().len() * LIMB_BITS;
 
         // out = 2**r - m where m = self.
-        let out = limb::limbs_negative_odd(out, m.limbs())?;
+        let out = limb::write_negative_assume_odd(out, m.limbs())?;
 
         let lg_m = m.len_bits().as_bits();
         let leading_zero_bits_in_m = r - lg_m;
@@ -101,7 +103,7 @@ impl<M> One<M, RR> {
         let m = &Mont::from_parts_unchecked_less_safe(m, &n0, cpu);
 
         let mut acc = out
-            .write_fully_with(|out| One::fillR(out, m))
+            .write_fully_with(|uninit| One::write_mont_identity(&mut uninit.into_cursor(), m))
             .map(Elem::<M, R>::assume_in_range_and_encoded_less_safe)?;
 
         // 2**t * R can be calculated by t doublings starting with R.

src/limb.rs

@@ -22,7 +22,7 @@ use crate::{
     bb, c,
     error::{self, LenMismatchError},
     polyfill::{
-        slice::{AliasingSlices, InOut, Uninit},
+        slice::{AliasingSlices, Cursor, InOut},
         sliceutil, usize_from_u32, ArrayFlatMap, StartMutPtr,
     },
     window5::Window5,
@@ -333,20 +333,22 @@ pub(crate) fn limbs_double_mod(r: &mut [Limb], m: &[Limb]) -> Result<(), LenMism
 }
 
 // *r = -a, assuming a is odd.
-pub(crate) fn limbs_negative_odd<'r>(
-    r: Uninit<'r, Limb>,
+pub(crate) fn write_negative_assume_odd<'r>(
+    r: &mut Cursor<'r, Limb>,
     a: &[Limb],
 ) -> Result<&'r mut [Limb], LenMismatchError> {
     // Two's complement step 1: flip all the bits.
     // The compiler should optimize this to vectorized (a ^ !0).
     let r = r
         .write_iter(a.iter().map(|&a| !a))
-        .uninit_empty()?
         .src_empty()?
         .into_written();
     // Two's complement step 2: Add one. Since `a` is odd, `r` is even. Thus we
     // can use a bitwise or for addition.
-    r[0] |= 1;
+    let Some(least_significant_limb) = r.get_mut(0) else {
+        return Err(LenMismatchError::new(a.len()));
+    };
+    *least_significant_limb |= 1;
     Ok(r)
 }
 

src/polyfill.rs

@@ -98,6 +98,7 @@ mod start_ptr;
 mod test;
 
 mod uninit_slice;
+mod uninit_slice_cursor;
 mod unwrap_const;
 
 pub use self::{

src/polyfill/slice.rs

@@ -27,6 +27,7 @@ use crate::polyfill::{StartMutPtr, StartPtr};
 pub(crate) use super::{
     aliasing_slices::{AliasSrc, AliasingSlices, InOut},
     uninit_slice::{AliasedUninit, Uninit},
+    uninit_slice_cursor::Cursor,
 };
 
 #[allow(dead_code)]

src/polyfill/uninit_slice.rs

@@ -15,7 +15,10 @@
 #[allow(unused_imports)]
 use crate::polyfill::prelude::*;
 
-use super::start_ptr::{StartMutPtr, StartPtr};
+use super::{
+    start_ptr::{StartMutPtr, StartPtr},
+    uninit_slice_cursor::Cursor,
+};
 use crate::{error::LenMismatchError, polyfill};
 use core::{
     marker::PhantomData,
@@ -53,11 +56,25 @@ impl<E> StartMutPtr for &mut Uninit<'_, E> {
 }
 
 impl<'target, E> Uninit<'target, E> {
+    pub fn into_cursor(self) -> Cursor<'target, E> {
+        Cursor::new(self)
+    }
+
     pub fn len(&self) -> usize {
         self.target.len()
     }
 
-    pub(super) fn split_off_mut(&mut self, range: RangeTo<usize>) -> Option<Uninit<'target, E>> {
+    #[allow(dead_code)]
+    pub fn reborrow_mut(&mut self) -> Uninit<'_, E> {
+        Uninit {
+            target: self.target,
+        }
+    }
+
+    pub(super) fn split_off_mut<'s>(
+        &'s mut self,
+        range: RangeTo<usize>,
+    ) -> Option<Uninit<'target, E>> {
         if self.target.len() < range.end {
             return None;
         }
@@ -74,7 +91,26 @@ impl<'target, E: Copy> Uninit<'target, E> {
         self,
         src: &[E],
     ) -> Result<WriteResult<'target, E, Self, ()>, LenMismatchError> {
-        self.write_iter(src.iter().copied()).src_empty()
+        self.write_copy_of_slice_(src)
+            .map_err(|uninit| LenMismatchError::new(uninit.len()))
+    }
+
+    pub fn write_copy_of_slice_(
+        mut self,
+        src: &[E],
+    ) -> Result<WriteResult<'target, E, Self, ()>, Self> {
+        let Some(mut dst) = self.split_off_mut(..src.len()) else {
+            return Err(self);
+        };
+        let written = unsafe {
+            ptr::copy_nonoverlapping(src.as_ptr(), dst.start_mut_ptr(), src.len());
+            dst.assume_init()
+        };
+        Ok(WriteResult {
+            written,
+            dst_leftover: self,
+            src_leftover: (),
+        })
     }
 
     pub fn write_iter<Src: IntoIterator<Item = E>>(
@@ -264,6 +300,7 @@ impl<'written, E, Dst, Src> WriteResult<'written, E, Dst, Src> {
 }
 
 impl<'written, E, Src> WriteResult<'written, E, Uninit<'written, E>, Src> {
+    #[allow(dead_code)]
     #[inline(always)]
     pub fn uninit_empty(self) -> Result<WriteResult<'written, E, (), Src>, LenMismatchError> {
         let (res, dst_leftover) = self.take_uninit();

src/polyfill/uninit_slice_cursor.rs

@@ -0,0 +1,68 @@
+// 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.
+
+#[allow(unused_imports)]
+use crate::polyfill::prelude::*;
+
+use super::uninit_slice::{Uninit, WriteResult};
+use crate::error::LenMismatchError;
+use core::mem;
+
+pub struct Cursor<'buf, E> {
+    uninit: Uninit<'buf, E>,
+}
+
+impl<'buf, E> Cursor<'buf, E> {
+    pub(super) const fn new(uninit: Uninit<'buf, E>) -> Self {
+        Self { uninit }
+    }
+
+    #[allow(dead_code)]
+    pub fn write_copy_of_slice(&mut self, src: &[E]) -> Result<&'buf mut [E], LenMismatchError>
+    where
+        E: Copy,
+    {
+        match mem::take(&mut self.uninit).write_copy_of_slice_(src) {
+            Ok(w) => {
+                let (res, uninit) = w.take_uninit();
+                self.uninit = uninit;
+                Ok(res.into_written())
+            }
+            Err(uninit) => {
+                self.uninit = uninit;
+                Err(LenMismatchError::new(self.uninit.len()))
+            }
+        }
+    }
+
+    pub fn write_iter<'s, Src: IntoIterator<Item = E>>(
+        &'s mut self,
+        src: Src,
+    ) -> WriteResult<'buf, E, (), Src::IntoIter>
+    where
+        E: Copy,
+    {
+        // TODO: Deal with panics.
+        let uninit = mem::replace(&mut self.uninit, Uninit::from([].as_mut_slice()));
+        let (res, uninit) = uninit.write_iter(src).take_uninit();
+        self.uninit = uninit;
+        res
+    }
+}
+
+impl<'buf, E> From<Uninit<'buf, E>> for Cursor<'buf, E> {
+    fn from(uninit: Uninit<'buf, E>) -> Self {
+        Self::new(uninit)
+    }
+}