openvm_recursion_circuit/primitives/pow/
air.rs

1use core::borrow::Borrow;
2
3use openvm_circuit_primitives::{ColumnsAir, StructReflection, StructReflectionHelper};
4use openvm_recursion_circuit_derive::AlignedBorrow;
5use openvm_stark_backend::{
6    interaction::InteractionBuilder, p3_util::log2_strict_usize, BaseAirWithPublicValues,
7    PartitionedBaseAir,
8};
9use p3_air::{Air, AirBuilder, BaseAir};
10use p3_field::PrimeCharacteristicRing;
11use p3_matrix::Matrix;
12
13use crate::primitives::bus::{
14    PowerCheckerBus, PowerCheckerBusMessage, RangeCheckerBus, RangeCheckerBusMessage,
15};
16
17#[repr(C)]
18#[derive(AlignedBorrow, Debug, StructReflection)]
19pub struct PowerCheckerCols<T> {
20    pub log: T,
21    pub pow: T,
22    pub mult_pow: T,
23    pub mult_range: T,
24}
25
26#[derive(Debug, ColumnsAir)]
27#[columns_via(PowerCheckerCols<u8>)]
28pub struct PowerCheckerAir<const BASE: usize, const N: usize> {
29    pub pow_bus: PowerCheckerBus,
30    pub range_bus: RangeCheckerBus,
31}
32
33impl<F, const BASE: usize, const N: usize> BaseAir<F> for PowerCheckerAir<BASE, N> {
34    fn width(&self) -> usize {
35        PowerCheckerCols::<F>::width()
36    }
37}
38impl<F, const B: usize, const N: usize> BaseAirWithPublicValues<F> for PowerCheckerAir<B, N> {}
39impl<F, const B: usize, const N: usize> PartitionedBaseAir<F> for PowerCheckerAir<B, N> {}
40
41impl<AB: AirBuilder + InteractionBuilder, const BASE: usize, const N: usize> Air<AB>
42    for PowerCheckerAir<BASE, N>
43{
44    fn eval(&self, builder: &mut AB) {
45        let main = builder.main();
46
47        let (local, next) = (
48            main.row_slice(0).expect("window should have two elements"),
49            main.row_slice(1).expect("window should have two elements"),
50        );
51        let local: &PowerCheckerCols<AB::Var> = (*local).borrow();
52        let next: &PowerCheckerCols<AB::Var> = (*next).borrow();
53
54        builder.when_first_row().assert_zero(local.log);
55        builder.when_first_row().assert_eq(local.pow, AB::F::ONE);
56
57        builder
58            .when_transition()
59            .assert_eq(local.log + AB::F::ONE, next.log);
60        builder
61            .when_transition()
62            .assert_eq(local.pow * AB::F::from_usize(BASE), next.pow);
63
64        builder
65            .when_last_row()
66            .assert_eq(local.log, AB::F::from_usize(N - 1));
67
68        self.pow_bus.add_key_with_lookups(
69            builder,
70            PowerCheckerBusMessage {
71                log: local.log,
72                exp: local.pow,
73            },
74            local.mult_pow,
75        );
76        self.range_bus.add_key_with_lookups(
77            builder,
78            RangeCheckerBusMessage {
79                value: local.log.into(),
80                max_bits: AB::Expr::from_usize(log2_strict_usize(N)),
81            },
82            local.mult_range,
83        );
84    }
85}