Column definitions for addition, range checks & lookups (#477)

* Column definitions for addition, range checks & lookups

I implemented addition (unsigned for now) as an example of how the arithmetic unit can interact with the 16-bit range check unit.

Range checks and lookups aren't implemented yet.

* Missing constraints

* Tweaks to get tests passing

* Reorg registers into files

* Minor

field/src/field_types.rs

@@ -278,6 +278,12 @@ pub trait Field:
         Self::from_canonical_u64(n as u64)
     }
 
+    /// Returns `n`. Assumes that `n` is already in canonical form, i.e. `n < Self::order()`.
+    // TODO: Should probably be unsafe.
+    fn from_canonical_u16(n: u16) -> Self {
+        Self::from_canonical_u64(n as u64)
+    }
+
     /// Returns `n`. Assumes that `n` is already in canonical form, i.e. `n < Self::order()`.
     // TODO: Should probably be unsafe.
     fn from_canonical_usize(n: usize) -> Self {

plonky2/src/plonk/plonk_common.rs

@@ -138,7 +138,7 @@ where
     sum
 }
 
-pub(crate) fn reduce_with_powers_ext_recursive<F: RichField + Extendable<D>, const D: usize>(
+pub fn reduce_with_powers_ext_recursive<F: RichField + Extendable<D>, const D: usize>(
     builder: &mut CircuitBuilder<F, D>,
     terms: &[ExtensionTarget<D>],
     alpha: Target,

starky/src/lib.rs

@@ -1,8 +1,6 @@
 // TODO: Remove these when crate is closer to being finished.
 #![allow(dead_code)]
 #![allow(unused_variables)]
-#![allow(unreachable_code)]
-#![allow(clippy::diverging_sub_expression)]
 #![allow(incomplete_features)]
 #![feature(generic_const_exprs)]
 

system_zero/src/arithmetic/addition.rs

@@ -0,0 +1,70 @@
+use plonky2::field::extension_field::Extendable;
+use plonky2::field::field_types::Field;
+use plonky2::field::packed_field::PackedField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::ext_target::ExtensionTarget;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::plonk_common::reduce_with_powers_ext_recursive;
+use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+
+use crate::registers::arithmetic::*;
+use crate::registers::NUM_COLUMNS;
+
+pub(crate) fn generate_addition<F: RichField>(values: &mut [F; NUM_COLUMNS]) {
+    let in_1 = values[COL_ADD_INPUT_1].to_canonical_u64();
+    let in_2 = values[COL_ADD_INPUT_2].to_canonical_u64();
+    let in_3 = values[COL_ADD_INPUT_3].to_canonical_u64();
+    let output = in_1 + in_2 + in_3;
+
+    values[COL_ADD_OUTPUT_1] = F::from_canonical_u16(output as u16);
+    values[COL_ADD_OUTPUT_2] = F::from_canonical_u16((output >> 16) as u16);
+    values[COL_ADD_OUTPUT_3] = F::from_canonical_u16((output >> 32) as u16);
+}
+
+pub(crate) fn eval_addition<F: Field, P: PackedField<Scalar = F>>(
+    local_values: &[P; NUM_COLUMNS],
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    let is_add = local_values[IS_ADD];
+    let in_1 = local_values[COL_ADD_INPUT_1];
+    let in_2 = local_values[COL_ADD_INPUT_2];
+    let in_3 = local_values[COL_ADD_INPUT_3];
+    let out_1 = local_values[COL_ADD_OUTPUT_1];
+    let out_2 = local_values[COL_ADD_OUTPUT_2];
+    let out_3 = local_values[COL_ADD_OUTPUT_3];
+
+    let weight_2 = F::from_canonical_u64(1 << 16);
+    let weight_3 = F::from_canonical_u64(1 << 32);
+    // Note that this can't overflow. Since each output limb has been range checked as 16-bits,
+    // this sum can be around 48 bits at most.
+    let out = out_1 + out_2 * weight_2 + out_3 * weight_3;
+
+    let computed_out = in_1 + in_2 + in_3;
+
+    yield_constr.constraint_wrapping(is_add * (out - computed_out));
+}
+
+pub(crate) fn eval_addition_recursively<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    local_values: &[ExtensionTarget<D>; NUM_COLUMNS],
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    let is_add = local_values[IS_ADD];
+    let in_1 = local_values[COL_ADD_INPUT_1];
+    let in_2 = local_values[COL_ADD_INPUT_2];
+    let in_3 = local_values[COL_ADD_INPUT_3];
+    let out_1 = local_values[COL_ADD_OUTPUT_1];
+    let out_2 = local_values[COL_ADD_OUTPUT_2];
+    let out_3 = local_values[COL_ADD_OUTPUT_3];
+
+    let limb_base = builder.constant(F::from_canonical_u64(1 << 16));
+    // Note that this can't overflow. Since each output limb has been range checked as 16-bits,
+    // this sum can be around 48 bits at most.
+    let out = reduce_with_powers_ext_recursive(builder, &[out_1, out_2, out_3], limb_base);
+
+    let computed_out = builder.add_many_extension(&[in_1, in_2, in_3]);
+
+    let diff = builder.sub_extension(out, computed_out);
+    let filtered_diff = builder.mul_extension(is_add, diff);
+    yield_constr.constraint_wrapping(builder, filtered_diff);
+}

system_zero/src/arithmetic/division.rs

@@ -0,0 +1,31 @@
+use plonky2::field::extension_field::Extendable;
+use plonky2::field::field_types::Field;
+use plonky2::field::packed_field::PackedField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::ext_target::ExtensionTarget;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+
+use crate::registers::arithmetic::*;
+use crate::registers::NUM_COLUMNS;
+
+pub(crate) fn generate_division<F: RichField>(values: &mut [F; NUM_COLUMNS]) {
+    // TODO
+}
+
+pub(crate) fn eval_division<F: Field, P: PackedField<Scalar = F>>(
+    local_values: &[P; NUM_COLUMNS],
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    let is_div = local_values[IS_DIV];
+    // TODO
+}
+
+pub(crate) fn eval_division_recursively<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    local_values: &[ExtensionTarget<D>; NUM_COLUMNS],
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    let is_div = local_values[IS_DIV];
+    // TODO
+}

