feat: add select_array_by_indicator to GateInstructions

halo2-base/src/gates/flex_gate.rs

@@ -727,6 +727,35 @@ pub trait GateInstructions<F: ScalarField> {
         self.select_by_indicator(ctx, cells, ind)
     }
 
+    /// `array2d` is an array of fixed length arrays.
+    /// Assumes:
+    /// * `array2d.len() == indicator.len()`
+    /// * `array2d[i].len() == array2d[j].len()` for all `i,j`.
+    /// * the values of `indicator` are boolean and that `indicator` has at most one `1` bit.
+    /// * the lengths of `array2d` and `indicator` are the same.
+    ///
+    /// Returns the "dot product" of `array2d` with `indicator` as a fixed length (1d) array of length `array2d[0].len()`.
+    fn select_array_by_indicator<AR, AV>(
+        &self,
+        ctx: &mut Context<F>,
+        array2d: &[AR],
+        indicator: &[AssignedValue<F>],
+    ) -> Vec<AssignedValue<F>>
+    where
+        AR: AsRef<[AV]>,
+        AV: AsRef<AssignedValue<F>>,
+    {
+        (0..array2d[0].as_ref().len())
+            .map(|j| {
+                self.select_by_indicator(
+                    ctx,
+                    array2d.iter().map(|array_i| *array_i.as_ref()[j].as_ref()),
+                    indicator.iter().copied(),
+                )
+            })
+            .collect()
+    }
+
     /// Constrains that a cell is equal to 0 and returns `1` if `a = 0`, otherwise `0`.
     ///
     /// Defines a vertical gate of form `| out | a | inv | 1 | 0 | a | out | 0 |`, where out = 1 if a = 0, otherwise out = 0.

halo2-base/src/gates/tests/flex_gate.rs

@@ -4,6 +4,7 @@ use crate::utils::biguint_to_fe;
 use crate::utils::testing::base_test;
 use crate::QuantumCell::Witness;
 use crate::{gates::flex_gate::GateInstructions, QuantumCell};
+use itertools::Itertools;
 use num_bigint::BigUint;
 use test_case::test_case;
 
@@ -156,6 +157,19 @@ pub fn test_select_from_idx(array: Vec<QuantumCell<Fr>>, idx: QuantumCell<Fr>) -
     base_test().run_gate(|ctx, chip| *chip.select_from_idx(ctx, array, idx).value())
 }
 
+#[test_case(vec![vec![1,2,3], vec![4,5,6], vec![7,8,9]].into_iter().map(|a| a.into_iter().map(Fr::from).collect_vec()).collect_vec(), 
+Fr::from(1) => 
+[4,5,6].map(Fr::from).to_vec(); 
+"select_array_by_indicator(1): [[1,2,3], [4,5,6], [7,8,9]] -> [4,5,6]")]
+pub fn test_select_array_by_indicator(array2d: Vec<Vec<Fr>>, idx: Fr) -> Vec<Fr> {
+    base_test().run_gate(|ctx, chip| {
+        let array2d = array2d.into_iter().map(|a| ctx.assign_witnesses(a)).collect_vec();
+        let idx = ctx.load_witness(idx);
+        let ind = chip.idx_to_indicator(ctx, idx, array2d.len());
+        chip.select_array_by_indicator(ctx, &array2d, &ind).iter().map(|a| *a.value()).collect()
+    })
+}
+
 #[test_case(Fr::zero() => Fr::from(1); "is_zero(): 0 -> 1")]
 pub fn test_is_zero(input: Fr) -> Fr {
     base_test().run_gate(|ctx, chip| {

halo2-base/src/lib.rs

@@ -135,6 +135,12 @@ impl<F: ScalarField> AssignedValue<F> {
     }
 }
 
+impl<F: ScalarField> AsRef<AssignedValue<F>> for AssignedValue<F> {
+    fn as_ref(&self) -> &AssignedValue<F> {
+        self
+    }
+}
+
 /// Represents a single thread of an execution trace.
 /// * We keep the naming [Context] for historical reasons.
 #[derive(Clone, Debug)]