Skip to content

Commit

Permalink
tk air boundary (#11)
Browse files Browse the repository at this point in the history
* impl: fraction field element

* air boundary tests in progress

* fft utils added to support list of cosets impl

* list of cosets .h implemented

* randomness - not necessary yet, abandon

* added notes to prime field element usage

* boundary constraints and test utils modified

* utility crate init

* fft-bases-in-progress-pause

* build fix

* intsrument fft-bases with test scaffolding

* tests pass
  • Loading branch information
thor314 committed Jan 10, 2024
1 parent 5b807ae commit 3d980dc
Show file tree
Hide file tree
Showing 24 changed files with 896 additions and 271 deletions.
2 changes: 2 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,11 @@ members=[
"air",
"algebra",
"composition_polynomial",
"fft_utils",
"math",
"proof_system",
"randomness",
"stark",
"statement",
"utils",
]
7 changes: 7 additions & 0 deletions air/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,10 @@ edition="2021"
ark-ff ="0.4.2"
algebra ={ path="../algebra" }
composition_polynomial={ path="../composition_polynomial" }
utils = { path = "../utils" }

#[dev-dependencies]
# todo: remove ark-test-curves eventually
ark-std ="0.4.0"
ark-test-curves="0.4.2"
rand ={ version="0.8.4", features=["std"] }
259 changes: 190 additions & 69 deletions air/src/boundary/mod.rs
Original file line number Diff line number Diff line change
@@ -1,18 +1,30 @@
//! https://github.com/starkware-libs/stone-prover/blob/a78ff37c1402dc9c3e3050a1090cd98b7ff123b3/src/starkware/air/boundary/boundary_air.h
//! Note: Using `stdVecDeque` in place of C++ `std::vector`, for efficient insertion and removal
//! from both ends.
// todo(tk): may need to write CompositionPolynomial first to implement:
// public:
// using FieldElementT_ = FieldElementT;
// using Builder = typename CompositionPolynomialImpl<BoundaryAir<FieldElementT>>::Builder;
//
// something like:
// type Builder = CompositionPolynomial<BoundaryAir<F::FieldExt>>;
use std::collections::{HashMap, VecDeque};
use std::{
collections::{HashMap, VecDeque},
vec,
};

use algebra::ConstFieldElementSpan;
use algebra::{ConstFieldElementSpan, FractionFieldElement};
use ark_ff::Field;
use composition_polynomial::CompositionPolynomial;
use utils::assert_on_release;

use crate::{assert_on_release, Air};
use crate::Air;

// todo: lift to crate root;
// todo: verify that VecDeque is optimal replacement
/// Annotation of type mirroring of `gsl::span` type in Stone Prover codebase.
pub type GslSpan<T> = VecDeque<T>;
// TODO(TK 2024-01-09): placeholder type

/// A simple AIR the describes the contraints:
/// (column_i(x) - y0_i) / (x - x0_i).
Expand All @@ -21,20 +33,21 @@ use crate::{assert_on_release, Air};
/// - trace_length, size of trace.
/// - n_columns, number of columns in trace.
/// - boundary_conditions, list of tuples (column, x, y) indicating the constraint that column(x)=y.
// todo: this method is difficult to read and likely implemented incorrectly
pub struct BoundaryAir<F: Field> {
pub trace_length: usize,
pub trace_length: u64,
pub num_columns: usize,
// todo(tk): check usage of VecDeq
constraints: VecDeque<ConstraintData<F>>,
/// Mask is the list of pairs (relative_row, col) that define the neighbors needed for the
/// constraint.
/// The mask touches each column once at the current row.
pub mask: Vec<(usize, usize)>,
pub mask: Vec<(isize, usize)>,
}

impl<F: Field> BoundaryAir<F> {
pub fn new(
trace_length: usize,
trace_length: u64,
n_columns: usize,
boundary_conditions: &[BoundaryCondition<F>],
) -> Self {
Expand Down Expand Up @@ -68,58 +81,63 @@ impl<F: Field> BoundaryAir<F> {
Self { trace_length, num_columns: n_columns, constraints, mask }
}

// todo(tk): find replacement for gsl::span
// pub fn precompute_domain_evals_on_coset(
// point: &F,
// generator: &F,
// point_exponents: gsl::span<usize>,
// shifts: gsl::span<F>,
// ) { // do nothing
// }

// pub fn constraints_eval(&self, ){
// gsl::span<const FieldElementT> neighbors,
// gsl::span<const FieldElementT> /* periodic_columns */,
// gsl::span<const FieldElementT> random_coefficients,
// gsl::span<const FieldElementT> point_powers, gsl::span<const FieldElementT> /*shifts*/,
// gsl::span<const FieldElementT> /* precomp_domains */) const {
// ASSERT_DEBUG(neighbors.size() == n_columns_, "Wrong number of neighbors");
// ASSERT_DEBUG(
// random_coefficients.size() == constraints_.size(), "Wrong number of random
// coefficients");

// const FieldElementT& point = point_powers[0];

// FractionFieldElement<FieldElementT> outer_sum(FieldElementT::Zero());
// FieldElementT inner_sum(FieldElementT::Zero());

// FieldElementT prev_x = constraints_[0].point_x;

// for (const ConstraintData& constraint : constraints_) {
// const FieldElementT constraint_value =
// random_coefficients[constraint.coeff_idx] *
// (neighbors[constraint.column_index] - constraint.point_y);
// if (prev_x == constraint.point_x) {
// // All constraints with the same constraint.point_x are summed into inner_sum.
// inner_sum += constraint_value;
// } else {
// // New constraint.point_x, add the old (inner_sum/prev_x) to the outer_sum
// // and start a new inner_sum.
// outer_sum += FractionFieldElement<FieldElementT>(inner_sum, point - prev_x);
// inner_sum = constraint_value;
// prev_x = constraint.point_x;
// }
// }
// outer_sum += FractionFieldElement<FieldElementT>(inner_sum, point - prev_x);

// return outer_sum;
// }}

// std::vector<FieldElementT> DomainEvalsAtPoint(
// gsl::span<const FieldElementT> /* point_powers */,
// gsl::span<const FieldElementT> /* shifts */) const {
// return {};
// }
// fn create_composition_polynomial in Air trait impl.

// TODO(TK 2024-01-09): should this be in Air trait?
pub fn precompute_domain_evals_on_coset(
point: &F,
generator: &F,
point_exponents: GslSpan<u64>,
shifts: GslSpan<F>,
) { // do nothing
}

// evaluate the set of constraints for the AIR
// TODO(TK 2024-01-09): almost certainly should be in Air trait
pub fn constraints_eval(
&self,
neighbors: GslSpan<F>,
periodic_columns: GslSpan<F>,
random_coefficients: GslSpan<F>,
point_powers: GslSpan<F>,
shifts: GslSpan<F>,
precomp_domains: GslSpan<F>,
) -> FractionFieldElement<F> /*outer_sum*/ {
assert_on_release(neighbors.len() == self.num_columns, "Wrong number of neighbors");
assert_on_release(
random_coefficients.len() == self.constraints.len(),
"Wrong number of coefficients",
);

let point = &point_powers[0];

let mut outer_sum = FractionFieldElement::<F>::zero();
let mut inner_sum = F::zero();
let mut prev_x = self.constraints.front().expect("Constraints cannot be empty").point_x;

for constraint in &self.constraints {
let constraint_value = random_coefficients[constraint.coeff_idx]
* (neighbors[constraint.column_index] - constraint.point_y);

if prev_x == constraint.point_x {
// All constraints with the same constraint.point_x are summed into inner_sum.
inner_sum += constraint_value;
} else {
// New constraint.point_x, add the old (inner_sum/prev_x) to the outer_sum
// and start a new inner_sum.
outer_sum += FractionFieldElement::new(inner_sum, *point - prev_x);
inner_sum = constraint_value;
prev_x = constraint.point_x;
}
}
outer_sum += FractionFieldElement::new(inner_sum, *point - prev_x);

outer_sum
}

fn domain_evals_at_point(point_powers: GslSpan<F>, shifts: GslSpan<F>) -> VecDeque<F> {
Default::default() // intentionally left blank
}
}

impl<F: Field> Air<F> for BoundaryAir<F> {
Expand All @@ -128,39 +146,142 @@ impl<F: Field> Air<F> for BoundaryAir<F> {
trace_generator: &F,
random_coefficients: &ConstFieldElementSpan<F>,
) -> Box<CompositionPolynomial<F>> {
// todo(tk): builder is a bad name for a builder of CompositionPolynomial, commenting until
// farther into refactor
//
// let builder = Builder::new(0);
// builder.build_unique_ptr(use_owned(&self), trace_generator, self.trace_length,
// random_coefficients, (), ())
todo!()
Box::new(CompositionPolynomial::new(
// TODO(TK 2024-01-09): placeholder impl
&F::one(),
*trace_generator,
self.trace_length,
vec![],
vec![],
// random_coefficients,
vec![],
vec![],
))
}

fn get_composition_polynomial_degree_bound(&self) -> usize { self.trace_length }
fn get_composition_polynomial_degree_bound(&self) -> u64 { self.trace_length }

fn num_random_coefficients(&self) -> usize { self.constraints.len() }

fn get_interaction_params(&self) -> Option<crate::InteractionParams> { None }

fn num_columns(&self) -> usize { self.num_columns }

fn trace_length(&self) -> usize { todo!() }
fn trace_length(&self) -> u64 { self.trace_length }
}

/// List of tuples (column, x, y) indicating the constraint that column(x)=y.
// todo(tk): does this match sig: `gsl::span<const std::tuple<size_t, FieldElement,
// FieldElement>> boundary_conditions`?
#[derive(Clone, Debug)]
pub struct BoundaryCondition<F: Field> {
pub column_idx: usize,
pub point_x: F,
pub point_y: F,
}

impl<F: Field> BoundaryCondition<F> {
pub fn new(column_idx: usize, point_x: F, point_y: F) -> Self {
Self { column_idx, point_x, point_y }
}
}

// doc(tk)
#[derive(Clone, Debug)]
struct ConstraintData<F: Field> {
coeff_idx: usize,
column_index: usize,
point_x: F,
point_y: F,
}

#[cfg(test)]
mod test {
// mirror: https://github.com/starkware-libs/stone-prover/blob/main/src/starkware/air/boundary/boundary_air_test.cc
// test for soundness and correctness of the boundary constraint.
use std::{iter::repeat_with, vec::Vec};

use ark_ff::{BigInteger, Field, PrimeField, UniformRand};
use ark_std::{test_rng, One, Zero};
use ark_test_curves::fp128::Fq;
use rand::Rng;

use super::*;
use crate::test_utils;

// use crate::{
// boundary_air::BoundaryAir,
// utils::{FieldElementVector, Prng},
// };

const N_COLUMNS: usize = 10;
const N_CONDITIONS: usize = 20;
const TRACE_LENGTH: u64 = 1024;

// This test builds a random trace. It then generate random points to sample on random
// columns, and creates boundary constraints for them. It then tests that the resulting
// composition polynomial is indeed of expected low degree., which means the constraints
// hold.
#[test]
fn correctness() {
let mut rng = test_rng();
let trace: Vec<Vec<_>> = repeat_with(|| {
repeat_with(|| <Fq as UniformRand>::rand(&mut rng)).take(N_COLUMNS).collect::<Vec<_>>()
})
.take(TRACE_LENGTH as usize)
.collect();

// Compute correct boundary conditions.
let boundary_conditions: Vec<BoundaryCondition<Fq>> = repeat_with(|| {
let column_index = rng.gen_range(0..N_COLUMNS) as usize;
let point_x = <Fq as UniformRand>::rand(&mut rng);
let point_y = <Fq as UniformRand>::rand(&mut rng);

BoundaryCondition::new(column_index, point_x, point_y)
})
.take(N_CONDITIONS)
.collect();

let air = BoundaryAir::new(TRACE_LENGTH, N_COLUMNS, &boundary_conditions);
let random_coefficients = repeat_with(|| <Fq as UniformRand>::rand(&mut rng))
.take(air.num_random_coefficients())
.collect::<Vec<_>>();
// let actual_degree =
// test_utils::compute_composition_degree(&air, trace, &random_coefficients);

// // Degree is expected to be trace_length - 2.
// assert_eq!(TRACE_LENGTH - 2, actual_degree);
// assert_eq!(air.get_composition_polynomial_degree_bound() - 2, actual_degree);
}

#[test]
fn soundness() {
// let mut prng = Prng::new();
// let mut trace = Vec::new();

// // Generate random trace.
// for _ in 0..n_columns {
// trace.push(prng.random_field_element_vector(trace_length));
// }

// // Compute incorrect boundary conditions.
// let mut boundary_conditions = Vec::new();
// for _ in 0..n_conditions {
// let column_index = prng.uniform_int(0, n_columns - 1);
// let point_x = FieldElement::random(&mut prng);
// let point_y = FieldElement::random(&mut prng);

// boundary_conditions.push((column_index, point_x, point_y));
// }

// let air = BoundaryAir::new(trace_length, n_columns, &boundary_conditions);

// let random_coefficients =
// prng.random_field_element_vector(air.num_random_coefficients());

// let num_of_cosets = 2;
// let actual_degree =
// compute_composition_degree(&air, trace, &random_coefficients, num_of_cosets);

// assert_eq!(num_of_cosets * air.get_composition_polynomial_degree_bound() - 1,
// actual_degree);
}
}
12 changes: 6 additions & 6 deletions air/src/boundary_constraints/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
use ark_ff::Field;
use composition_polynomial::PeriodicColumn;

use crate::TempGslSpan; // Using ark-ff crate for field generics
use crate::GslSpan; // Using ark-ff crate for field generics

/// Creates a periodic column
/// which satisfies:
Expand All @@ -23,8 +23,8 @@ use crate::TempGslSpan; // Using ark-ff crate for field generics
/// - `rows`: span of row indices.
/// - `values`: span of field element values.
pub fn create_boundary_periodic_column<F: Field>(
rows: TempGslSpan<usize>,
values: TempGslSpan<F>,
rows: GslSpan<usize>,
values: GslSpan<F>,
trace_length: usize,
trace_generator: F,
trace_offset: F,
Expand All @@ -34,7 +34,7 @@ pub fn create_boundary_periodic_column<F: Field>(

/// Creates a periodic column with y-values all set to 1.
pub fn create_base_boundary_periodic_column<F: Field>(
rows: TempGslSpan<usize>,
rows: GslSpan<usize>,
trace_length: u64,
trace_generator: F,
trace_offset: F,
Expand All @@ -44,7 +44,7 @@ pub fn create_base_boundary_periodic_column<F: Field>(

/// Creates a periodic column that is zero at the specified rows and invertible elsewhere.
pub fn create_vanishing_periodic_column<F: Field>(
rows: TempGslSpan<usize>,
rows: GslSpan<usize>,
trace_length: u64,
trace_generator: F,
trace_offset: F,
Expand All @@ -54,7 +54,7 @@ pub fn create_vanishing_periodic_column<F: Field>(

/// Creates a periodic column that is zero on rows {0, step, 2*step, ...} except for the given rows.
pub fn create_complement_vanishing_periodic_column<F: Field>(
rows: TempGslSpan<usize>,
rows: GslSpan<usize>,
step: u64,
trace_length: u64,
trace_generator: F,
Expand Down
Loading

0 comments on commit 3d980dc

Please sign in to comment.