system_zero/src/arithmetic/mod.rs

@@ -0,0 +1,75 @@
+use plonky2::field::extension_field::Extendable;
+use plonky2::field::field_types::Field;
+use plonky2::field::packed_field::PackedField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+use starky::vars::StarkEvaluationTargets;
+use starky::vars::StarkEvaluationVars;
+
+use crate::arithmetic::addition::{eval_addition, eval_addition_recursively, generate_addition};
+use crate::arithmetic::division::{eval_division, eval_division_recursively, generate_division};
+use crate::arithmetic::multiplication::{
+    eval_multiplication, eval_multiplication_recursively, generate_multiplication,
+};
+use crate::arithmetic::subtraction::{
+    eval_subtraction, eval_subtraction_recursively, generate_subtraction,
+};
+use crate::public_input_layout::NUM_PUBLIC_INPUTS;
+use crate::registers::arithmetic::*;
+use crate::registers::NUM_COLUMNS;
+
+mod addition;
+mod division;
+mod multiplication;
+mod subtraction;
+
+pub(crate) fn generate_arithmetic_unit<F: RichField>(values: &mut [F; NUM_COLUMNS]) {
+    if values[IS_ADD].is_one() {
+        generate_addition(values);
+    } else if values[IS_SUB].is_one() {
+        generate_subtraction(values);
+    } else if values[IS_MUL].is_one() {
+        generate_multiplication(values);
+    } else if values[IS_DIV].is_one() {
+        generate_division(values);
+    }
+}
+
+pub(crate) fn eval_arithmetic_unit<F: Field, P: PackedField<Scalar = F>>(
+    vars: StarkEvaluationVars<F, P, NUM_COLUMNS, NUM_PUBLIC_INPUTS>,
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    let local_values = &vars.local_values;
+
+    // Check that the operation flag values are binary.
+    for col in [IS_ADD, IS_SUB, IS_MUL, IS_DIV] {
+        let val = local_values[col];
+        yield_constr.constraint_wrapping(val * val - val);
+    }
+
+    eval_addition(local_values, yield_constr);
+    eval_subtraction(local_values, yield_constr);
+    eval_multiplication(local_values, yield_constr);
+    eval_division(local_values, yield_constr);
+}
+
+pub(crate) fn eval_arithmetic_unit_recursively<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    vars: StarkEvaluationTargets<D, NUM_COLUMNS, NUM_PUBLIC_INPUTS>,
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    let local_values = &vars.local_values;
+
+    // Check that the operation flag values are binary.
+    for col in [IS_ADD, IS_SUB, IS_MUL, IS_DIV] {
+        let val = local_values[col];
+        let constraint = builder.mul_add_extension(val, val, val);
+        yield_constr.constraint_wrapping(builder, constraint);
+    }
+
+    eval_addition_recursively(builder, local_values, yield_constr);
+    eval_subtraction_recursively(builder, local_values, yield_constr);
+    eval_multiplication_recursively(builder, local_values, yield_constr);
+    eval_division_recursively(builder, local_values, yield_constr);
+}

