QOL improvements with register file structure

starky/src/arithmetic/bool.rs

@@ -8,8 +8,9 @@ use plonky2::plonk::circuit_builder::CircuitBuilder;
 use super::builder::ChipBuilder;
 use super::chip::ChipParameters;
 use super::instruction::Instruction;
-use super::register::{BitRegister, MemorySlice, Register, WitnessData};
+use super::register::{BitRegister, MemorySlice, Register};
 use super::trace::TraceHandle;
+use crate::arithmetic::register::WitnessData;
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
 #[derive(Clone, Debug, Copy, PartialEq, Eq, PartialOrd, Ord)]

starky/src/arithmetic/ec/mod.rs

@@ -5,10 +5,11 @@ use num::BigUint;
 use plonky2::field::extension::Extendable;
 use plonky2::hash::hash_types::RichField;
 
-use super::register::Register;
+use super::register::FieldRegister;
 use crate::arithmetic::builder::ChipBuilder;
 use crate::arithmetic::chip::ChipParameters;
-use crate::arithmetic::field::{FieldParameters, FieldRegister};
+use crate::arithmetic::field::FieldParameters;
+use crate::arithmetic::register::Register;
 use crate::arithmetic::trace::TraceHandle;
 
 pub const LIMB: u32 = 2u32.pow(16);

starky/src/arithmetic/field/add.rs

@@ -10,7 +10,7 @@ use crate::arithmetic::builder::ChipBuilder;
 use crate::arithmetic::chip::ChipParameters;
 use crate::arithmetic::instruction::Instruction;
 use crate::arithmetic::polynomial::{Polynomial, PolynomialGadget, PolynomialOps};
-use crate::arithmetic::register::{MemorySlice, Register, WitnessData};
+use crate::arithmetic::register::{FieldRegister, MemorySlice, Register, WitnessData};
 use crate::arithmetic::trace::TraceHandle;
 use crate::arithmetic::util::{extract_witness_and_shift, split_digits, to_field_iter};
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};

starky/src/arithmetic/field/mod.rs

@@ -14,8 +14,9 @@ use self::mul::{FpMul, FpMulConst};
 use self::quad::FpQuad;
 use super::instruction::Instruction;
 use super::polynomial::Polynomial;
+use super::register::FieldRegister;
 use super::trace::TraceHandle;
-use crate::arithmetic::register::{CellType, MemorySlice, Register, U16Array};
+use crate::arithmetic::register::MemorySlice;
 
 pub const MAX_NB_LIMBS: usize = 32;
 pub const LIMB: u32 = 2u32.pow(16);
@@ -54,35 +55,6 @@ pub fn modulus_field_iter<F: Field, P: FieldParameters>() -> impl Iterator<Item
         .take(P::NB_LIMBS)
 }
 
-#[derive(Debug, Clone, Copy)]
-pub struct FieldRegister<P: FieldParameters> {
-    array: U16Array,
-    _marker: core::marker::PhantomData<P>,
-}
-
-impl<P: FieldParameters> Register for FieldRegister<P> {
-    const CELL: Option<CellType> = Some(CellType::U16);
-
-    fn from_raw_register(register: MemorySlice) -> Self {
-        Self {
-            array: U16Array::from_raw_register(register),
-            _marker: core::marker::PhantomData,
-        }
-    }
-
-    fn register(&self) -> &MemorySlice {
-        self.array.register()
-    }
-
-    fn size_of() -> usize {
-        P::NB_LIMBS
-    }
-
-    fn into_raw_register(self) -> MemorySlice {
-        self.array.into_raw_register()
-    }
-}
-
 /// The parameters for the Fp25519 field of modulues 2^255-19.
 #[derive(Debug, Clone, Copy)]
 pub struct Fp25519Param;

starky/src/arithmetic/instruction.rs

@@ -10,7 +10,8 @@ use plonky2::hash::hash_types::RichField;
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 
 use super::bool::ConstraintBool;
-use super::register::{MemorySlice, WitnessData};
+use super::register::MemorySlice;
+use crate::arithmetic::register::WitnessData;
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
 pub trait Instruction<F: RichField + Extendable<D>, const D: usize>:

