chore!: Move recover_evaluations_in_domain_order into cosets.rs (#175)

* move `recover_evaluations_in_domain_order` into cosets.rs

* refactor code

cryptography/kzg_multi_open/src/fk20.rs

@@ -10,6 +10,7 @@ mod prover;
 mod toeplitz;
 mod verifier;
 
+pub use cosets::recover_evaluations_in_domain_order;
 pub use errors::VerifierError;
 pub use prover::{FK20Prover as Prover, Input as ProverInput};
 pub use verifier::FK20Verifier as Verifier;

cryptography/kzg_multi_open/src/fk20/cosets.rs

@@ -64,6 +64,93 @@ pub fn coset_gens(num_points: usize, num_cosets: usize, bit_reversed: bool) -> V
     coset_gens
 }
 
+/// Given a group of coset evaluations, this method will return/reorder the evaluations as if
+/// we evaluated them on the relevant extended domain.
+/// The coset indices returned can be used to locate the coset_evaluations in the new flattened order:
+///   - The idea is that a particular coset evaluation is evenly distributed across the set of flattened
+///     evaluations.
+///   Example:
+///     - Let's say we have `k` cosets. Each coset holds `m` values. Each coset will have an associated index.
+///     - Once this method has completed, we will be given a flattened set of evaluations where the
+///       `m` values in each coset are now a distance of `k` values apart from each other.
+///     - The first value that was in the first coset, will be in position `0`.
+///     - The second value that was in the first coset, will be in position `k`
+///     - The third value that was in the first coset, will be in position `2k`
+///     - The first value that was in the second coset, will NOT be in position `1`
+///        Instead it will be in position `t = reverse_bit_order(1, k)`.
+///     - This value of `t` is what the function returns alongside the flattened evaluations,
+///       allowing the caller to deduce the other positions.
+///
+//
+// Note: For evaluations that are missing, this method will fill these in with zeroes.
+//
+// Note: It is the callers responsibility to ensure that there are no duplicate
+// coset indices.
+pub fn recover_evaluations_in_domain_order(
+    domain_size: usize,
+    coset_indices: Vec<usize>,
+    coset_evaluations: Vec<Vec<Scalar>>,
+) -> Option<(Vec<usize>, Vec<Scalar>)> {
+    assert_eq!(coset_indices.len(), coset_evaluations.len());
+
+    if coset_indices.is_empty() {
+        return None;
+    }
+
+    let mut elements = vec![Scalar::from(0u64); domain_size];
+
+    // Check that each coset has the same size
+    let coset_len = coset_evaluations[0].len();
+    let same_len = coset_evaluations
+        .iter()
+        .all(|coset| coset.len() == coset_len);
+    if !same_len {
+        return None;
+    }
+
+    // Check that none of the indices are "out of bounds"
+    // This would result in the subsequent indexing operations to panic
+    //
+    // The greatest index we will be using is:
+    // `t = coset_index * coset_len`
+    // Let's denote the returned vectors length as `k`
+    // We want t < k
+    // => coset_index * coset_len < k
+    // => coset_index < k / coset_len
+    let index_bound = domain_size / coset_len;
+    let all_coset_indices_within_bound = coset_indices
+        .iter()
+        .all(|coset_index| *coset_index < index_bound);
+    if !all_coset_indices_within_bound {
+        return None;
+    }
+
+    // Iterate over each coset evaluation set and place the evaluations in the correct locations
+    for (&coset_index, coset_evals) in coset_indices.iter().zip(coset_evaluations) {
+        let start = coset_index * coset_len;
+        let end = start + coset_len;
+
+        elements[start..end].copy_from_slice(&coset_evals);
+    }
+
+    // Now bit reverse the result, so we get the evaluations as if we had just done
+    // and FFT on them. ie we computed the evaluation set and did not do a reverse bit order.
+    reverse_bit_order(&mut elements);
+
+    // The order of the coset indices in the returned vector will be different.
+    // The new indices of the cosets can be figured out by reverse bit ordering
+    // the existing indices.
+    let cosets_per_full_domain = domain_size / coset_len;
+    let num_bits_coset_per_full_domain = log2(cosets_per_full_domain as u32);
+
+    let new_coset_indices: Vec<_> = coset_indices
+        .into_iter()
+        .map(|rbo_coset_index| reverse_bits(rbo_coset_index, num_bits_coset_per_full_domain))
+        .collect();
+
+    Some((new_coset_indices, elements))
+}
+
 /// Generate k = `num_points / points_per_coset` amount of cosets, each containing `points_per_coset` points.
 /// The points in each coset will be roots of unity.
 /// For FK20, this is a hard requirement for efficient proof generation.