system_zero/src/arithmetic/multiplication.rs

@@ -0,0 +1,31 @@
+use plonky2::field::extension_field::Extendable;
+use plonky2::field::field_types::Field;
+use plonky2::field::packed_field::PackedField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::ext_target::ExtensionTarget;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+
+use crate::registers::arithmetic::*;
+use crate::registers::NUM_COLUMNS;
+
+pub(crate) fn generate_multiplication<F: RichField>(values: &mut [F; NUM_COLUMNS]) {
+    // TODO
+}
+
+pub(crate) fn eval_multiplication<F: Field, P: PackedField<Scalar = F>>(
+    local_values: &[P; NUM_COLUMNS],
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    let is_mul = local_values[IS_MUL];
+    // TODO
+}
+
+pub(crate) fn eval_multiplication_recursively<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    local_values: &[ExtensionTarget<D>; NUM_COLUMNS],
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    let is_mul = local_values[IS_MUL];
+    // TODO
+}

system_zero/src/arithmetic/subtraction.rs

@@ -0,0 +1,31 @@
+use plonky2::field::extension_field::Extendable;
+use plonky2::field::field_types::Field;
+use plonky2::field::packed_field::PackedField;
+use plonky2::hash::hash_types::RichField;
+use plonky2::iop::ext_target::ExtensionTarget;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
+
+use crate::registers::arithmetic::*;
+use crate::registers::NUM_COLUMNS;
+
+pub(crate) fn generate_subtraction<F: RichField>(values: &mut [F; NUM_COLUMNS]) {
+    // TODO
+}
+
+pub(crate) fn eval_subtraction<F: Field, P: PackedField<Scalar = F>>(
+    local_values: &[P; NUM_COLUMNS],
+    yield_constr: &mut ConstraintConsumer<P>,
+) {
+    let is_sub = local_values[IS_SUB];
+    // TODO
+}
+
+pub(crate) fn eval_subtraction_recursively<F: RichField + Extendable<D>, const D: usize>(
+    builder: &mut CircuitBuilder<F, D>,
+    local_values: &[ExtensionTarget<D>; NUM_COLUMNS],
+    yield_constr: &mut RecursiveConstraintConsumer<F, D>,
+) {
+    let is_sub = local_values[IS_SUB];
+    // TODO
+}

system_zero/src/column_layout.rs

