remainder/utils/

mle.rs

//! Module for generating and manipulating mles.

use std::iter::repeat_with;

use itertools::{repeat_n, FoldWhile, Itertools};
use rand::Rng;
use rayon::prelude::*;
use shared_types::Field;

use crate::{
    claims::RawClaim,
    layer::LayerId,
    mle::{betavalues::BetaValues, dense::DenseMle, evals::MultilinearExtension, MleIndex},
};

/// Return a vector containing a padded version of the input data, with the
/// padding value at the end of the vector, such that the length is
/// `data.len().next_power_of_two()`.  This is a no-op if the length is already
/// a power of two.
/// # Examples:
/// ```
/// use remainder::utils::mle::pad_with;
/// let data = vec![1, 2, 3];
/// let padded_data = pad_with(0, &data);
/// assert_eq!(padded_data, vec![1, 2, 3, 0]);
/// assert_eq!(pad_with(0, &padded_data), vec![1, 2, 3, 0]); // length is already a power of two.
/// ```
pub fn pad_with<F: Clone>(padding_value: F, data: &[F]) -> Vec<F> {
    let padded_length = data.len().checked_next_power_of_two().unwrap();
    let mut padded_data = Vec::with_capacity(padded_length);
    padded_data.extend_from_slice(data);
    padded_data.extend(std::iter::repeat_n(
        padding_value,
        padded_length - data.len(),
    ));
    padded_data
}

/// Returns the argsort (i.e. indices) of the given vector slice.
///
/// ## Example:
/// ```
/// use remainder::utils::mle::argsort;
/// let data = vec![3, 1, 4, 1, 5, 9, 2];
///
/// // Ascending order
/// let indices = argsort(&data, false);
/// assert_eq!(indices, vec![1, 3, 6, 0, 2, 4, 5]);
/// ```
pub fn argsort<T: Ord>(slice: &[T], invert: bool) -> Vec<usize> {
    let mut indices: Vec<usize> = (0..slice.len()).collect();

    indices.sort_by(|&i, &j| {
        if invert {
            slice[j].cmp(&slice[i])
        } else {
            slice[i].cmp(&slice[j])
        }
    });

    indices
}

/// Helper function to create random MLE with specific number of vars
pub fn get_random_mle<F: Field>(num_vars: usize, rng: &mut impl Rng) -> DenseMle<F> {
    let capacity = 2_u32.pow(num_vars as u32);
    let bookkeeping_table = repeat_with(|| F::from(rng.gen::<u64>()) * F::from(rng.gen::<u64>()))
        .take(capacity as usize)
        .collect_vec();
    DenseMle::new_from_raw(bookkeeping_table, LayerId::Input(0))
}

/// Helper function to create random MLE with specific number of vars
pub fn get_random_mle_from_capacity<F: Field>(capacity: usize, rng: &mut impl Rng) -> DenseMle<F> {
    let bookkeeping_table = repeat_with(|| F::from(rng.gen::<u64>()) * F::from(rng.gen::<u64>()))
        .take(capacity)
        .collect_vec();
    DenseMle::new_from_raw(bookkeeping_table, LayerId::Input(0))
}

/// Returns a vector of MLEs for dataparallel testing according to the number of
/// variables and number of dataparallel bits.
pub fn get_dummy_random_mle_vec<F: Field>(
    num_vars: usize,
    num_dataparallel_bits: usize,
    rng: &mut impl Rng,
) -> Vec<DenseMle<F>> {
    (0..(1 << num_dataparallel_bits))
        .map(|_| {
            let mle_vec = (0..(1 << num_vars))
                .map(|_| F::from(rng.gen::<u64>()))
                .collect_vec();
            DenseMle::new_from_raw(mle_vec, LayerId::Input(0))
        })
        .collect_vec()
}

/// Returns the specific bit decomp for a given index, using `num_bits` bits.
/// Note that this returns the decomposition in BIG ENDIAN!
pub fn get_mle_idx_decomp_for_idx<F: Field>(idx: usize, num_bits: usize) -> Vec<MleIndex<F>> {
    (0..(num_bits))
        .rev()
        .map(|cur_num_bits| {
            let is_one =
                (idx % 2_usize.pow(cur_num_bits as u32 + 1)) >= 2_usize.pow(cur_num_bits as u32);
            MleIndex::Fixed(is_one)
        })
        .collect_vec()
}