@@ -107,7 +194,10 @@ mod tests {
 
     use crate::fk20::{
         batch_toeplitz::transpose,
-        cosets::{generate_cosets, log2, reverse_bit_order, reverse_bits},
+        cosets::{
+            generate_cosets, log2, recover_evaluations_in_domain_order, reverse_bit_order,
+            reverse_bits,
+        },
         h_poly::take_every_nth,
     };
 
@@ -291,6 +381,87 @@ mod tests {
         assert_eq!(full_subgroup_set, cosets_flattened_set)
     }
 
+    #[test]
+    fn show_data_distribution_on_recover_evaluations_in_domain_order() {
+        use bls12_381::ff::Field;
+
+        const DOMAIN_SIZE: usize = 32;
+        const POINTS_PER_COSET: usize = 4;
+        const NUM_COSETS: usize = 8;
+
+        // Let's pretend that we've generated the coset_evaluations in bit-reversed order
+        let bit_reversed_evaluations: Vec<_> = (0..DOMAIN_SIZE)
+            .map(|i| Scalar::from((i + 1) as u64))
+            .collect();
+        let mut bit_reversed_coset_evaluations: Vec<Vec<Scalar>> = bit_reversed_evaluations
+            .chunks(POINTS_PER_COSET)
+            .map(|chunk| chunk.to_vec())
+            .collect();
+
+        // We have 32 values and 4 points per coset, so we have 8 cosets.
+        let coset_indices: Vec<_> = (0..NUM_COSETS).collect();
+
+        // Zero out the first coset
+        let first_coset = &mut bit_reversed_coset_evaluations[0];
+        for evaluation in first_coset {
+            *evaluation = Scalar::ZERO
+        }
+        // Zero out the 4th coset
+        let fourth_coset = &mut bit_reversed_coset_evaluations[3];
+        for evaluation in fourth_coset {
+            *evaluation = Scalar::ZERO
+        }
+
+        // Now let's simulate the first and fourth coset missing
+        let coset_evaluations_missing: Vec<_> = bit_reversed_coset_evaluations
+            .into_iter()
+            .enumerate()
+            .filter(|(i, _)| *i != 0 && *i != 3)
+            .map(|(_, coset)| coset)
+            .collect();
+        let coset_indices_missing: Vec<_> = coset_indices
+            .into_iter()
+            .enumerate()
+            .filter(|(i, _)| *i != 0 && *i != 3)
+            .map(|(_, coset)| coset)
+            .collect();
+
+        let (coset_indices_normal_order, coset_evaluations_normal_order) =
+            recover_evaluations_in_domain_order(
+                DOMAIN_SIZE,
+                coset_indices_missing,
+                coset_evaluations_missing,
+            )
+            .unwrap();
+
+        let missing_coset_index_0 = reverse_bits(0, log2(NUM_COSETS as u32));
+        let missing_coset_index_3 = reverse_bits(3, log2(NUM_COSETS as u32));
+
+        // Let's show what happened to the evaluations in the first and fourth cosets which were missing
+        //
+        // It was in the first coset, so the idea is that there will be zeroes in every `rbo(0) + NUM_COSET * i` position
+        // where i ranges from 0 to NUM_COSET.
+        //
+        // The same is also the case for the fourth missing coset, ie we would also have 0s in every `rbo(4) + NUM_COSET * i` position.
+        //
+        // In general, if the `k`th coset is missing, then this function will return the evaluations with 0s
+        // in the `rbo(k) + NUM_COSET  * i`'th positions.
+        for block in coset_evaluations_normal_order.chunks(8) {
+            for (index, element) in block.into_iter().enumerate() {
+                if index == missing_coset_index_0 || index == missing_coset_index_3 {
+                    assert_eq!(*element, Scalar::ZERO)
+                } else {
+                    assert_ne!(*element, Scalar::ZERO)
+                }
+            }
+        }
+
+        // We also note that the coset indices that are returned will not have `missing_coset_index_3` or
+        // missing_coset_index_0
+        assert!(!coset_indices_normal_order.contains(&missing_coset_index_0));
+        assert!(!coset_indices_normal_order.contains(&missing_coset_index_3));
+    }
+
     #[test]
     fn bit_reverse_fuzz() {
         fn naive_bit_reverse(n: u32, l: u32) -> u32 {