lookup argument for range check

evm/src/memory/memory_stark.rs

@@ -8,12 +8,14 @@ use plonky2::hash::hash_types::RichField;
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 use crate::memory::registers::{
     memory_value_limb, sorted_memory_value_limb, MEMORY_ADDR_CONTEXT, MEMORY_ADDR_SEGMENT,
-    MEMORY_ADDR_VIRTUAL, MEMORY_CONTEXT_FIRST_CHANGE, MEMORY_COUNTER, MEMORY_IS_READ,
-    MEMORY_RANGE_CHECK, MEMORY_SEGMENT_FIRST_CHANGE, MEMORY_TIMESTAMP, MEMORY_VIRTUAL_FIRST_CHANGE,
-    NUM_REGISTERS, SORTED_MEMORY_ADDR_CONTEXT, SORTED_MEMORY_ADDR_SEGMENT,
-    SORTED_MEMORY_ADDR_VIRTUAL, SORTED_MEMORY_IS_READ, SORTED_MEMORY_TIMESTAMP,
+    MEMORY_ADDR_VIRTUAL, MEMORY_CONTEXT_FIRST_CHANGE, MEMORY_COUNTER, MEMORY_COUNTER_PERMUTED,
+    MEMORY_IS_READ, MEMORY_RANGE_CHECK, MEMORY_RANGE_CHECK_PERMUTED, MEMORY_SEGMENT_FIRST_CHANGE,
+    MEMORY_TIMESTAMP, MEMORY_VIRTUAL_FIRST_CHANGE, NUM_REGISTERS, SORTED_MEMORY_ADDR_CONTEXT,
+    SORTED_MEMORY_ADDR_SEGMENT, SORTED_MEMORY_ADDR_VIRTUAL, SORTED_MEMORY_IS_READ,
+    SORTED_MEMORY_TIMESTAMP,
 };
 use crate::stark::Stark;
+use crate::util::permuted_cols;
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
 #[derive(Default)]
@@ -197,7 +199,14 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
         trace_cols[MEMORY_VIRTUAL_FIRST_CHANGE] = virtual_first_change;
 
         trace_cols[MEMORY_RANGE_CHECK] = range_check_value;
-        trace_cols[MEMORY_COUNTER] = (0..trace_cols.len()).map(|i| F::from_canonical_usize(i)).collect();
+        trace_cols[MEMORY_COUNTER] = (0..trace_cols.len())
+            .map(|i| F::from_canonical_usize(i))
+            .collect();
+
+        let (permuted_inputs, permuted_table) =
+            permuted_cols(&trace_cols[MEMORY_RANGE_CHECK], &trace_cols[MEMORY_COUNTER]);
+        trace_cols[MEMORY_RANGE_CHECK_PERMUTED] = permuted_inputs;
+        trace_cols[MEMORY_COUNTER_PERMUTED] = permuted_table;
     }
 }
 
@@ -271,6 +280,23 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for MemoryStark<F
             yield_constr
                 .constraint(next_is_read * timestamp_first_change * (next_values[i] - values[i]));
         }
+
+        // Lookup argument for the range check.
+        let local_perm_input = vars.local_values[MEMORY_RANGE_CHECK_PERMUTED];
+        let next_perm_table = vars.next_values[MEMORY_COUNTER_PERMUTED];
+        let next_perm_input = vars.next_values[MEMORY_COUNTER_PERMUTED];
+
+        // A "vertical" diff between the local and next permuted inputs.
+        let diff_input_prev = next_perm_input - local_perm_input;
+        // A "horizontal" diff between the next permuted input and permuted table value.
+        let diff_input_table = next_perm_input - next_perm_table;
+
+        yield_constr.constraint(diff_input_prev * diff_input_table);
+
+        // This is actually constraining the first row, as per the spec, since `diff_input_table`
+        // is a diff of the next row's values. In the context of `constraint_last_row`, the next
+        // row is the first row.
+        yield_constr.constraint_last_row(diff_input_table);
     }
 
     fn eval_ext_circuit(
@@ -388,6 +414,24 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for MemoryStark<F
             let read_constraint = builder.mul_extension(next_is_read, zero_if_read);
             yield_constr.constraint(builder, read_constraint);
         }