/// Returns the total MLE indices given a `Vec<bool>`. for the prefix bits and
/// then the number of free bits after.
pub fn get_total_mle_indices<F: Field>(
    prefix_bits: &[bool],
    num_free_bits: usize,
) -> Vec<MleIndex<F>> {
    prefix_bits
        .iter()
        .map(|bit| MleIndex::Fixed(*bit))
        .chain(repeat_n(MleIndex::Free, num_free_bits))
        .collect()
}

/// Construct a parent MLE for the given MLEs and prefix bits, where the prefix
/// bits of each MLE specify how it should be inserted into the parent. Entries
/// left unspecified are filled with `F::ZERO`.
/// # Requires:
/// * that the number of variables in each MLE, plus the number of its prefix
///   bits, is the same across all pairs; this will be the number of variables
///   in the returned MLE.
/// * the slice is non-empty.
///
/// # Example:
/// ```
/// use remainder::utils::mle::build_composite_mle;
/// use remainder::mle::evals::MultilinearExtension;
/// use shared_types::Fr;
/// use itertools::{Itertools};
/// let mle1 = MultilinearExtension::new(vec![Fr::from(2)]);
/// let mle2 = MultilinearExtension::new(vec![Fr::from(1), Fr::from(3)]);
/// let result = build_composite_mle(&[(&mle1, vec![false, true]), (&mle2, vec![true])]);
/// assert_eq!(*result.f.iter().collect_vec().clone(), vec![Fr::from(0), Fr::from(2), Fr::from(1), Fr::from(3)]);
/// ```
pub fn build_composite_mle<F: Field>(
    mles: &[(&MultilinearExtension<F>, Vec<bool>)],
) -> MultilinearExtension<F> {
    assert!(!mles.is_empty());
    let out_num_vars = mles[0].0.num_vars() + mles[0].1.len();
    // Check that all (MLE, prefix bit) pairs require the same total number of
    // variables.

    mles.iter().for_each(|(mle, prefix_bits)| {
        assert_eq!(mle.num_vars() + prefix_bits.len(), out_num_vars);
    });
    let mut out = vec![F::ZERO; 1 << out_num_vars];
    for (mle, prefix_bits) in mles {
        let mut current_window = 1 << out_num_vars;
        let starting_index = prefix_bits.iter().fold(0, |acc_index, bit| {
            let starting_index_acc = if *bit {
                acc_index + current_window / 2
            } else {
                acc_index
            };
            current_window /= 2;
            starting_index_acc
        });
        // REMARK: Modifying this check so that the input mles can be non-
        // powers of 2.
        assert_eq!(current_window, mle.len().next_power_of_two());
        (starting_index..(starting_index + current_window))
            .enumerate()
            .for_each(|(mle_idx, out_idx)| {
                out[out_idx] = mle.get(mle_idx).unwrap_or(F::ZERO);
            });
    }
    MultilinearExtension::new(out)
}

/// Verifies a claim by evaluating the MLE at the challenge point and checking
/// that the result.
pub fn verify_claim<F: Field>(mle_unpadded_evaluations: &[F], claim: &RawClaim<F>) {
    let mle_evaluations = claim
        .get_point()
        .iter()
        .fold_while(mle_unpadded_evaluations, |acc, elem| {
            if elem == &F::ZERO {
                let sliced_acc = &acc[..(acc.len() / 2)];
                FoldWhile::Continue(sliced_acc)
            } else if elem == &F::ONE {
                let sliced_acc = &acc[(acc.len() / 2)..];
                FoldWhile::Continue(sliced_acc)
            } else {
                FoldWhile::Done(acc)
            }
        })
        .into_inner();
    let filtered_claim = claim
        .get_point()
        .iter()
        .skip_while(|x| x == &&F::ZERO || x == &&F::ONE)
        .copied()
        .collect_vec();
    let mle = MultilinearExtension::new(mle_evaluations.to_vec());
    assert_eq!(mle.num_vars(), filtered_claim.len());
    let eval = evaluate_mle_at_a_point_gray_codes(&mle, &filtered_claim);
    assert_eq!(eval, claim.get_eval());
}

/// A struct representing an iterator that iterates through the range
/// (1..2^{`num_bits`}) but in the ordering of a Gray Code, which means that the
/// Hamming distance between the bit representation of any consecutive indices
/// is only 1.
///
/// The iterator is of the type (u32, (u32, bool)) which represents: (index,
/// (index of the bit that was flipped, the previous value of the flipped bit.))

