From 0802d6c0211cb9873bbf33bde06bb077a28c02b4 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Thu, 14 Jul 2022 11:31:47 -0700
Subject: [PATCH 1/7] Continue work on bootstrapping

The kernel is hashed using a Keccak based sponge for now. We could switch to Poseidon later if our kernel grows too large.

Note that we use simple zero-padding (pad0*) instead of the standard pad10* rule. It's simpler, and we don't care that the prover can add extra 0s at the end of the code. The program counter can never reach those bytes, and even if it could, they'd be 0 anyway given the EVM's zero-initialization rule.

In one CPU row, we can do a whole Keccak hash (via the CTL), absorbing 136 bytes. But we can't actually bootstrap that many bytes of kernel code in one row, because we're also limited by memory bandwidth. Currently we can write 4 bytes of the kernel to memory in one row.

So we treat the `keccak_input_limbs` columns as a buffer. We gradually fill up this buffer, 4 bytes (one `u32` word) at a time. Every `136 / 4 = 34` rows, the buffer will be full, so at that point we activate the Keccak CTL to absorb the buffer.
---
 evm/Cargo.toml                    |   1 +
 evm/src/cpu/bootstrap_kernel.rs   | 110 ++++++++++++++++++++++++------
 evm/src/cpu/columns.rs            |   1 +
 evm/src/cpu/kernel/aggregator.rs  |   3 +
 evm/src/cpu/kernel/assembler.rs   |  27 +++++---
 evm/src/cpu/kernel/keccak_util.rs |  14 ++++
 evm/src/cpu/kernel/mod.rs         |   1 +
 evm/src/generation/mod.rs         |   2 +-
 evm/src/generation/state.rs       |  29 ++++----
 9 files changed, 144 insertions(+), 44 deletions(-)
 create mode 100644 evm/src/cpu/kernel/keccak_util.rs

diff --git a/evm/Cargo.toml b/evm/Cargo.toml
index facf300b..e8dcfdfe 100644
--- a/evm/Cargo.toml
+++ b/evm/Cargo.toml
@@ -13,6 +13,7 @@ ethereum-types = "0.13.1"
 hex = { version = "0.4.3", optional = true }
 itertools = "0.10.3"
 log = "0.4.14"
+once_cell = "1.13.0"
 pest = "2.1.3"
 pest_derive = "2.1.0"
 rayon = "1.5.1"
diff --git a/evm/src/cpu/bootstrap_kernel.rs b/evm/src/cpu/bootstrap_kernel.rs
index ba10b70e..3d5d4f1c 100644
--- a/evm/src/cpu/bootstrap_kernel.rs
+++ b/evm/src/cpu/bootstrap_kernel.rs
@@ -1,33 +1,74 @@
 //! The initial phase of execution, where the kernel code is hashed while being written to memory.
 //! The hash is then checked against a precomputed kernel hash.
 
+use std::borrow::Borrow;
+
 use itertools::Itertools;
 use plonky2::field::extension::Extendable;
 use plonky2::field::packed::PackedField;
 use plonky2::field::types::Field;
 use plonky2::hash::hash_types::RichField;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2_util::ceil_div_usize;
 
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
-use crate::cpu::columns::{COL_MAP, NUM_CPU_COLUMNS};
+use crate::cpu::columns::{CpuColumnsView, NUM_CPU_COLUMNS};
+use crate::cpu::kernel::aggregator::KERNEL;
+use crate::cpu::kernel::keccak_util::keccakf_u32s;
 use crate::cpu::public_inputs::NUM_PUBLIC_INPUTS;
 use crate::generation::state::GenerationState;
 use crate::memory;
