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)]
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 prefix_residue_sums<F: HardwareField>(
600 r: &[Flat<F>],
601 limit: usize,
602 stride: usize,
603) -> Vec<Flat<F>> {
604 let one = Flat::from_raw(F::ONE);
605 let zero = Flat::from_raw(F::ZERO);
606 let nv = r.len();
607
608 let full = match 1usize.checked_shl(nv as u32) {
609 Some(cap) => limit >= cap,
610 None => false,
611 };
612
613 let mut pow2_mod = Vec::with_capacity(nv);
614 let mut p2 = 1usize % stride;
615
616 for _ in 0..nv {
617 pow2_mod.push(p2);
618 p2 = (p2 * 2) % stride;
619 }
620
621 let mut free = alloc::vec![zero; stride];
622 let mut next_free = alloc::vec![zero; stride];
623 let mut tight = if full { zero } else { one };
624 let mut tight_residue = 0usize;
625
626 if full {
627 free[0] = one;
628 }
629
630 for k in (0..nv).rev() {
631 let p2k = pow2_mod[k];
632 for slot in next_free.iter_mut() {
633 *slot = zero;
634 }
635
636 for c in 0..stride {
637 let w = free[c];
638
639 next_free[c] += w * (one - r[k]);
640 next_free[(c + p2k) % stride] += w * r[k];
641 }
642
643 if (limit >> k) & 1 == 1 {
644 next_free[tight_residue] += tight * (one - r[k]);
645 tight *= r[k];
646 tight_residue = (tight_residue + p2k) % stride;
647 } else {
648 tight *= one - r[k];
649 }
650
651 core::mem::swap(&mut free, &mut next_free);
652 }
653
654 free
655}
656
657fn cadence_mle<F: HardwareField>(
658 r: &[Flat<F>],
659 stride: usize,
660 count: usize,
661 origin: usize,
662 values: &[F],
663) -> Flat<F> {
664 if stride == 0 || values.len() != stride {
665 return Flat::from_raw(F::ZERO);
666 }
667
668 let end = origin.saturating_add(stride.saturating_mul(count));
669
670 let mut class_sums = prefix_residue_sums(r, end, stride);
671 let start_sums = prefix_residue_sums(r, origin, stride);
672
673 for (c, s) in start_sums.iter().enumerate() {
674 class_sums[c] += *s;
675 }
676
677 let mut acc = Flat::from_raw(F::ZERO);
678 for (c, sum) in class_sums.iter().enumerate() {
679 let offset = (c + stride - origin % stride) % stride;
680 acc += values[offset].to_hardware() * *sum;
681 }
682
683 acc
684}
685
686fn cadence_value_at<F: HardwareField>(
687 row: usize,
688 stride: usize,
689 count: usize,
690 origin: usize,
691 values: &[F],
692) -> Flat<F> {
693 if stride == 0 || values.len() != stride {
694 return Flat::from_raw(F::ZERO);
695 }
696
697 let end = origin.saturating_add(stride.saturating_mul(count));
698
699 if row < origin || row >= end {
700 Flat::from_raw(F::ZERO)
701 } else {
702 values[(row - origin) % stride].to_hardware()
703 }
704}
705
706fn eq_index<F: HardwareField>(r: &[Flat<F>], index: usize) -> Flat<F> {
707 let one = Flat::from_raw(F::ONE);
708 let mut prod = one;
709
710 for (k, &r_k) in r.iter().enumerate() {
711 let factor = if (index >> k) & 1 == 1 {
712 r_k
713 } else {
714 one - r_k
715 };
716
717 prod *= factor;
718 }
719
720 prod
721}
722
723fn validate_shape<F: TowerField>(
724 shape: &FixedShape<F>,
725 col_type: ColumnType,
726 num_vars: Option<usize>,
727) -> errors::Result<()> {
728 match shape {
729 FixedShape::LastRow | FixedShape::FirstRow => Ok(()),
730 FixedShape::Custom(bits) => match num_vars {
731 Some(nv) if bits.len() != nv => Err(errors::Error::Protocol {
732 protocol: "fixed_column",
733 message: "Custom point bit width != num_vars",
734 }),
735 _ => Ok(()),
736 },
737 FixedShape::Periodic { period, values } => {
738 if !period.is_power_of_two() {
739 return Err(errors::Error::Protocol {
740 protocol: "fixed_column",
741 message: "Periodic period must be a power of two",
742 });
743 }
744
745 if values.len() != *period {
746 return Err(errors::Error::Protocol {
747 protocol: "fixed_column",
748 message: "Periodic values length != period",
749 });
750 }
751
752 if let Some(nv) = num_vars
753 && *period > (1usize << nv)
754 {
755 return Err(errors::Error::Protocol {
756 protocol: "fixed_column",
757 message: "Periodic period exceeds trace height",
758 });
759 }
760
761 check_bit_domain(values.iter().copied(), col_type)
762 }
763 FixedShape::Sparse(entries) => {
764 if let Some(nv) = num_vars {
765 let n = 1usize << nv;
766 for &(row, _) in entries {
767 if row >= n {
768 return Err(errors::Error::Protocol {
769 protocol: "fixed_column",
770 message: "Sparse row index exceeds trace height",
771 });
772 }
773 }
774 }
775
776 for (i, &(row, _)) in entries.iter().enumerate() {
777 if entries[..i].iter().any(|&(prior, _)| prior == row) {
778 return Err(errors::Error::Protocol {
779 protocol: "fixed_column",
780 message: "duplicate Sparse row",
781 });
782 }
783 }
784
785 check_bit_domain(entries.iter().map(|&(_, v)| v), col_type)
786 }
787 FixedShape::Dense(values) => {
788 if let Some(nv) = num_vars
789 && values.len() != (1usize << nv)
790 {
791 return Err(errors::Error::Protocol {
792 protocol: "fixed_column",
793 message: "Dense values length != trace height",
794 });
795 }
796
797 check_bit_domain(values.iter().copied(), col_type)
798 }
799 FixedShape::Cadence {
800 stride,
801 count,
802 origin,
803 values,
804 } => {
805 if *stride == 0 {
806 return Err(errors::Error::Protocol {
807 protocol: "fixed_column",
808 message: "Cadence stride must be non-zero",
809 });
810 }
811
812 if values.len() != *stride {
813 return Err(errors::Error::Protocol {
814 protocol: "fixed_column",
815 message: "Cadence values length != stride",
816 });
817 }
818
819 let end = stride
820 .checked_mul(*count)
821 .and_then(|span| span.checked_add(*origin))
822 .ok_or(errors::Error::Protocol {
823 protocol: "fixed_column",
824 message: "Cadence active range overflows",
825 })?;
826
827 if let Some(nv) = num_vars
828 && end > (1usize << nv)
829 {
830 return Err(errors::Error::Protocol {
831 protocol: "fixed_column",
832 message: "Cadence active range exceeds trace height",
833 });
834 }
835
836 check_bit_domain(values.iter().copied(), col_type)
837 }
838 FixedShape::Segments(segments) => {
839 if segments.is_empty() {
840 return Err(errors::Error::Protocol {
841 protocol: "fixed_column",
842 message: "Segments list is empty",
843 });
844 }
845
846 let mut prev_end = 0usize;
847 for seg in segments {
848 if seg.stride == 0 {
849 return Err(errors::Error::Protocol {
850 protocol: "fixed_column",
851 message: "Segments stride must be non-zero",
852 });
853 }
854
855 if seg.count == 0 {
856 return Err(errors::Error::Protocol {
857 protocol: "fixed_column",
858 message: "Segments count must be non-zero",
859 });
860 }
861
862 if seg.values.len() != seg.stride {
863 return Err(errors::Error::Protocol {
864 protocol: "fixed_column",
865 message: "Segments values length != stride",
866 });
867 }
868
869 let end = seg
870 .stride
871 .checked_mul(seg.count)
872 .and_then(|span| span.checked_add(seg.origin))
873 .ok_or(errors::Error::Protocol {
874 protocol: "fixed_column",
875 message: "Segments active range overflows",
876 })?;
877
878 if seg.origin < prev_end {
879 return Err(errors::Error::Protocol {
880 protocol: "fixed_column",
881 message: "Segments must be sorted by origin and disjoint",
882 });
883 }
884
885 prev_end = end;
886
887 check_bit_domain(seg.values.iter().copied(), col_type)?;
888 }
889
890 if let Some(nv) = num_vars
891 && prev_end > (1usize << nv)
892 {
893 return Err(errors::Error::Protocol {
894 protocol: "fixed_column",
895 message: "Segments active range exceeds trace height",
896 });
897 }
898
899 Ok(())
900 }
901 }
902}
903
904fn check_bit_domain<F: TowerField>(
905 values: impl Iterator<Item = F>,
906 col_type: ColumnType,
907) -> errors::Result<()> {
908 if col_type != ColumnType::Bit {
909 return Ok(());
910 }
911
912 for v in values {
913 if v != F::ZERO && v != F::ONE {
914 return Err(errors::Error::Protocol {
915 protocol: "fixed_column",
916 message: "Bit fixed column value not in {0,1}",
917 });
918 }
919 }
920
921 Ok(())
922}