From 6db3291aa267d4f30ebcae18cae9c0f170a06654 Mon Sep 17 00:00:00 2001
From: Matt Walker <matt.g.d.walker@gmail.com>
Date: Thu, 17 Oct 2024 12:45:54 -0400
Subject: [PATCH] o1vm/pickles: Don't forget the extra columns

---
 o1vm/src/pickles/column_env.rs | 4 ++--
 o1vm/src/pickles/tests.rs      | 8 +-------
 2 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/o1vm/src/pickles/column_env.rs b/o1vm/src/pickles/column_env.rs
index 1055a4dd42..c66f68160f 100644
--- a/o1vm/src/pickles/column_env.rs
+++ b/o1vm/src/pickles/column_env.rs
@@ -36,8 +36,8 @@ pub struct ColumnEnvironment<'a, F: FftField> {
 }
 
 pub fn get_all_columns() -> Vec<Column> {
-    let mut cols = Vec::<Column>::with_capacity(SCRATCH_SIZE + N_MIPS_SEL_COLS);
-    for i in 0..SCRATCH_SIZE {
+    let mut cols = Vec::<Column>::with_capacity(SCRATCH_SIZE + 2 + N_MIPS_SEL_COLS);
+    for i in 0..SCRATCH_SIZE + 2 {
         cols.push(Column::Relation(i));
     }
     for i in 0..N_MIPS_SEL_COLS {
diff --git a/o1vm/src/pickles/tests.rs b/o1vm/src/pickles/tests.rs
index 4d668dde11..734ffca8aa 100644
--- a/o1vm/src/pickles/tests.rs
+++ b/o1vm/src/pickles/tests.rs
@@ -74,12 +74,6 @@ fn test_regression_selectors_for_instructions() {
 
 #[test]
 fn test_small_circuit() {
-    /*   domain: EvaluationDomains<G::ScalarField>,
-       srs: &SRS<G>,
-       inputs: ProofInputs<G>,
-       constraints: &[E<G::ScalarField>],
-       rng: &mut RNG,
-    */
     debug!("0");
     let domain = EvaluationDomains::<Fq>::create(8).unwrap();
     let srs = SRS::create(8);
@@ -131,7 +125,7 @@ fn test_small_circuit() {
     };
     debug!("1");
     let mut expr = Expr::literal(Fq::zero());
-    for i in 0..SCRATCH_SIZE {
+    for i in 0..SCRATCH_SIZE + 2 {
         expr += Expr::cell(Column::Relation(i), CurrOrNext::Curr);
     }
     expr *= Expr::cell(Column::DynamicSelector(0), CurrOrNext::Curr);