p3_air/
virtual_column.rs
1use alloc::vec;
2use alloc::vec::Vec;
3use core::ops::Mul;
4
5use p3_field::{Field, FieldAlgebra};
6
7#[derive(Clone, Debug)]
9pub struct VirtualPairCol<F: Field> {
10 column_weights: Vec<(PairCol, F)>,
11 constant: F,
12}
13
14#[derive(Clone, Copy, Debug)]
16pub enum PairCol {
17 Preprocessed(usize),
18 Main(usize),
19}
20
21impl PairCol {
22 pub const fn get<T: Copy>(&self, preprocessed: &[T], main: &[T]) -> T {
23 match self {
24 PairCol::Preprocessed(i) => preprocessed[*i],
25 PairCol::Main(i) => main[*i],
26 }
27 }
28}
29
30impl<F: Field> VirtualPairCol<F> {
31 pub const fn new(column_weights: Vec<(PairCol, F)>, constant: F) -> Self {
32 Self {
33 column_weights,
34 constant,
35 }
36 }
37
38 pub fn new_preprocessed(column_weights: Vec<(usize, F)>, constant: F) -> Self {
39 Self::new(
40 column_weights
41 .into_iter()
42 .map(|(i, w)| (PairCol::Preprocessed(i), w))
43 .collect(),
44 constant,
45 )
46 }
47
48 pub fn new_main(column_weights: Vec<(usize, F)>, constant: F) -> Self {
49 Self::new(
50 column_weights
51 .into_iter()
52 .map(|(i, w)| (PairCol::Main(i), w))
53 .collect(),
54 constant,
55 )
56 }
57
58 pub const ONE: Self = Self::constant(F::ONE);
59
60 #[must_use]
61 pub const fn constant(x: F) -> Self {
62 Self {
63 column_weights: vec![],
64 constant: x,
65 }
66 }
67
68 #[must_use]
69 pub fn single(column: PairCol) -> Self {
70 Self {
71 column_weights: vec![(column, F::ONE)],
72 constant: F::ZERO,
73 }
74 }
75
76 #[must_use]
77 pub fn single_preprocessed(column: usize) -> Self {
78 Self::single(PairCol::Preprocessed(column))
79 }
80
81 #[must_use]
82 pub fn single_main(column: usize) -> Self {
83 Self::single(PairCol::Main(column))
84 }
85
86 #[must_use]
87 pub fn sum_main(columns: Vec<usize>) -> Self {
88 let column_weights = columns.into_iter().map(|col| (col, F::ONE)).collect();
89 Self::new_main(column_weights, F::ZERO)
90 }
91
92 #[must_use]
93 pub fn sum_preprocessed(columns: Vec<usize>) -> Self {
94 let column_weights = columns.into_iter().map(|col| (col, F::ONE)).collect();
95 Self::new_preprocessed(column_weights, F::ZERO)
96 }
97
98 #[must_use]
100 pub fn diff_preprocessed(a_col: usize, b_col: usize) -> Self {
101 Self::new_preprocessed(vec![(a_col, F::ONE), (b_col, F::NEG_ONE)], F::ZERO)
102 }
103
104 #[must_use]
106 pub fn diff_main(a_col: usize, b_col: usize) -> Self {
107 Self::new_main(vec![(a_col, F::ONE), (b_col, F::NEG_ONE)], F::ZERO)
108 }
109
110 pub fn apply<Expr, Var>(&self, preprocessed: &[Var], main: &[Var]) -> Expr
111 where
112 F: Into<Expr>,
113 Expr: FieldAlgebra + Mul<F, Output = Expr>,
114 Var: Into<Expr> + Copy,
115 {
116 let mut result = self.constant.into();
117 for (column, weight) in self.column_weights.iter() {
118 result += column.get(preprocessed, main).into() * *weight;
119 }
120 result
121 }
122}