bigint/x86_64: Eliminate most unnecessary chunking.

We don't actually need to pass arguments as `&[[Limb; 8]]` instead of
`&[Limb]` in most cases. As long as one argument has a type like that,
then we've ensured that the length is a multiple of 512 bits. In most
cases we were calling `as_chunks[_mut]` on something and then
immediately calling `as_flattened[_mut]` on it.

Author: briansmith

src/arithmetic/bigint.rs

@@ -637,9 +637,7 @@ fn elem_exp_consttime_inner<N, M, const STORAGE_LIMBS: usize>(
         .ok_or_else(|| LenMismatchError::new(m.limbs().len()))?;
     let cpe = m_original.len(); // 512-bit chunks per entry
 
-    let oneRRR = &oneRRR.as_ref().limbs;
-    let oneRRR =
-        as_chunks_exact(oneRRR.as_ref()).ok_or_else(|| LenMismatchError::new(oneRRR.len()))?;
+    let oneRRR = oneRRR.as_ref().limbs.as_ref();
 
     // The x86_64 assembly was written under the assumption that the input data
     // is aligned to `MOD_EXP_CTIME_ALIGN` bytes, which was/is 64 in OpenSSL.
@@ -669,15 +667,16 @@ fn elem_exp_consttime_inner<N, M, const STORAGE_LIMBS: usize>(
     // These are named `(tmp, am, np)` in BoringSSL.
     let (acc, rest) = state.split_at_mut(cpe);
     let (base_cached, m_cached) = rest.split_at_mut(cpe);
+    let base_cached = base_cached.as_flattened_mut();
+    let acc = acc.as_flattened_mut();
 
     // "To improve cache locality" according to upstream.
     m_cached
         .as_flattened_mut()
         .copy_from_slice(m_original.as_flattened());
 
     let out: Elem<M, RInverse> = elem_reduced(out, base_mod_n, m, other_prime_len_bits);
-    let base_rinverse = as_chunks_exact(out.limbs.as_ref())
-        .ok_or_else(|| LenMismatchError::new(out.limbs.len()))?;
+    let base_rinverse = out.limbs.as_ref();
 
     // base_cached = base*R == (base/R * RRR)/R
     let _: &[Limb] = mul_mont5(base_cached, base_rinverse, oneRRR, m_cached, n0, cpu2)?;
@@ -687,7 +686,7 @@ fn elem_exp_consttime_inner<N, M, const STORAGE_LIMBS: usize>(
     // gathering, storing the last calculated power into `acc`.
     fn scatter_powers_of_2(
         table: &mut [[Limb; 8]],
-        acc: &mut [[Limb; 8]],
+        mut acc: &mut [Limb],
         m_cached: &[[Limb; 8]],
         n0: &N0,
         mut i: LeakyWindow5,
@@ -699,20 +698,19 @@ fn elem_exp_consttime_inner<N, M, const STORAGE_LIMBS: usize>(
                 Some(i) => i,
                 None => break,
             };
-            let _: &[Limb] = sqr_mont5(acc.as_flattened_mut(), m_cached, n0, cpu)?;
+            acc = sqr_mont5(acc, m_cached, n0, cpu)?;
         }
         Ok(())
     }
 
     // All entries in `table` will be Montgomery encoded.
 
     // acc = table[0] = base**0 (i.e. 1).
-    let _: &[Limb] = m.oneR(polyfill::slice::Uninit::from_mut(acc.as_flattened_mut()))?;
+    let _: &[Limb] = m.oneR(polyfill::slice::Uninit::from_mut(acc))?;
     scatter5(acc, table, LeakyWindow5::_0)?;
 
     // acc = base**1 (i.e. base).
-    acc.as_flattened_mut()
-        .copy_from_slice(base_cached.as_flattened());
+    acc.copy_from_slice(base_cached);
 
     // Fill in entries 1, 2, 4, 8, 16.
     scatter_powers_of_2(table, acc, m_cached, n0, LeakyWindow5::_1, cpu2)?;
@@ -740,7 +738,7 @@ fn elem_exp_consttime_inner<N, M, const STORAGE_LIMBS: usize>(
         },
     );
 
-    out.as_mut().copy_from_slice(acc.as_flattened());
+    out.as_mut().copy_from_slice(acc);
     Ok(from_montgomery_amm(out, m))
 }
 

src/arithmetic/limbs/x86_64/mont.rs

@@ -38,19 +38,15 @@ const _512_IS_LIMB_BITS_TIMES_8: () = assert!(8 * Limb::BITS == 512);
 
 #[inline]
 pub(in super::super::super) fn mul_mont5<'o>(
