1use ark_ff::PrimeField;
 2use ark_std::log2;
 3
 4/// Decompose an integer into a binary vector in little endian.
 5pub fn bit_decompose(input: u64, num_var: usize) -> Vec<bool> {
 6    let mut res = Vec::with_capacity(num_var);
 7    let mut i = input;
 8    for _ in 0..num_var {
 9        res.push(i & 1 == 1);
10        i >>= 1;
11    }
12    res
13}
14
15/// given the evaluation input `point` of the `index`-th polynomial,
16/// obtain the evaluation point in the merged polynomial
17pub fn gen_eval_point<F: PrimeField>(index: usize, index_len: usize, point: &[F]) -> Vec<F> {
18    let index_vec: Vec<F> = bit_decompose(index as u64, index_len)
19        .into_iter()
20        .map(|x| F::from(x))
21        .collect();
22    [point, &index_vec].concat()
23}
24
25/// Return the number of variables that one need for an MLE to
26/// batch the list of MLEs
27#[inline]
28pub fn get_batched_nv(num_var: usize, polynomials_len: usize) -> usize {
29    num_var + log2(polynomials_len) as usize
30}
31
32// Input index
33// - `i := (i_0, ...i_{n-1})`,
34// - `num_vars := n`
35// return three elements:
36// - `x0 := (i_1, ..., i_{n-1}, 0)`
37// - `x1 := (i_1, ..., i_{n-1}, 1)`
38// - `sign := i_0`
39#[inline]
40pub fn get_index(i: usize, num_vars: usize) -> (usize, usize, bool) {
41    let bit_sequence = bit_decompose(i as u64, num_vars);
42
43    // the last bit comes first here because of LE encoding
44    let x0 = project(&[[false].as_ref(), bit_sequence[..num_vars - 1].as_ref()].concat()) as usize;
45    let x1 = project(&[[true].as_ref(), bit_sequence[..num_vars - 1].as_ref()].concat()) as usize;
46
47    (x0, x1, bit_sequence[num_vars - 1])
48}
49
50/// Project a little endian binary vector into an integer.
51#[inline]
52pub(crate) fn project(input: &[bool]) -> u64 {
53    let mut res = 0;
54    for &e in input.iter().rev() {
55        res <<= 1;
56        res += e as u64;
57    }
58    res
59}
60
61#[cfg(test)]
62mod test {
63    use super::{bit_decompose, get_index, project};
64    use ark_std::{rand::RngCore, test_rng};
65
66    #[test]
67    fn test_decomposition() {
68        let mut rng = test_rng();
69        for _ in 0..100 {
70            let t = rng.next_u64();
71            let b = bit_decompose(t, 64);
72            let r = project(&b);
73            assert_eq!(t, r)
74        }
75    }
76
77    #[test]
78    fn test_get_index() {
79        let a = 0b1010;
80        let (x0, x1, sign) = get_index(a, 4);
81        assert_eq!(x0, 0b0100);
82        assert_eq!(x1, 0b0101);
83        assert!(sign);
84
85        let (x0, x1, sign) = get_index(a, 5);
86        assert_eq!(x0, 0b10100);
87        assert_eq!(x1, 0b10101);
88        assert!(!sign);
89
90        let a = 0b1111;
91        let (x0, x1, sign) = get_index(a, 4);
92        assert_eq!(x0, 0b1110);
93        assert_eq!(x1, 0b1111);
94        assert!(sign);
95    }
96}