#[derive(Debug)]
pub struct GrayCodeIterator {
    num_bits: usize,
    current_iteration: u32,
    end_iteration: Option<u32>,
}

impl GrayCodeIterator {
    /// Note: `num_bits` cannot be more than 31 because we work with u32s in
    /// this iterator.
    pub fn new(num_bits: usize) -> Self {
        assert!(num_bits < 32);
        Self {
            num_bits,
            current_iteration: 0,
            end_iteration: None,
        }
    }

    pub(crate) fn new_at_index(
        num_bits: usize,
        current_iteration: u32,
        end_iteration: Option<u32>,
    ) -> Self {
        Self {
            num_bits,
            current_iteration,
            end_iteration,
        }
    }

    pub(crate) fn get_gray_index(num_bits: usize, index: u32) -> u32 {
        let mask = (1 << num_bits) - 1;

        (index ^ (index >> 1)) & mask
    }
}

impl Iterator for GrayCodeIterator {
    type Item = (u32, (u32, bool));

    fn next(&mut self) -> Option<Self::Item> {
        if self.current_iteration >= ((1 << self.num_bits) - 1) {
            return None;
        }

        if self.end_iteration.is_some()
            && self.current_iteration >= (self.end_iteration.unwrap() - 1)
        {
            return None;
        }

        if self.end_iteration.is_some()
            && self.current_iteration >= (self.end_iteration.unwrap() - 1)
        {
            return None;
        }

        // Mask current value to ensure we only get num_bits number of bits per
        // result.
        let mask = (1 << self.num_bits) - 1;

        // Because we don't store the previous value, just calculate it using
        // the current_val that's stored.
        let prev_gray = (self.current_iteration ^ (self.current_iteration >> 1)) & mask;
        // The next value is simply XOR of the counter incremented by 1 and
        // itself right-shifted. (source: algorithm in Wikipedia and ChatGPT
        // hehe)
        self.current_iteration += 1;
        let new_gray = (self.current_iteration ^ (self.current_iteration >> 1)) & mask;

        // Internally, the bits are stored in little-endian. NOTE: Our
        // bookkeeping tables are in "big-endian", so we need to take this into
        // account when evaluating an MLE.
        Some((
            new_gray,
            compute_flipped_bit_idx_and_value_graycode(prev_gray, new_gray),
        ))
    }
}

/// A struct representing lexicographic bit-order in little-endian.
///
/// The iterator returns elements of the form (u32, Vec<(u32, bool)>) where the
/// first u32 is the current index, and the Vec<(u32, bool)> represents the bits
/// that flipped and what they used to be. As opposed to [GrayCodeIterator],
/// these don't have Hamming distance 1, so the flipped bits go in a Vec.
pub struct LexicographicLE {
    num_bits: usize,
    current_val: u32,
}

impl LexicographicLE {
    fn new(num_bits: usize) -> Self {
        Self {
            num_bits,
            current_val: 0,
        }
    }
}

impl Iterator for LexicographicLE {
    type Item = (u32, Vec<(u32, bool)>);

    fn next(&mut self) -> Option<Self::Item> {
        if self.current_val >= ((1 << self.num_bits) - 1) {
            return None;
        }

        let prev_val = self.current_val;
        self.current_val += 1;

        let flipped_bit_idx_and_values =
            compute_flipped_bit_idx_and_values_lexicographic(prev_val, self.current_val);

        Some((self.current_val, flipped_bit_idx_and_values))
    }
}

/// Compute the single flipped bit and its previous value for the gray codes
/// iterator.
pub fn compute_flipped_bit_idx_and_value_graycode(curr_val: u32, next_val: u32) -> (u32, bool) {
    let flipped_bit = (curr_val ^ next_val).trailing_zeros();
    let previous_value = (curr_val & (1 << flipped_bit)) != 0;
    (flipped_bit, previous_value)
}

/// Compute the inverses and one minus the elem inverted for a vec of claim challenges.
pub fn compute_inverses_vec_and_one_minus_inverted_vec<F: Field>(
    claim_points: &[&[F]],
) -> (Vec<Vec<F>>, Vec<Vec<F>>) {
    let inverses_vec = claim_points
        .iter()
        .map(|claim_point| {
            claim_point
                .iter()
                .map(|elem| elem.invert().unwrap())
                .collect_vec()
        })
        .collect_vec();
    let one_minus_inverses_vec = claim_points
        .iter()
        .map(|claim_point| {
            claim_point
                .iter()
                .map(|elem| (F::ONE - elem).invert().unwrap())
                .collect_vec()
        })
        .collect_vec();
    (inverses_vec, one_minus_inverses_vec)
}