@@ -1,108 +0,0 @@
-//// CORE REGISTERS
-
-/// A cycle counter. Starts at 0; increments by 1.
-pub(crate) const COL_CLOCK: usize = 0;
-
-/// A column which contains the values `[0, ... 2^16 - 1]`, potentially with duplicates. Used for
-/// 16-bit range checks.
-///
-/// For ease of verification, we enforce that it must begin with 0 and end with `2^16 - 1`, and each
-/// delta must be either 0 or 1.
-pub(crate) const COL_RANGE_16: usize = COL_CLOCK + 1;
-
-/// Pointer to the current instruction.
-pub(crate) const COL_INSTRUCTION_PTR: usize = COL_RANGE_16 + 1;
-/// Pointer to the base of the current call's stack frame.
-pub(crate) const COL_FRAME_PTR: usize = COL_INSTRUCTION_PTR + 1;
-/// Pointer to the tip of the current call's stack frame.
-pub(crate) const COL_STACK_PTR: usize = COL_FRAME_PTR + 1;
-
-//// PERMUTATION UNIT
-pub(crate) mod permutation {
-    use plonky2::hash::hashing::SPONGE_WIDTH;
-    use plonky2::hash::poseidon;
-
-    const START_UNIT: usize = super::COL_STACK_PTR + 1;
-
-    const START_FULL_FIRST: usize = START_UNIT + SPONGE_WIDTH;
-
-    pub const fn col_full_first_mid_sbox(round: usize, i: usize) -> usize {
-        debug_assert!(round < poseidon::HALF_N_FULL_ROUNDS);
-        debug_assert!(i < SPONGE_WIDTH);
-        START_FULL_FIRST + 2 * round * SPONGE_WIDTH + i
-    }
-
-    pub const fn col_full_first_after_mds(round: usize, i: usize) -> usize {
-        debug_assert!(round < poseidon::HALF_N_FULL_ROUNDS);
-        debug_assert!(i < SPONGE_WIDTH);
-        START_FULL_FIRST + (2 * round + 1) * SPONGE_WIDTH + i
-    }
-
-    const START_PARTIAL: usize =
-        col_full_first_after_mds(poseidon::HALF_N_FULL_ROUNDS - 1, SPONGE_WIDTH - 1) + 1;
-
-    pub const fn col_partial_mid_sbox(round: usize) -> usize {
-        debug_assert!(round < poseidon::N_PARTIAL_ROUNDS);
-        START_PARTIAL + 2 * round
-    }
-
-    pub const fn col_partial_after_sbox(round: usize) -> usize {
-        debug_assert!(round < poseidon::N_PARTIAL_ROUNDS);
-        START_PARTIAL + 2 * round + 1
-    }
-
-    const START_FULL_SECOND: usize = col_partial_after_sbox(poseidon::N_PARTIAL_ROUNDS - 1) + 1;
-
-    pub const fn col_full_second_mid_sbox(round: usize, i: usize) -> usize {
-        debug_assert!(round <= poseidon::HALF_N_FULL_ROUNDS);
-        debug_assert!(i < SPONGE_WIDTH);
-        START_FULL_SECOND + 2 * round * SPONGE_WIDTH + i
-    }
-
-    pub const fn col_full_second_after_mds(round: usize, i: usize) -> usize {
-        debug_assert!(round <= poseidon::HALF_N_FULL_ROUNDS);
-        debug_assert!(i < SPONGE_WIDTH);
-        START_FULL_SECOND + (2 * round + 1) * SPONGE_WIDTH + i
-    }
-
-    pub const fn col_input(i: usize) -> usize {
-        debug_assert!(i < SPONGE_WIDTH);
-        START_UNIT + i
-    }
-
-    pub const fn col_output(i: usize) -> usize {
-        debug_assert!(i < SPONGE_WIDTH);
-        col_full_second_after_mds(poseidon::HALF_N_FULL_ROUNDS - 1, i)
-    }
-
-    pub(super) const END_UNIT: usize = col_output(SPONGE_WIDTH - 1);
-}
-
-//// MEMORY UNITS
-
-//// DECOMPOSITION UNITS
-pub(crate) mod decomposition {
-
-    const START_UNITS: usize = super::permutation::END_UNIT + 1;
-
-    const NUM_UNITS: usize = 4;
-    /// The number of bits associated with a single decomposition unit.
-    const UNIT_BITS: usize = 32;
-    /// One column for the value being decomposed, plus one column per bit.
-    const UNIT_COLS: usize = 1 + UNIT_BITS;
-
-    pub const fn col_input(unit: usize) -> usize {
-        debug_assert!(unit < NUM_UNITS);
-        START_UNITS + unit * UNIT_COLS
-    }
-
-    pub const fn col_bit(unit: usize, bit: usize) -> usize {
-        debug_assert!(unit < NUM_UNITS);
-        debug_assert!(bit < UNIT_BITS);
-        START_UNITS + unit * UNIT_COLS + 1 + bit
-    }
-
-    pub(super) const END_UNITS: usize = START_UNITS + UNIT_COLS * NUM_UNITS;
-}
-
-pub(crate) const NUM_COLUMNS: usize = decomposition::END_UNITS;

system_zero/src/core_registers.rs

@@ -6,10 +6,9 @@ use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsume
 use starky::vars::StarkEvaluationTargets;
 use starky::vars::StarkEvaluationVars;
 
-use crate::column_layout::{
-    COL_CLOCK, COL_FRAME_PTR, COL_INSTRUCTION_PTR, COL_RANGE_16, COL_STACK_PTR, NUM_COLUMNS,
-};
 use crate::public_input_layout::NUM_PUBLIC_INPUTS;
