fix: Revert some memory chip constraints

The MemoryProgram chip was modified to only be included in the first
shard. However, the rust out-of-circuit verifier was not updated to
account for this since it still expects the chip to be present in every
shard, and the in-circuit recursive verifier was updated with
constraints that are only valid if the core proof has a single shard.

This commit temporarily comments out these additional checks, but this
should be reverted when the verifier and Memory chip-related fixes are
integrated. This is related to the batch of PRs for issue #38

It might be necessary to revert commit
6aea5dc
after the proper fixes are incorporated.

core/src/memory/program.rs

@@ -1,6 +1,7 @@
+// TODO(wwared): Update this file to include MemoryProgram only in the first shard, see issue #38
 use core::borrow::{Borrow, BorrowMut};
 use core::mem::size_of;
-use p3_air::{Air, AirBuilderWithPublicValues, BaseAir, PairBuilder};
+use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir, PairBuilder};
 use p3_field::AbstractField;
 use p3_field::PrimeField;
 use p3_matrix::dense::RowMajorMatrix;
@@ -11,6 +12,7 @@ use sphinx_derive::AlignedBorrow;
 
 use crate::air::{AirInteraction, BaseAirBuilder, EventLens, PublicValues, WithEvents};
 use crate::air::{MachineAir, Word};
+use crate::operations::IsZeroOperation;
 use crate::runtime::{ExecutionRecord, Program};
 use crate::utils::pad_to_power_of_two;
 
@@ -35,6 +37,9 @@ pub struct MemoryProgramMultCols<T> {
     ///
     /// This column is technically redundant with `is_real`, but it's included for clarity.
     pub multiplicity: T,
+
+    /// Columns to see if current shard is 1.
+    pub is_first_shard: IsZeroOperation<T>,
 }
 
 /// Chip that initializes memory that is provided from the program. The table is preprocessed and
@@ -123,6 +128,7 @@ impl<F: PrimeField> MachineAir<F> for MemoryProgramChip {
                 let mut row = [F::zero(); NUM_MEMORY_PROGRAM_MULT_COLS];
                 let cols: &mut MemoryProgramMultCols<F> = row.as_mut_slice().borrow_mut();
                 cols.multiplicity = mult;
+                IsZeroOperation::populate(&mut cols.is_first_shard, input.index() - 1);
                 row
             })
             .collect::<Vec<_>>();
@@ -139,8 +145,8 @@ impl<F: PrimeField> MachineAir<F> for MemoryProgramChip {
         trace
     }
 
-    fn included(&self, record: &Self::Record) -> bool {
-        record.index == 1
+    fn included(&self, _record: &Self::Record) -> bool {
+        true
     }
 }
 
@@ -173,14 +179,26 @@ where
                 .collect::<Vec<_>>(),
         );
 
+        IsZeroOperation::<AB::F>::eval(
+            builder,
+            public_values.shard - AB::Expr::one(),
+            mult_local.is_first_shard,
+            prep_local.is_real.into(),
+        );
+        let is_first_shard = mult_local.is_first_shard.result;
+
         // Multiplicity must be either 0 or 1.
         builder.assert_bool(mult_local.multiplicity);
 
         // If first shard and preprocessed is real, multiplicity must be one.
-        builder.assert_eq(mult_local.multiplicity, prep_local.is_real.into());
-
-        // The shard this chip is contained in must be one.
-        builder.assert_one(public_values.shard);
+        builder
+            .when(is_first_shard * prep_local.is_real)
+            .assert_one(mult_local.multiplicity);
+
+        // If not first shard or preprocessed is not real, multiplicity must be zero.
+        builder
+            .when((AB::Expr::one() - is_first_shard) + (AB::Expr::one() - prep_local.is_real))
+            .assert_zero(mult_local.multiplicity);
 
         let mut values = vec![AB::Expr::zero(), AB::Expr::zero(), prep_local.addr.into()];
         values.extend(prep_local.value.map(Into::into));

prover/src/verify.rs

@@ -133,11 +133,12 @@ impl SphinxProver {
                     "memory should exist in the first chip".to_string(),
                 ));
             }