-use crate::memory::segments;
+use crate::memory::{segments, NUM_CHANNELS};
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
+/// The Keccak rate (1088 bits), measured in bytes.
+const KECCAK_RATE_BYTES: usize = 1088 / 8;
+
+/// The Keccak rate (1088 bits), measured in u32 limbs.
+const KECCAK_RATE_LIMBS: usize = 1088 / 32;
+
+/// We can't process more than `NUM_CHANNELS` bytes per row, since that's all the memory bandwidth
+/// we have. We also can't process more than 4 bytes (or the number of bytes in a `u32`), since we
+/// want them to fit in a single limb of Keccak input.
+const BYTES_PER_ROW: usize = 4;
+
 pub(crate) fn generate_bootstrap_kernel<F: Field>(state: &mut GenerationState<F>) {
-    for chunk in &state.kernel.code.clone().into_iter().enumerate().chunks(4) {
+    let mut code = KERNEL.code.clone();
+
+    // Zero-pad the code such that its size is a multiple of the Keccak rate.
+    let padded_size = ceil_div_usize(code.len(), KECCAK_RATE_BYTES) * KECCAK_RATE_BYTES;
+    code.resize(padded_size, 0);
+
+    let mut sponge_state = [0u32; 50];
+    let mut sponge_input_pos: usize = 0;
+
+    // Iterate through chunks of the code, such that we can write one chunk to memory per row.
+    for chunk in &code.into_iter().enumerate().chunks(BYTES_PER_ROW) {
+        state.current_cpu_row.is_bootstrap_kernel = F::ONE;
+
+        // Write this chunk to memory, while simultaneously packing its bytes into a u32 word.
+        let mut packed_bytes: u32 = 0;
         for (addr, byte) in chunk {
             let mut value = [F::ZERO; memory::VALUE_LIMBS];
             value[0] = F::from_canonical_u8(byte);
 
-            let channel = addr % memory::NUM_CHANNELS;
+            let channel = addr % NUM_CHANNELS;
             state.set_mem_current(channel, segments::CODE, addr, value);
 
-            // TODO: Set other registers.
+            packed_bytes = (packed_bytes << 8) | byte as u32;
+        }
 
-            state.commit_cpu_row();
+        sponge_state[sponge_input_pos] = packed_bytes;
+        state.current_cpu_row.keccak_input_limbs = sponge_state.map(F::from_canonical_u32);
+        state.commit_cpu_row();
+
+        sponge_input_pos = (sponge_input_pos + 1) % KECCAK_RATE_LIMBS;
+        // If we just crossed a multiple of KECCAK_RATE_LIMBS, then we've filled the Keccak input
+        // buffer, so it's time to absorb.
+        if sponge_input_pos == 0 {
+            state.current_cpu_row.is_keccak = F::ONE;
+            keccakf_u32s(&mut sponge_state);
+            state.current_cpu_row.keccak_output_limbs = sponge_state.map(F::from_canonical_u32);
         }
     }
 }
@@ -36,18 +77,34 @@ pub(crate) fn eval_bootstrap_kernel<F: Field, P: PackedField<Scalar = F>>(
     vars: StarkEvaluationVars<F, P, NUM_CPU_COLUMNS, NUM_PUBLIC_INPUTS>,
     yield_constr: &mut ConstraintConsumer<P>,
 ) {
+    let local_values: &CpuColumnsView<_> = vars.local_values.borrow();
+    let next_values: &CpuColumnsView<_> = vars.next_values.borrow();
+
     // IS_BOOTSTRAP_KERNEL must have an init value of 1, a final value of 0, and a delta in {0, -1}.
-    let local_is_bootstrap = vars.local_values[COL_MAP.is_bootstrap_kernel];
-    let next_is_bootstrap = vars.next_values[COL_MAP.is_bootstrap_kernel];
+    let local_is_bootstrap = local_values.is_bootstrap_kernel;
+    let next_is_bootstrap = next_values.is_bootstrap_kernel;
     yield_constr.constraint_first_row(local_is_bootstrap - P::ONES);
     yield_constr.constraint_last_row(local_is_bootstrap);
     let delta_is_bootstrap = next_is_bootstrap - local_is_bootstrap;
     yield_constr.constraint_transition(delta_is_bootstrap * (delta_is_bootstrap + P::ONES));
 
-    // If IS_BOOTSTRAP_KERNEL changed (from 1 to 0), check that the current kernel hash matches a
-    // precomputed one.
-    let hash_diff = F::ZERO; // TODO
-    yield_constr.constraint_transition(delta_is_bootstrap * hash_diff)
+    // TODO: Constraints to enforce that, if IS_BOOTSTRAP_KERNEL,
+    // - If CLOCK is a multiple of KECCAK_RATE_LIMBS, activate the Keccak CTL, and ensure the output
+    //   is copied to the next row (besides the first limb which will immediately be overwritten).
+    // - Otherwise, ensure that the Keccak input is copied to the next row (besides the next limb).
+
+    // If IS_BOOTSTRAP_KERNEL changed (from 1 to 0), check that
+    // - the clock is a multiple of KECCAK_RATE_LIMBS (TODO)
+    // - the current kernel hash matches a precomputed one
+    for (&expected, actual) in KERNEL
+        .code_hash
+        .iter()
+        .zip(local_values.keccak_output_limbs)
+    {
+        let expected = P::from(F::from_canonical_u32(expected));
+        let diff = expected - actual;
+        yield_constr.constraint_transition(delta_is_bootstrap * diff);
+    }
 }
 
 pub(crate) fn eval_bootstrap_kernel_circuit<F: RichField + Extendable<D>, const D: usize>(
@@ -55,11 +112,13 @@ pub(crate) fn eval_bootstrap_kernel_circuit<F: RichField + Extendable<D>, const
     vars: StarkEvaluationTargets<D, NUM_CPU_COLUMNS, NUM_PUBLIC_INPUTS>,
     yield_constr: &mut RecursiveConstraintConsumer<F, D>,
 ) {
+    let local_values: &CpuColumnsView<_> = vars.local_values.borrow();
+    let next_values: &CpuColumnsView<_> = vars.next_values.borrow();
     let one = builder.one_extension();
 
     // IS_BOOTSTRAP_KERNEL must have an init value of 1, a final value of 0, and a delta in {0, -1}.
-    let local_is_bootstrap = vars.local_values[COL_MAP.is_bootstrap_kernel];
-    let next_is_bootstrap = vars.next_values[COL_MAP.is_bootstrap_kernel];
+    let local_is_bootstrap = local_values.is_bootstrap_kernel;
+    let next_is_bootstrap = next_values.is_bootstrap_kernel;
     let constraint = builder.sub_extension(local_is_bootstrap, one);
     yield_constr.constraint_first_row(builder, constraint);
     yield_constr.constraint_last_row(builder, local_is_bootstrap);
@@ -68,9 +127,22 @@ pub(crate) fn eval_bootstrap_kernel_circuit<F: RichField + Extendable<D>, const
         builder.mul_add_extension(delta_is_bootstrap, delta_is_bootstrap, delta_is_bootstrap);
     yield_constr.constraint_transition(builder, constraint);
 
-    // If IS_BOOTSTRAP_KERNEL changed (from 1 to 0), check that the current kernel hash matches a
-    // precomputed one.
-    let hash_diff = builder.zero_extension(); // TODO
-    let constraint = builder.mul_extension(delta_is_bootstrap, hash_diff);
-    yield_constr.constraint_transition(builder, constraint)
+    // TODO: Constraints to enforce that, if IS_BOOTSTRAP_KERNEL,
+    // - If CLOCK is a multiple of KECCAK_RATE_LIMBS, activate the Keccak CTL, and ensure the output
+    //   is copied to the next row (besides the first limb which will immediately be overwritten).
+    // - Otherwise, ensure that the Keccak input is copied to the next row (besides the next limb).
+
+    // If IS_BOOTSTRAP_KERNEL changed (from 1 to 0), check that
+    // - the clock is a multiple of KECCAK_RATE_LIMBS (TODO)
+    // - the current kernel hash matches a precomputed one
+    for (&expected, actual) in KERNEL
+        .code_hash
+        .iter()
+        .zip(local_values.keccak_output_limbs)
+    {
+        let expected = builder.constant_extension(F::Extension::from_canonical_u32(expected));
+        let diff = builder.sub_extension(expected, actual);
+        let constraint = builder.mul_extension(delta_is_bootstrap, diff);
+        yield_constr.constraint_transition(builder, constraint);
+    }
 }
diff --git a/evm/src/cpu/columns.rs b/evm/src/cpu/columns.rs
index 1fce6cb7..f3a400c6 100644
--- a/evm/src/cpu/columns.rs
+++ b/evm/src/cpu/columns.rs
@@ -8,6 +8,7 @@ use std::ops::{Index, IndexMut};
 use crate::memory;
 
 #[repr(C)]
+#[derive(Eq, PartialEq, Debug)]
 pub struct CpuColumnsView<T> {
     /// Filter. 1 if the row is part of bootstrapping the kernel code, 0 otherwise.
     pub is_bootstrap_kernel: T,
diff --git a/evm/src/cpu/kernel/aggregator.rs b/evm/src/cpu/kernel/aggregator.rs
index 04d57e75..6ca88ba1 100644
--- a/evm/src/cpu/kernel/aggregator.rs
+++ b/evm/src/cpu/kernel/aggregator.rs
@@ -4,10 +4,13 @@ use std::collections::HashMap;
 
 use ethereum_types::U256;
 use itertools::Itertools;
+use once_cell::sync::Lazy;
 
 use super::assembler::{assemble, Kernel};
 use crate::cpu::kernel::parser::parse;
 
+pub static KERNEL: Lazy<Kernel> = Lazy::new(combined_kernel);
+
 pub fn evm_constants() -> HashMap<String, U256> {
     let mut c = HashMap::new();
     c.insert("SEGMENT_ID_TXN_DATA".into(), 0.into()); // TODO: Replace with actual segment ID.
diff --git a/evm/src/cpu/kernel/assembler.rs b/evm/src/cpu/kernel/assembler.rs
index bdc8ded4..c3ccea7a 100644
--- a/evm/src/cpu/kernel/assembler.rs
+++ b/evm/src/cpu/kernel/assembler.rs
@@ -6,6 +6,7 @@ use log::debug;
 
 use super::ast::PushTarget;
 use crate::cpu::kernel::ast::Literal;
+use crate::cpu::kernel::keccak_util::hash_kernel;
 use crate::cpu::kernel::{
     ast::{File, Item},
     opcodes::{get_opcode, get_push_opcode},
@@ -19,9 +20,25 @@ const BYTES_PER_OFFSET: u8 = 3;
 #[derive(PartialEq, Eq, Debug)]
 pub struct Kernel {
     pub(crate) code: Vec<u8>,
+
+    /// Computed using `hash_kernel`. It is encoded as `u32` limbs for convenience, since we deal
+    /// with `u32` limbs in our Keccak table.
+    pub(crate) code_hash: [u32; 8],
+
     pub(crate) global_labels: HashMap<String, usize>,
 }
 
+impl Kernel {
+    fn new(code: Vec<u8>, global_labels: HashMap<String, usize>) -> Self {
+        let code_hash = hash_kernel(&code);
+        Self {
+            code,
+            code_hash,
+            global_labels,
+        }
+    }
+}
+
 struct Macro {
     params: Vec<String>,
     items: Vec<Item>,
@@ -56,10 +73,7 @@ pub(crate) fn assemble(files: Vec<File>, constants: HashMap<String, U256>) -> Ke
         debug!("Assembled file size: {} bytes", file_len);
     }
     assert_eq!(code.len(), offset, "Code length doesn't match offset.");
-    Kernel {
-        code,
-        global_labels,
-    }
+    Kernel::new(code, global_labels)
 }
 
 fn find_macros(files: &[File]) -> HashMap<String, Macro> {
@@ -286,10 +300,7 @@ mod tests {
         expected_global_labels.insert("function_1".to_string(), 0);
         expected_global_labels.insert("function_2".to_string(), 3);
 
-        let expected_kernel = Kernel {
-            code: expected_code,
-            global_labels: expected_global_labels,
-        };
+        let expected_kernel = Kernel::new(expected_code, expected_global_labels);
 
         let program = vec![file_1, file_2];
         assert_eq!(assemble(program, HashMap::new()), expected_kernel);
diff --git a/evm/src/cpu/kernel/keccak_util.rs b/evm/src/cpu/kernel/keccak_util.rs
new file mode 100644
index 00000000..1498ba08
--- /dev/null
+++ b/evm/src/cpu/kernel/keccak_util.rs
@@ -0,0 +1,14 @@
+/// A Keccak-f based hash.
+///
+/// This hash does not use standard Keccak padding, since we don't care about extra zeros at the
+/// end of the code.
+pub(crate) fn hash_kernel(_code: &[u8]) -> [u32; 8] {
+    let state = [0u32; 50];
+    // TODO: absorb code
+    state[0..8].try_into().unwrap()
+}
+
+/// Like tiny-keccak's `keccakf`, but deals with `u32` limbs instead of `u64` limbs.
+pub(crate) fn keccakf_u32s(_state: &mut [u32; 50]) {
+    // TODO: Implement
+}
diff --git a/evm/src/cpu/kernel/mod.rs b/evm/src/cpu/kernel/mod.rs
index ae765b99..bc7608dd 100644
--- a/evm/src/cpu/kernel/mod.rs
+++ b/evm/src/cpu/kernel/mod.rs
@@ -1,6 +1,7 @@
 pub mod aggregator;
 pub mod assembler;
 mod ast;
+pub(crate) mod keccak_util;
 mod opcodes;
 mod parser;
 
diff --git a/evm/src/generation/mod.rs b/evm/src/generation/mod.rs
index 4eb3cbbf..edca27dd 100644
--- a/evm/src/generation/mod.rs
+++ b/evm/src/generation/mod.rs
@@ -46,7 +46,7 @@ pub fn generate_traces<F: RichField + Extendable<D>, const D: usize>(
         logic_ops: logic_inputs,
         ..
     } = state;
-    assert_eq!(current_cpu_row, [F::ZERO; NUM_CPU_COLUMNS]);
+    assert_eq!(current_cpu_row, [F::ZERO; NUM_CPU_COLUMNS].into());
 
     let cpu_trace = trace_rows_to_poly_values(cpu_rows);
     let keccak_trace = all_stark.keccak_stark.generate_trace(keccak_inputs);
diff --git a/evm/src/generation/state.rs b/evm/src/generation/state.rs
index 1b26a3c4..7a95f7e4 100644
--- a/evm/src/generation/state.rs
+++ b/evm/src/generation/state.rs
@@ -1,20 +1,17 @@
+use std::mem;
+
 use ethereum_types::U256;
 use plonky2::field::types::Field;
 
-use crate::cpu::columns::NUM_CPU_COLUMNS;
-use crate::cpu::kernel::aggregator::combined_kernel;
-use crate::cpu::kernel::assembler::Kernel;
+use crate::cpu::columns::{CpuColumnsView, NUM_CPU_COLUMNS};
 use crate::generation::memory::MemoryState;
-use crate::logic::{Op, Operation};
 use crate::memory::memory_stark::MemoryOp;
 use crate::{keccak, logic};
 
 #[derive(Debug)]
 pub(crate) struct GenerationState<F: Field> {
-    pub(crate) kernel: Kernel,
-
     pub(crate) cpu_rows: Vec<[F; NUM_CPU_COLUMNS]>,
-    pub(crate) current_cpu_row: [F; NUM_CPU_COLUMNS],
+    pub(crate) current_cpu_row: CpuColumnsView<F>,
 
     pub(crate) current_context: usize,
     pub(crate) memory: MemoryState<F>,
@@ -27,24 +24,24 @@ impl<F: Field> GenerationState<F> {
     /// Compute logical AND, and record the operation to be added in the logic table later.
     #[allow(unused)] // TODO: Should be used soon.
     pub(crate) fn and(&mut self, input0: U256, input1: U256) -> U256 {
-        self.logic_op(Op::And, input0, input1)
+        self.logic_op(logic::Op::And, input0, input1)
     }
 
     /// Compute logical OR, and record the operation to be added in the logic table later.
     #[allow(unused)] // TODO: Should be used soon.
     pub(crate) fn or(&mut self, input0: U256, input1: U256) -> U256 {
-        self.logic_op(Op::Or, input0, input1)
+        self.logic_op(logic::Op::Or, input0, input1)
     }
 
     /// Compute logical XOR, and record the operation to be added in the logic table later.
     #[allow(unused)] // TODO: Should be used soon.
     pub(crate) fn xor(&mut self, input0: U256, input1: U256) -> U256 {
-        self.logic_op(Op::Xor, input0, input1)
+        self.logic_op(logic::Op::Xor, input0, input1)
     }
 
     /// Compute logical AND, and record the operation to be added in the logic table later.
-    pub(crate) fn logic_op(&mut self, op: Op, input0: U256, input1: U256) -> U256 {
-        let operation = Operation::new(op, input0, input1);
+    pub(crate) fn logic_op(&mut self, op: logic::Op, input0: U256, input1: U256) -> U256 {
+        let operation = logic::Operation::new(op, input0, input1);
         let result = operation.result;
         self.logic_ops.push(operation);
         result
@@ -96,8 +93,9 @@ impl<F: Field> GenerationState<F> {
     }
 
     pub(crate) fn commit_cpu_row(&mut self) {
-        self.cpu_rows.push(self.current_cpu_row);
-        self.current_cpu_row = [F::ZERO; NUM_CPU_COLUMNS];
+        let mut swapped_row = [F::ZERO; NUM_CPU_COLUMNS].into();
+        mem::swap(&mut self.current_cpu_row, &mut swapped_row);
+        self.cpu_rows.push(swapped_row.into());
     }
 }
 
@@ -106,9 +104,8 @@ impl<F: Field> GenerationState<F> {
 impl<F: Field> Default for GenerationState<F> {
     fn default() -> Self {
         Self {
-            kernel: combined_kernel(),
             cpu_rows: vec![],
-            current_cpu_row: [F::ZERO; NUM_CPU_COLUMNS],
+            current_cpu_row: [F::ZERO; NUM_CPU_COLUMNS].into(),
             current_context: 0,
             memory: MemoryState::default(),
             keccak_inputs: vec![],

From 6d69e14a895caae590fcb457fa69ba64834a7d7a Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Thu, 14 Jul 2022 14:55:17 -0700
Subject: [PATCH 2/7] Add `%rep` syntax for repeating a block

Same syntax as NASM.
---
 evm/src/cpu/kernel/asm/basic_macros.asm | 15 +++++++-----
 evm/src/cpu/kernel/assembler.rs         | 31 +++++++++++++++++++++----
 evm/src/cpu/kernel/ast.rs               |  2 ++
 evm/src/cpu/kernel/evm_asm.pest         |  5 ++--
 evm/src/cpu/kernel/parser.rs            |  8 +++++++
 5 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/evm/src/cpu/kernel/asm/basic_macros.asm b/evm/src/cpu/kernel/asm/basic_macros.asm
index 200aeea0..9e884fea 100644
--- a/evm/src/cpu/kernel/asm/basic_macros.asm
+++ b/evm/src/cpu/kernel/asm/basic_macros.asm
@@ -9,18 +9,21 @@
 %endmacro
 
 %macro pop2
-    pop
-    pop
+    %rep 2
+        pop
+    %endrep
 %endmacro
 
 %macro pop3
-    pop
-    %pop2
+    %rep 3
+        pop
+    %endrep
 %endmacro
 
 %macro pop4
-    %pop2
-    %pop2
+    %rep 4
+        pop
+    %endrep
 %endmacro
 
 // If pred is zero, yields z; otherwise, yields nz
diff --git a/evm/src/cpu/kernel/assembler.rs b/evm/src/cpu/kernel/assembler.rs
index bdc8ded4..d1eaf752 100644
--- a/evm/src/cpu/kernel/assembler.rs
+++ b/evm/src/cpu/kernel/assembler.rs
@@ -44,6 +44,7 @@ pub(crate) fn assemble(files: Vec<File>, constants: HashMap<String, U256>) -> Ke
     let mut local_labels = Vec::with_capacity(files.len());
     for file in files {
         let expanded_file = expand_macros(file.body, &macros);
+        let expanded_file = expand_repeats(expanded_file);
         let expanded_file = inline_constants(expanded_file, &constants);
         local_labels.push(find_labels(&expanded_file, &mut offset, &mut global_labels));
         expanded_files.push(expanded_file);
@@ -132,6 +133,21 @@ fn expand_macro_call(
     expand_macros(expanded_item, macros)
 }
 
+fn expand_repeats(body: Vec<Item>) -> Vec<Item> {
+    let mut expanded = vec![];
+    for item in body {
+        if let Item::Repeat(count, block) = item {
+            let reps = count.to_u256().as_usize();
+            for _ in 0..reps {
+                expanded.extend(block.clone());
+            }
+        } else {
+            expanded.push(item);
+        }
+    }
+    expanded
+}
+
 fn inline_constants(body: Vec<Item>, constants: &HashMap<String, U256>) -> Vec<Item> {
     body.into_iter()
         .map(|item| {
@@ -157,8 +173,8 @@ fn find_labels(
     let mut local_labels = HashMap::<String, usize>::new();
     for item in body {
         match item {
-            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) => {
-                panic!("Macros should have been expanded already")
+            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) | Item::Repeat(_, _) => {
+                panic!("Macros and repeats should have been expanded already")
             }
             Item::GlobalLabelDeclaration(label) => {
                 let old = global_labels.insert(label.clone(), *offset);
@@ -185,8 +201,8 @@ fn assemble_file(
     // Assemble the file.
     for item in body {
         match item {
-            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) => {
-                panic!("Macros should have been expanded already")
+            Item::MacroDef(_, _, _) | Item::MacroCall(_, _) | Item::Repeat(_, _) => {
+                panic!("Macros and repeats should have been expanded already")
             }
             Item::GlobalLabelDeclaration(_) | Item::LocalLabelDeclaration(_) => {
                 // Nothing to do; we processed labels in the prior phase.
@@ -393,6 +409,13 @@ mod tests {
         assert_eq!(kernel.code, vec![push4, 0xDE, 0xAD, 0xBE, 0xEF]);
     }
 
+    #[test]
+    fn repeat() {
+        let kernel = parse_and_assemble(&["%rep 3 ADD %endrep"]);
+        let add = get_opcode("ADD");
+        assert_eq!(kernel.code, vec![add, add, add]);
+    }
+
     fn parse_and_assemble(files: &[&str]) -> Kernel {
         parse_and_assemble_with_constants(files, HashMap::new())
     }
diff --git a/evm/src/cpu/kernel/ast.rs b/evm/src/cpu/kernel/ast.rs
index 5025cd99..9bb315ff 100644
--- a/evm/src/cpu/kernel/ast.rs
+++ b/evm/src/cpu/kernel/ast.rs
@@ -12,6 +12,8 @@ pub(crate) enum Item {
     MacroDef(String, Vec<String>, Vec<Item>),
     /// Calls a macro: name, args.
     MacroCall(String, Vec<PushTarget>),
+    /// Repetition, like `%rep` in NASM.
+    Repeat(Literal, Vec<Item>),
     /// Declares a global label.
     GlobalLabelDeclaration(String),
     /// Declares a label that is local to the current file.
diff --git a/evm/src/cpu/kernel/evm_asm.pest b/evm/src/cpu/kernel/evm_asm.pest
index d7f4629a..d5a89d99 100644
--- a/evm/src/cpu/kernel/evm_asm.pest
+++ b/evm/src/cpu/kernel/evm_asm.pest
@@ -15,9 +15,10 @@ literal = { literal_hex | literal_decimal }
 variable = ${ "$" ~ identifier }
 constant = ${ "@" ~ identifier }
 
-item = { macro_def | macro_call | global_label | local_label | bytes_item | push_instruction | nullary_instruction }
+item = { macro_def | macro_call | repeat | global_label | local_label | bytes_item | push_instruction | nullary_instruction }
 macro_def = { ^"%macro" ~ identifier ~ macro_paramlist? ~ item* ~ ^"%endmacro" }
-macro_call = ${ "%" ~ !(^"macro" | ^"endmacro") ~ identifier ~ macro_arglist? }
+macro_call = ${ "%" ~ !(^"macro" | ^"endmacro" | ^"rep" | ^"endrep") ~ identifier ~ macro_arglist? }
+repeat = { ^"%rep" ~ literal ~ item* ~ ^"%endrep" }
 macro_paramlist = { "(" ~ identifier ~ ("," ~ identifier)* ~ ")" }
 macro_arglist = !{ "(" ~ push_target ~ ("," ~ push_target)* ~ ")" }
 global_label = { ^"GLOBAL " ~ identifier ~ ":" }
diff --git a/evm/src/cpu/kernel/parser.rs b/evm/src/cpu/kernel/parser.rs
index c6dd8392..b8ac3f40 100644
--- a/evm/src/cpu/kernel/parser.rs
+++ b/evm/src/cpu/kernel/parser.rs
@@ -23,6 +23,7 @@ fn parse_item(item: Pair<Rule>) -> Item {
     match item.as_rule() {
         Rule::macro_def => parse_macro_def(item),
         Rule::macro_call => parse_macro_call(item),
+        Rule::repeat => parse_repeat(item),
         Rule::global_label => {
             Item::GlobalLabelDeclaration(item.into_inner().next().unwrap().as_str().into())
         }
@@ -70,6 +71,13 @@ fn parse_macro_call(item: Pair<Rule>) -> Item {
     Item::MacroCall(name, args)
 }
 
+fn parse_repeat(item: Pair<Rule>) -> Item {
+    assert_eq!(item.as_rule(), Rule::repeat);
+    let mut inner = item.into_inner().peekable();
+    let count = parse_literal(inner.next().unwrap());
+    Item::Repeat(count, inner.map(parse_item).collect())
+}
+
 fn parse_push_target(target: Pair<Rule>) -> PushTarget {
     assert_eq!(target.as_rule(), Rule::push_target);
     let inner = target.into_inner().next().unwrap();

From 292bb4a02405922bc2d68926a442a5fa1f54357c Mon Sep 17 00:00:00 2001
From: wborgeaud <williamborgeaud@gmail.com>
Date: Fri, 15 Jul 2022 11:10:10 +0200
Subject: [PATCH 3/7] Implement memory for the interpreter

---
 evm/src/cpu/kernel/aggregator.rs  |  48 +++++++-------
 evm/src/cpu/kernel/interpreter.rs | 104 +++++++++++++++++++++++++++---
 2 files changed, 119 insertions(+), 33 deletions(-)

diff --git a/evm/src/cpu/kernel/aggregator.rs b/evm/src/cpu/kernel/aggregator.rs
index 04d57e75..53c13cf3 100644
--- a/evm/src/cpu/kernel/aggregator.rs
+++ b/evm/src/cpu/kernel/aggregator.rs
@@ -70,26 +70,26 @@ mod tests {
 
         // Random input
         let initial_stack = vec![U256::from_str("0xdeadbeef")?, b, a];
-        let stack_with_kernel = run(&kernel.code, exp, initial_stack);
+        let stack_with_kernel = run(&kernel.code, exp, initial_stack).stack;
         let initial_stack = vec![b, a];
         let code = [0xa, 0x63, 0xde, 0xad, 0xbe, 0xef, 0x56]; // EXP, PUSH4 deadbeef, JUMP
-        let stack_with_opcode = run(&code, 0, initial_stack);
+        let stack_with_opcode = run(&code, 0, initial_stack).stack;
         assert_eq!(stack_with_kernel, stack_with_opcode);
 
         // 0 base
         let initial_stack = vec![U256::from_str("0xdeadbeef")?, b, U256::zero()];
-        let stack_with_kernel = run(&kernel.code, exp, initial_stack);
+        let stack_with_kernel = run(&kernel.code, exp, initial_stack).stack;
         let initial_stack = vec![b, U256::zero()];
         let code = [0xa, 0x63, 0xde, 0xad, 0xbe, 0xef, 0x56]; // EXP, PUSH4 deadbeef, JUMP
-        let stack_with_opcode = run(&code, 0, initial_stack);
+        let stack_with_opcode = run(&code, 0, initial_stack).stack;
         assert_eq!(stack_with_kernel, stack_with_opcode);
 
         // 0 exponent
         let initial_stack = vec![U256::from_str("0xdeadbeef")?, U256::zero(), a];
-        let stack_with_kernel = run(&kernel.code, exp, initial_stack);
+        let stack_with_kernel = run(&kernel.code, exp, initial_stack).stack;
         let initial_stack = vec![U256::zero(), a];
         let code = [0xa, 0x63, 0xde, 0xad, 0xbe, 0xef, 0x56]; // EXP, PUSH4 deadbeef, JUMP
-        let stack_with_opcode = run(&code, 0, initial_stack);
+        let stack_with_opcode = run(&code, 0, initial_stack).stack;
         assert_eq!(stack_with_kernel, stack_with_opcode);
 
         Ok(())
@@ -131,76 +131,76 @@ mod tests {
 
         // Standard addition #1
         let initial_stack = u256ify(["0xdeadbeef", point0.1, point0.0, point1.1, point1.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, u256ify([point2.1, point2.0])?);
         // Standard addition #2
         let initial_stack = u256ify(["0xdeadbeef", point1.1, point1.0, point0.1, point0.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, u256ify([point2.1, point2.0])?);
 
         // Standard doubling #1
         let initial_stack = u256ify(["0xdeadbeef", point0.1, point0.0, point0.1, point0.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, u256ify([point3.1, point3.0])?);
         // Standard doubling #2
         let initial_stack = u256ify(["0xdeadbeef", point0.1, point0.0])?;
-        let stack = run(&kernel.code, ec_double, initial_stack);
+        let stack = run(&kernel.code, ec_double, initial_stack).stack;
         assert_eq!(stack, u256ify([point3.1, point3.0])?);
         // Standard doubling #3
         let initial_stack = u256ify(["0xdeadbeef", "0x2", point0.1, point0.0])?;
-        let stack = run(&kernel.code, ec_mul, initial_stack);
+        let stack = run(&kernel.code, ec_mul, initial_stack).stack;
         assert_eq!(stack, u256ify([point3.1, point3.0])?);
 
         // Addition with identity #1
         let initial_stack = u256ify(["0xdeadbeef", identity.1, identity.0, point1.1, point1.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, u256ify([point1.1, point1.0])?);
         // Addition with identity #2
         let initial_stack = u256ify(["0xdeadbeef", point1.1, point1.0, identity.1, identity.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, u256ify([point1.1, point1.0])?);
         // Addition with identity #3
         let initial_stack =
             u256ify(["0xdeadbeef", identity.1, identity.0, identity.1, identity.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, u256ify([identity.1, identity.0])?);
 
         // Addition with invalid point(s) #1
         let initial_stack = u256ify(["0xdeadbeef", point0.1, point0.0, invalid.1, invalid.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, vec![U256::MAX, U256::MAX]);
         // Addition with invalid point(s) #2
         let initial_stack = u256ify(["0xdeadbeef", invalid.1, invalid.0, point0.1, point0.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, vec![U256::MAX, U256::MAX]);
         // Addition with invalid point(s) #3
         let initial_stack = u256ify(["0xdeadbeef", invalid.1, invalid.0, identity.1, identity.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, vec![U256::MAX, U256::MAX]);
         // Addition with invalid point(s) #4
         let initial_stack = u256ify(["0xdeadbeef", invalid.1, invalid.0, invalid.1, invalid.0])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, vec![U256::MAX, U256::MAX]);
 
         // Scalar multiplication #1
         let initial_stack = u256ify(["0xdeadbeef", s, point0.1, point0.0])?;
-        let stack = run(&kernel.code, ec_mul, initial_stack);
+        let stack = run(&kernel.code, ec_mul, initial_stack).stack;
         assert_eq!(stack, u256ify([point4.1, point4.0])?);
         // Scalar multiplication #2
         let initial_stack = u256ify(["0xdeadbeef", "0x0", point0.1, point0.0])?;
-        let stack = run(&kernel.code, ec_mul, initial_stack);
+        let stack = run(&kernel.code, ec_mul, initial_stack).stack;
         assert_eq!(stack, u256ify([identity.1, identity.0])?);
         // Scalar multiplication #3
         let initial_stack = u256ify(["0xdeadbeef", "0x1", point0.1, point0.0])?;
-        let stack = run(&kernel.code, ec_mul, initial_stack);
+        let stack = run(&kernel.code, ec_mul, initial_stack).stack;
         assert_eq!(stack, u256ify([point0.1, point0.0])?);
         // Scalar multiplication #4
         let initial_stack = u256ify(["0xdeadbeef", s, identity.1, identity.0])?;
-        let stack = run(&kernel.code, ec_mul, initial_stack);
+        let stack = run(&kernel.code, ec_mul, initial_stack).stack;
         assert_eq!(stack, u256ify([identity.1, identity.0])?);
         // Scalar multiplication #5
         let initial_stack = u256ify(["0xdeadbeef", s, invalid.1, invalid.0])?;
-        let stack = run(&kernel.code, ec_mul, initial_stack);
+        let stack = run(&kernel.code, ec_mul, initial_stack).stack;
         assert_eq!(stack, vec![U256::MAX, U256::MAX]);
 
         // Multiple calls
@@ -214,7 +214,7 @@ mod tests {
             point0.1,
             point0.0,
         ])?;
-        let stack = run(&kernel.code, ec_add, initial_stack);
+        let stack = run(&kernel.code, ec_add, initial_stack).stack;
         assert_eq!(stack, u256ify([point4.1, point4.0])?);
 
         Ok(())
diff --git a/evm/src/cpu/kernel/interpreter.rs b/evm/src/cpu/kernel/interpreter.rs
index 09e493b9..81ce287b 100644
--- a/evm/src/cpu/kernel/interpreter.rs
+++ b/evm/src/cpu/kernel/interpreter.rs
@@ -3,20 +3,61 @@ use ethereum_types::{U256, U512};
 /// Halt interpreter execution whenever a jump to this offset is done.
 const HALT_OFFSET: usize = 0xdeadbeef;
 
-struct Interpreter<'a> {
+#[derive(Debug, Default)]
+pub(crate) struct EvmMemory {
+    memory: Vec<u8>,
+}
+
+impl EvmMemory {
+    fn len(&self) -> usize {
+        self.memory.len()
+    }
+
+    /// Expand memory until `self.len() >= offset`.
+    fn expand(&mut self, offset: usize) {
+        while self.len() < offset {
+            self.memory.extend([0; 32]);
+        }
+    }
+
+    fn mload(&mut self, offset: usize) -> U256 {
+        self.expand(offset + 32);
+        U256::from_big_endian(&self.memory[offset..offset + 32])
+    }
+
+    fn mstore(&mut self, offset: usize, value: U256) {
+        self.expand(offset + 32);
+        let value_be = {
+            let mut tmp = [0; 32];
+            value.to_big_endian(&mut tmp);
+            tmp
+        };
+        self.memory[offset..offset + 32].copy_from_slice(&value_be);
+    }
+
+    fn mstore8(&mut self, offset: usize, value: U256) {
+        self.expand(offset + 1);
+        let value_byte = value.0[0] as u8;
+        self.memory[offset] = value_byte;
+    }
+}
+
+pub(crate) struct Interpreter<'a> {
     code: &'a [u8],
     jumpdests: Vec<usize>,
     offset: usize,
-    stack: Vec<U256>,
+    pub(crate) stack: Vec<U256>,
+    pub(crate) memory: EvmMemory,
     running: bool,
 }
 
-pub fn run(code: &[u8], initial_offset: usize, initial_stack: Vec<U256>) -> Vec<U256> {
+pub(crate) fn run(code: &[u8], initial_offset: usize, initial_stack: Vec<U256>) -> Interpreter {
     let mut interpreter = Interpreter {
         code,
         jumpdests: find_jumpdests(code),
         offset: initial_offset,
         stack: initial_stack,
+        memory: EvmMemory::default(),
         running: true,
     };
 
@@ -24,7 +65,7 @@ pub fn run(code: &[u8], initial_offset: usize, initial_stack: Vec<U256>) -> Vec<
         interpreter.run_opcode();
     }
 
-    interpreter.stack
+    interpreter
 }
 
 impl<'a> Interpreter<'a> {
@@ -104,9 +145,9 @@ impl<'a> Interpreter<'a> {
             0x46 => todo!(),                                           // "CHAINID",
             0x48 => todo!(),                                           // "BASEFEE",
             0x50 => self.run_pop(),                                    // "POP",
-            0x51 => todo!(),                                           // "MLOAD",
-            0x52 => todo!(),                                           // "MSTORE",
-            0x53 => todo!(),                                           // "MSTORE8",
+            0x51 => self.run_mload(),                                  // "MLOAD",
+            0x52 => self.run_mstore(),                                 // "MSTORE",
+            0x53 => self.run_mstore8(),                                // "MSTORE8",
             0x54 => todo!(),                                           // "SLOAD",
             0x55 => todo!(),                                           // "SSTORE",
             0x56 => self.run_jump(),                                   // "JUMP",
@@ -249,6 +290,24 @@ impl<'a> Interpreter<'a> {
         self.pop();
     }
 
+    fn run_mload(&mut self) {
+        let offset = self.pop();
+        let value = self.memory.mload(offset.as_usize());
+        self.push(value);
+    }
+
+    fn run_mstore(&mut self) {
+        let offset = self.pop();
+        let value = self.pop();
+        self.memory.mstore(offset.as_usize(), value);
+    }
+
+    fn run_mstore8(&mut self) {
+        let offset = self.pop();
+        let value = self.pop();
+        self.memory.mstore8(offset.as_usize(), value);
+    }
+
     fn run_jump(&mut self) {
         let x = self.pop().as_usize();
         self.offset = x;
@@ -306,13 +365,40 @@ fn find_jumpdests(code: &[u8]) -> Vec<usize> {
 
 #[cfg(test)]
 mod tests {
-    use crate::cpu::kernel::interpreter::run;
+    use hex_literal::hex;
+
+    use crate::cpu::kernel::interpreter::{run, Interpreter};
 
     #[test]
     fn test_run() {
         let code = vec![
             0x60, 0x1, 0x60, 0x2, 0x1, 0x63, 0xde, 0xad, 0xbe, 0xef, 0x56,
         ]; // PUSH1, 1, PUSH1, 2, ADD, PUSH4 deadbeef, JUMP
-        assert_eq!(run(&code, 0, vec![]), vec![0x3.into()]);
+        assert_eq!(run(&code, 0, vec![]).stack, vec![0x3.into()]);
+    }
+
+    #[test]
+    fn test_run_with_memory() {
+        //         PUSH1 0xff
+        //         PUSH1 0
+        //         MSTORE
+
+        //         PUSH1 0
+        //         MLOAD
+
+        //         PUSH1 1
+        //         MLOAD
+
+        //         PUSH1 0x42
+        //         PUSH1 0x27
+        //         MSTORE8
+        let code = vec![
+            0x60, 0xff, 0x60, 0x0, 0x52, 0x60, 0, 0x51, 0x60, 0x1, 0x51, 0x60, 0x42, 0x60, 0x27,
+            0x53,
+        ];
+        let run = run(&code, 0, vec![]);
+        let Interpreter { stack, memory, .. } = run;
+        assert_eq!(stack, vec![0xff.into(), 0xff00.into()]);
+        assert_eq!(&memory.memory, &hex!("00000000000000000000000000000000000000000000000000000000000000ff0000000000000042000000000000000000000000000000000000000000000000"));
     }
 }

From 134c66b37d14ea7c8b29714cd2688dfc6181ecd9 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Fri, 15 Jul 2022 13:02:56 -0700
Subject: [PATCH 4/7] Missing TODO

---
 evm/src/cpu/bootstrap_kernel.rs | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/evm/src/cpu/bootstrap_kernel.rs b/evm/src/cpu/bootstrap_kernel.rs
index 3d5d4f1c..7b11bf60 100644
--- a/evm/src/cpu/bootstrap_kernel.rs
+++ b/evm/src/cpu/bootstrap_kernel.rs
@@ -92,6 +92,7 @@ pub(crate) fn eval_bootstrap_kernel<F: Field, P: PackedField<Scalar = F>>(
     // - If CLOCK is a multiple of KECCAK_RATE_LIMBS, activate the Keccak CTL, and ensure the output
     //   is copied to the next row (besides the first limb which will immediately be overwritten).
     // - Otherwise, ensure that the Keccak input is copied to the next row (besides the next limb).
+    // - The next limb we add to the buffer is also written to memory.
 
     // If IS_BOOTSTRAP_KERNEL changed (from 1 to 0), check that
     // - the clock is a multiple of KECCAK_RATE_LIMBS (TODO)
@@ -131,6 +132,7 @@ pub(crate) fn eval_bootstrap_kernel_circuit<F: RichField + Extendable<D>, const
     // - If CLOCK is a multiple of KECCAK_RATE_LIMBS, activate the Keccak CTL, and ensure the output
     //   is copied to the next row (besides the first limb which will immediately be overwritten).
     // - Otherwise, ensure that the Keccak input is copied to the next row (besides the next limb).
+    // - The next limb we add to the buffer is also written to memory.
 
     // If IS_BOOTSTRAP_KERNEL changed (from 1 to 0), check that
     // - the clock is a multiple of KECCAK_RATE_LIMBS (TODO)

From ab5abc391d308cada7f634b50aa201451bde034d Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sat, 16 Jul 2022 09:59:23 -0700
Subject: [PATCH 5/7] Organize segments in an enum

It's a bit more type-safe (can't mix up segment with context or virtual addr), and this way uniqueness of ordinals is enforced, partially addressing a concern raised in #591.

To avoid making `Segment` public (which I don't think would be appropriate), I had to make some other visibility changes, and had to move `generate_random_memory_ops` into the test module.
---
 evm/src/all_stark.rs             |   3 +-
 evm/src/cpu/bootstrap_kernel.rs  |   5 +-
 evm/src/cpu/kernel/aggregator.rs |   5 +-
 evm/src/generation/memory.rs     |   4 +-
 evm/src/generation/state.rs      |   9 +-
 evm/src/memory/memory_stark.rs   | 175 ++++++++++++++++---------------
 evm/src/memory/segments.rs       |  77 ++++++++++----
 7 files changed, 165 insertions(+), 113 deletions(-)

diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs
index f02e0202..ba157fc0 100644
--- a/evm/src/all_stark.rs
+++ b/evm/src/all_stark.rs
@@ -146,7 +146,8 @@ mod tests {
     use crate::cross_table_lookup::testutils::check_ctls;
     use crate::keccak::keccak_stark::{KeccakStark, NUM_INPUTS, NUM_ROUNDS};
     use crate::logic::{self, LogicStark, Operation};
-    use crate::memory::memory_stark::{generate_random_memory_ops, MemoryStark};
+    use crate::memory::memory_stark::tests::generate_random_memory_ops;
+    use crate::memory::memory_stark::MemoryStark;
     use crate::memory::NUM_CHANNELS;
     use crate::proof::AllProof;
     use crate::prover::prove;
diff --git a/evm/src/cpu/bootstrap_kernel.rs b/evm/src/cpu/bootstrap_kernel.rs
index 7b11bf60..bb0e2be9 100644
--- a/evm/src/cpu/bootstrap_kernel.rs
+++ b/evm/src/cpu/bootstrap_kernel.rs
@@ -18,7 +18,8 @@ use crate::cpu::kernel::keccak_util::keccakf_u32s;
 use crate::cpu::public_inputs::NUM_PUBLIC_INPUTS;
 use crate::generation::state::GenerationState;
 use crate::memory;
-use crate::memory::{segments, NUM_CHANNELS};
+use crate::memory::segments::Segment;
+use crate::memory::NUM_CHANNELS;
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
 /// The Keccak rate (1088 bits), measured in bytes.
@@ -53,7 +54,7 @@ pub(crate) fn generate_bootstrap_kernel<F: Field>(state: &mut GenerationState<F>
             value[0] = F::from_canonical_u8(byte);
 
             let channel = addr % NUM_CHANNELS;
-            state.set_mem_current(channel, segments::CODE, addr, value);
+            state.set_mem_current(channel, Segment::Code, addr, value);
 
             packed_bytes = (packed_bytes << 8) | byte as u32;
         }
diff --git a/evm/src/cpu/kernel/aggregator.rs b/evm/src/cpu/kernel/aggregator.rs
index 6ca88ba1..28a9c597 100644
--- a/evm/src/cpu/kernel/aggregator.rs
+++ b/evm/src/cpu/kernel/aggregator.rs
@@ -8,12 +8,15 @@ use once_cell::sync::Lazy;
 
 use super::assembler::{assemble, Kernel};
 use crate::cpu::kernel::parser::parse;
+use crate::memory::segments::Segment;
 
 pub static KERNEL: Lazy<Kernel> = Lazy::new(combined_kernel);
 
 pub fn evm_constants() -> HashMap<String, U256> {
     let mut c = HashMap::new();
-    c.insert("SEGMENT_ID_TXN_DATA".into(), 0.into()); // TODO: Replace with actual segment ID.
+    for segment in Segment::all() {
+        c.insert(segment.var_name().into(), (segment as u32).into());
+    }
     c
 }
 
diff --git a/evm/src/generation/memory.rs b/evm/src/generation/memory.rs
index dfff4388..2ef46d15 100644
--- a/evm/src/generation/memory.rs
+++ b/evm/src/generation/memory.rs
@@ -1,7 +1,7 @@
 use plonky2::field::types::Field;
 
 use crate::memory::memory_stark::MemoryOp;
-use crate::memory::segments::NUM_SEGMENTS;
+use crate::memory::segments::Segment;
 use crate::memory::VALUE_LIMBS;
 
 #[allow(unused)] // TODO: Should be used soon.
@@ -26,7 +26,7 @@ impl<F: Field> Default for MemoryState<F> {
 #[derive(Default, Debug)]
 pub(crate) struct MemoryContextState<F: Field> {
     /// The content of each memory segment.
-    pub segments: [MemorySegmentState<F>; NUM_SEGMENTS],
+    pub segments: [MemorySegmentState<F>; Segment::COUNT],
 }
 
 #[derive(Default, Debug)]
diff --git a/evm/src/generation/state.rs b/evm/src/generation/state.rs
index 7a95f7e4..c5a6bbc4 100644
--- a/evm/src/generation/state.rs
+++ b/evm/src/generation/state.rs
@@ -6,6 +6,7 @@ use plonky2::field::types::Field;
 use crate::cpu::columns::{CpuColumnsView, NUM_CPU_COLUMNS};
 use crate::generation::memory::MemoryState;
 use crate::memory::memory_stark::MemoryOp;
+use crate::memory::segments::Segment;
 use crate::{keccak, logic};
 
 #[derive(Debug)]
@@ -52,12 +53,12 @@ impl<F: Field> GenerationState<F> {
     pub(crate) fn get_mem_current(
         &mut self,
         channel_index: usize,
-        segment: usize,
+        segment: Segment,
         virt: usize,
     ) -> [F; crate::memory::VALUE_LIMBS] {
         let timestamp = self.cpu_rows.len();
         let context = self.current_context;
-        let value = self.memory.contexts[context].segments[segment].get(virt);
+        let value = self.memory.contexts[context].segments[segment as usize].get(virt);
         self.memory.log.push(MemoryOp {
             channel_index: Some(channel_index),
             timestamp,
@@ -74,7 +75,7 @@ impl<F: Field> GenerationState<F> {
     pub(crate) fn set_mem_current(
         &mut self,
         channel_index: usize,
-        segment: usize,
+        segment: Segment,
         virt: usize,
         value: [F; crate::memory::VALUE_LIMBS],
     ) {
@@ -89,7 +90,7 @@ impl<F: Field> GenerationState<F> {
             virt,
             value,
         });
-        self.memory.contexts[context].segments[segment].set(virt, value)
+        self.memory.contexts[context].segments[segment as usize].set(virt, value)
     }
 
     pub(crate) fn commit_cpu_row(&mut self) {
diff --git a/evm/src/memory/memory_stark.rs b/evm/src/memory/memory_stark.rs
index f150323b..843dfc2f 100644
--- a/evm/src/memory/memory_stark.rs
+++ b/evm/src/memory/memory_stark.rs
@@ -1,4 +1,3 @@
-use std::collections::{HashMap, HashSet};
 use std::marker::PhantomData;
 
 use itertools::Itertools;
@@ -10,7 +9,6 @@ use plonky2::hash::hash_types::RichField;
 use plonky2::timed;
 use plonky2::util::timing::TimingTree;
 use plonky2::util::transpose;
-use rand::Rng;
 use rayon::prelude::*;
 
 use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
@@ -21,6 +19,7 @@ use crate::memory::columns::{
     COUNTER, COUNTER_PERMUTED, IS_READ, NUM_COLUMNS, RANGE_CHECK, RANGE_CHECK_PERMUTED,
     SEGMENT_FIRST_CHANGE, TIMESTAMP, VIRTUAL_FIRST_CHANGE,
 };
+use crate::memory::segments::Segment;
 use crate::memory::{NUM_CHANNELS, VALUE_LIMBS};
 use crate::permutation::PermutationPair;
 use crate::stark::Stark;
@@ -46,13 +45,13 @@ pub struct MemoryStark<F, const D: usize> {
 }
 
 #[derive(Clone, Debug)]
-pub struct MemoryOp<F> {
+pub(crate) struct MemoryOp<F> {
     /// The channel this operation came from, or `None` if it's a dummy operation for padding.
     pub channel_index: Option<usize>,
     pub timestamp: usize,
     pub is_read: bool,
     pub context: usize,
-    pub segment: usize,
+    pub segment: Segment,
     pub virt: usize,
     pub value: [F; 8],
 }
@@ -70,7 +69,7 @@ impl<F: Field> MemoryOp<F> {
         row[TIMESTAMP] = F::from_canonical_usize(self.timestamp);
         row[IS_READ] = F::from_bool(self.is_read);
         row[ADDR_CONTEXT] = F::from_canonical_usize(self.context);
-        row[ADDR_SEGMENT] = F::from_canonical_usize(self.segment);
+        row[ADDR_SEGMENT] = F::from_canonical_usize(self.segment as usize);
         row[ADDR_VIRTUAL] = F::from_canonical_usize(self.virt);
         for j in 0..VALUE_LIMBS {
             row[value_limb(j)] = self.value[j];
@@ -79,79 +78,6 @@ impl<F: Field> MemoryOp<F> {
     }
 }
 
-pub fn generate_random_memory_ops<F: RichField, R: Rng>(
-    num_ops: usize,
-    rng: &mut R,
-) -> Vec<MemoryOp<F>> {
-    let mut memory_ops = Vec::new();
-
-    let mut current_memory_values: HashMap<(usize, usize, usize), [F; 8]> = HashMap::new();
-    let num_cycles = num_ops / 2;
-    for clock in 0..num_cycles {
-        let mut used_indices = HashSet::new();
-        let mut new_writes_this_cycle = HashMap::new();
-        let mut has_read = false;
-        for _ in 0..2 {
-            let mut channel_index = rng.gen_range(0..NUM_CHANNELS);
-            while used_indices.contains(&channel_index) {
-                channel_index = rng.gen_range(0..NUM_CHANNELS);
-            }
-            used_indices.insert(channel_index);
-
-            let is_read = if clock == 0 {
-                false
-            } else {
-                !has_read && rng.gen()
-            };
-            has_read = has_read || is_read;
-
-            let (context, segment, virt, vals) = if is_read {
-                let written: Vec<_> = current_memory_values.keys().collect();
-                let &(context, segment, virt) = written[rng.gen_range(0..written.len())];
-                let &vals = current_memory_values
-                    .get(&(context, segment, virt))
-                    .unwrap();
-
-                (context, segment, virt, vals)
-            } else {
-                // TODO: with taller memory table or more padding (to enable range-checking bigger diffs),
-                // test larger address values.
-                let mut context = rng.gen_range(0..40);
-                let mut segment = rng.gen_range(0..8);
-                let mut virt = rng.gen_range(0..20);
-                while new_writes_this_cycle.contains_key(&(context, segment, virt)) {
-                    context = rng.gen_range(0..40);
-                    segment = rng.gen_range(0..8);
-                    virt = rng.gen_range(0..20);
-                }
-
-                let val: [u32; 8] = rng.gen();
-                let vals: [F; 8] = val.map(F::from_canonical_u32);
-
-                new_writes_this_cycle.insert((context, segment, virt), vals);
-
-                (context, segment, virt, vals)
-            };
-
-            let timestamp = clock * NUM_CHANNELS + channel_index;
-            memory_ops.push(MemoryOp {
-                channel_index: Some(channel_index),
-                timestamp,
-                is_read,
-                context,
-                segment,
-                virt,
-                value: vals,
-            });
-        }
-        for (k, v) in new_writes_this_cycle {
-            current_memory_values.insert(k, v);
-        }
-    }
-
-    memory_ops
-}
-
 fn get_max_range_check<F: Field>(memory_ops: &[MemoryOp<F>]) -> usize {
     memory_ops
         .iter()
@@ -160,7 +86,7 @@ fn get_max_range_check<F: Field>(memory_ops: &[MemoryOp<F>]) -> usize {
             if curr.context != next.context {
                 next.context - curr.context - 1
             } else if curr.segment != next.segment {
-                next.segment - curr.segment - 1
+                next.segment as usize - curr.segment as usize - 1
             } else if curr.virt != next.virt {
                 next.virt - curr.virt - 1
             } else {
@@ -264,7 +190,7 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
         }
     }
 
-    pub fn generate_trace(&self, memory_ops: Vec<MemoryOp<F>>) -> Vec<PolynomialValues<F>> {
+    pub(crate) fn generate_trace(&self, memory_ops: Vec<MemoryOp<F>>) -> Vec<PolynomialValues<F>> {
         let mut timing = TimingTree::new("generate trace", log::Level::Debug);
 
         // Generate most of the trace in row-major form.
@@ -533,13 +459,94 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for MemoryStark<F
 }
 
 #[cfg(test)]
-mod tests {
-    use anyhow::Result;
-    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+pub(crate) mod tests {
+    use std::collections::{HashMap, HashSet};
 
-    use crate::memory::memory_stark::MemoryStark;
+    use anyhow::Result;
+    use plonky2::hash::hash_types::RichField;
+    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+    use rand::prelude::SliceRandom;
+    use rand::Rng;
+
+    use crate::memory::memory_stark::{MemoryOp, MemoryStark};
+    use crate::memory::segments::Segment;
+    use crate::memory::NUM_CHANNELS;
     use crate::stark_testing::{test_stark_circuit_constraints, test_stark_low_degree};
 
+    pub(crate) fn generate_random_memory_ops<F: RichField, R: Rng>(
+        num_ops: usize,
+        rng: &mut R,
+    ) -> Vec<MemoryOp<F>> {
+        let mut memory_ops = Vec::new();
+
+        let mut current_memory_values: HashMap<(usize, Segment, usize), [F; 8]> = HashMap::new();
+        let num_cycles = num_ops / 2;
+        for clock in 0..num_cycles {
+            let mut used_indices = HashSet::new();
+            let mut new_writes_this_cycle = HashMap::new();
+            let mut has_read = false;
+            for _ in 0..2 {
+                let mut channel_index = rng.gen_range(0..NUM_CHANNELS);
+                while used_indices.contains(&channel_index) {
+                    channel_index = rng.gen_range(0..NUM_CHANNELS);
+                }
+                used_indices.insert(channel_index);
+
+                let is_read = if clock == 0 {
+                    false
+                } else {
+                    !has_read && rng.gen()
+                };
+                has_read = has_read || is_read;
+
+                let (context, segment, virt, vals) = if is_read {
+                    let written: Vec<_> = current_memory_values.keys().collect();
+                    let &(context, segment, virt) = written[rng.gen_range(0..written.len())];
+                    let &vals = current_memory_values
+                        .get(&(context, segment, virt))
+                        .unwrap();
+
+                    (context, segment, virt, vals)
+                } else {
+                    // TODO: with taller memory table or more padding (to enable range-checking bigger diffs),
+                    // test larger address values.
+                    let mut context = rng.gen_range(0..40);
+                    let segments = [Segment::Code, Segment::Stack, Segment::MainMemory];
+                    let mut segment = *segments.choose(rng).unwrap();
+                    let mut virt = rng.gen_range(0..20);
+                    while new_writes_this_cycle.contains_key(&(context, segment, virt)) {
+                        context = rng.gen_range(0..40);
+                        segment = *segments.choose(rng).unwrap();
+                        virt = rng.gen_range(0..20);
+                    }
+
+                    let val: [u32; 8] = rng.gen();
+                    let vals: [F; 8] = val.map(F::from_canonical_u32);
+
+                    new_writes_this_cycle.insert((context, segment, virt), vals);
+
+                    (context, segment, virt, vals)
+                };
+
+                let timestamp = clock * NUM_CHANNELS + channel_index;
+                memory_ops.push(MemoryOp {
+                    channel_index: Some(channel_index),
+                    timestamp,
+                    is_read,
+                    context,
+                    segment,
+                    virt,
+                    value: vals,
+                });
+            }
+            for (k, v) in new_writes_this_cycle {
+                current_memory_values.insert(k, v);
+            }
+        }
+
+        memory_ops
+    }
+
     #[test]
     fn test_stark_degree() -> Result<()> {
         const D: usize = 2;
diff --git a/evm/src/memory/segments.rs b/evm/src/memory/segments.rs
index d20cb037..f1b92dfc 100644
--- a/evm/src/memory/segments.rs
+++ b/evm/src/memory/segments.rs
@@ -1,22 +1,61 @@
-/// Contains EVM bytecode.
-pub const CODE: usize = 0;
+#[allow(dead_code)] // TODO: Not all segments are used yet.
+#[derive(Copy, Clone, Eq, PartialEq, Hash, Ord, PartialOrd, Debug)]
+pub(crate) enum Segment {
+    /// Contains EVM bytecode.
+    Code = 0,
+    /// The program stack.
+    Stack = 1,
+    /// Main memory, owned by the contract code.
+    MainMemory = 2,
+    /// Data passed to the current context by its caller.
+    Calldata = 3,
+    /// Data returned to the current context by its latest callee.
+    Returndata = 4,
+    /// A segment which contains a few fixed-size metadata fields, such as the caller's context, or the
+    /// size of `CALLDATA` and `RETURNDATA`.
+    Metadata = 5,
+    /// General purpose kernel memory, used by various kernel functions.
+    /// In general, calling a helper function can result in this memory being clobbered.
+    KernelGeneral = 6,
+    /// Contains transaction data (after it's parsed and converted to a standard format).
+    TxnData = 7,
+    /// Raw RLP data.
+    RlpRaw = 8,
+    /// RLP data that has been parsed and converted to a more "friendly" format.
+    RlpParsed = 9,
+}
 
-pub const STACK: usize = 1;
+impl Segment {
+    pub(crate) const COUNT: usize = 10;
 
-/// Main memory, owned by the contract code.
-pub const MAIN_MEM: usize = 2;
+    pub(crate) fn all() -> [Self; Self::COUNT] {
+        [
+            Self::Code,
+            Self::Stack,
+            Self::MainMemory,
+            Self::Calldata,
+            Self::Returndata,
+            Self::Metadata,
+            Self::KernelGeneral,
+            Self::TxnData,
+            Self::RlpRaw,
+            Self::RlpParsed,
+        ]
+    }
 
-/// Memory owned by the kernel.
-pub const KERNEL_MEM: usize = 3;
-
-/// Data passed to the current context by its caller.
-pub const CALLDATA: usize = 4;
-
-/// Data returned to the current context by its latest callee.
-pub const RETURNDATA: usize = 5;
-
-/// A segment which contains a few fixed-size metadata fields, such as the caller's context, or the
-/// size of `CALLDATA` and `RETURNDATA`.
-pub const METADATA: usize = 6;
-
-pub const NUM_SEGMENTS: usize = 7;
+    /// The variable name that gets passed into kernel assembly code.
+    pub(crate) fn var_name(&self) -> &'static str {
+        match self {
+            Segment::Code => "SEGMENT_CODE",
+            Segment::Stack => "SEGMENT_STACK",
+            Segment::MainMemory => "SEGMENT_MAIN_MEMORY",
+            Segment::Calldata => "SEGMENT_CALLDATA",
+            Segment::Returndata => "SEGMENT_RETURNDATA",
+            Segment::Metadata => "SEGMENT_METADATA",
+            Segment::KernelGeneral => "SEGMENT_KERNEL_GENERAL",
+            Segment::TxnData => "SEGMENT_TXN_DATA",
+            Segment::RlpRaw => "SEGMENT_RLP_RAW",
+            Segment::RlpParsed => "SEGMENT_RLP_PARSED",
+        }
+    }
+}

From 997453237f4e5ab1cd34c22d488c50f5211b8224 Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Sat, 16 Jul 2022 09:14:51 -0700
Subject: [PATCH 6/7] Store memory values as `U256`s

Ultimately they're encoded as `[F; 8]`s in the table, but I don't anticipate that we'll have any use cases where we want to store more than 256 bits. Might as well store `U256` until we actually build the table since they're more compact.
---
 evm/src/cpu/bootstrap_kernel.rs |  6 +-----
 evm/src/generation/memory.rs    | 30 +++++++++++++---------------
 evm/src/generation/state.rs     |  6 +++---
 evm/src/memory/columns.rs       |  3 ++-
 evm/src/memory/memory_stark.rs  | 35 +++++++++++++++------------------
 5 files changed, 36 insertions(+), 44 deletions(-)

diff --git a/evm/src/cpu/bootstrap_kernel.rs b/evm/src/cpu/bootstrap_kernel.rs
index bb0e2be9..af307a28 100644
--- a/evm/src/cpu/bootstrap_kernel.rs
+++ b/evm/src/cpu/bootstrap_kernel.rs
@@ -17,7 +17,6 @@ use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::keccak_util::keccakf_u32s;
 use crate::cpu::public_inputs::NUM_PUBLIC_INPUTS;
 use crate::generation::state::GenerationState;
-use crate::memory;
 use crate::memory::segments::Segment;
 use crate::memory::NUM_CHANNELS;
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
@@ -50,11 +49,8 @@ pub(crate) fn generate_bootstrap_kernel<F: Field>(state: &mut GenerationState<F>
         // Write this chunk to memory, while simultaneously packing its bytes into a u32 word.
         let mut packed_bytes: u32 = 0;
         for (addr, byte) in chunk {
-            let mut value = [F::ZERO; memory::VALUE_LIMBS];
-            value[0] = F::from_canonical_u8(byte);
-
             let channel = addr % NUM_CHANNELS;
-            state.set_mem_current(channel, Segment::Code, addr, value);
+            state.set_mem_current(channel, Segment::Code, addr, byte.into());
 
             packed_bytes = (packed_bytes << 8) | byte as u32;
         }
diff --git a/evm/src/generation/memory.rs b/evm/src/generation/memory.rs
index 2ef46d15..0b80756a 100644
--- a/evm/src/generation/memory.rs
+++ b/evm/src/generation/memory.rs
@@ -1,19 +1,18 @@
-use plonky2::field::types::Field;
+use ethereum_types::U256;
 
 use crate::memory::memory_stark::MemoryOp;
 use crate::memory::segments::Segment;
-use crate::memory::VALUE_LIMBS;
 
 #[allow(unused)] // TODO: Should be used soon.
 #[derive(Debug)]
-pub(crate) struct MemoryState<F: Field> {
+pub(crate) struct MemoryState {
     /// A log of each memory operation, in the order that it occurred.
-    pub log: Vec<MemoryOp<F>>,
+    pub log: Vec<MemoryOp>,
 
-    pub contexts: Vec<MemoryContextState<F>>,
+    pub contexts: Vec<MemoryContextState>,
 }
 
-impl<F: Field> Default for MemoryState<F> {
+impl Default for MemoryState {
     fn default() -> Self {
         Self {
             log: vec![],
@@ -24,28 +23,27 @@ impl<F: Field> Default for MemoryState<F> {
 }
 
 #[derive(Default, Debug)]
-pub(crate) struct MemoryContextState<F: Field> {
+pub(crate) struct MemoryContextState {
     /// The content of each memory segment.
-    pub segments: [MemorySegmentState<F>; Segment::COUNT],
+    pub segments: [MemorySegmentState; Segment::COUNT],
 }
 
 #[derive(Default, Debug)]
-pub(crate) struct MemorySegmentState<F: Field> {
-    pub content: Vec<[F; VALUE_LIMBS]>,
+pub(crate) struct MemorySegmentState {
+    pub content: Vec<U256>,
 }
 
-impl<F: Field> MemorySegmentState<F> {
-    pub(super) fn get(&self, virtual_addr: usize) -> [F; VALUE_LIMBS] {
+impl MemorySegmentState {
+    pub(super) fn get(&self, virtual_addr: usize) -> U256 {
         self.content
             .get(virtual_addr)
             .copied()
-            .unwrap_or([F::ZERO; VALUE_LIMBS])
+            .unwrap_or(U256::zero())
     }
 
-    pub(super) fn set(&mut self, virtual_addr: usize, value: [F; VALUE_LIMBS]) {
+    pub(super) fn set(&mut self, virtual_addr: usize, value: U256) {
         if virtual_addr + 1 > self.content.len() {
-            self.content
-                .resize(virtual_addr + 1, [F::ZERO; VALUE_LIMBS]);
+            self.content.resize(virtual_addr + 1, U256::zero());
         }
         self.content[virtual_addr] = value;
     }
diff --git a/evm/src/generation/state.rs b/evm/src/generation/state.rs
index c5a6bbc4..46ccc4e3 100644
--- a/evm/src/generation/state.rs
+++ b/evm/src/generation/state.rs
@@ -15,7 +15,7 @@ pub(crate) struct GenerationState<F: Field> {
     pub(crate) current_cpu_row: CpuColumnsView<F>,
 
     pub(crate) current_context: usize,
-    pub(crate) memory: MemoryState<F>,
+    pub(crate) memory: MemoryState,
 
     pub(crate) keccak_inputs: Vec<[u64; keccak::keccak_stark::NUM_INPUTS]>,
     pub(crate) logic_ops: Vec<logic::Operation>,
@@ -55,7 +55,7 @@ impl<F: Field> GenerationState<F> {
         channel_index: usize,
         segment: Segment,
         virt: usize,
-    ) -> [F; crate::memory::VALUE_LIMBS] {
+    ) -> U256 {
         let timestamp = self.cpu_rows.len();
         let context = self.current_context;
         let value = self.memory.contexts[context].segments[segment as usize].get(virt);
@@ -77,7 +77,7 @@ impl<F: Field> GenerationState<F> {
         channel_index: usize,
         segment: Segment,
         virt: usize,
-        value: [F; crate::memory::VALUE_LIMBS],
+        value: U256,
     ) {
         let timestamp = self.cpu_rows.len();
         let context = self.current_context;
diff --git a/evm/src/memory/columns.rs b/evm/src/memory/columns.rs
index 5f6c3911..7229a834 100644
--- a/evm/src/memory/columns.rs
+++ b/evm/src/memory/columns.rs
@@ -9,7 +9,8 @@ pub(crate) const ADDR_CONTEXT: usize = IS_READ + 1;
 pub(crate) const ADDR_SEGMENT: usize = ADDR_CONTEXT + 1;
 pub(crate) const ADDR_VIRTUAL: usize = ADDR_SEGMENT + 1;
 
-// Eight limbs to hold up to a 256-bit value.
+// Eight 32-bit limbs hold a total of 256 bits.
+// If a value represents an integer, it is little-endian encoded.
 const VALUE_START: usize = ADDR_VIRTUAL + 1;
 pub(crate) const fn value_limb(i: usize) -> usize {
     debug_assert!(i < VALUE_LIMBS);
diff --git a/evm/src/memory/memory_stark.rs b/evm/src/memory/memory_stark.rs
index 843dfc2f..82e10869 100644
--- a/evm/src/memory/memory_stark.rs
+++ b/evm/src/memory/memory_stark.rs
@@ -1,5 +1,6 @@
 use std::marker::PhantomData;
 
+use ethereum_types::U256;
 use itertools::Itertools;
 use plonky2::field::extension::{Extendable, FieldExtension};
 use plonky2::field::packed::PackedField;
@@ -45,7 +46,7 @@ pub struct MemoryStark<F, const D: usize> {
 }
 
 #[derive(Clone, Debug)]
-pub(crate) struct MemoryOp<F> {
+pub(crate) struct MemoryOp {
     /// The channel this operation came from, or `None` if it's a dummy operation for padding.
     pub channel_index: Option<usize>,
     pub timestamp: usize,
@@ -53,15 +54,15 @@ pub(crate) struct MemoryOp<F> {
     pub context: usize,
     pub segment: Segment,
     pub virt: usize,
-    pub value: [F; 8],
+    pub value: U256,
 }
 
-impl<F: Field> MemoryOp<F> {
+impl MemoryOp {
     /// Generate a row for a given memory operation. Note that this does not generate columns which
     /// depend on the next operation, such as `CONTEXT_FIRST_CHANGE`; those are generated later.
     /// It also does not generate columns such as `COUNTER`, which are generated later, after the
     /// trace has been transposed into column-major form.
-    fn to_row(&self) -> [F; NUM_COLUMNS] {
+    fn to_row<F: Field>(&self) -> [F; NUM_COLUMNS] {
         let mut row = [F::ZERO; NUM_COLUMNS];
         if let Some(channel) = self.channel_index {
             row[is_channel(channel)] = F::ONE;
@@ -72,13 +73,13 @@ impl<F: Field> MemoryOp<F> {
         row[ADDR_SEGMENT] = F::from_canonical_usize(self.segment as usize);
         row[ADDR_VIRTUAL] = F::from_canonical_usize(self.virt);
         for j in 0..VALUE_LIMBS {
-            row[value_limb(j)] = self.value[j];
+            row[value_limb(j)] = F::from_canonical_u32((self.value >> (j * 32)).low_u32());
         }
         row
     }
 }
 
-fn get_max_range_check<F: Field>(memory_ops: &[MemoryOp<F>]) -> usize {
+fn get_max_range_check(memory_ops: &[MemoryOp]) -> usize {
     memory_ops
         .iter()
         .tuple_windows()
@@ -142,7 +143,7 @@ pub fn generate_first_change_flags_and_rc<F: RichField>(trace_rows: &mut [[F; NU
 impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
     /// Generate most of the trace rows. Excludes a few columns like `COUNTER`, which are generated
     /// later, after transposing to column-major form.
-    fn generate_trace_row_major(&self, mut memory_ops: Vec<MemoryOp<F>>) -> Vec<[F; NUM_COLUMNS]> {
+    fn generate_trace_row_major(&self, mut memory_ops: Vec<MemoryOp>) -> Vec<[F; NUM_COLUMNS]> {
         memory_ops.sort_by_key(|op| (op.context, op.segment, op.virt, op.timestamp));
 
         Self::pad_memory_ops(&mut memory_ops);
@@ -167,7 +168,7 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
         trace_col_vecs[COUNTER_PERMUTED] = permuted_table;
     }
 
-    fn pad_memory_ops(memory_ops: &mut Vec<MemoryOp<F>>) {
+    fn pad_memory_ops(memory_ops: &mut Vec<MemoryOp>) {
         let num_ops = memory_ops.len();
         let max_range_check = get_max_range_check(memory_ops);
         let num_ops_padded = num_ops.max(max_range_check + 1).next_power_of_two();
@@ -190,7 +191,7 @@ impl<F: RichField + Extendable<D>, const D: usize> MemoryStark<F, D> {
         }
     }
 
-    pub(crate) fn generate_trace(&self, memory_ops: Vec<MemoryOp<F>>) -> Vec<PolynomialValues<F>> {
+    pub(crate) fn generate_trace(&self, memory_ops: Vec<MemoryOp>) -> Vec<PolynomialValues<F>> {
         let mut timing = TimingTree::new("generate trace", log::Level::Debug);
 
         // Generate most of the trace in row-major form.
@@ -463,7 +464,7 @@ pub(crate) mod tests {
     use std::collections::{HashMap, HashSet};
 
     use anyhow::Result;
-    use plonky2::hash::hash_types::RichField;
+    use ethereum_types::U256;
     use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
     use rand::prelude::SliceRandom;
     use rand::Rng;
@@ -473,13 +474,10 @@ pub(crate) mod tests {
     use crate::memory::NUM_CHANNELS;
     use crate::stark_testing::{test_stark_circuit_constraints, test_stark_low_degree};
 
-    pub(crate) fn generate_random_memory_ops<F: RichField, R: Rng>(
-        num_ops: usize,
-        rng: &mut R,
-    ) -> Vec<MemoryOp<F>> {
+    pub(crate) fn generate_random_memory_ops<R: Rng>(num_ops: usize, rng: &mut R) -> Vec<MemoryOp> {
         let mut memory_ops = Vec::new();
 
-        let mut current_memory_values: HashMap<(usize, Segment, usize), [F; 8]> = HashMap::new();
+        let mut current_memory_values: HashMap<(usize, Segment, usize), U256> = HashMap::new();
         let num_cycles = num_ops / 2;
         for clock in 0..num_cycles {
             let mut used_indices = HashSet::new();
@@ -520,12 +518,11 @@ pub(crate) mod tests {
                         virt = rng.gen_range(0..20);
                     }
 
-                    let val: [u32; 8] = rng.gen();
-                    let vals: [F; 8] = val.map(F::from_canonical_u32);
+                    let val = U256(rng.gen());
 
-                    new_writes_this_cycle.insert((context, segment, virt), vals);
+                    new_writes_this_cycle.insert((context, segment, virt), val);
 
-                    (context, segment, virt, vals)
+                    (context, segment, virt, val)
                 };
 
                 let timestamp = clock * NUM_CHANNELS + channel_index;

From ef842b03c8cef953494ab56aba25757c33e5a51a Mon Sep 17 00:00:00 2001
From: Daniel Lubarov <daniel@lubarov.com>
Date: Wed, 13 Jul 2022 16:45:21 -0700
Subject: [PATCH 7/7] Address some feedback on #591

---
 evm/src/generation/memory.rs | 2 +-
 evm/src/generation/mod.rs    | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/evm/src/generation/memory.rs b/evm/src/generation/memory.rs
index 0b80756a..60bfe794 100644
--- a/evm/src/generation/memory.rs
+++ b/evm/src/generation/memory.rs
@@ -42,7 +42,7 @@ impl MemorySegmentState {
     }
 
     pub(super) fn set(&mut self, virtual_addr: usize, value: U256) {
-        if virtual_addr + 1 > self.content.len() {
+        if virtual_addr >= self.content.len() {
             self.content.resize(virtual_addr + 1, U256::zero());
         }
         self.content[virtual_addr] = value;
diff --git a/evm/src/generation/mod.rs b/evm/src/generation/mod.rs
index edca27dd..b73270db 100644
--- a/evm/src/generation/mod.rs
+++ b/evm/src/generation/mod.rs
@@ -12,6 +12,8 @@ use crate::util::trace_rows_to_poly_values;
 mod memory;
 pub(crate) mod state;
 
+/// A piece of data which has been encoded using Recursive Length Prefix (RLP) serialization.
+/// See https://ethereum.org/en/developers/docs/data-structures-and-encoding/rlp/
 pub type RlpBlob = Vec<u8>;
 
 /// Merkle proofs are encoded using an RLP blob for each node in the path.