fix: removed array_contains and handled partial subarray case for get_fixed_subarray

plonky2x/core/src/frontend/vars/array.rs

@@ -9,7 +9,7 @@ use plonky2::iop::challenger::RecursiveChallenger;
 use plonky2::iop::target::Target;
 use serde::{Deserialize, Serialize};
 
-use super::{BoolVariable, ByteVariable, CircuitVariable, ValueStream, Variable, VariableStream};
+use super::{ByteVariable, CircuitVariable, ValueStream, Variable, VariableStream};
 use crate::backend::circuit::PlonkParameters;
 use crate::frontend::builder::CircuitBuilder;
 use crate::frontend::hint::simple::hint::Hint;
@@ -180,7 +180,8 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
     /// Given an `array` of variables, and a dynamic `index` start_idx, returns
     /// `array[start_idx..start_idx+sub_array_size]` as an `array`.
     ///
-    /// `seed` is used to generate randomness for the proof.
+    /// `seed` is used to generate randomness for the proof.  Note that the seed should be the
+    /// Fiat-Shamir of the entire array and claimed subarray.
     ///
     /// The security of each challenge is log2(field_size) - log2(array_size), so the total security
     /// is (log2(field_size) - log2(array_size)) * num_loops.
@@ -252,6 +253,7 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
             let one: Variable = self.one();
 
             let mut accumulator1 = self.zero::<Variable>();
+            let mut subarray_size: Variable = self.zero();
 
             // r is the source of randomness from the challenger for this loop.
             let mut r = one;
@@ -266,6 +268,8 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
                 let at_end_idx = self.is_equal(idx, end_idx);
                 within_sub_array = self.select(at_end_idx, false_v, within_sub_array);
 
+                subarray_size = self.add(subarray_size, within_sub_array.variable);
+
                 // If within the subarray, multiply the current r by the challenge.
                 let multiplier = self.select(within_sub_array, challenges[i], one);
                 // For subarray[i], the multiplier should be r^i. i is the index within the subarray.
@@ -280,6 +284,11 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
                 accumulator1 = self.add(accumulator1, temp_accum);
             }
 
+            // Assert that the returned subarray's length is == SUB_ARRAY_SIZE.
+            let expected_subarray_size =
+                self.constant(L::Field::from_canonical_usize(SUB_ARRAY_SIZE));
+            self.assert_is_equal(subarray_size, expected_subarray_size);
+
             let mut accumulator2 = self.zero();
             let mut r = one;
             for j in 0..SUB_ARRAY_SIZE {
@@ -293,19 +302,6 @@ impl<L: PlonkParameters<D>, const D: usize> CircuitBuilder<L, D> {
 
         sub_array
     }
-
-    pub fn array_contains<V: CircuitVariable>(&mut self, array: &[V], element: V) -> BoolVariable {
-        assert!(array.len() < 1 << 16);
-        let mut accumulator = self.constant::<Variable>(L::Field::from_canonical_usize(0));
-
-        for i in 0..array.len() {
-            let element_equal = self.is_equal(array[i].clone(), element.clone());
-            accumulator = self.add(accumulator, element_equal.variable);
-        }
-
-        let one = self.constant::<Variable>(L::Field::from_canonical_usize(1));
-        self.is_equal(one, accumulator)
-    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -511,4 +507,49 @@ mod tests {
             expected_sub_array
         );
     }
+
+    #[test]
+    #[should_panic]
+    fn test_get_fixed_subarray_bad_case() {
+        utils::setup_logger();
+        type F = GoldilocksField;
+        const ARRAY_SIZE: usize = 12800;
+        const SUB_ARRAY_SIZE: usize = 3200;
+        const START_IDX: usize = 12000;
+
+        let mut builder = DefaultBuilder::new();
+
+        let array = builder.read::<ArrayVariable<Variable, ARRAY_SIZE>>();
+        let start_idx = builder.constant(F::from_canonical_usize(START_IDX));
+        let seed = builder.read::<Bytes32Variable>();
+        let result = builder.get_fixed_subarray::<ARRAY_SIZE, SUB_ARRAY_SIZE>(
+            &array,
+            start_idx,
+            &seed.as_bytes(),
+        );
+        builder.write(result);
+
+        let circuit = builder.build();
+
+        // The last 20 elements are dummy
+        let mut rng = OsRng;
+        let mut array_input = [F::default(); ARRAY_SIZE];
+        for elem in array_input.iter_mut() {
+            *elem = F::from_canonical_u64(rng.gen());
+        }
+
+        let mut seed_input = [0u8; 15];
+        for elem in seed_input.iter_mut() {
+            *elem = rng.gen();
+        }
+
+        let mut input = circuit.input();
+        input.write::<ArrayVariable<Variable, ARRAY_SIZE>>(array_input.to_vec());
+        input.write::<Bytes32Variable>(bytes32!(
+            "0x7c38fc8356aa20394c7f538e3cee3f924e6d9252494c8138d1a6aabfc253118f"
+        ));
+
+        let (proof, output) = circuit.prove(&input);
+        circuit.verify(&proof, &input, &output);
+    }
 }