#![cfg_attr(not(feature = "std"), no_std)]
extern crate alloc;
extern crate core;
use alloc::collections::BTreeMap;
use alloc::string::{String, ToString};
use alloc::vec;
use alloc::vec::Vec;
use constraint::{BoundaryConstraint, Constraint, ConstraintAst};
use core::marker::PhantomData;
use expander::VirtualExpander;
use hekate_core::errors;
use hekate_core::trace::{ColumnTrace, ColumnType, Trace, TraceCompatibleField};
use hekate_math::{Flat, HardwareField, TowerField};
use permutation::PermutationCheckSpec;
pub mod chiplet;
pub mod circuit;
pub mod constraint;
pub mod digest;
pub mod expander;
pub mod linearized;
pub mod outer;
pub mod permutation;
pub mod predicate;
pub mod schema;
pub trait Air<F: TowerField>: Sized + Clone + Sync {
fn name(&self) -> String {
"HekateAir".to_string()
}
fn num_columns(&self) -> usize {
self.virtual_column_layout().len()
}
fn constraints(&self) -> Vec<Constraint<F>> {
self.constraint_ast().to_constraints()
}
fn boundary_constraints(&self) -> Vec<BoundaryConstraint<F>> {
Vec::new()
}
fn column_layout(&self) -> &[ColumnType];
fn virtual_column_layout(&self) -> &[ColumnType] {
match self.virtual_expander() {
Some(e) => e.virtual_layout(),
None => self.column_layout(),
}
}
fn permutation_checks(&self) -> Vec<(String, PermutationCheckSpec)> {
Vec::new()
}
fn fixed_columns(&self) -> Vec<FixedColumn<F>> {
Vec::new()
}
fn virtual_expander(&self) -> Option<&VirtualExpander> {
None
}
fn parse_virtual_row(&self, bytes: &[u8], res: &mut Vec<Flat<F>>)
where
F: TraceCompatibleField,
{
res.clear();
if let Some(e) = self.virtual_expander() {
e.parse_row(bytes, res)
.expect("committed row byte length must match physical_row_bytes");
return;
}
let mut offset = 0;
for col_type in self.column_layout() {
let size = col_type.byte_size();
if offset + size <= bytes.len() {
res.push(col_type.parse_from_bytes(&bytes[offset..offset + size]));
offset += size;
}
}
}
fn constraint_ast(&self) -> ConstraintAst<F>;
fn inline_chiplets(&self) -> errors::Result<Vec<chiplet::ChipletDef<F>>> {
Ok(Vec::new())
}
fn inline_chiplet_kernels(&self) -> Vec<InlineKernelHint> {
Vec::new()
}
}
pub trait Program<F: TowerField>: Air<F> {
fn num_public_inputs(&self) -> usize {
0
}
fn chiplet_defs(&self) -> errors::Result<Vec<chiplet::ChipletDef<F>>> {
Ok(Vec::new())
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct ProgramCell {
pub col_idx: usize,
pub next_row: bool,
}
impl ProgramCell {
pub fn current(col_idx: usize) -> Self {
Self {
col_idx,
next_row: false,
}
}
pub fn next(col_idx: usize) -> Self {
Self {
col_idx,
next_row: true,
}
}
}
#[derive(Clone, Debug)]
pub struct ProgramInstance<F: TowerField> {
num_rows: usize,
public_inputs: Vec<F>,
}
impl<F: TowerField> ProgramInstance<F> {
pub fn new(num_rows: usize, public_inputs: Vec<F>) -> Self {
assert!(
num_rows.is_power_of_two(),
"Program trace height must be power of 2"
);
Self {
num_rows,
public_inputs,
}
}
#[inline(always)]
pub fn num_rows(&self) -> usize {
self.num_rows
}
#[inline(always)]
pub fn public_inputs(&self) -> &[F] {
&self.public_inputs
}
#[inline(always)]
pub fn public_input(&self, idx: usize) -> Option<F> {
self.public_inputs.get(idx).copied()
}
}
pub struct ProgramWitness<F: TowerField, T: Trace = ColumnTrace> {
pub trace: T,
pub chiplet_traces: Vec<ColumnTrace>,
_marker: PhantomData<F>,
}
impl<F: TowerField, T: Trace> ProgramWitness<F, T> {
pub fn new(trace: T) -> Self {
Self {
trace,
chiplet_traces: Vec::new(),
_marker: PhantomData,
}
}
pub fn with_chiplets(mut self, chiplet_traces: Vec<ColumnTrace>) -> Self {
self.chiplet_traces = chiplet_traces;
self
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct InlineKernelHint {
pub chiplet_idx: usize,
pub root_offset: usize,
pub column_offset: usize,
}
const EQ_TABLE_MAX_VARS: usize = 32;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct CadenceSegment<F> {
pub stride: usize,
pub count: usize,
pub origin: usize,
pub values: Vec<F>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum FixedShape<F> {
LastRow,
FirstRow,
Custom(Vec<bool>),
Periodic {
period: usize,
values: Vec<F>,
},
Sparse(Vec<(usize, F)>),
Dense(Vec<F>),
Cadence {
stride: usize,
count: usize,
origin: usize,
values: Vec<F>,
},
Segments(Vec<CadenceSegment<F>>),
}
impl<F> FixedShape<F> {
pub fn is_overlay(&self) -> bool {
match self {
Self::Periodic { .. }
| Self::Sparse(..)
| Self::Dense(..)
| Self::Cadence { .. }
| Self::Segments(..) => true,
Self::LastRow | Self::FirstRow | Self::Custom(..) => false,
}
}
}
impl<F: HardwareField> FixedShape<F> {
pub fn evaluate(&self, r: &[Flat<F>]) -> Flat<F> {
let one = Flat::from_raw(F::ONE);
match self {
FixedShape::LastRow => {
let mut prod = one;
for &r_k in r {
prod *= r_k;
}
one - prod
}
FixedShape::FirstRow => {
let mut prod = one;
for &r_k in r {
prod *= one - r_k;
}
prod
}
FixedShape::Custom(bits) => {
debug_assert_eq!(bits.len(), r.len(), "Custom point bit width != r.len()");
let mut prod = one;
for (k, &b) in bits.iter().enumerate() {
let factor = if b { r[k] } else { one - r[k] };
prod *= factor;
}
prod
}
FixedShape::Periodic { period, values } => {
let p = period.trailing_zeros() as usize;
let mut acc = Flat::from_raw(F::ZERO);
for (j, &v) in values.iter().enumerate() {
acc += v.to_hardware() * eq_index(&r[..p], j);
}
acc
}
FixedShape::Sparse(entries) => {
let mut acc = Flat::from_raw(F::ZERO);
for &(row, v) in entries {
acc += v.to_hardware() * eq_index(r, row);
}
acc
}
FixedShape::Dense(values) => {
let mut acc = Flat::from_raw(F::ZERO);
for (i, &v) in values.iter().enumerate() {
acc += v.to_hardware() * eq_index(r, i);
}
acc
}
FixedShape::Cadence {
stride,
count,
origin,
values,
} => cadence_mle(r, *stride, *count, *origin, values),
FixedShape::Segments(segments) => {
let mut acc = Flat::from_raw(F::ZERO);
for seg in segments {
acc += cadence_mle(r, seg.stride, seg.count, seg.origin, &seg.values);
}
acc
}
}
}
pub fn value_at_row(&self, row: usize, num_vars: usize) -> Flat<F> {
let one = Flat::from_raw(F::ONE);
let zero = Flat::from_raw(F::ZERO);
match self {
FixedShape::FirstRow => {
if row == 0 {
one
} else {
zero
}
}
FixedShape::LastRow => {
if row == (1usize << num_vars) - 1 {
zero
} else {
one
}
}
FixedShape::Custom(bits) => {
let target = bits
.iter()
.enumerate()
.fold(0usize, |acc, (k, &b)| acc | ((b as usize) << k));
if row == target { one } else { zero }
}
FixedShape::Periodic { period, values } => values[row % period].to_hardware(),
FixedShape::Sparse(entries) => {
let mut acc = zero;
for &(r, v) in entries {
if r == row {
acc += v.to_hardware();
}
}
acc
}
FixedShape::Dense(values) => values[row].to_hardware(),
FixedShape::Cadence {
stride,
count,
origin,
values,
} => cadence_value_at(row, *stride, *count, *origin, values),
FixedShape::Segments(segments) => {
let at = segments.partition_point(|seg| seg.origin <= row);
match at.checked_sub(1) {
Some(i) => {
let seg = &segments[i];
cadence_value_at(row, seg.stride, seg.count, seg.origin, &seg.values)
}
None => zero,
}
}
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct FixedColumn<F> {
pub col_idx: usize,
pub shape: FixedShape<F>,
}
impl<F> FixedColumn<F> {
pub fn last_row(col_idx: usize) -> Self {
Self {
col_idx,
shape: FixedShape::LastRow,
}
}
pub fn first_row(col_idx: usize) -> Self {
Self {
col_idx,
shape: FixedShape::FirstRow,
}
}
pub fn prefix(col_idx: usize, count: usize) -> Self
where
F: TowerField,
{
Self {
col_idx,
shape: FixedShape::Cadence {
stride: 1,
count,
origin: 0,
values: vec![F::ONE],
},
}
}
pub fn custom(col_idx: usize, bits: Vec<bool>) -> Self {
Self {
col_idx,
shape: FixedShape::Custom(bits),
}
}
pub fn periodic(col_idx: usize, period: usize, values: Vec<F>) -> Self {
Self {
col_idx,
shape: FixedShape::Periodic { period, values },
}
}
pub fn sparse(col_idx: usize, entries: Vec<(usize, F)>) -> Self {
Self {
col_idx,
shape: FixedShape::Sparse(entries),
}
}
pub fn dense(col_idx: usize, values: Vec<F>) -> Self {
Self {
col_idx,
shape: FixedShape::Dense(values),
}
}
pub fn segments(col_idx: usize, segments: Vec<CadenceSegment<F>>) -> Self {
Self {
col_idx,
shape: FixedShape::Segments(segments),
}
}
}
pub struct ShapeEvaluator<'a, F> {
r: &'a [Flat<F>],
eq: Option<EqTables<F>>,
sums: BTreeMap<(usize, usize), Vec<Flat<F>>>,
}
impl<'a, F: HardwareField> ShapeEvaluator<'a, F> {
pub fn new(r: &'a [Flat<F>]) -> Self {
Self {
r,
eq: (r.len() <= EQ_TABLE_MAX_VARS).then(|| EqTables::new(r)),
sums: BTreeMap::new(),
}
}
pub fn evaluate(&mut self, shape: &FixedShape<F>) -> Flat<F> {
let zero = Flat::from_raw(F::ZERO);
match shape {
FixedShape::Sparse(entries) => {
let mut acc = zero;
for &(row, value) in entries {
acc += scaled(value, self.eq(row));
}
acc
}
FixedShape::Dense(values) => {
let mut acc = zero;
for (row, &value) in values.iter().enumerate() {
acc += scaled(value, self.eq(row));
}
acc
}
FixedShape::Cadence {
stride,
count,
origin,
values,
} => self.cadence(*stride, *count, *origin, values),
FixedShape::Segments(segments) => {
let mut acc = zero;
for seg in segments {
acc += self.cadence(seg.stride, seg.count, seg.origin, &seg.values);
}
acc
}
FixedShape::LastRow
| FixedShape::FirstRow
| FixedShape::Custom(_)
| FixedShape::Periodic { .. } => shape.evaluate(self.r),
}
}
fn eq(&self, row: usize) -> Flat<F> {
match &self.eq {
Some(tables) => tables.eq(row),
None => eq_index(self.r, row),
}
}
fn cadence(&mut self, stride: usize, count: usize, origin: usize, values: &[F]) -> Flat<F> {
let zero = Flat::from_raw(F::ZERO);
if stride == 0 || values.len() != stride {
return zero;
}
let end = origin.saturating_add(stride.saturating_mul(count));
let Self { r, eq, sums } = self;
match (stride, eq) {
(1, Some(tables)) => scaled(values[0], tables.below(end) + tables.below(origin)),
_ => {
for limit in [end, origin] {
sums.entry((stride, limit))
.or_insert_with(|| residue_sums(r, limit, stride, None));
}
let upper = &sums[&(stride, end)];
let lower = &sums[&(stride, origin)];
let shift = stride - origin % stride;
let mut acc = zero;
for (c, (&u, &l)) in upper.iter().zip(lower).enumerate() {
acc += scaled(values[(c + shift) % stride], u + l);
}
acc
}
}
}
}
struct EqTables<F> {
low_bits: usize,
eq_low: Vec<Flat<F>>,
eq_high: Vec<Flat<F>>,
below_low: Vec<Flat<F>>,
below_high: Vec<Flat<F>>,
}
impl<F: HardwareField> EqTables<F> {
fn new(r: &[Flat<F>]) -> Self {
let low_bits = r.len() / 2;
let eq_low = eq_table(&r[..low_bits]);
let eq_high = eq_table(&r[low_bits..]);
Self {
low_bits,
below_low: running_sums(&eq_low),
below_high: running_sums(&eq_high),
eq_low,
eq_high,
}
}
fn eq(&self, row: usize) -> Flat<F> {
let low = row & (self.eq_low.len() - 1);
let high = (row >> self.low_bits) & (self.eq_high.len() - 1);
self.eq_low[low] * self.eq_high[high]
}
fn below(&self, limit: usize) -> Flat<F> {
let high = limit >> self.low_bits;
if high >= self.eq_high.len() {
return Flat::from_raw(F::ONE);
}
let low = limit & (self.eq_low.len() - 1);
self.below_high[high] + self.eq_high[high] * self.below_low[low]
}
}
pub fn fix<F>(col_idx: usize, shape: FixedShape<F>) -> FixedColumn<F> {
FixedColumn { col_idx, shape }
}
pub fn validate_fixed_columns<F: TowerField>(
fixed: &[FixedColumn<F>],
layout: &[ColumnType],
num_vars: Option<usize>,
) -> errors::Result<()> {
for (i, fc) in fixed.iter().enumerate() {
if fc.col_idx >= layout.len() {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "col_idx out of range",
});
}
validate_shape(&fc.shape, layout[fc.col_idx], num_vars)?;
for prior in &fixed[..i] {
if prior.col_idx == fc.col_idx {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "duplicate pin on same column",
});
}
}
}
Ok(())
}
fn residue_sums<F: HardwareField>(
r: &[Flat<F>],
limit: usize,
stride: usize,
steps: Option<&[Flat<F>]>,
) -> Vec<Flat<F>> {
let one = Flat::from_raw(F::ONE);
let zero = Flat::from_raw(F::ZERO);
let full = match 1usize.checked_shl(r.len() as u32) {
Some(cap) => limit >= cap,
None => false,
};
let mut free = vec![zero; stride];
let mut next = vec![zero; stride];
let mut tight = if full { zero } else { one };
let mut tight_residue = 0usize;
if full {
free[0] = one;
}
for t in (0..r.len()).rev() {
let r_t = r[t];
let step = steps.map(|steps| steps[t]);
next.fill(zero);
for (residue, &weight) in free.iter().enumerate() {
let low = 2 * residue;
let hi = weight * r_t;
let (at, w) = division_step(low, weight + hi, stride, step);
next[at] += w;
let (at, w) = division_step(low + 1, hi, stride, step);
next[at] += w;
}
let low = 2 * tight_residue;
let hi = tight * r_t;
(tight_residue, tight) = match (limit >> t) & 1 {
1 => {
let (at, w) = division_step(low, tight + hi, stride, step);
next[at] += w;
division_step(low + 1, hi, stride, step)
}
_ => division_step(low, tight + hi, stride, step),
};
core::mem::swap(&mut free, &mut next);
}
free
}
fn division_step<F: HardwareField>(
value: usize,
weight: Flat<F>,
stride: usize,
step: Option<Flat<F>>,
) -> (usize, Flat<F>) {
match (value >= stride, step) {
(true, Some(step)) => (value - stride, weight * step),
(true, None) => (value - stride, weight),
(false, _) => (value, weight),
}
}
fn cadence_mle<F: HardwareField>(
r: &[Flat<F>],
stride: usize,
count: usize,
origin: usize,
values: &[F],
) -> Flat<F> {
if stride == 0 || values.len() != stride {
return Flat::from_raw(F::ZERO);
}
let end = origin.saturating_add(stride.saturating_mul(count));
let mut class_sums = residue_sums(r, end, stride, None);
let start_sums = residue_sums(r, origin, stride, None);
for (c, s) in start_sums.iter().enumerate() {
class_sums[c] += *s;
}
let mut acc = Flat::from_raw(F::ZERO);
for (c, sum) in class_sums.iter().enumerate() {
let offset = (c + stride - origin % stride) % stride;
acc += values[offset].to_hardware() * *sum;
}
acc
}
fn cadence_value_at<F: HardwareField>(
row: usize,
stride: usize,
count: usize,
origin: usize,
values: &[F],
) -> Flat<F> {
if stride == 0 || values.len() != stride {
return Flat::from_raw(F::ZERO);
}
let end = origin.saturating_add(stride.saturating_mul(count));
if row < origin || row >= end {
Flat::from_raw(F::ZERO)
} else {
values[(row - origin) % stride].to_hardware()
}
}
fn eq_index<F: HardwareField>(r: &[Flat<F>], index: usize) -> Flat<F> {
let one = Flat::from_raw(F::ONE);
let mut prod = one;
for (k, &r_k) in r.iter().enumerate() {
let factor = if (index >> k) & 1 == 1 {
r_k
} else {
one - r_k
};
prod *= factor;
}
prod
}
fn eq_table<F: HardwareField>(r: &[Flat<F>]) -> Vec<Flat<F>> {
let mut table = Vec::with_capacity(1 << r.len());
table.push(Flat::from_raw(F::ONE));
for &r_k in r {
for i in 0..table.len() {
let hi = table[i] * r_k;
table[i] += hi;
table.push(hi);
}
}
table
}
fn running_sums<F: HardwareField>(values: &[Flat<F>]) -> Vec<Flat<F>> {
let mut sums = Vec::with_capacity(values.len() + 1);
let mut acc = Flat::from_raw(F::ZERO);
sums.push(acc);
for &v in values {
acc += v;
sums.push(acc);
}
sums
}
fn scaled<F: HardwareField>(value: F, x: Flat<F>) -> Flat<F> {
if value == F::ZERO {
Flat::from_raw(F::ZERO)
} else if value == F::ONE {
x
} else {
value.to_hardware() * x
}
}
fn validate_shape<F: TowerField>(
shape: &FixedShape<F>,
col_type: ColumnType,
num_vars: Option<usize>,
) -> errors::Result<()> {
match shape {
FixedShape::LastRow | FixedShape::FirstRow => Ok(()),
FixedShape::Custom(bits) => match num_vars {
Some(nv) if bits.len() != nv => Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Custom point bit width != num_vars",
}),
_ => Ok(()),
},
FixedShape::Periodic { period, values } => {
if !period.is_power_of_two() {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Periodic period must be a power of two",
});
}
if values.len() != *period {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Periodic values length != period",
});
}
if let Some(nv) = num_vars
&& *period > (1usize << nv)
{
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Periodic period exceeds trace height",
});
}
check_bit_domain(values.iter().copied(), col_type)
}
FixedShape::Sparse(entries) => {
if let Some(nv) = num_vars {
let n = 1usize << nv;
for &(row, _) in entries {
if row >= n {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Sparse row index exceeds trace height",
});
}
}
}
let mut rows: Vec<usize> = entries.iter().map(|&(row, _)| row).collect();
rows.sort_unstable();
if rows.windows(2).any(|pair| pair[0] == pair[1]) {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "duplicate Sparse row",
});
}
check_bit_domain(entries.iter().map(|&(_, v)| v), col_type)
}
FixedShape::Dense(values) => {
if let Some(nv) = num_vars
&& values.len() != (1usize << nv)
{
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Dense values length != trace height",
});
}
check_bit_domain(values.iter().copied(), col_type)
}
FixedShape::Cadence {
stride,
count,
origin,
values,
} => {
if *stride == 0 {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Cadence stride must be non-zero",
});
}
if values.len() != *stride {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Cadence values length != stride",
});
}
let end = stride
.checked_mul(*count)
.and_then(|span| span.checked_add(*origin))
.ok_or(errors::Error::Protocol {
protocol: "fixed_column",
message: "Cadence active range overflows",
})?;
if let Some(nv) = num_vars
&& end > (1usize << nv)
{
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Cadence active range exceeds trace height",
});
}
check_bit_domain(values.iter().copied(), col_type)
}
FixedShape::Segments(segments) => {
if segments.is_empty() {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Segments list is empty",
});
}
let mut prev_end = 0usize;
for seg in segments {
if seg.stride == 0 {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Segments stride must be non-zero",
});
}
if seg.count == 0 {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Segments count must be non-zero",
});
}
if seg.values.len() != seg.stride {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Segments values length != stride",
});
}
let end = seg
.stride
.checked_mul(seg.count)
.and_then(|span| span.checked_add(seg.origin))
.ok_or(errors::Error::Protocol {
protocol: "fixed_column",
message: "Segments active range overflows",
})?;
if seg.origin < prev_end {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Segments must be sorted by origin and disjoint",
});
}
prev_end = end;
check_bit_domain(seg.values.iter().copied(), col_type)?;
}
if let Some(nv) = num_vars
&& prev_end > (1usize << nv)
{
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Segments active range exceeds trace height",
});
}
Ok(())
}
}
}
fn check_bit_domain<F: TowerField>(
values: impl Iterator<Item = F>,
col_type: ColumnType,
) -> errors::Result<()> {
if col_type != ColumnType::Bit {
return Ok(());
}
for v in values {
if v != F::ZERO && v != F::ONE {
return Err(errors::Error::Protocol {
protocol: "fixed_column",
message: "Bit fixed column value not in {0,1}",
});
}
}
Ok(())
}