openvm_circuit_primitives/xor/lookup/
mod.rs

1//! A chip which uses preprocessed trace to provide a lookup table for XOR operations
2//! between two numbers `x` and `y` of at most `M` bits.
3//! It generates a preprocessed table with a row for each possible triple `(x, y, x^y)`
4//! and keeps count of the number of times each triple is requested.
5
6use std::{
7    borrow::Borrow,
8    mem::size_of,
9    sync::atomic::{self, AtomicU32},
10};
11
12use openvm_circuit_primitives_derive::AlignedBorrow;
13use openvm_cpu_backend::CpuBackend;
14use openvm_stark_backend::{
15    interaction::{BusIndex, InteractionBuilder, LookupBus},
16    p3_air::{Air, BaseAir, PairBuilder},
17    p3_field::Field,
18    p3_matrix::{dense::RowMajorMatrix, Matrix},
19    prover::AirProvingContext,
20    BaseAirWithPublicValues, PartitionedBaseAir, StarkProtocolConfig, Val,
21};
22
23use super::bus::XorBus;
24use crate::{Chip, ColumnsAir, StructReflection, StructReflectionHelper};
25
26#[cfg(test)]
27mod tests;
28
29/// Columns for the main trace of the XOR lookup
30#[repr(C)]
31#[derive(Copy, Clone, Debug, AlignedBorrow, StructReflection)]
32pub struct XorLookupCols<T> {
33    /// Multiplicity counter tracking the number of XOR operations requested for each triple
34    pub mult: T,
35}
36
37/// Columns for the preprocessed table of the XOR lookup
38#[repr(C)]
39#[derive(Copy, Clone, Debug, AlignedBorrow, StructReflection)]
40pub struct XorLookupPreprocessedCols<T> {
41    pub x: T,
42    pub y: T,
43    /// XOR result (x ⊕ y)
44    pub z: T,
45}
46
47pub const NUM_XOR_LOOKUP_COLS: usize = size_of::<XorLookupCols<u8>>();
48pub const NUM_XOR_LOOKUP_PREPROCESSED_COLS: usize = size_of::<XorLookupPreprocessedCols<u8>>();
49
50/// Xor via preprocessed lookup table. Can only be used if inputs have less than approximately
51/// 10-bits.
52#[derive(Clone, Copy, Debug, derive_new::new, ColumnsAir)]
53#[columns_via(XorLookupCols<u8>)]
54pub struct XorLookupAir<const M: usize> {
55    pub bus: XorBus,
56}
57
58impl<F: Field, const M: usize> BaseAirWithPublicValues<F> for XorLookupAir<M> {}
59impl<F: Field, const M: usize> PartitionedBaseAir<F> for XorLookupAir<M> {}
60impl<F: Field, const M: usize> BaseAir<F> for XorLookupAir<M> {
61    fn width(&self) -> usize {
62        NUM_XOR_LOOKUP_COLS
63    }
64
65    /// Generates a preprocessed table with a row for each possible triple (x, y, x^y)
66    fn preprocessed_trace(&self) -> Option<RowMajorMatrix<F>> {
67        let rows: Vec<_> = (0..(1 << M) * (1 << M))
68            .flat_map(|i| {
69                let x = i / (1 << M);
70                let y = i % (1 << M);
71                let z = x ^ y;
72                [x, y, z].map(F::from_u32)
73            })
74            .collect();
75
76        Some(RowMajorMatrix::new(rows, NUM_XOR_LOOKUP_PREPROCESSED_COLS))
77    }
78}
79
80impl<AB, const M: usize> Air<AB> for XorLookupAir<M>
81where
82    AB: InteractionBuilder + PairBuilder,
83{
84    fn eval(&self, builder: &mut AB) {
85        let main = builder.main();
86        let preprocessed = builder.preprocessed();
87
88        let prep_local = preprocessed.row_slice(0).unwrap();
89        let prep_local: &XorLookupPreprocessedCols<AB::Var> = (*prep_local).borrow();
90        let local = main.row_slice(0).expect("window should have two elements");
91        let local: &XorLookupCols<AB::Var> = (*local).borrow();
92
93        self.bus
94            .receive(prep_local.x, prep_local.y, prep_local.z)
95            .eval(builder, local.mult);
96    }
97}
98
99/// This chip gets requests to compute the xor of two numbers x and y of at most M bits.
100/// It generates a preprocessed table with a row for each possible triple (x, y, x^y)
101/// and keeps count of the number of times each triple is requested for the single main trace
102/// column.
103#[derive(Debug)]
104pub struct XorLookupChip<const M: usize> {
105    pub air: XorLookupAir<M>,
106    /// Tracks the count of each (x,y) pair requested
107    pub count: Vec<Vec<AtomicU32>>,
108}
109
110impl<const M: usize> XorLookupChip<M> {
111    pub fn new(bus: BusIndex) -> Self {
112        let mut count = vec![];
113        for _ in 0..(1 << M) {
114            let mut row = vec![];
115            for _ in 0..(1 << M) {
116                row.push(AtomicU32::new(0));
117            }
118            count.push(row);
119        }
120        Self {
121            air: XorLookupAir::new(XorBus(LookupBus::new(bus))),
122            count,
123        }
124    }
125
126    /// The xor bus this chip interacts with
127    pub fn bus(&self) -> XorBus {
128        self.air.bus
129    }
130
131    fn calc_xor(&self, x: u32, y: u32) -> u32 {
132        x ^ y
133    }
134
135    /// Request an XOR operation for inputs x and y
136    /// Increments the count for this (x,y) pair and returns x ⊕ y
137    pub fn request(&self, x: u32, y: u32) -> u32 {
138        let val_atomic = &self.count[x as usize][y as usize];
139        val_atomic.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
140
141        self.calc_xor(x, y)
142    }
143
144    /// Resets all request counters to zero
145    pub fn clear(&self) {
146        for i in 0..(1 << M) {
147            for j in 0..(1 << M) {
148                self.count[i][j].store(0, std::sync::atomic::Ordering::Relaxed);
149            }
150        }
151    }
152
153    /// Generates the multiplicity trace based on requests
154    pub fn generate_trace<F: Field>(&self) -> RowMajorMatrix<F> {
155        debug_assert_eq!(self.count.len(), 1 << M);
156        let multiplicities: Vec<_> = self
157            .count
158            .iter()
159            .flat_map(|count_x| {
160                debug_assert_eq!(count_x.len(), 1 << M);
161                count_x
162                    .iter()
163                    .map(|count_xy| F::from_u32(count_xy.load(atomic::Ordering::SeqCst)))
164            })
165            .collect();
166
167        RowMajorMatrix::new_col(multiplicities)
168    }
169}
170
171impl<R, SC: StarkProtocolConfig, const M: usize> Chip<R, CpuBackend<SC>> for XorLookupChip<M> {
172    fn generate_proving_ctx(&self, _: R) -> AirProvingContext<CpuBackend<SC>> {
173        let trace_row_maj = self.generate_trace::<Val<SC>>();
174        AirProvingContext::simple_no_pis(trace_row_maj)
175    }
176}
openvm_circuit_primitives/xor/lookup/mod.rs

openvm_circuit_primitives/xor/lookup/
mod.rs
