From ceb29181c395c08a2c4590e9fa63740882309a2a Mon Sep 17 00:00:00 2001 From: John Guibas Date: Thu, 13 Apr 2023 12:42:23 -0700 Subject: [PATCH 1/3] stable --- starky/src/arithmetic/bool.rs | 5 +- starky/src/arithmetic/builder.rs | 16 +- starky/src/arithmetic/ec/edwards/den.rs | 3 +- .../src/arithmetic/ec/edwards/instructions.rs | 2 +- .../src/arithmetic/ec/edwards/scalar_mul.rs | 2 +- starky/src/arithmetic/ec/mod.rs | 5 +- starky/src/arithmetic/field/add.rs | 3 +- starky/src/arithmetic/field/mod.rs | 32 +- starky/src/arithmetic/field/mul.rs | 3 +- starky/src/arithmetic/field/quad.rs | 3 +- starky/src/arithmetic/instruction.rs | 3 +- starky/src/arithmetic/mod.rs | 1 + starky/src/arithmetic/register.rs | 313 +----------------- starky/src/arithmetic/register2/bit.rs | 27 ++ starky/src/arithmetic/register2/element.rs | 26 ++ starky/src/arithmetic/register2/field.rs | 34 ++ starky/src/arithmetic/register2/memory.rs | 171 ++++++++++ starky/src/arithmetic/register2/mod.rs | 14 + starky/src/arithmetic/register2/register.rs | 39 +++ starky/src/arithmetic/register2/u16.rs | 28 ++ starky/src/arithmetic/trace.rs | 2 +- 21 files changed, 374 insertions(+), 358 deletions(-) create mode 100644 starky/src/arithmetic/register2/bit.rs create mode 100644 starky/src/arithmetic/register2/element.rs create mode 100644 starky/src/arithmetic/register2/field.rs create mode 100644 starky/src/arithmetic/register2/memory.rs create mode 100644 starky/src/arithmetic/register2/mod.rs create mode 100644 starky/src/arithmetic/register2/register.rs create mode 100644 starky/src/arithmetic/register2/u16.rs diff --git a/starky/src/arithmetic/bool.rs b/starky/src/arithmetic/bool.rs index 3bac2664e..8e74587da 100644 --- a/starky/src/arithmetic/bool.rs +++ b/starky/src/arithmetic/bool.rs @@ -8,7 +8,8 @@ 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::WitnessData; +use super::register2::{BitRegister, MemorySlice, Register}; use super::trace::TraceHandle; use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars}; @@ -194,7 +195,7 @@ mod tests { use super::*; use crate::arithmetic::builder::ChipBuilder; use crate::arithmetic::chip::{ChipParameters, TestStark}; - use crate::arithmetic::register::BitRegister; + use crate::arithmetic::register2::BitRegister; use crate::arithmetic::trace::trace; use crate::config::StarkConfig; use crate::prover::prove; diff --git a/starky/src/arithmetic/builder.rs b/starky/src/arithmetic/builder.rs index c94bd4251..d1bca6e57 100644 --- a/starky/src/arithmetic/builder.rs +++ b/starky/src/arithmetic/builder.rs @@ -12,7 +12,7 @@ use plonky2::hash::hash_types::RichField; use super::bool::ConstraintBool; use super::chip::{Chip, ChipParameters}; use super::instruction::{EqualityConstraint, Instruction, StandardInstruction, WriteInstruction}; -use super::register::{CellType, MemorySlice, Register}; +use super::register2::{MemorySlice, Register, RegisterType}; #[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord)] pub enum InsID { @@ -102,8 +102,8 @@ impl, F: RichField + Extendable, const D: usize> Chip /// Allocates a new local row register and returns it pub fn alloc_local(&mut self) -> Result { let register = match T::CELL { - Some(CellType::U16) => self.get_local_u16_memory(T::size_of())?, - Some(CellType::Bit) => { + Some(RegisterType::U16) => self.get_local_u16_memory(T::size_of())?, + Some(RegisterType::Bit) => { let reg = self.get_local_memory(T::size_of())?; let consr = EqualityConstraint::Bool(ConstraintBool(reg)); self.constraints.push(consr); @@ -117,8 +117,8 @@ impl, F: RichField + Extendable, const D: usize> Chip /// Allocates a new next row register and returns it pub fn alloc_next(&mut self) -> Result { let register = match T::CELL { - Some(CellType::U16) => self.get_next_u16_memory(T::size_of())?, - Some(CellType::Bit) => { + Some(RegisterType::U16) => self.get_next_u16_memory(T::size_of())?, + Some(RegisterType::Bit) => { let reg = self.get_next_memory(T::size_of())?; let consr = EqualityConstraint::Bool(ConstraintBool(reg)); self.constraints.push(consr); @@ -172,8 +172,8 @@ impl, F: RichField + Extendable, const D: usize> Chip if let Some(data) = inst.witness_data() { let (size, cell_type) = data.destruct(); let register = match cell_type { - Some(CellType::U16) => self.get_local_u16_memory(size)?, - Some(CellType::Bit) => { + Some(RegisterType::U16) => self.get_local_u16_memory(size)?, + Some(RegisterType::Bit) => { unimplemented!("Bit cells are not supported yet"); //self.get_local_memory(size)? } @@ -273,7 +273,7 @@ mod tests { use crate::arithmetic::chip::{ChipParameters, TestStark}; use crate::arithmetic::field::mul::FpMul; use crate::arithmetic::field::Fp25519Param; - use crate::arithmetic::register::ElementRegister; + use crate::arithmetic::register2::ElementRegister; use crate::arithmetic::trace::trace; use crate::config::StarkConfig; use crate::prover::prove; diff --git a/starky/src/arithmetic/ec/edwards/den.rs b/starky/src/arithmetic/ec/edwards/den.rs index 1c9f23df6..ed212e30d 100644 --- a/starky/src/arithmetic/ec/edwards/den.rs +++ b/starky/src/arithmetic/ec/edwards/den.rs @@ -12,7 +12,8 @@ use crate::arithmetic::chip::ChipParameters; use crate::arithmetic::field::modulus_field_iter; use crate::arithmetic::instruction::Instruction; use crate::arithmetic::polynomial::{Polynomial, PolynomialGadget, PolynomialOps}; -use crate::arithmetic::register::{MemorySlice, Register, WitnessData}; +use crate::arithmetic::register::WitnessData; +use crate::arithmetic::register2::{MemorySlice, Register}; use crate::arithmetic::trace::TraceHandle; use crate::arithmetic::util::{extract_witness_and_shift, split_digits, to_field_iter}; use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars}; diff --git a/starky/src/arithmetic/ec/edwards/instructions.rs b/starky/src/arithmetic/ec/edwards/instructions.rs index ca0ac1f78..96227f316 100644 --- a/starky/src/arithmetic/ec/edwards/instructions.rs +++ b/starky/src/arithmetic/ec/edwards/instructions.rs @@ -7,7 +7,7 @@ use crate::arithmetic::field::add::FpAdd; use crate::arithmetic::field::mul::{FpMul, FpMulConst}; use crate::arithmetic::field::quad::FpQuad; use crate::arithmetic::instruction::Instruction; -use crate::arithmetic::register::MemorySlice; +use crate::arithmetic::register2::MemorySlice; #[derive(Debug, Clone, Copy)] pub enum EdWardsMicroInstruction { diff --git a/starky/src/arithmetic/ec/edwards/scalar_mul.rs b/starky/src/arithmetic/ec/edwards/scalar_mul.rs index a24064e83..c12f12b3e 100644 --- a/starky/src/arithmetic/ec/edwards/scalar_mul.rs +++ b/starky/src/arithmetic/ec/edwards/scalar_mul.rs @@ -4,7 +4,7 @@ use crate::arithmetic::bool::Selector; use crate::arithmetic::builder::ChipBuilder; use crate::arithmetic::chip::ChipParameters; use crate::arithmetic::polynomial::Polynomial; -use crate::arithmetic::register::BitRegister; +use crate::arithmetic::register2::BitRegister; use crate::arithmetic::util::biguint_to_bits_le; #[derive(Clone, Copy)] diff --git a/starky/src/arithmetic/ec/mod.rs b/starky/src/arithmetic/ec/mod.rs index 26e3e169e..c8a3c5b57 100644 --- a/starky/src/arithmetic/ec/mod.rs +++ b/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::register2::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::register2::Register; use crate::arithmetic::trace::TraceHandle; pub const LIMB: u32 = 2u32.pow(16); diff --git a/starky/src/arithmetic/field/add.rs b/starky/src/arithmetic/field/add.rs index b7e43147f..401416055 100644 --- a/starky/src/arithmetic/field/add.rs +++ b/starky/src/arithmetic/field/add.rs @@ -10,7 +10,8 @@ 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::WitnessData; +use crate::arithmetic::register2::{FieldRegister, MemorySlice, Register}; use crate::arithmetic::trace::TraceHandle; use crate::arithmetic::util::{extract_witness_and_shift, split_digits, to_field_iter}; use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars}; diff --git a/starky/src/arithmetic/field/mod.rs b/starky/src/arithmetic/field/mod.rs index a5a6a61fe..112157892 100644 --- a/starky/src/arithmetic/field/mod.rs +++ b/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::register2::FieldRegister; use super::trace::TraceHandle; -use crate::arithmetic::register::{CellType, MemorySlice, Register, U16Array}; +use crate::arithmetic::register2::{MemorySlice, Register, RegisterType}; pub const MAX_NB_LIMBS: usize = 32; pub const LIMB: u32 = 2u32.pow(16); @@ -54,35 +55,6 @@ pub fn modulus_field_iter() -> impl Iterator { - array: U16Array, - _marker: core::marker::PhantomData

, -} - -impl Register for FieldRegister

