feat: new trace writing API (#122)

README.md

@@ -4,7 +4,9 @@
 
 
 ## Introduction
+Curta is a library for writing AIR constraints with composable and typed constraints. Currently, the main purpose of the library is to prove STARK-based accelerators for [Plonky2](https://github.com/mir-protocol/plonky2) circuits.
 
+Stark proofs and verification are implemented via [Starky](https://github.com/0xPolygonZero/plonky2/tree/main/starky). This repository contains a modified fork of Starky to enable the support a 1-round AIR with random challenges and using the Curta AIR constraints. All the cryptographic primitives are imported from the [Plonky2](https://github.com/mir-protocol/plonky2) proving system.  
 
 ## Building
 Curta requires the use of the nightly Rust toolchain. To use it by default, run the following command:
@@ -25,6 +27,8 @@ cargo test --release
 ## Building an AIR computation using Curta
 
 
+## Creating a STARK proof for an AIR computation
+
 
 ## Integrating into a Plonky2 circuit
 Curta starks can be integrated into a [Plonky2](https://github.com/mir-protocol/plonky2) circuit

curta/src/chip/arithmetic/expression_slice.rs

@@ -4,7 +4,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::air::parser::AirParser;
 use crate::chip::register::memory::MemorySlice;
-use crate::chip::register::Register;
+use crate::chip::trace::writer::AirWriter;
 use crate::math::prelude::*;
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
@@ -38,22 +38,10 @@ pub enum ArithmeticExpressionSlice<F> {
 }
 
 impl<F: Field> ArithmeticExpressionSlice<F> {
-    pub fn new<T: Register>(input: &T) -> Self {
-        ArithmeticExpressionSlice::Input(*input.register())
-    }
-
     pub fn from_raw_register(input: MemorySlice) -> Self {
         ArithmeticExpressionSlice::Input(input)
     }
 
-    pub fn from_constant(constant: F) -> Self {
-        ArithmeticExpressionSlice::Const(vec![constant])
-    }
-
-    pub fn from_constant_vec(constants: Vec<F>) -> Self {
-        ArithmeticExpressionSlice::Const(constants)
-    }
-
     pub fn registers(&self) -> Vec<MemorySlice> {
         match self {
             ArithmeticExpressionSlice::Input(input) => vec![*input],
@@ -175,4 +163,39 @@ impl<F: Field> ArithmeticExpressionSlice<F> {
             }
         }
     }
+
+    pub fn eval_writer(&self, writer: &impl AirWriter<Field = F>) -> Vec<F> {
+        match self {
+            ArithmeticExpressionSlice::Input(input) => writer.read_slice(input).to_vec(),
+            ArithmeticExpressionSlice::Const(constants) => constants.to_vec(),
+            ArithmeticExpressionSlice::Add(left, right) => {
+                let left = left.eval_writer(writer);
+                let right = right.eval_writer(writer);
+                left.into_iter().zip(right).map(|(l, r)| l + r).collect()
+            }
+            ArithmeticExpressionSlice::Sub(left, right) => {
+                let left = left.eval_writer(writer);
+                let right = right.eval_writer(writer);
+                left.into_iter().zip(right).map(|(l, r)| l - r).collect()
+            }
+            ArithmeticExpressionSlice::ConstMul(scalar, expr) => {
+                let expr_val = expr.eval_writer(writer);
+                expr_val.into_iter().map(|x| x * *scalar).collect()
+            }
+            ArithmeticExpressionSlice::ScalarMul(scalar, expr) => {
+                let scalar_val = scalar.eval_writer(writer)[0];
+                let expr_val = expr.eval_writer(writer);
+                expr_val.into_iter().map(|x| x * scalar_val).collect()
+            }
+            ArithmeticExpressionSlice::Mul(left, right) => {
+                let left_vals = left.eval_writer(writer);
+                let right_vals = right.eval_writer(writer);
+                left_vals
+                    .into_iter()
+                    .zip(right_vals)
+                    .map(|(l, r)| l * r)
+                    .collect()
+            }
+        }
+    }
 }

curta/src/chip/arithmetic/mod.rs

@@ -5,7 +5,7 @@ use crate::air::parser::AirParser;
 use crate::air::AirConstraint;
 
 pub mod expression;
-pub mod expression_slice;
+pub(crate) mod expression_slice;
 
 use crate::math::prelude::*;
 

curta/src/chip/bool.rs

@@ -10,7 +10,7 @@ use super::register::bit::BitRegister;
 use super::register::element::ElementRegister;
 use super::register::memory::MemorySlice;
 use super::register::{Register, RegisterSerializable};
-use super::trace::writer::TraceWriter;
+use super::trace::writer::{AirWriter, TraceWriter};
 use super::AirParameters;
 use crate::air::parser::AirParser;
 use crate::air::AirConstraint;
@@ -108,6 +108,22 @@ impl<F: Field> Instruction<F> for SelectInstruction {
             writer.write_slice(&self.result, &false_value, row_index);
         }
     }
+
+    fn write_to_air(&self, writer: &mut impl AirWriter<Field = F>) {
+        let bit = writer.read(&self.bit);
+        let true_value_array =
+            ArrayRegister::<ElementRegister>::from_register_unsafe(self.true_value);
+        let false_value_array =
+            ArrayRegister::<ElementRegister>::from_register_unsafe(self.false_value);
+        let true_value = writer.read_vec(&true_value_array);
+        let false_value = writer.read_vec(&false_value_array);
+
+        if bit == F::ONE {
+            writer.write_slice(&self.result, &true_value);
+        } else {
+            writer.write_slice(&self.result, &false_value);
+        }
+    }
 }
 
 // #[cfg(test)]

curta/src/chip/builder/mod.rs

@@ -8,20 +8,20 @@ use core::cmp::Ordering;
 use self::shared_memory::SharedMemory;
 use super::arithmetic::expression::ArithmeticExpression;
 use super::constraint::Constraint;
+use super::instruction::clock::ClockInstruction;
 use super::instruction::set::AirInstruction;
 use super::memory::pointer::accumulate::PointerAccumulator;
 use super::register::array::ArrayRegister;
 use super::register::cubic::CubicRegister;
 use super::register::element::ElementRegister;
-use super::register::{Register, RegisterSerializable};
+use super::register::Register;
 use super::table::accumulator::Accumulator;
 use super::table::bus::channel::BusChannel;
 use super::table::bus::global::Bus;
 use super::table::lookup::table::LookupTable;
 use super::table::lookup::values::LookupValues;
 use super::trace::data::AirTraceData;
 use super::{AirParameters, Chip};
-use crate::math::prelude::*;
 
 #[derive(Debug, Clone)]
 #[allow(clippy::type_complexity)]
@@ -179,9 +179,8 @@ impl<L: AirParameters> AirBuilder<L> {
     pub fn clock(&mut self) -> ElementRegister {
         let clk = self.alloc::<ElementRegister>();
 
-        self.set_to_expression_first_row(&clk, ArithmeticExpression::zero());
-        self.set_to_expression_transition(&clk.next(), clk.expr() + L::Field::ONE);
-
+        let instruction = AirInstruction::clock(ClockInstruction { clk });
+        self.register_air_instruction_internal(instruction);
         clk
     }
 
@@ -269,6 +268,7 @@ impl<L: AirParameters> AirBuilder<L> {
                 num_challenges: self.shared_memory.challenge_index(),
                 num_public_inputs: self.shared_memory.public_index(),
                 num_global_values: self.shared_memory.global_index(),
+                execution_trace_length,
                 instructions: self.instructions,
                 global_instructions: self.global_instructions,
                 accumulators: self.accumulators,
@@ -301,6 +301,7 @@ pub(crate) mod tests {
     pub use crate::chip::register::RegisterSerializable;
     pub use crate::chip::trace::generator::ArithmeticGenerator;
     pub use crate::math::goldilocks::cubic::GoldilocksCubicParameters;
+    use crate::math::prelude::*;
     pub use crate::maybe_rayon::*;
     pub use crate::plonky2::stark::config::PoseidonGoldilocksStarkConfig;
     pub(crate) use crate::plonky2::stark::tests::{test_recursive_starky, test_starky};

curta/src/chip/ec/edwards/assert_valid.rs

@@ -1,5 +1,3 @@
-use itertools::Itertools;
-
 use super::{EdwardsCurve, EdwardsParameters};
 use crate::chip::builder::AirBuilder;
 use crate::chip::ec::point::AffinePointRegister;
@@ -29,15 +27,15 @@ impl<L: AirParameters> AirBuilder<L> {
             one_limbs
                 .iter()
                 .map(|x| L::Field::from_canonical_u16(*x))
-                .collect_vec(),
+                .collect::<Vec<_>>(),
         );
         let one = self.constant(&one_p);
 
         let d_p = Polynomial::<L::Field>::from_coefficients(
             E::D[0..num_limbs]
                 .iter()
                 .map(|x| L::Field::from_canonical_u16(*x))
-                .collect_vec(),
+                .collect::<Vec<_>>(),
         );
 
         let d = self.constant(&d_p);

curta/src/chip/ec/edwards/ed25519/decompress.rs

@@ -1,5 +1,4 @@
 use curve25519_dalek::edwards::CompressedEdwardsY;
-use itertools::Itertools;
 use num::{BigUint, One};
 
 use super::params::{Ed25519, Ed25519BaseField, Ed25519Parameters};
@@ -35,15 +34,15 @@ impl<L: AirParameters> AirBuilder<L> {
             one_limbs
                 .iter()
                 .map(|x| L::Field::from_canonical_u16(*x))
-                .collect_vec(),
+                .collect::<Vec<_>>(),
         );
         let one = self.constant(&one_p);
 
         let d_p = Polynomial::<L::Field>::from_coefficients(
             Ed25519Parameters::D[0..num_limbs]
                 .iter()
                 .map(|x| L::Field::from_canonical_u16(*x))
-                .collect_vec(),
+                .collect::<Vec<_>>(),
         );
         let d: FieldRegister<Ed25519BaseField> = self.constant(&d_p);
 
@@ -52,7 +51,7 @@ impl<L: AirParameters> AirBuilder<L> {
             zero_limbs
                 .iter()
                 .map(|x| L::Field::from_canonical_u16(*x))
-                .collect_vec(),
+                .collect::<Vec<_>>(),
         );
         let zero: FieldRegister<Ed25519BaseField> = self.constant(&zero_p);
 

curta/src/chip/ec/edwards/ed25519/gadget.rs

@@ -6,7 +6,7 @@ use super::point::CompressedPointRegister;
 use crate::chip::builder::AirBuilder;
 use crate::chip::field::register::FieldRegister;
 use crate::chip::register::bit::BitRegister;
-use crate::chip::trace::writer::TraceWriter;
+use crate::chip::trace::writer::{AirWriter, TraceWriter};
 use crate::chip::AirParameters;
 use crate::math::prelude::*;
 use crate::polynomial::to_u16_le_limbs_polynomial;
@@ -32,6 +32,31 @@ pub trait CompressedPointWriter {
     );
 }
 
+pub trait CompressedPointAirWriter: AirWriter {
+    fn write_ec_compressed_point(
+        &mut self,
+        data: &CompressedPointRegister,
+        value: &CompressedEdwardsY,
+    ) {
+        let mut value_bytes = *value.as_bytes();
+        let compressed_sign_bit = Self::Field::from_canonical_u8(value_bytes[31] >> 7);
+
+        //println!("compressed_sign_bit is {:?}", compressed_sign_bit);
+
+        self.write(&data.sign, &compressed_sign_bit);
+
+        // mask the most significant bit
+        value_bytes[31] &= 0x7f;
+
+        let y = BigUint::from_bytes_le(&value_bytes);
+
+        let value_y = to_u16_le_limbs_polynomial::<Self::Field, Ed25519BaseField>(&y);
+        self.write(&data.y, &value_y);
+    }
+}
+
+impl<W: AirWriter> CompressedPointAirWriter for W {}
+
 impl<L: AirParameters> CompressedPointGadget for AirBuilder<L> {
     fn alloc_local_ec_compressed_point(&mut self) -> CompressedPointRegister {
         let y = self.alloc::<FieldRegister<Ed25519BaseField>>();

curta/src/chip/ec/edwards/ed25519/instruction.rs

@@ -14,7 +14,7 @@ use crate::chip::field::mul::FpMulInstruction;
 use crate::chip::field::mul_const::FpMulConstInstruction;
 use crate::chip::field::sub::FpSubInstruction;
 use crate::chip::instruction::Instruction;
-use crate::chip::trace::writer::TraceWriter;
+use crate::chip::trace::writer::{AirWriter, TraceWriter};
 use crate::math::field::PrimeField64;
 use crate::polynomial::parser::PolynomialParser;
 
@@ -55,6 +55,17 @@ impl<F: PrimeField64> Instruction<F> for Ed25519FpInstruction {
             }
         }
     }
+
+    fn write_to_air(&self, writer: &mut impl AirWriter<Field = F>) {
+        match self {
+            Ed25519FpInstruction::EC(instruction) => {
+                Instruction::<F>::write_to_air(instruction, writer)
+            }
+            Ed25519FpInstruction::Sqrt(instruction) => {
+                Instruction::<F>::write_to_air(instruction, writer)
+            }
+        }
+    }
 }
 
 impl From<LimbBitInstruction> for Ed25519FpInstruction {

curta/src/chip/ec/edwards/ed25519/sqrt.rs

@@ -14,7 +14,7 @@ use crate::chip::register::array::ArrayRegister;
 use crate::chip::register::bit::BitRegister;
 use crate::chip::register::u16::U16Register;
 use crate::chip::register::{Register, RegisterSerializable};
-use crate::chip::trace::writer::TraceWriter;
+use crate::chip::trace::writer::{AirWriter, TraceWriter};
 use crate::chip::utils::digits_to_biguint;
 use crate::chip::AirParameters;
 use crate::math::prelude::*;
@@ -151,6 +151,32 @@ impl<F: PrimeField64> Instruction<F> for Ed25519FpSqrtInstruction {
 
         self.square.write(writer, row_index);
     }
+
+    fn write_to_air(&self, writer: &mut impl AirWriter<Field = F>) {
+        let p_a = writer.read(&self.square.result);
+
+        let a_digits = p_a
+            .coefficients
+            .iter()
+            .map(|x| x.as_canonical_u64() as u16)
+            .collect::<Vec<_>>();
+
+        let a = digits_to_biguint(&a_digits);
+
+        let beta = sqrt(a);
+        let p_beta = to_u16_le_limbs_polynomial::<F, Ed25519BaseField>(&beta);
+        let a = &self.square.a;
+
+        let limb = p_beta.coefficients[0].as_canonical_u64();
+        let limb_bits = (0..Ed25519BaseField::NB_BITS_PER_LIMB)
+            .map(|i| F::from_canonical_u64((limb >> i) & 1))
+            .skip(1);
+
+        writer.write(a, &p_beta);
+        writer.write_array(&self.limb_witness, limb_bits);
+
+        self.square.write_to_air(writer);
+    }
 }
 
 pub fn sqrt(a: BigUint) -> BigUint {

curta/src/chip/ec/gadget.rs

@@ -4,7 +4,7 @@ use crate::chip::builder::AirBuilder;
 use crate::chip::field::parameters::FieldParameters;
 use crate::chip::field::register::FieldRegister;
 use crate::chip::register::Register;
-use crate::chip::trace::writer::TraceWriter;
+use crate::chip::trace::writer::{AirWriter, TraceWriter};
 use crate::chip::utils::field_limbs_to_biguint;
 use crate::chip::AirParameters;
 use crate::math::prelude::*;
@@ -35,6 +35,30 @@ pub trait EllipticCurveWriter<E: EllipticCurve> {
     );
 }
 
+pub trait EllipticCurveAirWriter<E: EllipticCurve>: AirWriter {
+    fn read_ec_point(&self, data: &AffinePointRegister<E>) -> AffinePoint<E>
+    where
+        Self::Field: PrimeField64,
+    {
+        let p_x = self.read(&data.x);
+        let p_y = self.read(&data.y);
+
+        let x = field_limbs_to_biguint(p_x.coefficients());
+        let y = field_limbs_to_biguint(p_y.coefficients());
+
+        AffinePoint::<E>::new(x, y)
+    }
+
+    fn write_ec_point(&mut self, data: &AffinePointRegister<E>, value: &AffinePoint<E>) {
+        let value_x = to_u16_le_limbs_polynomial::<Self::Field, E::BaseField>(&value.x);
+        let value_y = to_u16_le_limbs_polynomial::<Self::Field, E::BaseField>(&value.y);
+        self.write(&data.x, &value_x);
+        self.write(&data.y, &value_y);
+    }
+}
+
+impl<W: AirWriter, E: EllipticCurve> EllipticCurveAirWriter<E> for W {}
+
 impl<L: AirParameters, E: EllipticCurve> EllipticCurveGadget<E> for AirBuilder<L> {
     /// Allocates registers for a next affine elliptic curve point without range-checking.
     fn alloc_unchecked_ec_point(&mut self) -> AffinePointRegister<E> {