starky/src/arithmetic/register/bit.rs

@@ -0,0 +1,24 @@
+use super::CellType;
+use crate::arithmetic::register::memory::MemorySlice;
+use crate::arithmetic::register::register::Register;
+
+/// A register for a single element/column in the trace that is supposed to represent a bit. The
+/// value is automatically constrained to be 0 or 1 via the quadratic constraint x * (x - 1) == 0.
+#[derive(Debug, Clone, Copy)]
+pub struct BitRegister(MemorySlice);
+
+impl Register for BitRegister {
+    const CELL: Option<CellType> = Some(CellType::Bit);
+
+    fn register(&self) -> &MemorySlice {
+        &self.0
+    }
+
+    fn size_of() -> usize {
+        1
+    }
+
+    fn from_raw_register(register: MemorySlice) -> Self {
+        Self(register)
+    }
+}

starky/src/arithmetic/register/cell.rs

@@ -0,0 +1,5 @@
+#[derive(Debug, Clone, Copy)]
+pub enum CellType {
+    U16,
+    Bit,
+}

starky/src/arithmetic/register/element.rs

@@ -0,0 +1,23 @@
+use super::CellType;
+use crate::arithmetic::register::memory::MemorySlice;
+use crate::arithmetic::register::register::Register;
+
+/// A register for a single element/column in the trace. The value is not constrainted.
+#[derive(Debug, Clone, Copy)]
+pub struct ElementRegister(MemorySlice);
+
+impl Register for ElementRegister {
+    const CELL: Option<CellType> = None;
+
+    fn register(&self) -> &MemorySlice {
+        &self.0
+    }
+
+    fn size_of() -> usize {
+        1
+    }
+
+    fn from_raw_register(register: MemorySlice) -> Self {
+        ElementRegister(register)
+    }
+}

starky/src/arithmetic/register/field.rs

@@ -0,0 +1,33 @@
+use std::marker::PhantomData;
+
+use super::cell::CellType;
+use crate::arithmetic::field::FieldParameters;
+use crate::arithmetic::register::memory::MemorySlice;
+use crate::arithmetic::register::register::Register;
+
+/// A register for representing a field element. The value is decomposed into a series of U16 limbs
+/// which is controlled by `NB_LIMBS` in FieldParameters. Each limb is range checked using a lookup.
+#[derive(Debug, Clone, Copy)]
+pub struct FieldRegister<P: FieldParameters> {
+    register: MemorySlice,
+    _marker: PhantomData<P>,
+}
+
+impl<P: FieldParameters> Register for FieldRegister<P> {
+    const CELL: Option<CellType> = Some(CellType::U16);
+
+    fn register(&self) -> &MemorySlice {
+        &self.register
+    }
+
+    fn size_of() -> usize {
+        P::NB_LIMBS
+    }
+
+    fn from_raw_register(register: MemorySlice) -> Self {
+        Self {
+            register,
+            _marker: core::marker::PhantomData,
+        }
+    }
+}

starky/src/arithmetic/register.rs → starky/src/arithmetic/register/memory.rs

@@ -1,10 +1,11 @@
-use anyhow::{anyhow, Result};
 use plonky2::field::extension::FieldExtension;
 use plonky2::field::packed::PackedField;
 use plonky2::iop::ext_target::ExtensionTarget;
 
 use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