/// Compute the flipped bits between a previous value and a current value, and
/// return each of the flipped bits' indices and previous value.
pub fn compute_flipped_bit_idx_and_values_lexicographic(
    curr_val: u32,
    next_val: u32,
) -> Vec<(u32, bool)> {
    let flipped_bits = curr_val ^ next_val;
    let mut flipped_bit_idx_and_values = Vec::<(u32, bool)>::new();
    (0..32).for_each(|idx| {
        if (flipped_bits & (1 << idx)) != 0 {
            // NOTE: Our bookkeeping tables are in "big-endian", so we need
            // to take this into account when evaluating an MLE.
            flipped_bit_idx_and_values.push((idx, (curr_val & (1 << idx)) != 0))
        }
    });
    flipped_bit_idx_and_values
}

/// Compute the next beta values from the previous by multiplying by appropriate
/// inverses and challenge points given the flipped bits and their previous
/// values. This is for when we have multiple claims to compute the beta over.
pub fn compute_next_beta_values_vec_from_current<F: Field>(
    current_beta_values: &[F],
    inverses_vec: &[Vec<F>],
    one_minus_elem_inverted_vec: &[Vec<F>],
    claim_points: &[&[F]],
    flipped_bit_idx_and_values: &[(u32, bool)],
) -> Vec<F> {
    current_beta_values
        .iter()
        .zip(inverses_vec.iter().zip(one_minus_elem_inverted_vec))
        .zip(claim_points)
        .map(
            |((current_beta_value, (inverses, one_minus_inverses)), claim_point)| {
                compute_next_beta_value_from_current(
                    current_beta_value,
                    inverses,
                    one_minus_inverses,
                    claim_point,
                    flipped_bit_idx_and_values,
                )
            },
        )
        .collect_vec()
}

/// Compute the next beta value from the previous by multiplying by appropriate
/// inverses and challenge points given the flipped bits and their previous
/// values. This is for when we have a single claim to compute the beta over.
pub fn compute_next_beta_value_from_current<F: Field>(
    current_beta_value: &F,
    inverses: &[F],
    one_minus_elem_inverted: &[F],
    claim_point: &[F],
    flipped_bit_idx_and_values: &[(u32, bool)],
) -> F {
    let n = claim_point.len();
    flipped_bit_idx_and_values.iter().fold(
        // For each of the flipped bits, multiply by the
        // appropriate inverse depending on if the value was
        // previously 0 or 1.
        *current_beta_value,
        |acc, (idx, value)| {
            if *value {
                acc * inverses[n - 1 - *idx as usize]
                    * (F::ONE - claim_point[n - 1 - *idx as usize])
            } else {
                acc * (one_minus_elem_inverted[n - 1 - *idx as usize])
                    * claim_point[n - 1 - *idx as usize]
            }
        },
    )
}

/// This function non-destructively evaluates an MLE at a given point using the
/// [LexicographicLE] iterator.
pub fn evaluate_mle_at_a_point_lexicographic_order<F: Field>(
    mle: &MultilinearExtension<F>,
    point: &[F],
) -> F {
    let n = point.len();
    let mle_num_vars = mle.num_vars();
    assert_eq!(n, mle_num_vars);

    let starting_beta_value =
        BetaValues::compute_beta_over_two_challenges(&vec![F::ZERO; mle_num_vars], point);

    let starting_evaluation_acc = starting_beta_value * mle.first();
    let lexicographic_le = LexicographicLE::new(mle_num_vars);
    let inverses = point
        .iter()
        .map(|elem| elem.invert().unwrap())
        .collect_vec();
    let one_minus_inverses = point
        .iter()
        .map(|elem| (F::ONE - elem).invert().unwrap())
        .collect_vec();

    let (_final_beta_value, evaluation) = lexicographic_le.fold(
        (starting_beta_value, starting_evaluation_acc),
        |(prev_beta_value, evaluation_acc), (index, flipped_bit_indices_and_values)| {
            let next_beta_value = flipped_bit_indices_and_values.iter().fold(
                prev_beta_value,
                |acc, (flipped_bit_index, flipped_bit_value)| {
                    // For every bit i that is flipped, if it used to be a 1,
                    // then we multiply by r_i^{-1} and multiply by (1 - r_i) to
                    // account for this bit flip. NOTE: we subtract from n - 1
                    // to account for the fact that internally, these u32s are
                    // stored in little endian, but our bookkeeping tables are
                    // stored in "big endian" indexing.
                    if *flipped_bit_value {
                        acc * inverses[n - 1 - *flipped_bit_index as usize]
                            * (F::ONE - point[n - 1 - *flipped_bit_index as usize])
                    }
                    // For every bit i that is flipped, if it used to be a 0,
                    // then we multiply by (1 - r_i)^{-1} and multiply by r_i to
                    // account for this bit flip.
                    else {
                        acc * (one_minus_inverses[n - 1 - *flipped_bit_index as usize])
                            * point[n - 1 - *flipped_bit_index as usize]
                    }
                },
            );

            // Multiply this by the appropriate MLE coefficient.
            let next_evaluation_acc = next_beta_value * mle.get(index as usize).unwrap();
            (next_beta_value, evaluation_acc + next_evaluation_acc)
        },
    );
    evaluation
}