{ - const CELL: Option = 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; diff --git a/starky/src/arithmetic/field/mul.rs b/starky/src/arithmetic/field/mul.rs index 9aa3519b6..4abd96fb0 100644 --- a/starky/src/arithmetic/field/mul.rs +++ b/starky/src/arithmetic/field/mul.rs @@ -10,7 +10,8 @@ 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::WitnessData; +use crate::arithmetic::register2::{MemorySlice, Register}; use crate::arithmetic::trace::TraceHandle; use crate::arithmetic::util::{extract_witness_and_shift, split_digits, to_field_iter}; use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars}; diff --git a/starky/src/arithmetic/field/quad.rs b/starky/src/arithmetic/field/quad.rs index 8dc092030..b9f675428 100644 --- a/starky/src/arithmetic/field/quad.rs +++ b/starky/src/arithmetic/field/quad.rs @@ -10,7 +10,8 @@ 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::WitnessData; +use crate::arithmetic::register2::{MemorySlice, Register}; use crate::arithmetic::trace::TraceHandle; use crate::arithmetic::util::{extract_witness_and_shift, split_digits, to_field_iter}; use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars}; diff --git a/starky/src/arithmetic/instruction.rs b/starky/src/arithmetic/instruction.rs index ed367ae74..66ee0806f 100644 --- a/starky/src/arithmetic/instruction.rs +++ b/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::WitnessData; +use super::register2::MemorySlice; use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars}; pub trait Instruction, const D: usize>: diff --git a/starky/src/arithmetic/mod.rs b/starky/src/arithmetic/mod.rs index c565a7472..39f25cfc8 100644 --- a/starky/src/arithmetic/mod.rs +++ b/starky/src/arithmetic/mod.rs @@ -7,5 +7,6 @@ pub mod field; pub mod instruction; pub mod polynomial; pub mod register; +pub mod register2; pub mod trace; pub(crate) mod util; diff --git a/starky/src/arithmetic/register.rs b/starky/src/arithmetic/register.rs index 99837790e..f70c64377 100644 --- a/starky/src/arithmetic/register.rs +++ b/starky/src/arithmetic/register.rs @@ -1,78 +1,22 @@ -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}; - -#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)] -pub enum MemorySlice { - Local(usize, usize), - Next(usize, usize), - - // Not sure if these are needed - First(usize, usize), - 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; - - /// 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 { - 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; -} +use super::register2::RegisterType; pub struct WitnessData { size: usize, - cell_type: Option, + cell_type: Option, } impl WitnessData { pub fn u16(size: usize) -> Self { WitnessData { size, - cell_type: Some(CellType::U16), + cell_type: Some(RegisterType::U16), } } pub fn bitarray(size: usize) -> Self { WitnessData { size, - cell_type: Some(CellType::Bit), + cell_type: Some(RegisterType::Bit), } } @@ -83,254 +27,7 @@ impl WitnessData { } } - pub fn destruct(self) -> (usize, Option) { + pub fn destruct(self) -> (usize, Option) { (self.size, self.cell_type) } } - -#[derive(Debug, Clone, Copy)] -pub struct ElementRegister(MemorySlice); - -impl Register for ElementRegister { - const CELL: Option = 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 { - register: MemorySlice, -} - -#[derive(Debug, Clone, Copy)] -pub struct BitRegister { - register: MemorySlice, -} - -impl MemorySlice { - #[inline] - pub fn is_next(&self) -> bool { - matches!(self, MemorySlice::Next(_, _)) - } - - pub fn next(&self) -> Self { - match self { - MemorySlice::Local(index, length) => MemorySlice::Next(*index, *length), - _ => panic!("Invalid register type for the next register"), - } - } - - #[inline] - pub const fn get_range(&self) -> (usize, usize) { - match self { - MemorySlice::Local(index, length) => (*index, *index + length), - MemorySlice::Next(index, length) => (*index, *index + length), - MemorySlice::First(index, length) => (*index, *index + length), - MemorySlice::Last(index, length) => (*index, *index + length), - } - } - - #[inline] - pub const fn index(&self) -> usize { - match self { - MemorySlice::Local(index, _) => *index, - MemorySlice::Next(index, _) => *index, - MemorySlice::First(index, _) => *index, - MemorySlice::Last(index, _) => *index, - } - } - - #[inline] - pub const fn len(&self) -> usize { - match self { - MemorySlice::Local(_, length) => *length, - MemorySlice::Next(_, length) => *length, - MemorySlice::First(_, length) => *length, - MemorySlice::Last(_, length) => *length, - } - } - - #[inline] - pub const fn is_empty(&self) -> bool { - self.len() == 0 - } - - #[inline] - pub fn read(&self, trace_rows: &mut [Vec], value: &mut [T], row_index: usize) { - match self { - MemorySlice::Local(index, length) => { - value.copy_from_slice(&trace_rows[row_index][*index..*index + length]); - } - MemorySlice::Next(index, length) => { - value.copy_from_slice(&trace_rows[row_index + 1][*index..*index + length]); - } - _ => panic!("Cannot read from a non-local register"), - } - } - - #[inline] - pub fn assign(&self, trace_rows: &mut [Vec], value: &mut [T], row_index: usize) { - match self { - MemorySlice::Local(index, length) => { - trace_rows[row_index][*index..*index + length].copy_from_slice(value); - } - MemorySlice::Next(index, length) => { - trace_rows[row_index + 1][*index..*index + length].copy_from_slice(value); - } - MemorySlice::First(index, length) => { - trace_rows[0][*index..*index + length].copy_from_slice(value); - } - MemorySlice::Last(index, length) => { - trace_rows[trace_rows.len() - 1][*index..*index + length].copy_from_slice(value); - } - } - } - - #[inline] - pub fn packed_entries_slice< - 'a, - F, - FE, - P, - const D2: usize, - const COLUMNS: usize, - const PUBLIC_INPUTS: usize, - >( - &self, - vars: &StarkEvaluationVars<'a, FE, P, { COLUMNS }, { PUBLIC_INPUTS }>, - ) -> &'a [P] - where - FE: FieldExtension, - P: PackedField, - { - match self { - MemorySlice::Local(index, length) => &vars.local_values[*index..*index + length], - MemorySlice::Next(index, length) => &vars.next_values[*index..*index + length], - _ => panic!("Cannot read from a non-local register"), - } - } - - #[inline] - pub fn packed_entries< - F, - FE, - P, - const D2: usize, - const COLUMNS: usize, - const PUBLIC_INPUTS: usize, - >( - &self, - vars: &StarkEvaluationVars, - ) -> Vec