+/// A row-wise contiguous chunk of memory in the trace. Corresponds to a slice in vars.local_values
+/// or vars.next_values.
 #[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
 pub enum MemorySlice {
     Local(usize, usize),
@@ -15,112 +16,6 @@ pub enum MemorySlice {
     Last(usize, usize),
 }
 
-#[derive(Debug, Clone, Copy)]
-pub struct U16Array {
-    register: MemorySlice,
-}
-
-#[derive(Debug, Clone, Copy)]
-pub enum CellType {
-    U16,
-    Bit,
-}
-
-pub trait Register: 'static + Sized + Clone + Send + Sync {
-    const CELL: Option<CellType>;
-
-    /// Returns an element of the field
-    ///
-    /// Assumes register is of the correct size
-    fn from_raw_register(register: MemorySlice) -> Self;
-
-    fn into_raw_register(self) -> MemorySlice;
-
-    fn register(&self) -> &MemorySlice;
-
-    /// Returns an element of the field
-    ///
-    /// Checks that the register is of the correct size
-    fn from_register(register: MemorySlice) -> Result<Self> {
-        if register.len() != Self::size_of() {
-            return Err(anyhow!("Invalid register length"));
-        }
-
-        Ok(Self::from_raw_register(register))
-    }
-
-    fn next(&self) -> Self {
-        Self::from_raw_register(self.register().next())
-    }
-
-    fn size_of() -> usize;
-}
-
-pub struct WitnessData {
-    size: usize,
-    cell_type: Option<CellType>,
-}
-
-impl WitnessData {
-    pub fn u16(size: usize) -> Self {
-        WitnessData {
-            size,
-            cell_type: Some(CellType::U16),
-        }
-    }
-
-    pub fn bitarray(size: usize) -> Self {
-        WitnessData {
-            size,
-            cell_type: Some(CellType::Bit),
-        }
-    }
-
-    pub fn untyped(size: usize) -> Self {
-        WitnessData {
-            size,
-            cell_type: None,
-        }
-    }
-
-    pub fn destruct(self) -> (usize, Option<CellType>) {
-        (self.size, self.cell_type)
-    }
-}
-
-#[derive(Debug, Clone, Copy)]
-pub struct ElementRegister(MemorySlice);
-
-impl Register for ElementRegister {
-    const CELL: Option<CellType> = None;
-
-    fn from_raw_register(register: MemorySlice) -> Self {
-        ElementRegister(register)
-    }
-
-    fn into_raw_register(self) -> MemorySlice {
-        self.0
-    }
-
-    fn register(&self) -> &MemorySlice {
-        &self.0
-    }
-
-    fn size_of() -> usize {
-        1
-    }
-}
-
-#[derive(Debug, Clone, Copy)]
-pub struct BitArray<const N: usize> {
-    register: MemorySlice,
-}
-
-#[derive(Debug, Clone, Copy)]
-pub struct BitRegister {
-    register: MemorySlice,
-}
-
 impl MemorySlice {
     #[inline]
     pub fn is_next(&self) -> bool {
@@ -274,63 +169,3 @@ impl MemorySlice {
         }
     }
 }
-
-impl Register for BitRegister {
-    const CELL: Option<CellType> = Some(CellType::Bit);
-
-    fn from_raw_register(register: MemorySlice) -> Self {
-        Self { register }
-    }
-
-    fn register(&self) -> &MemorySlice {
-        &self.register
-    }
-
-    fn size_of() -> usize {
-        1
-    }
-
-    fn into_raw_register(self) -> MemorySlice {
-        self.register
-    }
-}
-
-impl<const N: usize> Register for BitArray<N> {
-    const CELL: Option<CellType> = Some(CellType::Bit);
-
-    fn from_raw_register(register: MemorySlice) -> Self {
-        Self { register }
-    }
-
-    fn register(&self) -> &MemorySlice {
-        &self.register
-    }
-
-    fn size_of() -> usize {
-        N
-    }
-
-    fn into_raw_register(self) -> MemorySlice {
-        self.register
-    }
-}
-
-impl Register for U16Array {
-    const CELL: Option<CellType> = Some(CellType::U16);
-
-    fn from_raw_register(register: MemorySlice) -> Self {
-        Self { register }
-    }
-
-    fn register(&self) -> &MemorySlice {
-        &self.register
-    }
-
-    fn size_of() -> usize {
-        panic!("Cannot get size of U16Array")
-    }
-
-    fn into_raw_register(self) -> MemorySlice {
-        self.register
-    }
-}

starky/src/arithmetic/register/mod.rs

@@ -0,0 +1,18 @@
+mod bit;
+mod cell;
+mod element;
+mod field;
+mod memory;
+mod register;
+mod u16;
+mod witness;
+
+pub use bit::BitRegister;
+pub use cell::CellType;
+pub use element::ElementRegister;
+pub use field::FieldRegister;
+pub use memory::MemorySlice;
+pub use register::Register;
+pub use witness::WitnessData;
+
+pub use self::u16::U16Register;