1#![cfg_attr(not(feature = "std"), no_std)]
6
7extern crate alloc;
8extern crate core;
9
10use alloc::string::{String, ToString};
11use alloc::vec;
12use alloc::vec::Vec;
13use constraint::{BoundaryConstraint, Constraint, ConstraintAst};
14use core::marker::PhantomData;
15use expander::VirtualExpander;
16use hekate_core::errors;
17use hekate_core::trace::{ColumnTrace, ColumnType, Trace, TraceCompatibleField};
18use hekate_math::{Flat, HardwareField, TowerField};
19use permutation::PermutationCheckSpec;
20
21pub mod chiplet;
22pub mod circuit;
23pub mod constraint;
24pub mod digest;
25pub mod expander;
26pub mod linearized;
27pub mod outer;
28pub mod permutation;
29pub mod predicate;
30pub mod schema;
31
32pub trait Air<F: TowerField>: Sized + Clone + Sync {
42 fn name(&self) -> String {
43 "HekateAir".to_string()
44 }
45
46 fn num_columns(&self) -> usize {
47 self.virtual_column_layout().len()
48 }
49
50 fn constraints(&self) -> Vec<Constraint<F>> {
52 self.constraint_ast().to_constraints()
53 }
54
55 fn boundary_constraints(&self) -> Vec<BoundaryConstraint<F>> {
59 Vec::new()
60 }
61
62 fn column_layout(&self) -> &[ColumnType];
67
68 fn virtual_column_layout(&self) -> &[ColumnType] {
72 match self.virtual_expander() {
73 Some(e) => e.virtual_layout(),
74 None => self.column_layout(),
75 }
76 }
77
78 fn permutation_checks(&self) -> Vec<(String, PermutationCheckSpec)> {
90 Vec::new()
91 }
92
93 fn fixed_columns(&self) -> Vec<FixedColumn<F>> {
97 Vec::new()
98 }
99
100 fn virtual_expander(&self) -> Option<&VirtualExpander> {
105 None
106 }
107
108 fn parse_virtual_row(&self, bytes: &[u8], res: &mut Vec<Flat<F>>)
115 where
116 F: TraceCompatibleField,
117 {
118 res.clear();
119
120 if let Some(e) = self.virtual_expander() {
121 e.parse_row(bytes, res)
122 .expect("committed row byte length must match physical_row_bytes");
123 return;
124 }
125
126 let mut offset = 0;
127 for col_type in self.column_layout() {
128 let size = col_type.byte_size();
129 if offset + size <= bytes.len() {
130 res.push(col_type.parse_from_bytes(&bytes[offset..offset + size]));
131 offset += size;
132 }
133 }
134 }
135
136 fn constraint_ast(&self) -> ConstraintAst<F>;
138
139 fn inline_chiplets(&self) -> errors::Result<Vec<chiplet::ChipletDef<F>>> {
141 Ok(Vec::new())
142 }
143
144 fn inline_chiplet_kernels(&self) -> Vec<InlineKernelHint> {
146 Vec::new()
147 }
148}
149
150pub trait Program<F: TowerField>: Air<F> {
161 fn num_public_inputs(&self) -> usize {
163 0
164 }
165
166 fn chiplet_defs(&self) -> errors::Result<Vec<chiplet::ChipletDef<F>>> {
171 Ok(Vec::new())
172 }
173}
174
175#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
179pub struct ProgramCell {
180 pub col_idx: usize,
181
182 pub next_row: bool,
185}
186
187impl ProgramCell {
188 pub fn current(col_idx: usize) -> Self {
190 Self {
191 col_idx,
192 next_row: false,
193 }
194 }
195
196 pub fn next(col_idx: usize) -> Self {
198 Self {
199 col_idx,
200 next_row: true,
201 }
202 }
203}
204
205#[derive(Clone, Debug)]
212pub struct ProgramInstance<F: TowerField> {
213 num_rows: usize,
214 public_inputs: Vec<F>,
215}
216
217impl<F: TowerField> ProgramInstance<F> {
218 pub fn new(num_rows: usize, public_inputs: Vec<F>) -> Self {
219 assert!(
220 num_rows.is_power_of_two(),
221 "Program trace height must be power of 2"
222 );
223
224 Self {
225 num_rows,
226 public_inputs,
227 }
228 }
229
230 #[inline(always)]
231 pub fn num_rows(&self) -> usize {
232 self.num_rows
233 }
234
235 #[inline(always)]
237 pub fn public_inputs(&self) -> &[F] {
238 &self.public_inputs
239 }
240
241 #[inline(always)]
242 pub fn public_input(&self, idx: usize) -> Option<F> {
243 self.public_inputs.get(idx).copied()
244 }
245}
246
247pub struct ProgramWitness<F: TowerField, T: Trace = ColumnTrace> {
251 pub trace: T,
252 pub chiplet_traces: Vec<ColumnTrace>,
253 _marker: PhantomData<F>,
254}
255
256impl<F: TowerField, T: Trace> ProgramWitness<F, T> {
257 pub fn new(trace: T) -> Self {
258 Self {
259 trace,
260 chiplet_traces: Vec::new(),
261 _marker: PhantomData,
262 }
263 }
264
265 pub fn with_chiplets(mut self, chiplet_traces: Vec<ColumnTrace>) -> Self {
268 self.chiplet_traces = chiplet_traces;
269 self
270 }
271}
272
273#[derive(Clone, Copy, Debug, PartialEq, Eq)]
277pub struct InlineKernelHint {
278 pub chiplet_idx: usize,
280
281 pub root_offset: usize,
284
285 pub column_offset: usize,
288}
289
290#[derive(Clone, Debug, PartialEq, Eq)]
295pub struct CadenceSegment<F> {
296 pub stride: usize,
297 pub count: usize,
298 pub origin: usize,
299 pub values: Vec<F>,
300}
301
302#[derive(Clone, Debug, PartialEq, Eq)]
305pub enum FixedShape<F> {
306 LastRow,
307 FirstRow,
308 Custom(Vec<bool>),
309 Periodic {
310 period: usize,
311 values: Vec<F>,
312 },
313 Sparse(Vec<(usize, F)>),
314 Dense(Vec<F>),
315
316 Cadence {
319 stride: usize,
320 count: usize,
321 origin: usize,
322 values: Vec<F>,
323 },
324
325 Segments(Vec<CadenceSegment<F>>),
328}
329
330impl<F> FixedShape<F> {
331 pub fn is_overlay(&self) -> bool {
334 match self {
335 Self::Periodic { .. }
336 | Self::Sparse(..)
337 | Self::Dense(..)
338 | Self::Cadence { .. }
339 | Self::Segments(..) => true,
340 Self::LastRow | Self::FirstRow | Self::Custom(..) => false,
341 }
342 }
343}
344
345impl<F: HardwareField> FixedShape<F> {
346 pub fn evaluate(&self, r: &[Flat<F>]) -> Flat<F> {
349 let one = Flat::from_raw(F::ONE);
350 match self {
351 FixedShape::LastRow => {
352 let mut prod = one;
353 for &r_k in r {
354 prod *= r_k;
355 }
356
357 one - prod
358 }
359 FixedShape::FirstRow => {
360 let mut prod = one;
361 for &r_k in r {
362 prod *= one - r_k;
363 }
364
365 prod
366 }
367 FixedShape::Custom(bits) => {
368 debug_assert_eq!(bits.len(), r.len(), "Custom point bit width != r.len()");
369
370 let mut prod = one;
371 for (k, &b) in bits.iter().enumerate() {
372 let factor = if b { r[k] } else { one - r[k] };
373 prod *= factor;
374 }
375
376 prod
377 }
378 FixedShape::Periodic { period, values } => {
379 let p = period.trailing_zeros() as usize;
382
383 let mut acc = Flat::from_raw(F::ZERO);
384 for (j, &v) in values.iter().enumerate() {
385 acc += v.to_hardware() * eq_index(&r[..p], j);
386 }
387
388 acc
389 }
390 FixedShape::Sparse(entries) => {
391 let mut acc = Flat::from_raw(F::ZERO);
392 for &(row, v) in entries {
393 acc += v.to_hardware() * eq_index(r, row);
394 }
395
396 acc
397 }
398 FixedShape::Dense(values) => {
399 let mut acc = Flat::from_raw(F::ZERO);
400 for (i, &v) in values.iter().enumerate() {
401 acc += v.to_hardware() * eq_index(r, i);
402 }
403
404 acc
405 }
406 FixedShape::Cadence {
407 stride,
408 count,
409 origin,
410 values,
411 } => cadence_mle(r, *stride, *count, *origin, values),
412 FixedShape::Segments(segments) => {
413 let mut acc = Flat::from_raw(F::ZERO);
414 for seg in segments {
415 acc += cadence_mle(r, seg.stride, seg.count, seg.origin, &seg.values);
416 }
417
418 acc
419 }
420 }
421 }
422
423 pub fn value_at_row(&self, row: usize, num_vars: usize) -> Flat<F> {
427 let one = Flat::from_raw(F::ONE);
428 let zero = Flat::from_raw(F::ZERO);
429
430 match self {
431 FixedShape::FirstRow => {
432 if row == 0 {
433 one
434 } else {
435 zero
436 }
437 }
438 FixedShape::LastRow => {
439 if row == (1usize << num_vars) - 1 {
440 zero
441 } else {
442 one
443 }
444 }
445 FixedShape::Custom(bits) => {
446 let target = bits
447 .iter()
448 .enumerate()
449 .fold(0usize, |acc, (k, &b)| acc | ((b as usize) << k));
450
451 if row == target { one } else { zero }
452 }
453 FixedShape::Periodic { period, values } => values[row % period].to_hardware(),
454 FixedShape::Sparse(entries) => {
455 let mut acc = zero;
456 for &(r, v) in entries {
457 if r == row {
458 acc += v.to_hardware();
459 }
460 }
461
462 acc
463 }
464 FixedShape::Dense(values) => values[row].to_hardware(),
465 FixedShape::Cadence {
466 stride,
467 count,
468 origin,
469 values,
470 } => cadence_value_at(row, *stride, *count, *origin, values),
471 FixedShape::Segments(segments) => {
472 let at = segments.partition_point(|seg| seg.origin <= row);
475
476 match at.checked_sub(1) {
477 Some(i) => {
478 let seg = &segments[i];
479
480 cadence_value_at(row, seg.stride, seg.count, seg.origin, &seg.values)
481 }
482 None => zero,
483 }
484 }
485 }
486 }
487}
488
489#[derive(Clone, Debug, PartialEq, Eq)]
491pub struct FixedColumn<F> {
492 pub col_idx: usize,
493 pub shape: FixedShape<F>,
494}
495
496impl<F> FixedColumn<F> {
497 pub fn last_row(col_idx: usize) -> Self {
498 Self {
499 col_idx,
500 shape: FixedShape::LastRow,
501 }
502 }
503
504 pub fn first_row(col_idx: usize) -> Self {
505 Self {
506 col_idx,
507 shape: FixedShape::FirstRow,
508 }
509 }
510
511 pub fn prefix(col_idx: usize, count: usize) -> Self
512 where
513 F: TowerField,
514 {
515 Self {
516 col_idx,
517 shape: FixedShape::Cadence {
518 stride: 1,
519 count,
520 origin: 0,
521 values: vec![F::ONE],
522 },
523 }
524 }
525
526 pub fn custom(col_idx: usize, bits: Vec<bool>) -> Self {
527 Self {
528 col_idx,
529 shape: FixedShape::Custom(bits),
530 }
531 }
532
533 pub fn periodic(col_idx: usize, period: usize, values: Vec<F>) -> Self {
534 Self {
535 col_idx,
536 shape: FixedShape::Periodic { period, values },
537 }
538 }
539
540 pub fn sparse(col_idx: usize, entries: Vec<(usize, F)>) -> Self {
541 Self {
542 col_idx,
543 shape: FixedShape::Sparse(entries),
544 }
545 }
546
547 pub fn dense(col_idx: usize, values: Vec<F>) -> Self {
548 Self {
549 col_idx,
550 shape: FixedShape::Dense(values),
551 }
552 }
553
554 pub fn segments(col_idx: usize, segments: Vec<CadenceSegment<F>>) -> Self {
555 Self {
556 col_idx,
557 shape: FixedShape::Segments(segments),
558 }
559 }
560}
561
562pub fn fix<F>(col_idx: usize, shape: FixedShape<F>) -> FixedColumn<F> {
564 FixedColumn { col_idx, shape }
565}
566
567pub fn validate_fixed_columns<F: TowerField>(
571 fixed: &[FixedColumn<F>],
572 layout: &[ColumnType],
573 num_vars: Option<usize>,
574) -> errors::Result<()> {
575 for (i, fc) in fixed.iter().enumerate() {
576 if fc.col_idx >= layout.len() {
577 return Err(errors::Error::Protocol {
578 protocol: "fixed_column",
579 message: "col_idx out of range",
580 });
581 }
582
583 validate_shape(&fc.shape, layout[fc.col_idx], num_vars)?;
584
585 for prior in &fixed[..i] {
586 if prior.col_idx == fc.col_idx {
587 return Err(errors::Error::Protocol {
588 protocol: "fixed_column",
589 message: "duplicate pin on same column",
590 });
591 }
592 }
593 }
594
595 Ok(())
596}
597
598fn residue_sums<F: HardwareField>(
601 r: &[Flat<F>],
602 limit: usize,
603 stride: usize,
604 steps: &[Flat<F>],
605) -> Vec<Flat<F>> {
606 let one = Flat::from_raw(F::ONE);
607 let zero = Flat::from_raw(F::ZERO);
608
609 let full = match 1usize.checked_shl(r.len() as u32) {
610 Some(cap) => limit >= cap,
611 None => false,
612 };
613
614 let mut free = vec![zero; stride];
615 let mut next = vec![zero; stride];
616
617 let mut tight = if full { zero } else { one };
618 let mut tight_residue = 0usize;
619
620 if full {
621 free[0] = one;
622 }
623
624 for t in (0..r.len()).rev() {
625 let bit_one = r[t];
626 let bit_zero = one + r[t];
627 let step = steps[t];
628
629 next.fill(zero);
630
631 for (residue, &weight) in free.iter().enumerate() {
632 let low = 2 * residue;
633
634 let (at, w) = division_step(low, weight * bit_zero, stride, step);
635 next[at] += w;
636
637 let (at, w) = division_step(low + 1, weight * bit_one, stride, step);
638 next[at] += w;
639 }
640
641 let low = 2 * tight_residue;
642
643 (tight_residue, tight) = match (limit >> t) & 1 {
644 1 => {
645 let (at, w) = division_step(low, tight * bit_zero, stride, step);
646 next[at] += w;
647
648 division_step(low + 1, tight * bit_one, stride, step)
649 }
650 _ => division_step(low, tight * bit_zero, stride, step),
651 };
652
653 core::mem::swap(&mut free, &mut next);
654 }
655
656 free
657}
658
659fn division_step<F: HardwareField>(
660 value: usize,
661 weight: Flat<F>,
662 stride: usize,
663 step: Flat<F>,
664) -> (usize, Flat<F>) {
665 match value >= stride {
666 true => (value - stride, weight * step),
667 false => (value, weight),
668 }
669}
670
671fn cadence_mle<F: HardwareField>(
672 r: &[Flat<F>],
673 stride: usize,
674 count: usize,
675 origin: usize,
676 values: &[F],
677) -> Flat<F> {
678 if stride == 0 || values.len() != stride {
679 return Flat::from_raw(F::ZERO);
680 }
681
682 let end = origin.saturating_add(stride.saturating_mul(count));
683 let unit = vec![Flat::from_raw(F::ONE); r.len()];
684
685 let mut class_sums = residue_sums(r, end, stride, &unit);
686 let start_sums = residue_sums(r, origin, stride, &unit);
687
688 for (c, s) in start_sums.iter().enumerate() {
689 class_sums[c] += *s;
690 }
691
692 let mut acc = Flat::from_raw(F::ZERO);
693 for (c, sum) in class_sums.iter().enumerate() {
694 let offset = (c + stride - origin % stride) % stride;
695 acc += values[offset].to_hardware() * *sum;
696 }
697
698 acc
699}
700
701fn cadence_value_at<F: HardwareField>(
702 row: usize,
703 stride: usize,
704 count: usize,
705 origin: usize,
706 values: &[F],
707) -> Flat<F> {
708 if stride == 0 || values.len() != stride {
709 return Flat::from_raw(F::ZERO);
710 }
711
712 let end = origin.saturating_add(stride.saturating_mul(count));
713
714 if row < origin || row >= end {
715 Flat::from_raw(F::ZERO)
716 } else {
717 values[(row - origin) % stride].to_hardware()
718 }
719}
720
721fn eq_index<F: HardwareField>(r: &[Flat<F>], index: usize) -> Flat<F> {
722 let one = Flat::from_raw(F::ONE);
723 let mut prod = one;
724
725 for (k, &r_k) in r.iter().enumerate() {
726 let factor = if (index >> k) & 1 == 1 {
727 r_k
728 } else {
729 one - r_k
730 };
731
732 prod *= factor;
733 }
734
735 prod
736}
737
738fn validate_shape<F: TowerField>(
739 shape: &FixedShape<F>,
740 col_type: ColumnType,
741 num_vars: Option<usize>,
742) -> errors::Result<()> {
743 match shape {
744 FixedShape::LastRow | FixedShape::FirstRow => Ok(()),
745 FixedShape::Custom(bits) => match num_vars {
746 Some(nv) if bits.len() != nv => Err(errors::Error::Protocol {
747 protocol: "fixed_column",
748 message: "Custom point bit width != num_vars",
749 }),
750 _ => Ok(()),
751 },
752 FixedShape::Periodic { period, values } => {
753 if !period.is_power_of_two() {
754 return Err(errors::Error::Protocol {
755 protocol: "fixed_column",
756 message: "Periodic period must be a power of two",
757 });
758 }
759
760 if values.len() != *period {
761 return Err(errors::Error::Protocol {
762 protocol: "fixed_column",
763 message: "Periodic values length != period",
764 });
765 }
766
767 if let Some(nv) = num_vars
768 && *period > (1usize << nv)
769 {
770 return Err(errors::Error::Protocol {
771 protocol: "fixed_column",
772 message: "Periodic period exceeds trace height",
773 });
774 }
775
776 check_bit_domain(values.iter().copied(), col_type)
777 }
778 FixedShape::Sparse(entries) => {
779 if let Some(nv) = num_vars {
780 let n = 1usize << nv;
781 for &(row, _) in entries {
782 if row >= n {
783 return Err(errors::Error::Protocol {
784 protocol: "fixed_column",
785 message: "Sparse row index exceeds trace height",
786 });
787 }
788 }
789 }
790
791 for (i, &(row, _)) in entries.iter().enumerate() {
792 if entries[..i].iter().any(|&(prior, _)| prior == row) {
793 return Err(errors::Error::Protocol {
794 protocol: "fixed_column",
795 message: "duplicate Sparse row",
796 });
797 }
798 }
799
800 check_bit_domain(entries.iter().map(|&(_, v)| v), col_type)
801 }
802 FixedShape::Dense(values) => {
803 if let Some(nv) = num_vars
804 && values.len() != (1usize << nv)
805 {
806 return Err(errors::Error::Protocol {
807 protocol: "fixed_column",
808 message: "Dense values length != trace height",
809 });
810 }
811
812 check_bit_domain(values.iter().copied(), col_type)
813 }
814 FixedShape::Cadence {
815 stride,
816 count,
817 origin,
818 values,
819 } => {
820 if *stride == 0 {
821 return Err(errors::Error::Protocol {
822 protocol: "fixed_column",
823 message: "Cadence stride must be non-zero",
824 });
825 }
826
827 if values.len() != *stride {
828 return Err(errors::Error::Protocol {
829 protocol: "fixed_column",
830 message: "Cadence values length != stride",
831 });
832 }
833
834 let end = stride
835 .checked_mul(*count)
836 .and_then(|span| span.checked_add(*origin))
837 .ok_or(errors::Error::Protocol {
838 protocol: "fixed_column",
839 message: "Cadence active range overflows",
840 })?;
841
842 if let Some(nv) = num_vars
843 && end > (1usize << nv)
844 {
845 return Err(errors::Error::Protocol {
846 protocol: "fixed_column",
847 message: "Cadence active range exceeds trace height",
848 });
849 }
850
851 check_bit_domain(values.iter().copied(), col_type)
852 }
853 FixedShape::Segments(segments) => {
854 if segments.is_empty() {
855 return Err(errors::Error::Protocol {
856 protocol: "fixed_column",
857 message: "Segments list is empty",
858 });
859 }
860
861 let mut prev_end = 0usize;
862 for seg in segments {
863 if seg.stride == 0 {
864 return Err(errors::Error::Protocol {
865 protocol: "fixed_column",
866 message: "Segments stride must be non-zero",
867 });
868 }
869
870 if seg.count == 0 {
871 return Err(errors::Error::Protocol {
872 protocol: "fixed_column",
873 message: "Segments count must be non-zero",
874 });
875 }
876
877 if seg.values.len() != seg.stride {
878 return Err(errors::Error::Protocol {
879 protocol: "fixed_column",
880 message: "Segments values length != stride",
881 });
882 }
883
884 let end = seg
885 .stride
886 .checked_mul(seg.count)
887 .and_then(|span| span.checked_add(seg.origin))
888 .ok_or(errors::Error::Protocol {
889 protocol: "fixed_column",
890 message: "Segments active range overflows",
891 })?;
892
893 if seg.origin < prev_end {
894 return Err(errors::Error::Protocol {
895 protocol: "fixed_column",
896 message: "Segments must be sorted by origin and disjoint",
897 });
898 }
899
900 prev_end = end;
901
902 check_bit_domain(seg.values.iter().copied(), col_type)?;
903 }
904
905 if let Some(nv) = num_vars
906 && prev_end > (1usize << nv)
907 {
908 return Err(errors::Error::Protocol {
909 protocol: "fixed_column",
910 message: "Segments active range exceeds trace height",
911 });
912 }
913
914 Ok(())
915 }
916 }
917}
918
919fn check_bit_domain<F: TowerField>(
920 values: impl Iterator<Item = F>,
921 col_type: ColumnType,
922) -> errors::Result<()> {
923 if col_type != ColumnType::Bit {
924 return Ok(());
925 }
926
927 for v in values {
928 if v != F::ZERO && v != F::ONE {
929 return Err(errors::Error::Protocol {
930 protocol: "fixed_column",
931 message: "Bit fixed column value not in {0,1}",
932 });
933 }
934 }
935
936 Ok(())
937}