/// This function non-destructively evaluates an MLE at a given point using the
/// gray codes iterator. Optimized version that uses 2 multiplications instead
/// of 1.
///
/// Currently does not support for when the value in the point is either 0 or 1.
pub fn evaluate_mle_at_a_point_gray_codes<F: Field>(
    mle: &MultilinearExtension<F>,
    point: &[F],
) -> F {
    let n = point.len();
    let mle_num_vars = mle.num_vars();
    assert_eq!(n, mle_num_vars);
    // The gray codes start at index 1, so we start with the first value which
    // is \widetilde{\beta}(\vec{0}, point).
    let starting_beta_value =
        BetaValues::compute_beta_over_two_challenges(&vec![F::ZERO; mle_num_vars], point);
    // This is the value that gets multiplied to the first MLE coefficient,
    // which is (1 - r_1) * (1 - r_2) * ... * (1 - r_n) where (r_1, ..., r_n) is
    // the point.
    let starting_evaluation_acc = starting_beta_value * mle.first();
    let gray_code = GrayCodeIterator::new(mle_num_vars);
    let inverses = point
        .iter()
        .map(|elem| elem.invert().unwrap())
        .collect_vec();
    let one_minus_inverses = point
        .iter()
        .map(|elem| (F::ONE - elem).invert().unwrap())
        .collect_vec();

    // We simply compute the correct inverse and new multiplicative term for
    // each bit that is flipped in the beta value, and accumulate these by doing
    // an element-wise multiplication with the correct index of the MLE
    // coefficients. We precompute these so that during the scanning we only
    // need to do one multiplication instead of two
    let multiplier_if_flipped_bit_is_one = inverses
        .iter()
        .zip(point.iter())
        .map(|(inverse, point_elem)| *inverse * (F::ONE - point_elem))
        .collect_vec();
    let multiplier_if_flipped_bit_is_zero = one_minus_inverses
        .iter()
        .zip(point.iter())
        .map(|(one_minus_inverse, point_elem)| *one_minus_inverse * point_elem)
        .collect_vec();

    let (_final_beta_value, evaluation) = gray_code.fold(
        (starting_beta_value, starting_evaluation_acc),
        |(prev_beta_value, evaluation_acc), (index, (flipped_bit_index, flipped_bit_value))| {
            // For every bit i that is flipped, if it used to be a 1, then we
            // multiply by r_i^{-1} and multiply by (1 - r_i) to account for
            // this bit flip. NOTE: we subtract from n - 1 to account for the
            // fact that internally, these u32s are stored in little endian, but
            // our bookkeeping tables are stored in "big endian" indexing.
            let next_beta_value = if flipped_bit_value {
                prev_beta_value
                    * multiplier_if_flipped_bit_is_one[n - 1 - flipped_bit_index as usize]
            }
            // For every bit i that is flipped, if it used to be a 0, then we
            // multiply by (1 - r_i)^{-1} and multiply by r_i to account for
            // this bit flip.
            else {
                prev_beta_value
                    * multiplier_if_flipped_bit_is_zero[n - 1 - flipped_bit_index as usize]
            };
            // Multiply this by the appropriate MLE coefficient.
            let next_evaluation_acc = next_beta_value * mle.get(index as usize).unwrap();
            (next_beta_value, evaluation_acc + next_evaluation_acc)
        },
    );
    evaluation
}

