chore(backend): implement some circuit op logic (#165)

* Initial circuit op work

* Fix copy op

* Add more ops

* Fixes

* Code review

src/backends/plonky2/circuits/common.rs

@@ -1,8 +1,12 @@
 //! Common functionality to build Pod circuits with plonky2
 
+use crate::backends::plonky2::basetypes::D;
 use crate::backends::plonky2::mock::mainpod::Statement;
 use crate::backends::plonky2::mock::mainpod::{Operation, OperationArg};
-use crate::middleware::{Params, StatementArg, ToFields, Value, F, HASH_SIZE, VALUE_SIZE};
+use crate::middleware::{
+    NativeOperation, NativePredicate, Params, Predicate, StatementArg, ToFields, Value, F,
+    HASH_SIZE, VALUE_SIZE,
+};
 use crate::middleware::{OPERATION_ARG_F_LEN, STATEMENT_ARG_F_LEN};
 use anyhow::Result;
 use plonky2::field::extension::Extendable;
@@ -11,13 +15,42 @@ use plonky2::hash::hash_types::RichField;
 use plonky2::iop::target::{BoolTarget, Target};
 use plonky2::iop::witness::{PartialWitness, WitnessWrite};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
-use std::iter;
+use std::{array, iter};
+
+pub const CODE_SIZE: usize = HASH_SIZE + 2;
 
 #[derive(Copy, Clone)]
 pub struct ValueTarget {
     pub elements: [Target; VALUE_SIZE],
 }
 
+impl ValueTarget {
+    pub fn zero(builder: &mut CircuitBuilder<F, D>) -> Self {
+        Self {
+            elements: [builder.zero(); VALUE_SIZE],
+        }
+    }
+
+    pub fn one(builder: &mut CircuitBuilder<F, D>) -> Self {
+        Self {
+            elements: array::from_fn(|i| {
+                if i == 0 {
+                    builder.one()
+                } else {
+                    builder.zero()
+                }
+            }),
+        }
+    }
+
+    pub fn from_slice(xs: &[Target]) -> Self {
+        assert_eq!(xs.len(), VALUE_SIZE);
+        Self {
+            elements: array::from_fn(|i| xs[i]),
+        }
+    }
+}
+
 #[derive(Clone)]
 pub struct StatementTarget {
     pub predicate: [Target; Params::predicate_size()],
@@ -25,12 +58,24 @@ pub struct StatementTarget {
 }
 
 impl StatementTarget {
-    pub fn to_flattened(&self) -> Vec<Target> {
+    pub fn new_native(
-        self.predicate
+        builder: &mut CircuitBuilder<F, D>,
-            .iter()
+        params: &Params,
-            .chain(self.args.iter().flatten())
+        predicate: NativePredicate,
-            .cloned()
+        args: &[[Target; STATEMENT_ARG_F_LEN]],
-            .collect()
+    ) -> Self {
+        let predicate_vec = builder.constants(&Predicate::Native(predicate).to_fields(params));
+        Self {
+            predicate: array::from_fn(|i| predicate_vec[i]),
+            args: args
+                .iter()
+                .map(|arg| *arg)
+                .chain(
+                    iter::repeat([builder.zero(); STATEMENT_ARG_F_LEN])
+                        .take(params.max_statement_args - args.len()),
+                )
+                .collect(),
+        }
     }
 
     pub fn set_targets(
@@ -51,6 +96,16 @@ impl StatementTarget {
         }
         Ok(())
     }
+
+    pub fn has_native_type(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        params: &Params,
+        t: NativePredicate,
+    ) -> BoolTarget {
+        let st_code = builder.constants(&Predicate::Native(t).to_fields(params));
+        builder.is_equal_slice(&self.predicate, &st_code)
+    }
 }
 
 // TODO: Implement Operation::to_field to determine the size of each element
@@ -79,6 +134,49 @@ impl OperationTarget {
         }
         Ok(())
     }
+
+    pub fn has_native_type(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        t: NativeOperation,
+    ) -> BoolTarget {
+        let one = builder.one();
+        let op_is_native = builder.is_equal(self.op_type[0], one);
+        let op_code = builder.constant(F::from_canonical_u64(t as u64));
+        let op_code_matches = builder.is_equal(self.op_type[1], op_code);
+        builder.and(op_is_native, op_code_matches)
+    }
+}
+
+/// Trait for target structs that may be converted to and from vectors
+/// of targets.
+pub trait Flattenable {
+    fn flatten(&self) -> Vec<Target>;
+    fn from_flattened(vs: &[Target]) -> Self;
+}
+
+impl Flattenable for StatementTarget {
+    fn flatten(&self) -> Vec<Target> {
+        self.predicate
+            .iter()
+            .chain(self.args.iter().flatten())
+            .cloned()
+            .collect()
+    }
+
+    fn from_flattened(v: &[Target]) -> Self {
+        let num_args = (v.len() - Params::predicate_size()) / STATEMENT_ARG_F_LEN;
+        assert_eq!(
+            v.len(),
+            Params::predicate_size() + num_args * STATEMENT_ARG_F_LEN
+        );
+        let predicate: [Target; Params::predicate_size()] = array::from_fn(|i| v[i]);
+        let args = (0..num_args)
+            .map(|i| array::from_fn(|j| v[Params::predicate_size() + i * STATEMENT_ARG_F_LEN + j]))
+            .collect();
+
+        Self { predicate, args }
+    }
 }
 
 pub trait CircuitBuilderPod<F: RichField + Extendable<D>, const D: usize> {
@@ -91,11 +189,27 @@ pub trait CircuitBuilderPod<F: RichField + Extendable<D>, const D: usize> {
     fn select_bool(&mut self, b: BoolTarget, x: BoolTarget, y: BoolTarget) -> BoolTarget;
     fn constant_value(&mut self, v: Value) -> ValueTarget;
     fn is_equal_slice(&mut self, xs: &[Target], ys: &[Target]) -> BoolTarget;
+
+    // Convenience methods for checking values.
+    /// Checks whether `xs` is right-padded with 0s so as to represent a `Value`.
+    fn statement_arg_is_value(&mut self, xs: &[Target]) -> BoolTarget;
+    /// Checks whether `x < y` if `b` is true. This involves checking
+    /// that `x` and `y` each consist of two `u32` limbs.
+    fn assert_less_if(&mut self, b: BoolTarget, x: ValueTarget, y: ValueTarget);
+
+    // Convenience methods for accessing and connecting elements of
+    // (vectors of) flattenables.
+    fn vec_ref<T: Flattenable>(&mut self, ts: &[T], i: Target) -> T;
+    fn select_flattenable<T: Flattenable>(&mut self, b: BoolTarget, x: &T, y: &T) -> T;
+    fn connect_flattenable<T: Flattenable>(&mut self, xs: &T, ys: &T);
+    fn is_equal_flattenable<T: Flattenable>(&mut self, xs: &T, ys: &T) -> BoolTarget;
+
+    // Convenience methods for Boolean into-iters.
+    fn all(&mut self, xs: impl IntoIterator<Item = BoolTarget>) -> BoolTarget;
+    fn any(&mut self, xs: impl IntoIterator<Item = BoolTarget>) -> BoolTarget;
 }
 
-impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilderPod<F, D>
+impl CircuitBuilderPod<F, D> for CircuitBuilder<F, D> {
-    for CircuitBuilder<F, D>
-{
     fn connect_slice(&mut self, xs: &[Target], ys: &[Target]) {
         assert_eq!(xs.len(), ys.len());
         for (x, y) in xs.iter().zip(ys.iter()) {
@@ -157,4 +271,110 @@ impl<F: RichField + Extendable<D>, const D: usize> CircuitBuilderPod<F, D>
             self.and(ok, is_eq)
         })
     }
+
+    fn statement_arg_is_value(&mut self, xs: &[Target]) -> BoolTarget {
+        let zeros = iter::repeat(self.zero())
+            .take(STATEMENT_ARG_F_LEN - VALUE_SIZE)
+            .collect::<Vec<_>>();
+        self.is_equal_slice(&xs[VALUE_SIZE..], &zeros)
+    }
+
+    fn assert_less_if(&mut self, b: BoolTarget, x: ValueTarget, y: ValueTarget) {
+        const NUM_BITS: usize = 32;
+
+        // Lt assertion with 32-bit range check.
+        let assert_limb_lt = |builder: &mut Self, x, y| {
+            // Check that targets fit within `NUM_BITS` bits.
+            builder.range_check(x, NUM_BITS);
+            builder.range_check(y, NUM_BITS);
+            // Check that `y-1-x` fits within `NUM_BITS` bits.
+            let one = builder.one();
+            let y_minus_one = builder.sub(y, one);
+            let expr = builder.sub(y_minus_one, x);
+            builder.range_check(expr, NUM_BITS);
+        };
+
+        // If b is false, replace `x` and `y` with dummy values.
+        let zero = ValueTarget::zero(self);
+        let one = ValueTarget::one(self);
+        let x = self.select_value(b, x, zero);
+        let y = self.select_value(b, y, one);
+
+        // `x` and `y` should only have two limbs each.
+        x.elements
+            .into_iter()
+            .skip(2)
+            .for_each(|l| self.assert_zero(l));
+        y.elements
+            .into_iter()
+            .skip(2)
+            .for_each(|l| self.assert_zero(l));
+
+        let big_limbs_eq = self.is_equal(x.elements[1], y.elements[1]);
+        let lhs = self.select(big_limbs_eq, x.elements[0], x.elements[1]);
+        let rhs = self.select(big_limbs_eq, y.elements[0], y.elements[1]);
+        assert_limb_lt(self, lhs, rhs);
+    }
+
+    fn vec_ref<T: Flattenable>(&mut self, ts: &[T], i: Target) -> T {
+        // TODO: Revisit this when we need more than 64 statements.
+        let vector_ref = |builder: &mut CircuitBuilder<F, D>, v: &[Target], i| {
+            assert!(v.len() <= 64);
+            builder.random_access(i, v.to_vec())
+        };
+        let matrix_row_ref = |builder: &mut CircuitBuilder<F, D>, m: &[Vec<Target>], i| {
+            let num_rows = m.len();
+            let num_columns = m
+                .get(0)
+                .map(|row| {
+                    let row_len = row.len();
+                    assert!(m.iter().all(|row| row.len() == row_len));
+                    row_len
+                })
+                .unwrap_or(0);
+            (0..num_columns)
+                .map(|j| {
+                    vector_ref(
+                        builder,
+                        &(0..num_rows).map(|i| m[i][j]).collect::<Vec<_>>(),
+                        i,
+                    )
+                })
+                .collect::<Vec<_>>()
+        };
+
+        let flattened_ts = ts.iter().map(|t| t.flatten()).collect::<Vec<_>>();
+        T::from_flattened(&matrix_row_ref(self, &flattened_ts, i))
+    }
+
+    fn select_flattenable<T: Flattenable>(&mut self, b: BoolTarget, x: &T, y: &T) -> T {
+        let flattened_x = x.flatten();
+        let flattened_y = y.flatten();
+
+        T::from_flattened(
+            &iter::zip(flattened_x, flattened_y)
+                .map(|(x, y)| self.select(b, x, y))
+                .collect::<Vec<_>>(),
+        )
+    }
+
+    fn connect_flattenable<T: Flattenable>(&mut self, xs: &T, ys: &T) {
+        self.connect_slice(&xs.flatten(), &ys.flatten())
+    }
+
+    fn is_equal_flattenable<T: Flattenable>(&mut self, xs: &T, ys: &T) -> BoolTarget {
+        self.is_equal_slice(&xs.flatten(), &ys.flatten())
+    }
+
+    fn all(&mut self, xs: impl IntoIterator<Item = BoolTarget>) -> BoolTarget {
+        xs.into_iter()
+            .reduce(|a, b| self.and(a, b))
+            .unwrap_or(self._true())
+    }
+
+    fn any(&mut self, xs: impl IntoIterator<Item = BoolTarget>) -> BoolTarget {
+        xs.into_iter()
+            .reduce(|a, b| self.or(a, b))
+            .unwrap_or(self._false())
+    }
 }

src/backends/plonky2/circuits/mainpod.rs

@@ -19,15 +19,17 @@ use crate::backends::plonky2::basetypes::{Hash, Value, D, EMPTY_HASH, EMPTY_VALU
 use crate::backends::plonky2::circuits::common::{
     CircuitBuilderPod, OperationTarget, StatementTarget, ValueTarget,
 };
-use crate::backends::plonky2::primitives::merkletree::{MerkleProof, MerkleTree};
+use crate::backends::plonky2::primitives::merkletree::MerkleTree;
 use crate::backends::plonky2::primitives::merkletree::{
     MerkleProofExistenceGate, MerkleProofExistenceTarget,
 };
 use crate::middleware::{
-    hash_str, AnchoredKey, NativeOperation, NativePredicate, Params, PodType, Predicate, Statement,
+    hash_str, AnchoredKey, NativeOperation, NativePredicate, Params, PodType, Statement,
-    StatementArg, ToFields, KEY_TYPE, SELF, STATEMENT_ARG_F_LEN,
+    StatementArg, ToFields, KEY_TYPE, SELF,
 };
 
+use super::common::Flattenable;
+
 //
 // SignedPod verification
 //
@@ -137,91 +139,208 @@ impl OperationVerifyGate {
     ) -> Result<OperationVerifyTarget> {
         let _true = builder._true();
         let _false = builder._false();
-        let one = builder.constant(F::ONE);
 
         // Verify that the operation `op` correctly generates the statement `st`.  The operation
         // can reference any of the `prev_statements`.
+        // TODO: Clean this up.
+        let resolved_op_args = if prev_statements.len() == 0 {
+            vec![]
+        } else {
+            op.args
+                .iter()
+                .flatten()
+                .map(|&i| builder.vec_ref(prev_statements, i))
+                .collect::<Vec<_>>()
+        };
+
         // The verification may require aux data which needs to be stored in the
         // `OperationVerifyTarget` so that we can set during witness generation.
 
         // For now only support native operations
-        builder.connect(op.op_type[0], one);
+        // Op checks to carry out. Each 'eval_X' should be thought of
-        let native_op = op.op_type[1];
+        // as 'eval' restricted to the op of type X, where the
+        // returned target is `false` if the input targets lie outside
+        // of the domain.
+        let op_checks = vec![
+            vec![
+                self.eval_none(builder, st, op),
+                self.eval_new_entry(builder, st, op, prev_statements),
+            ],
+            // Skip these if there are no resolved op args
+            if resolved_op_args.len() == 0 {
+                vec![]
+            } else {
+                vec![
+                    self.eval_copy(builder, st, op, &resolved_op_args)?,
+                    self.eval_eq_from_entries(builder, st, op, &resolved_op_args),
+                    self.eval_lt_from_entries(builder, st, op, &resolved_op_args),
+                ]
+            },
+        ]
+        .concat();
 
-        let mut op_flags = Vec::new();
+        let ok = builder.any(op_checks);
-        let op_none = builder.constant(F::from_canonical_u64(NativeOperation::None as u64));
-        let is_none = builder.is_equal(native_op, op_none);
-        op_flags.push(is_none);
-        let op_new_entry =
-            builder.constant(F::from_canonical_u64(NativeOperation::NewEntry as u64));
-        let is_new_entry = builder.is_equal(native_op, op_new_entry);
-        op_flags.push(is_new_entry);
-        let op_copy_statement =
-            builder.constant(F::from_canonical_u64(NativeOperation::CopyStatement as u64));
-        let is_copy_statement = builder.is_equal(native_op, op_copy_statement);
-        op_flags.push(is_copy_statement);
-        let op_eq_from_entries = builder.constant(F::from_canonical_u64(
-            NativeOperation::EqualFromEntries as u64,
-        ));
-        let is_eq_from_entries = builder.is_equal(native_op, op_eq_from_entries);
-        op_flags.push(is_eq_from_entries);
-        let op_lt_from_entries =
-            builder.constant(F::from_canonical_u64(NativeOperation::LtFromEntries as u64));
-        let is_lt_from_entries = builder.is_equal(native_op, op_lt_from_entries);
-        op_flags.push(is_lt_from_entries);
-        let op_not_contains_from_entries = builder.constant(F::from_canonical_u64(
-            NativeOperation::NotContainsFromEntries as u64,
-        ));
-        let is_not_contains_from_entries =
-            builder.is_equal(native_op, op_not_contains_from_entries);
-        op_flags.push(is_not_contains_from_entries);
-
-        // One supported operation must be used.  We sum all operation flags and expect the result
-        // to be 1.  Since the flags are boolean and at most one of them is true the sum is
-        // equivalent to the OR.
-        let or_op_flags = op_flags
-            .iter()
-            .map(|b| b.target)
-            .fold(_false.target, |acc, x| builder.add(acc, x));
-        builder.connect(or_op_flags, _true.target);
-
-        let ok = builder._true();
-        let none_ok = self.eval_none(builder, st, op);
-        let ok = builder.select_bool(is_none, none_ok, ok);
-        let new_entry_ok = self.eval_new_entry(builder, st, op);
-        let ok = builder.select_bool(is_new_entry, new_entry_ok, ok);
 
         builder.connect(ok.target, _true.target);
 
         Ok(OperationVerifyTarget {})
     }
 
+    fn eval_eq_from_entries(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        st: &StatementTarget,
+        op: &OperationTarget,
+        resolved_op_args: &[StatementTarget],
+    ) -> BoolTarget {
+        let op_code_ok = op.has_native_type(builder, NativeOperation::EqualFromEntries);
+
+        // Expect 2 op args of type `ValueOf`.
+        let op_arg_type_checks = resolved_op_args
+            .iter()
+            .take(2)
+            .map(|op_arg| op_arg.has_native_type(builder, &self.params, NativePredicate::ValueOf))
+            .collect::<Vec<_>>();
+        let op_arg_types_ok = builder.all(op_arg_type_checks);
+
+        // The values embedded in the op args must match, the last
+        // `STATEMENT_ARG_F_LEN - VALUE_SIZE` slots of each being 0.
+        let arg1_value = resolved_op_args[0].args[1];
+        let arg2_value = resolved_op_args[1].args[1];
+        let op_arg_range_checks = [
+            builder.statement_arg_is_value(&arg1_value),
+            builder.statement_arg_is_value(&arg2_value),
+        ];
+        let op_arg_range_ok = builder.all(op_arg_range_checks);
+        let op_args_eq =
+            builder.is_equal_slice(&arg1_value[..VALUE_SIZE], &arg2_value[..VALUE_SIZE]);
+
+        let arg1_key = resolved_op_args[0].args[0];
+        let arg2_key = resolved_op_args[1].args[0];
+        let expected_statement = StatementTarget::new_native(
+            builder,
+            &self.params,
+            NativePredicate::Equal,
+            &[arg1_key, arg2_key],
+        );
+        let st_ok = builder.is_equal_flattenable(st, &expected_statement);
+
+        builder.all([
+            op_code_ok,
+            op_arg_types_ok,
+            op_arg_range_ok,
+            op_args_eq,
+            st_ok,
+        ])
+    }
+
+    fn eval_lt_from_entries(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        st: &StatementTarget,
+        op: &OperationTarget,
+        resolved_op_args: &[StatementTarget],
+    ) -> BoolTarget {
+        let op_code_ok = op.has_native_type(builder, NativeOperation::LtFromEntries);
+
+        // Expect 2 op args of type `ValueOf`.
+        let op_arg_type_checks = resolved_op_args
+            .iter()
+            .take(2)
+            .map(|op_arg| op_arg.has_native_type(builder, &self.params, NativePredicate::ValueOf))
+            .collect::<Vec<_>>();
+        let op_arg_types_ok = builder.all(op_arg_type_checks);
+
+        // The values embedded in the op args must satisfy `<`, the
+        // last `STATEMENT_ARG_F_LEN - VALUE_SIZE` slots of each being
+        // 0.
+        let arg1_value = resolved_op_args[0].args[1];
+        let arg2_value = resolved_op_args[1].args[1];
+        let op_arg_range_checks = [&arg1_value, &arg2_value]
+            .into_iter()
+            .map(|x| builder.statement_arg_is_value(x))
+            .collect::<Vec<_>>();
+        let op_arg_range_ok = builder.all(op_arg_range_checks);
+        builder.assert_less_if(
+            op_code_ok,
+            ValueTarget::from_slice(&arg1_value[..VALUE_SIZE]),
+            ValueTarget::from_slice(&arg2_value[..VALUE_SIZE]),
+        );
+
+        let arg1_key = resolved_op_args[0].args[0];
+        let arg2_key = resolved_op_args[1].args[0];
+        let expected_statement = StatementTarget::new_native(
+            builder,
+            &self.params,
+            NativePredicate::Lt,
+            &[arg1_key, arg2_key],
+        );
+        let st_ok = builder.is_equal_flattenable(st, &expected_statement);
+
+        builder.all([op_code_ok, op_arg_types_ok, op_arg_range_ok, st_ok])
+    }
+
     fn eval_none(
         &self,
         builder: &mut CircuitBuilder<F, D>,
         st: &StatementTarget,
-        _op: &OperationTarget,
+        op: &OperationTarget,
     ) -> BoolTarget {
-        let expected_statement_flattened =
+        let op_code_ok = op.has_native_type(builder, NativeOperation::None);
-            builder.constants(&Statement::None.to_fields(&self.params));
+
-        builder.is_equal_slice(&st.to_flattened(), &expected_statement_flattened)
+        let expected_statement =
+            StatementTarget::new_native(builder, &self.params, NativePredicate::None, &[]);
+        let st_ok = builder.is_equal_flattenable(st, &expected_statement);
+
+        builder.all([op_code_ok, st_ok])
     }
 
     fn eval_new_entry(
         &self,
         builder: &mut CircuitBuilder<F, D>,
         st: &StatementTarget,
-        _op: &OperationTarget,
+        op: &OperationTarget,
+        prev_statements: &[StatementTarget],
     ) -> BoolTarget {
-        let value_of_st = &Statement::ValueOf(AnchoredKey(SELF, EMPTY_HASH), EMPTY_VALUE);
+        let op_code_ok = op.has_native_type(builder, NativeOperation::NewEntry);
-        let expected_predicate =
+
-            builder.constants(&Predicate::Native(NativePredicate::ValueOf).to_fields(&self.params));
+        let st_code_ok = st.has_native_type(builder, &self.params, NativePredicate::ValueOf);
-        let predicate_ok = builder.is_equal_slice(&st.predicate, &expected_predicate);
+
         let expected_arg_prefix = builder.constants(
             &StatementArg::Key(AnchoredKey(SELF, EMPTY_HASH)).to_fields(&self.params)[..VALUE_SIZE],
         );
         let arg_prefix_ok = builder.is_equal_slice(&st.args[0][..VALUE_SIZE], &expected_arg_prefix);
-        builder.and(predicate_ok, arg_prefix_ok)
+
+        let dupe_check = {
+            let individual_checks = prev_statements
+                .into_iter()
+                .map(|ps| {
+                    let same_predicate = builder.is_equal_slice(&st.predicate, &ps.predicate);
+                    let same_anchored_key = builder.is_equal_slice(&st.args[0], &ps.args[0]);
+                    builder.and(same_predicate, same_anchored_key)
+                })
+                .collect::<Vec<_>>();
+            builder.any(individual_checks)
+        };
+
+        let no_dupes_ok = builder.not(dupe_check);
+
+        builder.all([op_code_ok, st_code_ok, arg_prefix_ok, no_dupes_ok])
+    }
+
+    fn eval_copy(
+        &self,
+        builder: &mut CircuitBuilder<F, D>,
+        st: &StatementTarget,
+        op: &OperationTarget,
+        resolved_op_args: &[StatementTarget],
+    ) -> Result<BoolTarget> {
+        let op_code_ok = op.has_native_type(builder, NativeOperation::CopyStatement);
+
+        let expected_statement = &resolved_op_args[0];
+        let st_ok = builder.is_equal_flattenable(st, expected_statement);
+
+        Ok(builder.all([op_code_ok, st_ok]))
     }
 }
 
@@ -290,14 +409,13 @@ impl MainPodVerifyGate {
         // TODO: Store this hash in a global static with lazy init so that we don't have to
         // compute it every time.
         let key_type = hash_str(KEY_TYPE);
-        let expected_type_statement_flattened = builder.constants(
+        let expected_type_statement = StatementTarget::from_flattened(
-            &Statement::ValueOf(AnchoredKey(SELF, key_type), Value::from(PodType::MockMain))
+            &builder.constants(
-                .to_fields(params),
+                &Statement::ValueOf(AnchoredKey(SELF, key_type), Value::from(PodType::MockMain))
-        );
+                    .to_fields(params),
-        builder.connect_slice(
+            ),
-            &type_statement.to_flattened(),
-            &expected_type_statement_flattened,
         );
+        builder.connect_flattenable(type_statement, &expected_type_statement);
 
         // 5. Verify input statements
         let mut op_verifications = Vec::new();
@@ -372,9 +490,9 @@ impl MainPodVerifyCircuit {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::backends::plonky2::basetypes::C;
     use crate::backends::plonky2::mock::mainpod;
-    use crate::middleware::OperationType;
+    use crate::backends::plonky2::{basetypes::C, mock::mainpod::OperationArg};
+    use crate::middleware::{OperationType, PodId};
     use plonky2::plonk::{circuit_builder::CircuitBuilder, circuit_data::CircuitConfig};
 
     #[test]
@@ -455,6 +573,55 @@ mod tests {
         let st: mainpod::Statement = Statement::None.into();
         let op = mainpod::Operation(OperationType::Native(NativeOperation::None), vec![]);
         let prev_statements = vec![Statement::None.into()];
+        operation_verify(st.clone(), op, prev_statements.clone())?;
+
+        // NewEntry
+        let st1: mainpod::Statement =
+            Statement::ValueOf(AnchoredKey(SELF, "hello".into()), 55.into()).into();
+        let op = mainpod::Operation(OperationType::Native(NativeOperation::NewEntry), vec![]);
+        operation_verify(st1.clone(), op, vec![])?;
+
+        // Copy
+        let op = mainpod::Operation(
+            OperationType::Native(NativeOperation::CopyStatement),
+            vec![OperationArg::Index(0)],
+        );
+        operation_verify(st, op, prev_statements)?;
+
+        // Eq
+        let st2: mainpod::Statement = Statement::ValueOf(
+            AnchoredKey(PodId(Value::from(75).into()), "hello".into()),
+            55.into(),
+        )
+        .into();
+        let st: mainpod::Statement = Statement::Equal(
+            AnchoredKey(SELF, "hello".into()),
+            AnchoredKey(PodId(Value::from(75).into()), "hello".into()),
+        )
+        .into();
+        let op = mainpod::Operation(
+            OperationType::Native(NativeOperation::EqualFromEntries),
+            vec![OperationArg::Index(0), OperationArg::Index(1)],
+        );
+        let prev_statements = vec![st1.clone(), st2];
+        operation_verify(st, op, prev_statements)?;
+
+        // Lt
+        let st2: mainpod::Statement = Statement::ValueOf(
+            AnchoredKey(PodId(Value::from(88).into()), "hello".into()),
+            56.into(),
+        )
+        .into();
+        let st: mainpod::Statement = Statement::Lt(
+            AnchoredKey(SELF, "hello".into()),
+            AnchoredKey(PodId(Value::from(88).into()), "hello".into()),
+        )
+        .into();
+        let op = mainpod::Operation(
+            OperationType::Native(NativeOperation::LtFromEntries),
+            vec![OperationArg::Index(0), OperationArg::Index(1)],
+        );
+        let prev_statements = vec![st1.clone(), st2];
         operation_verify(st, op, prev_statements)?;
 
         Ok(())