p3_air/
virtual_column.rs

1use alloc::vec;
2use alloc::vec::Vec;
3use core::ops::Mul;
4
5use p3_field::{Field, FieldAlgebra};
6
7/// An affine function over columns in a PAIR.
8#[derive(Clone, Debug)]
9pub struct VirtualPairCol<F: Field> {
10    column_weights: Vec<(PairCol, F)>,
11    constant: F,
12}
13
14/// A column in a PAIR, i.e. either a preprocessed column or a main trace column.
15#[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    /// `a - b`, where `a` and `b` are columns in the preprocessed trace.
99    #[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    /// `a - b`, where `a` and `b` are columns in the main trace.
105    #[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}