/// This function non-destructively evaluates an MLE at a given point using the
/// gray codes iterator. Parallelized version that uses K threads.
pub fn evaluate_mle_at_a_point_gray_codes_parallel<F: Field, const K: usize>(
    mle: &MultilinearExtension<F>,
    point: &[F],
) -> F {
    let n = point.len();
    let mle_num_vars = mle.num_vars();
    assert_eq!(n, mle_num_vars);
    assert!(
        (1 << mle_num_vars) >= K,
        "cannot have more partitions than the length of MLE"
    );

    let starting_indices = (0..K)
        .map(|partition| partition * ((1 << mle_num_vars) / K))
        .collect_vec();
    let starting_gray_code_indices = starting_indices
        .iter()
        .map(|idx| GrayCodeIterator::get_gray_index(mle_num_vars, *idx as u32))
        .collect_vec();
    let starting_beta_values = starting_gray_code_indices
        .iter()
        .map(|gray_code| {
            BetaValues::compute_beta_over_challenge_and_index(point, *gray_code as usize)
        })
        .collect_vec();

    // This is the value that gets multiplied to the first MLE coefficient,
    // which is (1 - r_1) * (1 - r_2) * ... * (1 - r_n) where (r_1, ..., r_n) is
    // the point.
    let starting_evaluation_accs = starting_beta_values
        .iter()
        .zip(starting_gray_code_indices.iter())
        .map(|(beta_value, gray_code)| *beta_value * mle.get(*gray_code as usize).unwrap())
        .collect_vec();

    let gray_codes = starting_indices
        .iter()
        .enumerate()
        .map(|(partition, &starting_index)| {
            let end_iteration = if partition == K - 1 {
                None
            } else {
                Some(starting_indices[partition + 1] as u32)
            };
            GrayCodeIterator::new_at_index(mle_num_vars, starting_index as u32, end_iteration)
        })
        .collect_vec();

    let inverses = point
        .iter()
        .map(|elem| elem.invert().unwrap())
        .collect_vec();

    let one_minus_inverses = point
        .iter()
        .map(|elem| (F::ONE - elem).invert().unwrap())
        .collect_vec();

    // We simply compute the correct inverse and new multiplicative term for
    // each bit that is flipped in the beta value, and accumulate these by doing
    // an element-wise multiplication with the correct index of the MLE
    // coefficients. We precompute these so that during the scanning we only
    // need to do one multiplication instead of two
    let multiplier_if_flipped_bit_is_one = inverses
        .iter()
        .zip(point.iter())
        .map(|(inverse, point_elem)| *inverse * (F::ONE - point_elem))
        .collect_vec();

    let multiplier_if_flipped_bit_is_zero = one_minus_inverses
        .iter()
        .zip(point.iter())
        .map(|(one_minus_inverse, point_elem)| *one_minus_inverse * point_elem)
        .collect_vec();

    (0..K)
        .into_par_iter()
        .zip(gray_codes.into_par_iter())
        .map(|(partition, gray_code)| {
            let starting_beta_value = starting_beta_values[partition];
            let starting_evaluation_acc = starting_evaluation_accs[partition];
            let (_final_beta_value, evaluation) = gray_code.fold(
                (starting_beta_value, starting_evaluation_acc),
                |(prev_beta_value, evaluation_acc),
                 (index, (flipped_bit_index, flipped_bit_value))| {
                    // For every bit i that is flipped, if it used to be a 1,
                    // then we multiply by r_i^{-1} and multiply by (1 - r_i) to
                    // account for this bit flip. NOTE: we subtract from n - 1
                    // to account for the fact that internally, these u32s are
                    // stored in little endian, but our bookkeeping tables are
                    // stored in "big endian" indexing.
                    let next_beta_value = if flipped_bit_value {
                        prev_beta_value
                            * multiplier_if_flipped_bit_is_one[n - 1 - flipped_bit_index as usize]
                    }
                    // For every bit i that is flipped, if it used to be a 0,
                    // then we multiply by (1 - r_i)^{-1} and multiply by r_i to
                    // account for this bit flip.
                    else {
                        prev_beta_value
                            * multiplier_if_flipped_bit_is_zero[n - 1 - flipped_bit_index as usize]
                    };
                    // Multiply this by the appropriate MLE coefficient.
                    let next_evaluation_acc = next_beta_value * mle.get(index as usize).unwrap();
                    (next_beta_value, evaluation_acc + next_evaluation_acc)
                },
            );
            evaluation
        })
        .reduce(|| F::ZERO, |a, b| a + b)
}