- where - FE: FieldExtension, - P: PackedField, - { - match self { - MemorySlice::Local(index, length) => { - vars.local_values[*index..*index + length].to_vec() - } - MemorySlice::Next(index, length) => vars.next_values[*index..*index + length].to_vec(), - MemorySlice::First(index, length) => vars.public_inputs[*index..*index + length] - .iter() - .map(|x| P::from(*x)) - .collect(), - MemorySlice::Last(index, length) => vars.public_inputs[*index..*index + length] - .iter() - .map(|x| P::from(*x)) - .collect(), - } - } - - #[inline] - pub fn evaluation_targets< - 'a, - const COLUMNS: usize, - const PUBLIC_INPUTS: usize, - const D: usize, - >( - &self, - vars: &StarkEvaluationTargets<'a, D, { COLUMNS }, { PUBLIC_INPUTS }>, - ) -> &'a [ExtensionTarget] { - match self { - MemorySlice::Local(index, length) => &vars.local_values[*index..*index + length], - MemorySlice::Next(index, length) => &vars.next_values[*index..*index + length], - MemorySlice::First(index, length) => &vars.public_inputs[*index..*index + length], - MemorySlice::Last(index, length) => &vars.public_inputs[*index..*index + length], - } - } -} - -impl Register for BitRegister { - const CELL: Option = 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 Register for BitArray { - const CELL: Option = 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 = 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 - } -} diff --git a/starky/src/arithmetic/register2/bit.rs b/starky/src/arithmetic/register2/bit.rs new file mode 100644 index 000000000..8df717471 --- /dev/null +++ b/starky/src/arithmetic/register2/bit.rs @@ -0,0 +1,27 @@ +use crate::arithmetic::register2::memory::MemorySlice; +use crate::arithmetic::register2::register::{Register, RegisterType}; + +/// 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 = Some(RegisterType::Bit); + + fn from_raw_register(register: MemorySlice) -> Self { + Self(register) + } + + fn register(&self) -> &MemorySlice { + &self.0 + } + + fn size_of() -> usize { + 1 + } + + fn into_raw_register(self) -> MemorySlice { + self.0 + } +} diff --git a/starky/src/arithmetic/register2/element.rs b/starky/src/arithmetic/register2/element.rs new file mode 100644 index 000000000..b7cc05003 --- /dev/null +++ b/starky/src/arithmetic/register2/element.rs @@ -0,0 +1,26 @@ +use crate::arithmetic::register2::memory::MemorySlice; +use crate::arithmetic::register2::register::{Register, RegisterType}; + +/// 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 = 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 + } +} diff --git a/starky/src/arithmetic/register2/field.rs b/starky/src/arithmetic/register2/field.rs new file mode 100644 index 000000000..d5260d12f --- /dev/null +++ b/starky/src/arithmetic/register2/field.rs @@ -0,0 +1,34 @@ +use std::marker::PhantomData; + +use crate::arithmetic::field::FieldParameters; +use crate::arithmetic::register2::memory::MemorySlice; +use crate::arithmetic::register2::register::{Register, RegisterType}; + +#[derive(Debug, Clone, Copy)] +pub struct FieldRegister { + register: MemorySlice, + _marker: PhantomData

, +} + +impl Register for FieldRegister