+use crate::registers::core::*;
+use crate::registers::NUM_COLUMNS;
 use crate::system_zero::SystemZero;
 
 impl<F: RichField + Extendable<D>, const D: usize> SystemZero<F, D> {
@@ -35,11 +34,11 @@ impl<F: RichField + Extendable<D>, const D: usize> SystemZero<F, D> {
         let next_range_16 = (prev_range_16 + 1).min((1 << 16) - 1);
         next_values[COL_RANGE_16] = F::from_canonical_u64(next_range_16);
 
-        next_values[COL_INSTRUCTION_PTR] = todo!();
+        // next_values[COL_INSTRUCTION_PTR] = todo!();
 
-        next_values[COL_FRAME_PTR] = todo!();
+        // next_values[COL_FRAME_PTR] = todo!();
 
-        next_values[COL_STACK_PTR] = todo!();
+        // next_values[COL_STACK_PTR] = todo!();
     }
 
     #[inline]
@@ -64,9 +63,9 @@ impl<F: RichField + Extendable<D>, const D: usize> SystemZero<F, D> {
         let delta_range_16 = next_range_16 - local_range_16;
         yield_constr.constraint_first_row(local_range_16);
         yield_constr.constraint_last_row(local_range_16 - FE::from_canonical_u64((1 << 16) - 1));
-        yield_constr.constraint(delta_range_16 * (delta_range_16 - FE::ONE));
+        yield_constr.constraint(delta_range_16 * delta_range_16 - delta_range_16);
 
-        todo!()
+        // TODO constraints for stack etc.
     }
 
     pub(crate) fn eval_core_registers_recursively(
@@ -75,6 +74,28 @@ impl<F: RichField + Extendable<D>, const D: usize> SystemZero<F, D> {
         vars: StarkEvaluationTargets<D, NUM_COLUMNS, NUM_PUBLIC_INPUTS>,
         yield_constr: &mut RecursiveConstraintConsumer<F, D>,
     ) {
-        todo!()
+        let one_ext = builder.one_extension();
+        let max_u16 = builder.constant(F::from_canonical_u64((1 << 16) - 1));
+        let max_u16_ext = builder.convert_to_ext(max_u16);
+
+        // The clock must start with 0, and increment by 1.
+        let local_clock = vars.local_values[COL_CLOCK];
+        let next_clock = vars.next_values[COL_CLOCK];
+        let delta_clock = builder.sub_extension(next_clock, local_clock);
+        yield_constr.constraint_first_row(builder, local_clock);
+        let constraint = builder.sub_extension(delta_clock, one_ext);
+        yield_constr.constraint(builder, constraint);
+
+        // The 16-bit table must start with 0, end with 2^16 - 1, and increment by 0 or 1.
+        let local_range_16 = vars.local_values[COL_RANGE_16];
+        let next_range_16 = vars.next_values[COL_RANGE_16];
+        let delta_range_16 = builder.sub_extension(next_range_16, local_range_16);
+        yield_constr.constraint_first_row(builder, local_range_16);
+        let constraint = builder.sub_extension(local_range_16, max_u16_ext);
+        yield_constr.constraint_last_row(builder, constraint);
+        let constraint = builder.mul_add_extension(delta_range_16, delta_range_16, delta_range_16);
+        yield_constr.constraint(builder, constraint);
+
+        // TODO constraints for stack etc.
     }
 }

system_zero/src/lib.rs

@@ -1,12 +1,11 @@
 // TODO: Remove these when crate is closer to being finished.
 #![allow(dead_code)]
 #![allow(unused_variables)]
-#![allow(unreachable_code)]
-#![allow(clippy::diverging_sub_expression)]
 
-mod column_layout;
+mod arithmetic;
 mod core_registers;
 mod memory;
 mod permutation_unit;
 mod public_input_layout;
+mod registers;
 pub mod system_zero;

system_zero/src/permutation_unit.rs

@@ -8,12 +8,9 @@ use starky::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsume
 use starky::vars::StarkEvaluationTargets;
 use starky::vars::StarkEvaluationVars;
 
-use crate::column_layout::permutation::{
-    col_full_first_after_mds, col_full_first_mid_sbox, col_full_second_after_mds,
-    col_full_second_mid_sbox, col_input, col_partial_after_sbox, col_partial_mid_sbox,
-};
-use crate::column_layout::NUM_COLUMNS;
 use crate::public_input_layout::NUM_PUBLIC_INPUTS;
+use crate::registers::permutation::*;
+use crate::registers::NUM_COLUMNS;
 use crate::system_zero::SystemZero;
 
 fn constant_layer<F, FE, P, const D2: usize>(
@@ -272,10 +269,10 @@ mod tests {
     use starky::constraint_consumer::ConstraintConsumer;
     use starky::vars::StarkEvaluationVars;
 
-    use crate::column_layout::permutation::{col_input, col_output};
-    use crate::column_layout::NUM_COLUMNS;
     use crate::permutation_unit::SPONGE_WIDTH;
     use crate::public_input_layout::NUM_PUBLIC_INPUTS;
+    use crate::registers::permutation::{col_input, col_output};
+    use crate::registers::NUM_COLUMNS;
     use crate::system_zero::SystemZero;
 
     #[test]

system_zero/src/registers/arithmetic.rs

@@ -0,0 +1,37 @@
+//! Arithmetic unit.
+
+pub(crate) const IS_ADD: usize = super::START_ARITHMETIC;
+pub(crate) const IS_SUB: usize = IS_ADD + 1;
+pub(crate) const IS_MUL: usize = IS_SUB + 1;
+pub(crate) const IS_DIV: usize = IS_MUL + 1;
+
+const START_SHARED_COLS: usize = IS_DIV + 1;
+
+/// Within the arithmetic unit, there are shared columns which can be used by any arithmetic
+/// circuit, depending on which one is active this cycle.
+// Can be increased as needed as other operations are implemented.
+const NUM_SHARED_COLS: usize = 3;
+
+const fn shared_col(i: usize) -> usize {
+    debug_assert!(i < NUM_SHARED_COLS);
+    START_SHARED_COLS + i
+}
+
+/// The first value to be added; treated as an unsigned u32.
+pub(crate) const COL_ADD_INPUT_1: usize = shared_col(0);
+/// The second value to be added; treated as an unsigned u32.
+pub(crate) const COL_ADD_INPUT_2: usize = shared_col(1);
+/// The third value to be added; treated as an unsigned u32.
+pub(crate) const COL_ADD_INPUT_3: usize = shared_col(2);
+
+// Note: Addition outputs three 16-bit chunks, and since these values need to be range-checked
+// anyway, we might as well use the range check unit's columns as our addition outputs. So the
+// three proceeding columns are basically aliases, not columns owned by the arithmetic unit.
+/// The first 16-bit chunk of the output, based on little-endian ordering.
+pub(crate) const COL_ADD_OUTPUT_1: usize = super::range_check_16::col_rc_16_input(0);
+/// The second 16-bit chunk of the output, based on little-endian ordering.
+pub(crate) const COL_ADD_OUTPUT_2: usize = super::range_check_16::col_rc_16_input(1);
+/// The third 16-bit chunk of the output, based on little-endian ordering.
+pub(crate) const COL_ADD_OUTPUT_3: usize = super::range_check_16::col_rc_16_input(2);
+
+pub(super) const END: usize = super::START_ARITHMETIC + NUM_SHARED_COLS;

system_zero/src/registers/boolean.rs

@@ -0,0 +1,10 @@
+//! Boolean unit. Contains columns whose values must be 0 or 1.
+
+const NUM_BITS: usize = 128;
+
+pub const fn col_bit(index: usize) -> usize {
+    debug_assert!(index < NUM_BITS);
+    super::START_BOOLEAN + index
+}
+
+pub(super) const END: usize = super::START_BOOLEAN + NUM_BITS;

system_zero/src/registers/core.rs

@@ -0,0 +1,20 @@
+//! Core registers.
+
+/// A cycle counter. Starts at 0; increments by 1.
+pub(crate) const COL_CLOCK: usize = super::START_CORE;
+
+/// A column which contains the values `[0, ... 2^16 - 1]`, potentially with duplicates. Used for
+/// 16-bit range checks.
+///
+/// For ease of verification, we enforce that it must begin with 0 and end with `2^16 - 1`, and each
+/// delta must be either 0 or 1.
+pub(crate) const COL_RANGE_16: usize = COL_CLOCK + 1;
+
+/// Pointer to the current instruction.
+pub(crate) const COL_INSTRUCTION_PTR: usize = COL_RANGE_16 + 1;
+/// Pointer to the base of the current call's stack frame.
+pub(crate) const COL_FRAME_PTR: usize = COL_INSTRUCTION_PTR + 1;
+/// Pointer to the tip of the current call's stack frame.
+pub(crate) const COL_STACK_PTR: usize = COL_FRAME_PTR + 1;
+
+pub(super) const END: usize = COL_STACK_PTR + 1;

system_zero/src/registers/logic.rs

@@ -0,0 +1,3 @@
+//! Logic unit.
+
+pub(super) const END: usize = super::START_LOGIC;

system_zero/src/registers/lookup.rs

@@ -0,0 +1,21 @@
+//! Lookup unit.
+//! See https://zcash.github.io/halo2/design/proving-system/lookup.html
+
+const START_UNIT: usize = super::START_LOOKUP;
+
+const NUM_LOOKUPS: usize =
+    super::range_check_16::NUM_RANGE_CHECKS + super::range_check_degree::NUM_RANGE_CHECKS;
+
+/// This column contains a permutation of the input values.
+const fn col_permuted_input(i: usize) -> usize {
+    debug_assert!(i < NUM_LOOKUPS);
+    START_UNIT + 2 * i
+}
+
+/// This column contains a permutation of the table values.
+const fn col_permuted_table(i: usize) -> usize {
+    debug_assert!(i < NUM_LOOKUPS);
+    START_UNIT + 2 * i + 1
+}
+
+pub(super) const END: usize = START_UNIT + NUM_LOOKUPS;

system_zero/src/registers/memory.rs

@@ -0,0 +1,3 @@
+//! Memory unit.
+
+pub(super) const END: usize = super::START_MEMORY;

system_zero/src/registers/mod.rs

@@ -0,0 +1,20 @@
+pub(crate) mod arithmetic;
+pub(crate) mod boolean;
+pub(crate) mod core;
+pub(crate) mod logic;
+pub(crate) mod lookup;
+pub(crate) mod memory;
+pub(crate) mod permutation;
+pub(crate) mod range_check_16;
+pub(crate) mod range_check_degree;
+
+const START_ARITHMETIC: usize = 0;
+const START_BOOLEAN: usize = arithmetic::END;
+const START_CORE: usize = boolean::END;
+const START_LOGIC: usize = core::END;
+const START_LOOKUP: usize = logic::END;
+const START_MEMORY: usize = lookup::END;
+const START_PERMUTATION: usize = memory::END;
+const START_RANGE_CHECK_16: usize = permutation::END;
+const START_RANGE_CHECK_DEGREE: usize = range_check_16::END;
+pub(crate) const NUM_COLUMNS: usize = range_check_degree::END;

system_zero/src/registers/permutation.rs

@@ -0,0 +1,57 @@
+//! Permutation unit.
+
+use plonky2::hash::hashing::SPONGE_WIDTH;
+use plonky2::hash::poseidon;
+
+const START_FULL_FIRST: usize = super::START_PERMUTATION + SPONGE_WIDTH;
+
+pub const fn col_full_first_mid_sbox(round: usize, i: usize) -> usize {
+    debug_assert!(round < poseidon::HALF_N_FULL_ROUNDS);
+    debug_assert!(i < SPONGE_WIDTH);
+    START_FULL_FIRST + 2 * round * SPONGE_WIDTH + i
+}
+
+pub const fn col_full_first_after_mds(round: usize, i: usize) -> usize {
+    debug_assert!(round < poseidon::HALF_N_FULL_ROUNDS);
+    debug_assert!(i < SPONGE_WIDTH);
+    START_FULL_FIRST + (2 * round + 1) * SPONGE_WIDTH + i
+}
+
+const START_PARTIAL: usize =
+    col_full_first_after_mds(poseidon::HALF_N_FULL_ROUNDS - 1, SPONGE_WIDTH - 1) + 1;
+
+pub const fn col_partial_mid_sbox(round: usize) -> usize {
+    debug_assert!(round < poseidon::N_PARTIAL_ROUNDS);
+    START_PARTIAL + 2 * round
+}
+
+pub const fn col_partial_after_sbox(round: usize) -> usize {
+    debug_assert!(round < poseidon::N_PARTIAL_ROUNDS);
+    START_PARTIAL + 2 * round + 1
+}
+
+const START_FULL_SECOND: usize = col_partial_after_sbox(poseidon::N_PARTIAL_ROUNDS - 1) + 1;
+
+pub const fn col_full_second_mid_sbox(round: usize, i: usize) -> usize {
+    debug_assert!(round <= poseidon::HALF_N_FULL_ROUNDS);
+    debug_assert!(i < SPONGE_WIDTH);
+    START_FULL_SECOND + 2 * round * SPONGE_WIDTH + i
+}
+
+pub const fn col_full_second_after_mds(round: usize, i: usize) -> usize {
+    debug_assert!(round <= poseidon::HALF_N_FULL_ROUNDS);
+    debug_assert!(i < SPONGE_WIDTH);
+    START_FULL_SECOND + (2 * round + 1) * SPONGE_WIDTH + i
+}
+
+pub const fn col_input(i: usize) -> usize {
+    debug_assert!(i < SPONGE_WIDTH);
+    super::START_PERMUTATION + i
+}
+
+pub const fn col_output(i: usize) -> usize {
+    debug_assert!(i < SPONGE_WIDTH);
+    col_full_second_after_mds(poseidon::HALF_N_FULL_ROUNDS - 1, i)
+}
+
+pub(super) const END: usize = col_output(SPONGE_WIDTH - 1) + 1;

system_zero/src/registers/range_check_16.rs

@@ -0,0 +1,11 @@
+//! Range check unit which checks that values are in `[0, 2^16)`.
+
+pub(super) const NUM_RANGE_CHECKS: usize = 5;
+
+/// The input of the `i`th range check, i.e. the value being range checked.
+pub(crate) const fn col_rc_16_input(i: usize) -> usize {
+    debug_assert!(i < NUM_RANGE_CHECKS);
+    super::START_RANGE_CHECK_16 + i
+}
+
+pub(super) const END: usize = super::START_RANGE_CHECK_16 + NUM_RANGE_CHECKS;

system_zero/src/registers/range_check_degree.rs

@@ -0,0 +1,11 @@
+//! Range check unit which checks that values are in `[0, degree)`.
+
+pub(super) const NUM_RANGE_CHECKS: usize = 5;
+
+/// The input of the `i`th range check, i.e. the value being range checked.
+pub(crate) const fn col_rc_degree_input(i: usize) -> usize {
+    debug_assert!(i < NUM_RANGE_CHECKS);
+    super::START_RANGE_CHECK_DEGREE + i
+}
+
+pub(super) const END: usize = super::START_RANGE_CHECK_DEGREE + NUM_RANGE_CHECKS;

system_zero/src/system_zero.rs

@@ -9,9 +9,12 @@ use starky::stark::Stark;
 use starky::vars::StarkEvaluationTargets;
 use starky::vars::StarkEvaluationVars;
 
-use crate::column_layout::NUM_COLUMNS;
+use crate::arithmetic::{
+    eval_arithmetic_unit, eval_arithmetic_unit_recursively, generate_arithmetic_unit,
+};
 use crate::memory::TransactionMemory;
 use crate::public_input_layout::NUM_PUBLIC_INPUTS;
+use crate::registers::NUM_COLUMNS;
 
 /// We require at least 2^16 rows as it helps support efficient 16-bit range checks.
 const MIN_TRACE_ROWS: usize = 1 << 16;
@@ -34,10 +37,16 @@ impl<F: RichField + Extendable<D>, const D: usize> SystemZero<F, D> {
         loop {
             let mut next_row = [F::ZERO; NUM_COLUMNS];
             self.generate_next_row_core_registers(&row, &mut next_row);
+            generate_arithmetic_unit(&mut next_row);
             Self::generate_permutation_unit(&mut next_row);
 
             trace.push(row);
             row = next_row;
+
+            // TODO: Replace with proper termination condition.
+            if trace.len() == (1 << 16) - 1 {
+                break;
+            }
         }
 
         trace.push(row);
@@ -66,8 +75,9 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for SystemZero<F,
         P: PackedField<Scalar = FE>,
     {
         self.eval_core_registers(vars, yield_constr);
+        eval_arithmetic_unit(vars, yield_constr);
         Self::eval_permutation_unit(vars, yield_constr);
-        todo!()
+        // TODO: Other units
     }
 
     fn eval_ext_recursively(
@@ -77,8 +87,9 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for SystemZero<F,
         yield_constr: &mut RecursiveConstraintConsumer<F, D>,
     ) {
         self.eval_core_registers_recursively(builder, vars, yield_constr);
+        eval_arithmetic_unit_recursively(builder, vars, yield_constr);
         Self::eval_permutation_unit_recursively(builder, vars, yield_constr);
-        todo!()
+        // TODO: Other units
     }
 
     fn constraint_degree(&self) -> usize {
@@ -103,7 +114,7 @@ mod tests {
     use crate::system_zero::SystemZero;
 
     #[test]
-    #[ignore] // TODO
+    #[ignore] // A bit slow.
     fn run() -> Result<()> {
         type F = GoldilocksField;
         type C = PoseidonGoldilocksConfig;
@@ -121,7 +132,6 @@ mod tests {
     }
 
     #[test]
-    #[ignore] // TODO
     fn degree() -> Result<()> {
         type F = GoldilocksField;
         type C = PoseidonGoldilocksConfig;