miri; fix unsafe ub

.github/workflows/ci.yaml

@@ -50,6 +50,20 @@ jobs:
       - uses: Swatinem/rust-cache@v2
       - run: cargo +nightly-2024-01-04 doc
 
+  ub-detection:
+    if: github.event.pull_request.draft == false
+    name: Check for undefined behaviour (UB)
+    timeout-minutes: 30
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: dtolnay/rust-toolchain@master
+        with:
+          toolchain: nightly-2024-01-04
+          components: miri
+      - uses: Swatinem/rust-cache@v2
+      - run: MIRIFLAGS="-Zmiri-strict-provenance" cargo miri test --lib
+
   run-wasm32-wasi-tests:
     runs-on: ubuntu-latest
     steps:

crates/prover/src/core/backend/simd/circle.rs

@@ -136,7 +136,7 @@ impl PolyOps for SimdBackend {
         values: Col<Self, BaseField>,
     ) -> CircleEvaluation<Self, BaseField, BitReversedOrder> {
         // TODO(Ohad): Optimize.
-        let eval = CpuBackend::new_canonical_ordered(coset, values.into_cpu_vec());
+        let eval = CpuBackend::new_canonical_ordered(coset, values.as_slice().to_vec());
         CircleEvaluation::new(
             eval.domain,
             Col::<SimdBackend, BaseField>::from_iter(eval.values),
@@ -421,6 +421,7 @@ mod tests {
     use crate::core::poly::{BitReversedOrder, NaturalOrder};
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_interpolate_and_eval() {
         for log_size in MIN_FFT_LOG_SIZE..CACHED_FFT_LOG_SIZE + 4 {
             let domain = CanonicCoset::new(log_size).circle_domain();
@@ -437,6 +438,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_eval_extension() {
         for log_size in MIN_FFT_LOG_SIZE..CACHED_FFT_LOG_SIZE + 2 {
             let domain = CanonicCoset::new(log_size).circle_domain();
@@ -457,6 +459,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_eval_at_point() {
         for log_size in MIN_FFT_LOG_SIZE + 1..CACHED_FFT_LOG_SIZE + 4 {
             let domain = CanonicCoset::new(log_size).circle_domain();
@@ -480,6 +483,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_circle_poly_extend() {
         for log_size in MIN_FFT_LOG_SIZE..CACHED_FFT_LOG_SIZE + 2 {
             let poly =
@@ -495,6 +499,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_eval_securefield() {
         let mut rng = SmallRng::seed_from_u64(0);
         for log_size in MIN_FFT_LOG_SIZE..CACHED_FFT_LOG_SIZE + 2 {
@@ -515,6 +520,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_optimized_precompute_twiddles() {
         let coset = CanonicCoset::new(10).half_coset();
         let twiddles = SimdBackend::precompute_twiddles(coset);

crates/prover/src/core/backend/simd/cm31.rs

@@ -205,6 +205,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn multiplication_works() {
         let mut rng = SmallRng::seed_from_u64(0);
         let lhs = rng.gen();

crates/prover/src/core/backend/simd/column.rs

@@ -1,5 +1,5 @@
+use std::array;
 use std::iter::zip;
-use std::{array, mem};
 
 use bytemuck::allocation::cast_vec;
 use bytemuck::{cast_slice, cast_slice_mut, Zeroable};
@@ -51,15 +51,6 @@ impl BaseColumn {
         &mut cast_slice_mut(&mut self.data)[..self.length]
     }
 
-    pub fn into_cpu_vec(mut self) -> Vec<BaseField> {
-        let capacity = self.data.capacity() * N_LANES;
-        let length = self.length;
-        let ptr = self.data.as_mut_ptr() as *mut BaseField;
-        let res = unsafe { Vec::from_raw_parts(ptr, length, capacity) };
-        mem::forget(self);
-        res
-    }
-
     pub fn from_cpu(values: Vec<BaseField>) -> Self {
         values.into_iter().collect()
     }

crates/prover/src/core/backend/simd/domain.rs

@@ -67,6 +67,7 @@ impl Iterator for CircleDomainBitRevIterator {
 }
 
 #[test]
+#[cfg_attr(miri, ignore)]
 fn test_circle_domain_bit_rev_iterator() {
     let domain = CircleDomain::new(crate::core::circle::Coset::new(
         crate::core::circle::CirclePointIndex::generator(),

crates/prover/src/core/backend/simd/fft/ifft.rs

@@ -567,6 +567,7 @@ mod tests {
     use crate::core::poly::circle::{CanonicCoset, CircleDomain};
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_ibutterfly() {
         let mut rng = SmallRng::seed_from_u64(0);
         let mut v0: [BaseField; N_LANES] = rng.gen();
@@ -585,6 +586,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_ifft3() {
         let mut rng = SmallRng::seed_from_u64(0);
         let values = rng.gen::<[BaseField; 8]>().map(PackedBaseField::broadcast);
@@ -645,6 +647,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_vecwise_ibutterflies() {
         let domain = CanonicCoset::new(5).circle_domain();
         let twiddle_dbls = get_itwiddle_dbls(domain.half_coset);
@@ -668,6 +671,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_ifft_lower_with_vecwise() {
         for log_size in 5..12 {
             let domain = CanonicCoset::new(log_size).circle_domain();
@@ -690,6 +694,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_ifft_full() {
         for log_size in CACHED_FFT_LOG_SIZE + 1..CACHED_FFT_LOG_SIZE + 3 {
             let domain = CanonicCoset::new(log_size).circle_domain();

crates/prover/src/core/backend/simd/fft/rfft.rs

@@ -593,6 +593,7 @@ mod tests {
     use crate::core::poly::circle::{CanonicCoset, CircleDomain};
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_butterfly() {
         let mut rng = SmallRng::seed_from_u64(0);
         let mut v0: [BaseField; N_LANES] = rng.gen();
@@ -611,6 +612,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_fft3() {
         let mut rng = SmallRng::seed_from_u64(0);
         let values = rng.gen::<[BaseField; 8]>().map(PackedBaseField::broadcast);
@@ -672,6 +674,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_vecwise_butterflies() {
         let domain = CanonicCoset::new(5).circle_domain();
         let twiddle_dbls = get_twiddle_dbls(domain.half_coset);
@@ -699,6 +702,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_fft_lower() {
         for log_size in 5..12 {
             let domain = CanonicCoset::new(log_size).circle_domain();
@@ -722,6 +726,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_fft_full() {
         for log_size in CACHED_FFT_LOG_SIZE + 1..CACHED_FFT_LOG_SIZE + 3 {
             let domain = CanonicCoset::new(log_size).circle_domain();

crates/prover/src/core/backend/simd/fri.rs

@@ -183,6 +183,7 @@ mod tests {
     use crate::qm31;
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_fold_line() {
         const LOG_SIZE: u32 = 7;
         let mut rng = SmallRng::seed_from_u64(0);
@@ -205,6 +206,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn test_fold_circle_into_line() {
         const LOG_SIZE: u32 = 7;
         let values: Vec<SecureField> = (0..(1 << LOG_SIZE))
@@ -239,6 +241,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn decomposition_test() {
         const DOMAIN_LOG_SIZE: u32 = 5;
         const DOMAIN_LOG_HALF_SIZE: u32 = DOMAIN_LOG_SIZE - 1;

crates/prover/src/core/backend/simd/lookups/gkr.rs

@@ -526,6 +526,7 @@ mod tests {
     use crate::core::test_utils::test_channel;
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn gen_eq_evals_matches_cpu() {
         let two = BaseField::from(2).into();
         let y = [7, 3, 5, 6, 1, 1, 9].map(|v| BaseField::from(v).into());
@@ -548,6 +549,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn grand_product_works() -> Result<(), GkrError> {
         const N: usize = 1 << 8;
         let values = test_channel().draw_felts(N);
@@ -571,6 +573,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn logup_with_generic_trace_works() -> Result<(), GkrError> {
         const N: usize = 1 << 8;
         let mut rng = SmallRng::seed_from_u64(0);
@@ -610,6 +613,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn logup_with_multiplicities_trace_works() -> Result<(), GkrError> {
         const N: usize = 1 << 8;
         let mut rng = SmallRng::seed_from_u64(0);
@@ -649,6 +653,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn logup_with_singles_trace_works() -> Result<(), GkrError> {
         const N: usize = 1 << 8;
         let mut rng = SmallRng::seed_from_u64(0);

crates/prover/src/core/backend/simd/lookups/mle.rs

@@ -103,6 +103,7 @@ mod tests {
     use crate::core::test_utils::test_channel;
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn fix_first_variable_with_secure_field_mle_matches_cpu() {
         const N_VARIABLES: u32 = 8;
         let values = test_channel().draw_felts(1 << N_VARIABLES);
@@ -117,6 +118,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn fix_first_variable_with_base_field_mle_matches_cpu() {
         const N_VARIABLES: u32 = 8;
         let values = (0..1 << N_VARIABLES).map(BaseField::from).collect_vec();

crates/prover/src/core/backend/simd/m31.rs

@@ -7,6 +7,7 @@ use std::simd::{u32x16, Simd, Swizzle};
 
 use bytemuck::{Pod, Zeroable};
 use num_traits::{One, Zero};
+#[cfg(test)]
 use rand::distributions::{Distribution, Standard};
 
 use super::qm31::PackedQM31;
@@ -274,6 +275,7 @@ impl From<BaseField> for PackedM31 {
     }
 }
 
+#[cfg(test)]
 impl Distribution<PackedM31> for Standard {
     fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> PackedM31 {
         PackedM31::from_array(rng.gen())
@@ -611,6 +613,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn multiplication_works() {
         let mut rng = SmallRng::seed_from_u64(0);
         let lhs = rng.gen();
@@ -654,6 +657,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn inverse_works() {
         let mut rng = SmallRng::seed_from_u64(0);
         let values = rng.gen();

crates/prover/src/core/backend/simd/prefix_sum.rs

@@ -121,12 +121,12 @@ fn down_sweep_val<F: Sub<Output = F> + Copy>(lo: &mut F, hi: &mut F) {
 }
 
 fn inclusive_prefix_sum_slow(
-    bit_rev_circle_domain_evals: Col<SimdBackend, BaseField>,
+    mut bit_rev_circle_domain_evals: Col<SimdBackend, BaseField>,
 ) -> Col<SimdBackend, BaseField> {
     // Obtain values in coset order.
-    let mut coset_order_eval = bit_rev_circle_domain_evals.into_cpu_vec();
-    bit_reverse(&mut coset_order_eval);
-    coset_order_eval = circle_domain_order_to_coset_order(&coset_order_eval);
+    let coset_order_eval = bit_rev_circle_domain_evals.as_mut_slice();
+    bit_reverse(coset_order_eval);
+    let coset_order_eval = circle_domain_order_to_coset_order(coset_order_eval);
     let coset_order_prefix_sum = coset_order_eval
         .into_iter()
         .scan(BaseField::zero(), |acc, v| {
@@ -150,9 +150,7 @@ mod tests {
     use crate::core::backend::simd::prefix_sum::inclusive_prefix_sum_slow;
     use crate::core::backend::Column;
 
-    #[test]
-    fn exclusive_prefix_sum_simd_with_log_size_3_works() {
-        const LOG_N: u32 = 3;
+    fn exclusive_prefix_sum_simd_with_log_size_n_works<const LOG_N: u32>() {
         let mut rng = SmallRng::seed_from_u64(0);
         let evals: BaseColumn = (0..1 << LOG_N).map(|_| rng.gen()).collect();
         let expected = inclusive_prefix_sum_slow(evals.clone());
@@ -163,26 +161,17 @@ mod tests {
     }
 
     #[test]
-    fn exclusive_prefix_sum_simd_with_log_size_6_works() {
-        const LOG_N: u32 = 6;
-        let mut rng = SmallRng::seed_from_u64(0);
-        let evals: BaseColumn = (0..1 << LOG_N).map(|_| rng.gen()).collect();
-        let expected = inclusive_prefix_sum_slow(evals.clone());
-
-        let res = inclusive_prefix_sum(evals);
+    fn exclusive_prefix_sum_simd_with_log_size_3_works() {
+        exclusive_prefix_sum_simd_with_log_size_n_works::<3>();
+    }
 
-        assert_eq!(res.to_cpu(), expected.to_cpu());
+    #[test]
+    fn exclusive_prefix_sum_simd_with_log_size_6_works() {
+        exclusive_prefix_sum_simd_with_log_size_n_works::<6>();
     }
 
     #[test]
     fn exclusive_prefix_sum_simd_with_log_size_8_works() {
-        const LOG_N: u32 = 8;
-        let mut rng = SmallRng::seed_from_u64(0);
-        let evals: BaseColumn = (0..1 << LOG_N).map(|_| rng.gen()).collect();
-        let expected = inclusive_prefix_sum_slow(evals.clone());
-
-        let res = inclusive_prefix_sum(evals);
-
-        assert_eq!(res.to_cpu(), expected.to_cpu());
+        exclusive_prefix_sum_simd_with_log_size_n_works::<8>();
     }
 }

crates/prover/src/core/backend/simd/qm31.rs

@@ -4,6 +4,7 @@ use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
 
 use bytemuck::{Pod, Zeroable};
 use num_traits::{One, Zero};
+#[cfg(test)]
 use rand::distributions::{Distribution, Standard};
 
 use super::cm31::PackedCM31;
@@ -292,6 +293,7 @@ impl Neg for PackedQM31 {
     }
 }
 
+#[cfg(test)]
 impl Distribution<PackedQM31> for Standard {
     fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> PackedQM31 {
         PackedQM31::from_array(rng.gen())
@@ -351,6 +353,7 @@ mod tests {
     }
 
     #[test]
+    #[cfg_attr(miri, ignore)]
     fn multiplication_works() {
         let mut rng = SmallRng::seed_from_u64(0);
         let lhs = rng.gen();