{ + const CELL: Option = Some(RegisterType::U16); + + fn from_raw_register(register: MemorySlice) -> Self { + Self { + register, + _marker: core::marker::PhantomData, + } + } + + fn register(&self) -> &MemorySlice { + &self.register + } + + fn size_of() -> usize { + P::NB_LIMBS + } + + fn into_raw_register(self) -> MemorySlice { + self.register + } +} diff --git a/starky/src/arithmetic/register2/memory.rs b/starky/src/arithmetic/register2/memory.rs new file mode 100644 index 000000000..c5e8c7a01 --- /dev/null +++ b/starky/src/arithmetic/register2/memory.rs @@ -0,0 +1,171 @@ +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), + Next(usize, usize), + + // Not sure if these are needed + First(usize, usize), + Last(usize, usize), +} + +impl MemorySlice { + #[inline] + pub fn is_next(&self) -> bool { + matches!(self, MemorySlice::Next(_, _)) + } + + pub fn next(&self) -> Self { + match self { + MemorySlice::Local(index, length) => MemorySlice::Next(*index, *length), + _ => panic!("Invalid register type for the next register"), + } + } + + #[inline] + pub const fn get_range(&self) -> (usize, usize) { + match self { + MemorySlice::Local(index, length) => (*index, *index + length), + MemorySlice::Next(index, length) => (*index, *index + length), + MemorySlice::First(index, length) => (*index, *index + length), + MemorySlice::Last(index, length) => (*index, *index + length), + } + } + + #[inline] + pub const fn index(&self) -> usize { + match self { + MemorySlice::Local(index, _) => *index, + MemorySlice::Next(index, _) => *index, + MemorySlice::First(index, _) => *index, + MemorySlice::Last(index, _) => *index, + } + } + + #[inline] + pub const fn len(&self) -> usize { + match self { + MemorySlice::Local(_, length) => *length, + MemorySlice::Next(_, length) => *length, + MemorySlice::First(_, length) => *length, + MemorySlice::Last(_, length) => *length, + } + } + + #[inline] + pub const fn is_empty(&self) -> bool { + self.len() == 0 + } + + #[inline] + pub fn read(&self, trace_rows: &mut [Vec], value: &mut [T], row_index: usize) { + match self { + MemorySlice::Local(index, length) => { + value.copy_from_slice(&trace_rows[row_index][*index..*index + length]); + } + MemorySlice::Next(index, length) => { + value.copy_from_slice(&trace_rows[row_index + 1][*index..*index + length]); + } + _ => panic!("Cannot read from a non-local register"), + } + } + + #[inline] + pub fn assign(&self, trace_rows: &mut [Vec], value: &mut [T], row_index: usize) { + match self { + MemorySlice::Local(index, length) => { + trace_rows[row_index][*index..*index + length].copy_from_slice(value); + } + MemorySlice::Next(index, length) => { + trace_rows[row_index + 1][*index..*index + length].copy_from_slice(value); + } + MemorySlice::First(index, length) => { + trace_rows[0][*index..*index + length].copy_from_slice(value); + } + MemorySlice::Last(index, length) => { + trace_rows[trace_rows.len() - 1][*index..*index + length].copy_from_slice(value); + } + } + } + + #[inline] + pub fn packed_entries_slice< + 'a, + F, + FE, + P, + const D2: usize, + const COLUMNS: usize, + const PUBLIC_INPUTS: usize, + >( + &self, + vars: &StarkEvaluationVars<'a, FE, P, { COLUMNS }, { PUBLIC_INPUTS }>, + ) -> &'a [P] + where + FE: FieldExtension, + P: PackedField, + { + match self { + MemorySlice::Local(index, length) => &vars.local_values[*index..*index + length], + MemorySlice::Next(index, length) => &vars.next_values[*index..*index + length], + _ => panic!("Cannot read from a non-local register"), + } + } + + #[inline] + pub fn packed_entries< + F, + FE, + P, + const D2: usize, + const COLUMNS: usize, + const PUBLIC_INPUTS: usize, + >( + &self, + vars: &StarkEvaluationVars, + ) -> Vec

