From ebdfca66ce4e5bd9674d0d40a0ca54e007af68fd Mon Sep 17 00:00:00 2001 From: Danny Willems Date: Tue, 17 Sep 2024 11:40:38 +0200 Subject: [PATCH] Arrabiata: interface for "next row" and keep "next state" in the env --- arrabiata/src/constraints.rs | 16 ++++++++++++++++ arrabiata/src/interpreter.rs | 7 +++++++ arrabiata/src/witness.rs | 36 ++++++++++++++++++++++++++++++++++-- 3 files changed, 57 insertions(+), 2 deletions(-) diff --git a/arrabiata/src/constraints.rs b/arrabiata/src/constraints.rs index 86b53dca07..dc8f60fca9 100644 --- a/arrabiata/src/constraints.rs +++ b/arrabiata/src/constraints.rs @@ -55,6 +55,13 @@ impl InterpreterEnv for Env { pos } + fn access_next_row(&self, pos: Self::Position) -> Self::Variable { + Expr::Atom(ExprInner::Cell(Variable { + col: pos, + row: CurrOrNext::Next, + })) + } + fn allocate_public_input(&mut self) -> Self::Position { assert!(self.idx_var_pi < NUMBER_OF_PUBLIC_INPUTS, "Maximum number of public inputs reached ({NUMBER_OF_PUBLIC_INPUTS}), increase the number of public inputs"); let pos = Column::PublicInput(self.idx_var_pi); @@ -87,6 +94,15 @@ impl InterpreterEnv for Env { res } + fn write_column_next_row(&mut self, col: Self::Position, v: Self::Variable) -> Self::Variable { + let res = Expr::Atom(ExprInner::Cell(Variable { + col, + row: CurrOrNext::Next, + })); + self.assert_equal(res.clone(), v); + res + } + fn activate_gadget(&mut self, _gadget: Gadget) { // Nothing to do. It is only useful for the witness. } diff --git a/arrabiata/src/interpreter.rs b/arrabiata/src/interpreter.rs index 219f476803..0bc407320b 100644 --- a/arrabiata/src/interpreter.rs +++ b/arrabiata/src/interpreter.rs @@ -386,10 +386,17 @@ pub trait InterpreterEnv { /// Allocate a new variable in the circuit fn allocate(&mut self) -> Self::Position; + /// Return the corresponding variable at the given position, for the next row. + fn access_next_row(&self, pos: Self::Position) -> Self::Variable; + fn allocate_public_input(&mut self) -> Self::Position; + /// Set the value of the variable at the given position for the current row fn write_column(&mut self, col: Self::Position, v: Self::Variable) -> Self::Variable; + /// Set the value of the variable at the given position for the next row + fn write_column_next_row(&mut self, col: Self::Position, v: Self::Variable) -> Self::Variable; + /// Write the corresponding public inputs. // FIXME: This design might not be the best. Feel free to come up with a // better solution. The PI should be static for all witnesses diff --git a/arrabiata/src/witness.rs b/arrabiata/src/witness.rs index 0bd73d51af..23773043d8 100644 --- a/arrabiata/src/witness.rs +++ b/arrabiata/src/witness.rs @@ -82,6 +82,11 @@ pub struct Env< /// State of the current row in the execution trace pub state: [BigInt; NUMBER_OF_COLUMNS], + /// Next row in the execution trace. It is useful when we deal with + /// polynomials accessing "the next row", i.e. witness columns where we do + /// evaluate at ζ and ζω. + pub next_state: [BigInt; NUMBER_OF_COLUMNS], + /// Contain the public state // FIXME: I don't like this design. Feel free to suggest a better solution pub public_state: [BigInt; NUMBER_OF_PUBLIC_INPUTS], @@ -203,6 +208,13 @@ where pos } + fn access_next_row(&self, pos: Self::Position) -> Self::Variable { + let Column::X(idx) = pos else { + unimplemented!("Only works for private inputs") + }; + self.witness[idx][self.current_row + 1].clone() + } + fn allocate_public_input(&mut self) -> Self::Position { assert!(self.idx_var_pi < NUMBER_OF_PUBLIC_INPUTS, "Maximum number of public inputs reached ({NUMBER_OF_PUBLIC_INPUTS}), increase the number of public inputs"); let pos = Column::PublicInput(self.idx_var_pi); @@ -224,6 +236,20 @@ where v } + fn write_column_next_row(&mut self, col: Self::Position, v: Self::Variable) -> Self::Variable { + let Column::X(idx) = col else { + unimplemented!("Only works for private inputs") + }; + let modulus: BigInt = if self.current_iteration % 2 == 0 { + Fp::modulus_biguint().into() + } else { + Fq::modulus_biguint().into() + }; + let v = v.mod_floor(&modulus); + self.next_state[idx] = v.clone(); + v + } + fn write_public_input(&mut self, col: Self::Position, v: BigInt) -> Self::Variable { let Column::PublicInput(idx) = col else { unimplemented!("Only works for public input columns") @@ -327,11 +353,16 @@ where self.state.iter().enumerate().for_each(|(i, x)| { self.witness[i][self.current_row] = x.clone(); }); + // We increment the row + // TODO: should we check that we are not going over the domain size? self.current_row += 1; + // We reset the indices for the variables self.idx_var = 0; self.idx_var_pi = 0; - // Rest the state for the next row - self.state = std::array::from_fn(|_| BigInt::from(0_usize)); + // We keep track of the values we already set. + self.state = self.next_state.clone(); + // And we reset the next state + self.next_state = std::array::from_fn(|_| BigInt::from(0_usize)); } /// FIXME: check if we need to pick the left or right sponge @@ -879,6 +910,7 @@ impl< idx_var_pi: 0, current_row: 0, state: std::array::from_fn(|_| BigInt::from(0_usize)), + next_state: std::array::from_fn(|_| BigInt::from(0_usize)), public_state: std::array::from_fn(|_| BigInt::from(0_usize)), selectors, challenges,