/// Destructively evaluate an MLE at a point by using the `fix_variable`
/// algorithm iteratively until all of the variables have been bound.
///
pub fn evaluate_mle_destructive<F: Field>(mle: &mut MultilinearExtension<F>, point: &[F]) -> F {
    point.iter().for_each(|challenge| {
        mle.fix_variable(*challenge);
    });
    assert!(mle.is_fully_bound());
    mle.first()
}

#[cfg(test)]
mod tests {
    use ark_std::test_rng;
    use itertools::Itertools;
    use shared_types::{ff_field, Fr};

    use crate::{
        mle::evals::MultilinearExtension,
        utils::mle::{
            evaluate_mle_at_a_point_gray_codes_parallel,
            evaluate_mle_at_a_point_lexicographic_order, evaluate_mle_destructive,
            GrayCodeIterator,
        },
    };

    use super::evaluate_mle_at_a_point_gray_codes;

    #[test]
    fn test_gray_code_0_vars() {
        let mut gray_code_iterator = GrayCodeIterator::new(0);

        assert_eq!(gray_code_iterator.next(), None);
    }

    #[test]
    fn test_gray_code_iterator_len() {
        for n in 1..16 {
            assert_eq!(GrayCodeIterator::new(n).count(), (1 << n) - 1);
        }
    }

    // Note that this is the only test we have here that verifies that we're
    // actually using Gray Codes and not any of the other codes that could
    // satisfy the properties listed on the `test_gray_code_property` test.
    #[test]
    fn test_gray_code_3_vars() {
        let mut gray_code_iterator = GrayCodeIterator::new(3);

        assert_eq!(gray_code_iterator.next(), Some((1, (0, false))));
        assert_eq!(gray_code_iterator.next(), Some((3, (1, false))));
        assert_eq!(gray_code_iterator.next(), Some((2, (0, true))));
        assert_eq!(gray_code_iterator.next(), Some((6, (2, false))));
        assert_eq!(gray_code_iterator.next(), Some((7, (0, false))));
        assert_eq!(gray_code_iterator.next(), Some((5, (1, true))));
        assert_eq!(gray_code_iterator.next(), Some((4, (0, true))));
        assert_eq!(gray_code_iterator.next(), None);
    }

    // Property testing of `GrayCode` for values of `num_bits` of up to 15 (for
    // efficient testing). This test ensures that:
    //   1. The Hamming distance between consecutive codes is exactly 1.
    //   2. The index of the flipped bit correct.
    //   3. Each of the `2^num_codes` appear exactly once.
    //
    // Note that there are multiple codes that satisfy the above properties. Any
    // code with those properties can be used for computing MLEs in linear time.
    #[test]
    fn test_gray_code_property() {
        for n in 1..16 {
            let gray_code_iterator = GrayCodeIterator::new(n);

            let mut seen: Vec<bool> = vec![false; 1 << n];

            // Assume we're starting from 0.
            seen[0] = true;

            gray_code_iterator.fold(0, |prev, (cur, (idx, val))| {
                // Ensure `cur` has NOT been seen before.
                assert!(!seen[cur as usize]);
                seen[cur as usize] = true;

                let mask: u32 = 1 << idx;

                // This ensures that:
                //   1. The Hamming Distance between `prev` and `cur` is exactly
                //      1, because `mask` contains exactly one bit set to 1,
                //      and,
                //   2. The flipped bit is indeed in the `idx`-th position.
                assert_eq!(prev ^ cur, mask);

                // Ensures `val` is the previous value of the flipped bit.
                assert_eq!((prev & mask) >> idx, val as u32);

                cur
            });

            // Ensure all codes have been encountered during the iteration.
            assert!(seen.iter().all(|x| *x))
        }
    }