+ where + FE: FieldExtension, + P: PackedField, + { + match self { + MemorySlice::Local(index, length) => { + vars.local_values[*index..*index + length].to_vec() + } + MemorySlice::Next(index, length) => vars.next_values[*index..*index + length].to_vec(), + MemorySlice::First(index, length) => vars.public_inputs[*index..*index + length] + .iter() + .map(|x| P::from(*x)) + .collect(), + MemorySlice::Last(index, length) => vars.public_inputs[*index..*index + length] + .iter() + .map(|x| P::from(*x)) + .collect(), + } + } + + #[inline] + pub fn evaluation_targets< + 'a, + const COLUMNS: usize, + const PUBLIC_INPUTS: usize, + const D: usize, + >( + &self, + vars: &StarkEvaluationTargets<'a, D, { COLUMNS }, { PUBLIC_INPUTS }>, + ) -> &'a [ExtensionTarget] { + match self { + MemorySlice::Local(index, length) => &vars.local_values[*index..*index + length], + MemorySlice::Next(index, length) => &vars.next_values[*index..*index + length], + MemorySlice::First(index, length) => &vars.public_inputs[*index..*index + length], + MemorySlice::Last(index, length) => &vars.public_inputs[*index..*index + length], + } + } +} diff --git a/starky/src/arithmetic/register2/mod.rs b/starky/src/arithmetic/register2/mod.rs new file mode 100644 index 000000000..201f89cbb --- /dev/null +++ b/starky/src/arithmetic/register2/mod.rs @@ -0,0 +1,14 @@ +mod bit; +mod element; +mod field; +mod memory; +mod register; +mod u16; + +pub use bit::BitRegister; +pub use element::ElementRegister; +pub use field::FieldRegister; +pub use memory::MemorySlice; +pub use register::{Register, RegisterType}; + +pub use self::u16::U16Register; diff --git a/starky/src/arithmetic/register2/register.rs b/starky/src/arithmetic/register2/register.rs new file mode 100644 index 000000000..baa8b03f4 --- /dev/null +++ b/starky/src/arithmetic/register2/register.rs @@ -0,0 +1,39 @@ +use anyhow::{anyhow, Result}; + +use crate::arithmetic::register2::memory::MemorySlice; + +#[derive(Debug, Clone, Copy)] +pub enum RegisterType { + U16, + Bit, +} + +pub trait Register: 'static + Sized + Clone + Send + Sync { + const CELL: Option; + + /// 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 { + 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; +} diff --git a/starky/src/arithmetic/register2/u16.rs b/starky/src/arithmetic/register2/u16.rs new file mode 100644 index 000000000..332f579f9 --- /dev/null +++ b/starky/src/arithmetic/register2/u16.rs @@ -0,0 +1,28 @@ +use crate::arithmetic::register2::memory::MemorySlice; +use crate::arithmetic::register2::register::{Register, RegisterType}; + +/// A register for a single element/column in the trace that is supposed to represent a u16. The +/// value is automatically range checked via the lookup table if the register is allocated through +/// the builder. +#[derive(Debug, Clone, Copy)] +pub struct U16Register(MemorySlice); + +impl Register for U16Register { + const CELL: Option = Some(RegisterType::U16); + + fn from_raw_register(register: MemorySlice) -> Self { + Self(register) + } + + fn register(&self) -> &MemorySlice { + &self.0 + } + + fn size_of() -> usize { + panic!("Cannot get size of U16Array") + } + + fn into_raw_register(self) -> MemorySlice { + self.0 + } +} diff --git a/starky/src/arithmetic/trace.rs b/starky/src/arithmetic/trace.rs index a27dbae48..e014c4a72 100644 --- a/starky/src/arithmetic/trace.rs +++ b/starky/src/arithmetic/trace.rs @@ -11,7 +11,7 @@ use plonky2_maybe_rayon::*; use super::builder::InsID; use super::chip::ChipParameters; use super::instruction::{Instruction, StandardInstruction}; -use super::register::{MemorySlice, Register}; +use super::register2::{MemorySlice, Register}; use crate::arithmetic::chip::Chip; use crate::lookup::permuted_cols; From 56a2a7902895812430d347bcb37490409743f3cc Mon Sep 17 00:00:00 2001 From: John Guibas Date: Thu, 13 Apr 2023 13:01:39 -0700 Subject: [PATCH 2/3] Register refactor --- starky/src/arithmetic/bool.rs | 6 +++--- starky/src/arithmetic/builder.rs | 16 ++++++++-------- starky/src/arithmetic/ec/edwards/den.rs | 3 +-- starky/src/arithmetic/ec/edwards/instructions.rs | 2 +- starky/src/arithmetic/ec/edwards/scalar_mul.rs | 2 +- starky/src/arithmetic/ec/mod.rs | 4 ++-- starky/src/arithmetic/field/add.rs | 3 +-- starky/src/arithmetic/field/mod.rs | 4 ++-- starky/src/arithmetic/field/mul.rs | 3 +-- starky/src/arithmetic/field/quad.rs | 3 +-- starky/src/arithmetic/instruction.rs | 4 ++-- starky/src/arithmetic/mod.rs | 1 - .../arithmetic/{register2 => register}/bit.rs | 7 ++++--- starky/src/arithmetic/register/cell.rs | 5 +++++ .../{register2 => register}/element.rs | 7 ++++--- .../arithmetic/{register2 => register}/field.rs | 9 ++++++--- .../arithmetic/{register2 => register}/memory.rs | 0 .../arithmetic/{register2 => register}/mod.rs | 6 +++++- .../{register2 => register}/register.rs | 11 +++-------- .../arithmetic/{register2 => register}/u16.rs | 7 ++++--- .../{register.rs => register/witness.rs} | 10 +++++----- starky/src/arithmetic/trace.rs | 2 +- 22 files changed, 60 insertions(+), 55 deletions(-) rename starky/src/arithmetic/{register2 => register}/bit.rs (74%) create mode 100644 starky/src/arithmetic/register/cell.rs rename starky/src/arithmetic/{register2 => register}/element.rs (73%) rename starky/src/arithmetic/{register2 => register}/field.rs (61%) rename starky/src/arithmetic/{register2 => register}/memory.rs (100%) rename starky/src/arithmetic/{register2 => register}/mod.rs (67%) rename starky/src/arithmetic/{register2 => register}/register.rs (82%) rename starky/src/arithmetic/{register2 => register}/u16.rs (76%) rename starky/src/arithmetic/{register.rs => register/witness.rs} (65%) diff --git a/starky/src/arithmetic/bool.rs b/starky/src/arithmetic/bool.rs index 8e74587da..003afc495 100644 --- a/starky/src/arithmetic/bool.rs +++ b/starky/src/arithmetic/bool.rs @@ -8,9 +8,9 @@ use plonky2::plonk::circuit_builder::CircuitBuilder; use super::builder::ChipBuilder; use super::chip::ChipParameters; use super::instruction::Instruction; -use super::register::WitnessData; -use super::register2::{BitRegister, MemorySlice, Register}; +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)] @@ -195,7 +195,7 @@ mod tests { use super::*; use crate::arithmetic::builder::ChipBuilder; use crate::arithmetic::chip::{ChipParameters, TestStark}; - use crate::arithmetic::register2::BitRegister; + use crate::arithmetic::register::BitRegister; use crate::arithmetic::trace::trace; use crate::config::StarkConfig; use crate::prover::prove; diff --git a/starky/src/arithmetic/builder.rs b/starky/src/arithmetic/builder.rs index d1bca6e57..c94bd4251 100644 --- a/starky/src/arithmetic/builder.rs +++ b/starky/src/arithmetic/builder.rs @@ -12,7 +12,7 @@ use plonky2::hash::hash_types::RichField; use super::bool::ConstraintBool; use super::chip::{Chip, ChipParameters}; use super::instruction::{EqualityConstraint, Instruction, StandardInstruction, WriteInstruction}; -use super::register2::{MemorySlice, Register, RegisterType}; +use super::register::{CellType, MemorySlice, Register}; #[derive(Debug, Clone, Eq, PartialEq, PartialOrd, Ord)] pub enum InsID { @@ -102,8 +102,8 @@ impl, F: RichField + Extendable, const D: usize> Chip /// Allocates a new local row register and returns it pub fn alloc_local(&mut self) -> Result { let register = match T::CELL { - Some(RegisterType::U16) => self.get_local_u16_memory(T::size_of())?, - Some(RegisterType::Bit) => { + Some(CellType::U16) => self.get_local_u16_memory(T::size_of())?, + Some(CellType::Bit) => { let reg = self.get_local_memory(T::size_of())?; let consr = EqualityConstraint::Bool(ConstraintBool(reg)); self.constraints.push(consr); @@ -117,8 +117,8 @@ impl, F: RichField + Extendable, const D: usize> Chip /// Allocates a new next row register and returns it pub fn alloc_next(&mut self) -> Result { let register = match T::CELL { - Some(RegisterType::U16) => self.get_next_u16_memory(T::size_of())?, - Some(RegisterType::Bit) => { + Some(CellType::U16) => self.get_next_u16_memory(T::size_of())?, + Some(CellType::Bit) => { let reg = self.get_next_memory(T::size_of())?; let consr = EqualityConstraint::Bool(ConstraintBool(reg)); self.constraints.push(consr); @@ -172,8 +172,8 @@ impl, F: RichField + Extendable, const D: usize> Chip if let Some(data) = inst.witness_data() { let (size, cell_type) = data.destruct(); let register = match cell_type { - Some(RegisterType::U16) => self.get_local_u16_memory(size)?, - Some(RegisterType::Bit) => { + Some(CellType::U16) => self.get_local_u16_memory(size)?, + Some(CellType::Bit) => { unimplemented!("Bit cells are not supported yet"); //self.get_local_memory(size)? } @@ -273,7 +273,7 @@ mod tests { use crate::arithmetic::chip::{ChipParameters, TestStark}; use crate::arithmetic::field::mul::FpMul; use crate::arithmetic::field::Fp25519Param; - use crate::arithmetic::register2::ElementRegister; + use crate::arithmetic::register::ElementRegister; use crate::arithmetic::trace::trace; use crate::config::StarkConfig; use crate::prover::prove; diff --git a/starky/src/arithmetic/ec/edwards/den.rs b/starky/src/arithmetic/ec/edwards/den.rs index ed212e30d..1c9f23df6 100644 --- a/starky/src/arithmetic/ec/edwards/den.rs +++ b/starky/src/arithmetic/ec/edwards/den.rs @@ -12,8 +12,7 @@ use crate::arithmetic::chip::ChipParameters; use crate::arithmetic::field::modulus_field_iter; use crate::arithmetic::instruction::Instruction; use crate::arithmetic::polynomial::{Polynomial, PolynomialGadget, PolynomialOps}; -use crate::arithmetic::register::WitnessData; -use crate::arithmetic::register2::{MemorySlice, Register}; +use crate::arithmetic::register::{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}; diff --git a/starky/src/arithmetic/ec/edwards/instructions.rs b/starky/src/arithmetic/ec/edwards/instructions.rs index 96227f316..ca0ac1f78 100644 --- a/starky/src/arithmetic/ec/edwards/instructions.rs +++ b/starky/src/arithmetic/ec/edwards/instructions.rs @@ -7,7 +7,7 @@ use crate::arithmetic::field::add::FpAdd; use crate::arithmetic::field::mul::{FpMul, FpMulConst}; use crate::arithmetic::field::quad::FpQuad; use crate::arithmetic::instruction::Instruction; -use crate::arithmetic::register2::MemorySlice; +use crate::arithmetic::register::MemorySlice; #[derive(Debug, Clone, Copy)] pub enum EdWardsMicroInstruction { diff --git a/starky/src/arithmetic/ec/edwards/scalar_mul.rs b/starky/src/arithmetic/ec/edwards/scalar_mul.rs index c12f12b3e..a24064e83 100644 --- a/starky/src/arithmetic/ec/edwards/scalar_mul.rs +++ b/starky/src/arithmetic/ec/edwards/scalar_mul.rs @@ -4,7 +4,7 @@ use crate::arithmetic::bool::Selector; use crate::arithmetic::builder::ChipBuilder; use crate::arithmetic::chip::ChipParameters; use crate::arithmetic::polynomial::Polynomial; -use crate::arithmetic::register2::BitRegister; +use crate::arithmetic::register::BitRegister; use crate::arithmetic::util::biguint_to_bits_le; #[derive(Clone, Copy)] diff --git a/starky/src/arithmetic/ec/mod.rs b/starky/src/arithmetic/ec/mod.rs index c8a3c5b57..8e6f95dfc 100644 --- a/starky/src/arithmetic/ec/mod.rs +++ b/starky/src/arithmetic/ec/mod.rs @@ -5,11 +5,11 @@ use num::BigUint; use plonky2::field::extension::Extendable; use plonky2::hash::hash_types::RichField; -use super::register2::FieldRegister; +use super::register::FieldRegister; use crate::arithmetic::builder::ChipBuilder; use crate::arithmetic::chip::ChipParameters; use crate::arithmetic::field::FieldParameters; -use crate::arithmetic::register2::Register; +use crate::arithmetic::register::Register; use crate::arithmetic::trace::TraceHandle; pub const LIMB: u32 = 2u32.pow(16); diff --git a/starky/src/arithmetic/field/add.rs b/starky/src/arithmetic/field/add.rs index 401416055..4f1f3ddbb 100644 --- a/starky/src/arithmetic/field/add.rs +++ b/starky/src/arithmetic/field/add.rs @@ -10,8 +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::WitnessData; -use crate::arithmetic::register2::{FieldRegister, MemorySlice, Register}; +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}; diff --git a/starky/src/arithmetic/field/mod.rs b/starky/src/arithmetic/field/mod.rs index 112157892..33fc11a69 100644 --- a/starky/src/arithmetic/field/mod.rs +++ b/starky/src/arithmetic/field/mod.rs @@ -14,9 +14,9 @@ use self::mul::{FpMul, FpMulConst}; use self::quad::FpQuad; use super::instruction::Instruction; use super::polynomial::Polynomial; -use super::register2::FieldRegister; +use super::register::FieldRegister; use super::trace::TraceHandle; -use crate::arithmetic::register2::{MemorySlice, Register, RegisterType}; +use crate::arithmetic::register::MemorySlice; pub const MAX_NB_LIMBS: usize = 32; pub const LIMB: u32 = 2u32.pow(16); diff --git a/starky/src/arithmetic/field/mul.rs b/starky/src/arithmetic/field/mul.rs index 4abd96fb0..9aa3519b6 100644 --- a/starky/src/arithmetic/field/mul.rs +++ b/starky/src/arithmetic/field/mul.rs @@ -10,8 +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::WitnessData; -use crate::arithmetic::register2::{MemorySlice, Register}; +use crate::arithmetic::register::{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}; diff --git a/starky/src/arithmetic/field/quad.rs b/starky/src/arithmetic/field/quad.rs index b9f675428..8dc092030 100644 --- a/starky/src/arithmetic/field/quad.rs +++ b/starky/src/arithmetic/field/quad.rs @@ -10,8 +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::WitnessData; -use crate::arithmetic::register2::{MemorySlice, Register}; +use crate::arithmetic::register::{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}; diff --git a/starky/src/arithmetic/instruction.rs b/starky/src/arithmetic/instruction.rs index 66ee0806f..bd1f00512 100644 --- a/starky/src/arithmetic/instruction.rs +++ b/starky/src/arithmetic/instruction.rs @@ -10,8 +10,8 @@ use plonky2::hash::hash_types::RichField; use plonky2::plonk::circuit_builder::CircuitBuilder; use super::bool::ConstraintBool; -use super::register::WitnessData; -use super::register2::MemorySlice; +use super::register::MemorySlice; +use crate::arithmetic::register::WitnessData; use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars}; pub trait Instruction, const D: usize>: diff --git a/starky/src/arithmetic/mod.rs b/starky/src/arithmetic/mod.rs index 39f25cfc8..c565a7472 100644 --- a/starky/src/arithmetic/mod.rs +++ b/starky/src/arithmetic/mod.rs @@ -7,6 +7,5 @@ pub mod field; pub mod instruction; pub mod polynomial; pub mod register; -pub mod register2; pub mod trace; pub(crate) mod util; diff --git a/starky/src/arithmetic/register2/bit.rs b/starky/src/arithmetic/register/bit.rs similarity index 74% rename from starky/src/arithmetic/register2/bit.rs rename to starky/src/arithmetic/register/bit.rs index 8df717471..fc6581924 100644 --- a/starky/src/arithmetic/register2/bit.rs +++ b/starky/src/arithmetic/register/bit.rs @@ -1,5 +1,6 @@ -use crate::arithmetic::register2::memory::MemorySlice; -use crate::arithmetic::register2::register::{Register, RegisterType}; +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. @@ -7,7 +8,7 @@ use crate::arithmetic::register2::register::{Register, RegisterType}; pub struct BitRegister(MemorySlice); impl Register for BitRegister { - const CELL: Option = Some(RegisterType::Bit); + const CELL: Option = Some(CellType::Bit); fn from_raw_register(register: MemorySlice) -> Self { Self(register) diff --git a/starky/src/arithmetic/register/cell.rs b/starky/src/arithmetic/register/cell.rs new file mode 100644 index 000000000..65fa90933 --- /dev/null +++ b/starky/src/arithmetic/register/cell.rs @@ -0,0 +1,5 @@ +#[derive(Debug, Clone, Copy)] +pub enum CellType { + U16, + Bit, +} diff --git a/starky/src/arithmetic/register2/element.rs b/starky/src/arithmetic/register/element.rs similarity index 73% rename from starky/src/arithmetic/register2/element.rs rename to starky/src/arithmetic/register/element.rs index b7cc05003..64e352a93 100644 --- a/starky/src/arithmetic/register2/element.rs +++ b/starky/src/arithmetic/register/element.rs @@ -1,12 +1,13 @@ -use crate::arithmetic::register2::memory::MemorySlice; -use crate::arithmetic::register2::register::{Register, RegisterType}; +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 = None; + const CELL: Option = None; fn from_raw_register(register: MemorySlice) -> Self { ElementRegister(register) diff --git a/starky/src/arithmetic/register2/field.rs b/starky/src/arithmetic/register/field.rs similarity index 61% rename from starky/src/arithmetic/register2/field.rs rename to starky/src/arithmetic/register/field.rs index d5260d12f..3af49e66b 100644 --- a/starky/src/arithmetic/register2/field.rs +++ b/starky/src/arithmetic/register/field.rs @@ -1,9 +1,12 @@ use std::marker::PhantomData; +use super::cell::CellType; use crate::arithmetic::field::FieldParameters; -use crate::arithmetic::register2::memory::MemorySlice; -use crate::arithmetic::register2::register::{Register, RegisterType}; +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 { register: MemorySlice, @@ -11,7 +14,7 @@ pub struct FieldRegister { } impl Register for FieldRegister