+
+        // Lookup argument for range check.
+        let local_perm_input = vars.local_values[MEMORY_RANGE_CHECK_PERMUTED];
+        let next_perm_table = vars.next_values[MEMORY_COUNTER_PERMUTED];
+        let next_perm_input = vars.next_values[MEMORY_COUNTER_PERMUTED];
+
+        // A "vertical" diff between the local and next permuted inputs.
+        let diff_input_prev = builder.sub_extension(next_perm_input, local_perm_input);
+        // A "horizontal" diff between the next permuted input and permuted table value.
+        let diff_input_table = builder.sub_extension(next_perm_input, next_perm_table);
+
+        let diff_product = builder.mul_extension(diff_input_prev, diff_input_table);
+        yield_constr.constraint(builder, diff_product);
+
+        // This is actually constraining the first row, as per the spec, since `diff_input_table`
+        // is a diff of the next row's values. In the context of `constraint_last_row`, the next
+        // row is the first row.
+        yield_constr.constraint_last_row(builder, diff_input_table);
     }
 
     fn constraint_degree(&self) -> usize {

evm/src/memory/registers.rs

@@ -30,5 +30,7 @@ pub(crate) const MEMORY_VIRTUAL_FIRST_CHANGE: usize = MEMORY_SEGMENT_FIRST_CHANG
 
 pub(crate) const MEMORY_RANGE_CHECK: usize = MEMORY_VIRTUAL_FIRST_CHANGE + 1;
 pub(crate) const MEMORY_COUNTER: usize = MEMORY_RANGE_CHECK + 1;
+pub(crate) const MEMORY_RANGE_CHECK_PERMUTED: usize = MEMORY_COUNTER + 1;
+pub(crate) const MEMORY_COUNTER_PERMUTED: usize = MEMORY_RANGE_CHECK_PERMUTED + 1;
 
-pub(crate) const NUM_REGISTERS: usize = MEMORY_COUNTER + 1;
+pub(crate) const NUM_REGISTERS: usize = MEMORY_COUNTER_PERMUTED + 1;

evm/src/util.rs

@@ -1,5 +1,7 @@
+use std::cmp::Ordering;
+
 use itertools::Itertools;
-use plonky2::field::field_types::Field;
+use plonky2::field::field_types::{Field, PrimeField64};
 use plonky2::field::polynomial::PolynomialValues;
 use plonky2::util::transpose;
 
@@ -14,3 +16,71 @@ pub fn trace_rows_to_poly_values<F: Field, const COLUMNS: usize>(
         .map(|column| PolynomialValues::new(column))
         .collect()
 }
+
+/// Given an input column and a table column, generate the permuted input and permuted table columns
+/// used in the Halo2 permutation argument.
+pub fn permuted_cols<F: PrimeField64>(inputs: &[F], table: &[F]) -> (Vec<F>, Vec<F>) {
+    let n = inputs.len();
+
+    // The permuted inputs do not have to be ordered, but we found that sorting was faster than
+    // hash-based grouping. We also sort the table, as this helps us identify "unused" table
+    // elements efficiently.
+
+    // To compare elements, e.g. for sorting, we first need them in canonical form. It would be
+    // wasteful to canonicalize in each comparison, as a single element may be involved in many
+    // comparisons. So we will canonicalize once upfront, then use `to_noncanonical_u64` when
+    // comparing elements.
+
+    let sorted_inputs = inputs
+        .iter()
+        .map(|x| x.to_canonical())
+        .sorted_unstable_by_key(|x| x.to_noncanonical_u64())
+        .collect_vec();
+    let sorted_table = table
+        .iter()
+        .map(|x| x.to_canonical())
+        .sorted_unstable_by_key(|x| x.to_noncanonical_u64())
+        .collect_vec();
+
+    let mut unused_table_inds = Vec::with_capacity(n);
+    let mut unused_table_vals = Vec::with_capacity(n);
+    let mut permuted_table = vec![F::ZERO; n];
+    let mut i = 0;
+    let mut j = 0;
+    while (j < n) && (i < n) {
+        let input_val = sorted_inputs[i].to_noncanonical_u64();
+        let table_val = sorted_table[j].to_noncanonical_u64();
+        match input_val.cmp(&table_val) {
+            Ordering::Greater => {
+                unused_table_vals.push(sorted_table[j]);
+                j += 1;
+            }
+            Ordering::Less => {
+                if let Some(x) = unused_table_vals.pop() {
+                    permuted_table[i] = x;
+                } else {
+                    unused_table_inds.push(i);
+                }
+                i += 1;
+            }
+            Ordering::Equal => {
+                permuted_table[i] = sorted_table[j];
+                i += 1;
+                j += 1;
+            }
+        }
+    }
+
+    #[allow(clippy::needless_range_loop)] // indexing is just more natural here
+    for jj in j..n {
+        unused_table_vals.push(sorted_table[jj]);
+    }
+    for ii in i..n {
+        unused_table_inds.push(ii);
+    }
+    for (ind, val) in unused_table_inds.into_iter().zip_eq(unused_table_vals) {
+        permuted_table[ind] = val;
+    }
+
+    (sorted_inputs, permuted_table)
+}