    /// there cannot be more threads than the length of the MLE
    #[test]
    #[should_panic]
    fn test_evaluate_mle_at_a_point_1_variable_gray_codes_parallel_more_threads_than_mle_length() {
        const K: usize = 3;
        let mut mle: MultilinearExtension<Fr> = vec![1, 2].into();
        let point = &[Fr::from(2)];
        let computed_evaluation = evaluate_mle_at_a_point_gray_codes(&mle, point);
        let expected_evaluation =
            evaluate_mle_at_a_point_gray_codes_parallel::<shared_types::Fr, K>(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_1_variable_gray_codes_parallel_1_thread() {
        const K: usize = 1;
        let mut mle: MultilinearExtension<Fr> = vec![1, 2].into();
        let point = &[Fr::from(2)];
        let computed_evaluation = evaluate_mle_at_a_point_gray_codes(&mle, point);
        let expected_evaluation =
            evaluate_mle_at_a_point_gray_codes_parallel::<shared_types::Fr, K>(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_1_variable_gray_codes_parallel_2_threads() {
        const K: usize = 2;
        let mut mle: MultilinearExtension<Fr> = vec![1, 2].into();
        let point = &[Fr::from(2)];
        let computed_evaluation = evaluate_mle_at_a_point_gray_codes(&mle, point);
        let expected_evaluation =
            evaluate_mle_at_a_point_gray_codes_parallel::<shared_types::Fr, K>(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_1_variable_gray_codes() {
        let mut mle: MultilinearExtension<Fr> = vec![1, 2].into();
        let point = &[Fr::from(2)];
        let computed_evaluation = evaluate_mle_at_a_point_gray_codes(&mle, point);
        let expected_evaluation = evaluate_mle_destructive(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_2_variable_gray_codes() {
        let mut mle: MultilinearExtension<Fr> = vec![1, 2, 1, 2].into();
        let point = &[Fr::from(2), Fr::from(3)];
        let computed_evaluation = evaluate_mle_at_a_point_gray_codes(&mle, point);
        let expected_evaluation = evaluate_mle_destructive(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_3_variable_gray_codes_random_parallel() {
        const K: usize = 5;
        let mut rng = test_rng();
        let mut mle = MultilinearExtension::new((0..8).map(|_| Fr::random(&mut rng)).collect());
        let point = &(0..3).map(|_| Fr::random(&mut rng)).collect_vec();
        let computed_evaluation =
            evaluate_mle_at_a_point_gray_codes_parallel::<shared_types::Fr, K>(&mle, point);
        let expected_evaluation = evaluate_mle_destructive(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_3_variable_gray_codes_random() {
        let mut rng = test_rng();
        let mut mle = MultilinearExtension::new((0..8).map(|_| Fr::random(&mut rng)).collect());
        let point = &(0..3).map(|_| Fr::random(&mut rng)).collect_vec();
        let computed_evaluation = evaluate_mle_at_a_point_gray_codes(&mle, point);
        let expected_evaluation = evaluate_mle_destructive(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_1_variable_lexicographic() {
        let mut mle: MultilinearExtension<Fr> = vec![1, 2].into();
        let point = &[Fr::from(2)];
        let computed_evaluation = evaluate_mle_at_a_point_lexicographic_order(&mle, point);
        let expected_evaluation = evaluate_mle_destructive(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_2_variable_lexicographic() {
        let mut mle: MultilinearExtension<Fr> = vec![1, 2, 1, 2].into();
        let point = &[Fr::from(2), Fr::from(3)];
        let computed_evaluation = evaluate_mle_at_a_point_lexicographic_order(&mle, point);
        let expected_evaluation = evaluate_mle_destructive(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    #[test]
    fn test_evaluate_mle_at_a_point_3_variable_lexicographic_random() {
        let mut rng = test_rng();
        let mut mle = MultilinearExtension::new((0..8).map(|_| Fr::random(&mut rng)).collect());
        let point = &(0..3).map(|_| Fr::random(&mut rng)).collect_vec();
        let computed_evaluation = evaluate_mle_at_a_point_lexicographic_order(&mle, point);
        let expected_evaluation = evaluate_mle_destructive(&mut mle, point);
        assert_eq!(computed_evaluation, expected_evaluation);
    }

    /// Ensure that all three methods of computing MLEs at a point produce the
    /// same result for random MLEs and points of sizes up to 15 bits.
    #[test]
    fn test_evaluation_equivalence() {
        for n in 1..16 {
            let num_vars = n;
            let num_evals = 1 << num_vars;

            let mut rng = test_rng();
            let mut mle =
                MultilinearExtension::new((0..num_evals).map(|_| Fr::random(&mut rng)).collect());
            let point = (0..num_vars).map(|_| Fr::random(&mut rng)).collect_vec();

            let gray_code_evaluation = evaluate_mle_at_a_point_gray_codes(&mle, &point);
            let lexicographic_evaluation =
                evaluate_mle_at_a_point_lexicographic_order(&mle, &point);
            let destructive_evaluation = evaluate_mle_destructive(&mut mle, &point);

            assert!(
                gray_code_evaluation == lexicographic_evaluation
                    && lexicographic_evaluation == destructive_evaluation
            );
        }
    }
}