{ - const CELL: Option = Some(RegisterType::U16); + const CELL: Option = Some(CellType::U16); fn from_raw_register(register: MemorySlice) -> Self { Self { diff --git a/starky/src/arithmetic/register2/memory.rs b/starky/src/arithmetic/register/memory.rs similarity index 100% rename from starky/src/arithmetic/register2/memory.rs rename to starky/src/arithmetic/register/memory.rs diff --git a/starky/src/arithmetic/register2/mod.rs b/starky/src/arithmetic/register/mod.rs similarity index 67% rename from starky/src/arithmetic/register2/mod.rs rename to starky/src/arithmetic/register/mod.rs index 201f89cbb..a986bf419 100644 --- a/starky/src/arithmetic/register2/mod.rs +++ b/starky/src/arithmetic/register/mod.rs @@ -1,14 +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, RegisterType}; +pub use register::Register; +pub use witness::WitnessData; pub use self::u16::U16Register; diff --git a/starky/src/arithmetic/register2/register.rs b/starky/src/arithmetic/register/register.rs similarity index 82% rename from starky/src/arithmetic/register2/register.rs rename to starky/src/arithmetic/register/register.rs index baa8b03f4..794d149ab 100644 --- a/starky/src/arithmetic/register2/register.rs +++ b/starky/src/arithmetic/register/register.rs @@ -1,15 +1,10 @@ use anyhow::{anyhow, Result}; -use crate::arithmetic::register2::memory::MemorySlice; - -#[derive(Debug, Clone, Copy)] -pub enum RegisterType { - U16, - Bit, -} +use super::cell::CellType; +use crate::arithmetic::register::memory::MemorySlice; pub trait Register: 'static + Sized + Clone + Send + Sync { - const CELL: Option; + const CELL: Option; /// Returns an element of the field /// diff --git a/starky/src/arithmetic/register2/u16.rs b/starky/src/arithmetic/register/u16.rs similarity index 76% rename from starky/src/arithmetic/register2/u16.rs rename to starky/src/arithmetic/register/u16.rs index 332f579f9..16a18d08f 100644 --- a/starky/src/arithmetic/register2/u16.rs +++ b/starky/src/arithmetic/register/u16.rs @@ -1,5 +1,6 @@ -use crate::arithmetic::register2::memory::MemorySlice; -use crate::arithmetic::register2::register::{Register, RegisterType}; +use super::cell::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 u16. The /// value is automatically range checked via the lookup table if the register is allocated through @@ -8,7 +9,7 @@ use crate::arithmetic::register2::register::{Register, RegisterType}; pub struct U16Register(MemorySlice); impl Register for U16Register { - const CELL: Option = Some(RegisterType::U16); + const CELL: Option = Some(CellType::U16); fn from_raw_register(register: MemorySlice) -> Self { Self(register) diff --git a/starky/src/arithmetic/register.rs b/starky/src/arithmetic/register/witness.rs similarity index 65% rename from starky/src/arithmetic/register.rs rename to starky/src/arithmetic/register/witness.rs index f70c64377..a318ee465 100644 --- a/starky/src/arithmetic/register.rs +++ b/starky/src/arithmetic/register/witness.rs @@ -1,22 +1,22 @@ -use super::register2::RegisterType; +use super::CellType; pub struct WitnessData { size: usize, - cell_type: Option, + cell_type: Option, } impl WitnessData { pub fn u16(size: usize) -> Self { WitnessData { size, - cell_type: Some(RegisterType::U16), + cell_type: Some(CellType::U16), } } pub fn bitarray(size: usize) -> Self { WitnessData { size, - cell_type: Some(RegisterType::Bit), + cell_type: Some(CellType::Bit), } } @@ -27,7 +27,7 @@ impl WitnessData { } } - pub fn destruct(self) -> (usize, Option) { + pub fn destruct(self) -> (usize, Option) { (self.size, self.cell_type) } } diff --git a/starky/src/arithmetic/trace.rs b/starky/src/arithmetic/trace.rs index e014c4a72..a27dbae48 100644 --- a/starky/src/arithmetic/trace.rs +++ b/starky/src/arithmetic/trace.rs @@ -11,7 +11,7 @@ use plonky2_maybe_rayon::*; use super::builder::InsID; use super::chip::ChipParameters; use super::instruction::{Instruction, StandardInstruction}; -use super::register2::{MemorySlice, Register}; +use super::register::{MemorySlice, Register}; use crate::arithmetic::chip::Chip; use crate::lookup::permuted_cols; From cf199fc05c356a38a29c2f8c1ac63c416487bdeb Mon Sep 17 00:00:00 2001 From: John Guibas Date: Thu, 13 Apr 2023 13:13:11 -0700 Subject: [PATCH 3/3] Make register trait simpler --- starky/src/arithmetic/register/bit.rs | 8 ++--- starky/src/arithmetic/register/element.rs | 12 +++----- starky/src/arithmetic/register/field.rs | 14 ++++----- starky/src/arithmetic/register/register.rs | 35 ++++++++++++---------- starky/src/arithmetic/register/u16.rs | 10 ++----- 5 files changed, 33 insertions(+), 46 deletions(-) diff --git a/starky/src/arithmetic/register/bit.rs b/starky/src/arithmetic/register/bit.rs index fc6581924..2fa3960f6 100644 --- a/starky/src/arithmetic/register/bit.rs +++ b/starky/src/arithmetic/register/bit.rs @@ -10,10 +10,6 @@ pub struct BitRegister(MemorySlice); impl Register for BitRegister { const CELL: Option = Some(CellType::Bit); - fn from_raw_register(register: MemorySlice) -> Self { - Self(register) - } - fn register(&self) -> &MemorySlice { &self.0 } @@ -22,7 +18,7 @@ impl Register for BitRegister { 1 } - fn into_raw_register(self) -> MemorySlice { - self.0 + fn from_raw_register(register: MemorySlice) -> Self { + Self(register) } } diff --git a/starky/src/arithmetic/register/element.rs b/starky/src/arithmetic/register/element.rs index 64e352a93..deec32c1d 100644 --- a/starky/src/arithmetic/register/element.rs +++ b/starky/src/arithmetic/register/element.rs @@ -9,14 +9,6 @@ pub struct ElementRegister(MemorySlice); impl Register for ElementRegister { const CELL: Option = None; - fn from_raw_register(register: MemorySlice) -> Self { - ElementRegister(register) - } - - fn into_raw_register(self) -> MemorySlice { - self.0 - } - fn register(&self) -> &MemorySlice { &self.0 } @@ -24,4 +16,8 @@ impl Register for ElementRegister { fn size_of() -> usize { 1 } + + fn from_raw_register(register: MemorySlice) -> Self { + ElementRegister(register) + } } diff --git a/starky/src/arithmetic/register/field.rs b/starky/src/arithmetic/register/field.rs index 3af49e66b..46f5129dd 100644 --- a/starky/src/arithmetic/register/field.rs +++ b/starky/src/arithmetic/register/field.rs @@ -16,13 +16,6 @@ pub struct FieldRegister { impl Register for FieldRegister

{ const CELL: Option = Some(CellType::U16); - fn from_raw_register(register: MemorySlice) -> Self { - Self { - register, - _marker: core::marker::PhantomData, - } - } - fn register(&self) -> &MemorySlice { &self.register } @@ -31,7 +24,10 @@ impl Register for FieldRegister

{ P::NB_LIMBS } - fn into_raw_register(self) -> MemorySlice { - self.register + fn from_raw_register(register: MemorySlice) -> Self { + Self { + register, + _marker: core::marker::PhantomData, + } } } diff --git a/starky/src/arithmetic/register/register.rs b/starky/src/arithmetic/register/register.rs index 794d149ab..ad86672a5 100644 --- a/starky/src/arithmetic/register/register.rs +++ b/starky/src/arithmetic/register/register.rs @@ -3,32 +3,35 @@ use anyhow::{anyhow, Result}; use super::cell::CellType; use crate::arithmetic::register::memory::MemorySlice; +/// A register is a slice of memory in the trace that is supposed to represent a specific type of +/// data. A register can be thought as a compiler provided type--it should not be necessary to +/// ever use the register type directly to access values. If you want to access the values, you +/// should instead compose multiple different register types into a struct. pub trait Register: 'static + Sized + Clone + Send + Sync { + /// The type of each cell in the register. Useful for specific constraints we want to apply + /// on cells such as them falling in a specific range. const CELL: Option; - /// Returns an element of the field - /// - /// Assumes register is of the correct size - fn from_raw_register(register: MemorySlice) -> Self; + /// Returns the memory slice of the register. + fn register(&self) -> &MemorySlice; + + /// Returns the register but in the next row. + fn next(&self) -> Self { + Self::from_raw_register(self.register().next()) + } - fn into_raw_register(self) -> MemorySlice; + /// Returns the expected size of the register in cells. + fn size_of() -> usize; - fn register(&self) -> &MemorySlice; + /// Initializes the register given a memory slice with no checks on length. Avoid using this + /// function unless you know what you are doing and use `from_register` instead. + fn from_raw_register(register: MemorySlice) -> Self; - /// Returns an element of the field - /// - /// Checks that the register is of the correct size + /// Initializes the register given a memory slice with checks on length. fn from_register(register: MemorySlice) -> Result { 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; } diff --git a/starky/src/arithmetic/register/u16.rs b/starky/src/arithmetic/register/u16.rs index 16a18d08f..4f3e03c1b 100644 --- a/starky/src/arithmetic/register/u16.rs +++ b/starky/src/arithmetic/register/u16.rs @@ -11,19 +11,15 @@ pub struct U16Register(MemorySlice); impl Register for U16Register { const CELL: Option = Some(CellType::U16); - fn from_raw_register(register: MemorySlice) -> Self { - Self(register) - } - fn register(&self) -> &MemorySlice { &self.0 } fn size_of() -> usize { - panic!("Cannot get size of U16Array") + 1 } - fn into_raw_register(self) -> MemorySlice { - self.0 + fn from_raw_register(register: MemorySlice) -> Self { + Self(register) } }