-    r: &'o mut [[Limb; 8]],
-    a: &[[Limb; 8]],
-    b: &[[Limb; 8]],
+    r: &'o mut [Limb],
+    a: &[Limb],
+    b: &[Limb],
     m: &[[Limb; 8]],
     n0: &N0,
     maybe_adx_bmi2: Option<(Adx, Bmi2)>,
 ) -> Result<&'o mut [Limb], LimbSliceError> {
     mul_mont5_4x(
-        (
-            Uninit::from_mut(r.as_flattened_mut()),
-            a.as_flattened(),
-            b.as_flattened(),
-        ),
+        (Uninit::from_mut(r), a, b),
         SmallerChunks::as_smaller_chunks(m),
         n0,
         maybe_adx_bmi2,
@@ -125,7 +121,7 @@ pub(in super::super::super) fn sqr_mont5<'o>(
 
 #[inline(always)]
 pub(in super::super::super) fn gather5(
-    r: &mut [[Limb; 8]],
+    r: &mut [Limb],
     table: &[[Limb; 8]],
     power: Window5,
 ) -> Result<(), LimbSliceError> {
@@ -139,16 +135,15 @@ pub(in super::super::super) fn gather5(
             power: Window5);
     }
     let num_limbs = check_common(r, table)?;
-    let r = r.as_flattened_mut();
     let table = table.as_flattened();
     unsafe { bn_gather5(r.as_mut_ptr(), num_limbs, table.as_ptr(), power) };
     Ok(())
 }
 
 #[inline(always)]
 pub(in super::super::super) fn mul_mont_gather5_amm(
-    r: &mut [[Limb; 8]],
-    a: &[[Limb; 8]],
+    r: &mut [Limb],
+    a: &[Limb],
     table: &[[Limb; 8]],
     n: &[[Limb; 8]],
     n0: &N0,
@@ -180,11 +175,9 @@ pub(in super::super::super) fn mul_mont_gather5_amm(
         );
     }
     let num_limbs = check_common_with_n(r, table, n)?;
-    let a = a.as_flattened();
     if a.len() != num_limbs.get() {
         Err(LenMismatchError::new(a.len()))?;
     }
-    let r = r.as_flattened_mut();
     let r = r.as_mut_ptr();
     let a = a.as_ptr();
     let table = table.as_flattened();
@@ -202,7 +195,7 @@ pub(in super::super::super) fn mul_mont_gather5_amm(
 // SAFETY: `power` must be less than 32.
 #[inline(always)]
 pub(in super::super::super) fn power5_amm(
-    in_out: &mut [[Limb; 8]],
+    in_out: &mut [Limb],
     table: &[[Limb; 8]],
     n: &[[Limb; 8]],
     n0: &N0,
@@ -234,7 +227,6 @@ pub(in super::super::super) fn power5_amm(
         );
     }
     let num_limbs = check_common_with_n(in_out, table, n)?;
-    let in_out = in_out.as_flattened_mut();
     let r = in_out.as_mut_ptr();
     let a = in_out.as_ptr();
     let table = table.as_flattened();

src/arithmetic/limbs512/scatter.rs

@@ -10,7 +10,7 @@ use {
 // `table` has space for 32 entries the same size as `a`. Instead of storing
 // entries consecutively row-wise, instead store them column-wise.
 pub(in super::super::super) fn scatter5(
-    a: &[[Limb; LIMBS_PER_CHUNK]],
+    a: &[Limb],
     table: &mut [[Limb; LIMBS_PER_CHUNK]],
     i: LeakyWindow5,
 ) -> Result<(), LimbSliceError> {
@@ -22,7 +22,7 @@ pub(in super::super::super) fn scatter5(
         .iter_mut()
         .skip(i)
         .step_by(32)
-        .zip(a.as_flattened())
+        .zip(a)
         .for_each(|(t, &a)| {
             *t = a;
         });

src/arithmetic/limbs512/storage.rs

@@ -71,11 +71,10 @@ impl<const N: usize> AlignedStorage<N> {
 // callers.
 #[inline(always)]
 pub(crate) fn check_common(
-    a: &[[Limb; LIMBS_PER_CHUNK]],
+    a: &[Limb],
     table: &[[Limb; LIMBS_PER_CHUNK]],
 ) -> Result<NonZeroUsize, LimbSliceError> {
     assert_eq!((table.as_ptr() as usize) % 16, 0); // According to BoringSSL.
-    let a = a.as_flattened();
     let table = table.as_flattened();
     let num_limbs = NonZeroUsize::new(a.len()).ok_or_else(|| LimbSliceError::too_short(a.len()))?;
     if num_limbs.get() > MAX_LIMBS {
@@ -90,7 +89,7 @@ pub(crate) fn check_common(
 #[cfg(target_arch = "x86_64")]
 #[inline(always)]
 pub(crate) fn check_common_with_n(
-    a: &[[Limb; LIMBS_PER_CHUNK]],
+    a: &[Limb],
     table: &[[Limb; LIMBS_PER_CHUNK]],
     n: &[[Limb; LIMBS_PER_CHUNK]],
 ) -> Result<NonZeroUsize, LimbSliceError> {