-            if i != 0 && program_memory_init_count > 0 {
-                return Err(MachineVerificationError::InvalidChipOccurence(
-                    "memory program should not exist in the first chip".to_string(),
-                ));
-            }
+            // TODO(wwared): Uncomment after upcoming ports update verifiers, see issue #38
+            // if i != 0 && program_memory_init_count > 0 {
+            //     return Err(MachineVerificationError::InvalidChipOccurence(
+            //         "memory program should not exist in the first chip".to_string(),
+            //     ));
+            // }
 
             // Assert that the `MemoryInit` and `MemoryFinalize` chips only exist in the last shard.
             if i != proof.0.len() - 1 && (memory_final_count > 0 || memory_init_count > 0) {

recursion/program/src/stark.rs

@@ -118,7 +118,7 @@ where
         machine: &StarkMachine<SC, A>,
         challenger: &mut DuplexChallengerVariable<C>,
         proof: &ShardProofVariable<C>,
-        shard_idx: Var<C::N>,
+        _shard_idx: Var<C::N>,
     ) where
         A: MachineAir<C::F> + for<'a> Air<RecursiveVerifierConstraintFolder<'a, C>>,
         C::F: TwoAdicField,
@@ -314,38 +314,40 @@ where
                 builder.assert_var_ne(index, C::N::from_canonical_usize(EMPTY));
             }
 
+            // TODO(wwared): Update and uncomment these constraints with future security ports, see issue #38
             if chip.as_ref().name() == "MemoryProgram" {
-                builder.if_eq(shard_idx, C::N::one()).then_or_else(
-                    |builder| {
-                        builder.assert_var_ne(index, C::N::from_canonical_usize(EMPTY));
-                    },
-                    |builder| {
-                        builder.assert_var_eq(index, C::N::from_canonical_usize(EMPTY));
-                    },
-                );
-            }
-
-            if chip.as_ref().name() == "MemoryInit" {
-                builder.if_eq(shard_idx, C::N::one()).then_or_else(
-                    |builder| {
-                        builder.assert_var_ne(index, C::N::from_canonical_usize(EMPTY));
-                    },
-                    |builder| {
-                        builder.assert_var_eq(index, C::N::from_canonical_usize(EMPTY));
-                    },
-                );
+                builder.assert_var_ne(index, C::N::from_canonical_usize(EMPTY));
+                // builder.if_eq(shard_idx, C::N::one()).then_or_else(
+                //     |builder| {
+                //         builder.assert_var_ne(index, C::N::from_canonical_usize(EMPTY));
+                //     },
+                //     |builder| {
+                //         builder.assert_var_eq(index, C::N::from_canonical_usize(EMPTY));
+                //     },
+                // );
             }
 
-            if chip.as_ref().name() == "MemoryFinalize" {
-                builder.if_eq(shard_idx, C::N::one()).then_or_else(
-                    |builder| {
-                        builder.assert_var_ne(index, C::N::from_canonical_usize(EMPTY));
-                    },
-                    |builder| {
-                        builder.assert_var_eq(index, C::N::from_canonical_usize(EMPTY));
-                    },
-                );
-            }
+            // if chip.as_ref().name() == "MemoryInit" {
+            //     builder.if_eq(shard_idx, C::N::one()).then_or_else(
+            //         |builder| {
+            //             builder.assert_var_ne(index, C::N::from_canonical_usize(EMPTY));
+            //         },
+            //         |builder| {
+            //             builder.assert_var_eq(index, C::N::from_canonical_usize(EMPTY));
+            //         },
+            //     );
+            // }
+
+            // if chip.as_ref().name() == "MemoryFinalize" {
+            //     builder.if_eq(shard_idx, C::N::one()).then_or_else(
+            //         |builder| {
+            //             builder.assert_var_ne(index, C::N::from_canonical_usize(EMPTY));
+            //         },
+            //         |builder| {
+            //             builder.assert_var_eq(index, C::N::from_canonical_usize(EMPTY));
+            //         },
+            //     );
+            // }
 
             builder
                 .if_ne(index, C::N::from_canonical_usize(EMPTY))