1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

use std::collections::HashMap;

use folding::expressions::FoldingColumnTrait;
use kimchi::circuits::expr::{CacheId, FormattedOutput};

/// Describe a generic indexed variable X_{i}.
#[derive(PartialEq, Eq, Clone, Copy, Debug, Hash)]
pub enum Column {
    /// Columns related to the relation encoded in the circuit
    Relation(usize),
    /// Columns related to dynamic selectors to indicate gate type
    DynamicSelector(usize),
    /// Constant column that is /always/ fixed for a given circuit.
    FixedSelector(usize),
    // Columns related to the lookup protocol
    /// Partial sums. This corresponds to the `h_i`.
    /// It is first indexed by the table ID, and after that internal index.
    LookupPartialSum((u32, usize)),
    /// Multiplicities, indexed. This corresponds to the `m_i`. First
    /// indexed by table ID, then internal index.
    LookupMultiplicity((u32, usize)),
    /// The lookup aggregation, i.e. `phi`
    LookupAggregation,
    /// The fixed tables. The parameter is considered to the indexed table.
    LookupFixedTable(u32),
}

impl Column {
    /// Adds offset if the column is `Relation`. Fails otherwise.
    pub fn add_rel_offset(self, offset: usize) -> Column {
        let Column::Relation(i) = self else {
            todo!("add_rel_offset is only implemented for the relation columns")
        };
        Column::Relation(offset + i)
    }
}

impl FormattedOutput for Column {
    fn latex(&self, _cache: &mut HashMap<CacheId, Self>) -> String {
        match self {
            Column::Relation(i) => format!("x_{{{i}}}"),
            Column::FixedSelector(i) => format!("fs_{{{i}}}"),
            Column::DynamicSelector(i) => format!("ds_{{{i}}}"),
            Column::LookupPartialSum((table_id, i)) => format!("h_{{{table_id}, {i}}}"),
            Column::LookupMultiplicity((table_id, i)) => format!("m_{{{table_id}, {i}}}"),
            Column::LookupFixedTable(i) => format!("t_{{{i}}}"),
            Column::LookupAggregation => String::from("φ"),
        }
    }

    fn text(&self, _cache: &mut HashMap<CacheId, Self>) -> String {
        match self {
            Column::Relation(i) => format!("x[{i}]"),
            Column::FixedSelector(i) => format!("fs[{i}]"),
            Column::DynamicSelector(i) => format!("ds[{i}]"),
            Column::LookupPartialSum((table_id, i)) => format!("h[{table_id}, {i}]"),
            Column::LookupMultiplicity((table_id, i)) => format!("m[{table_id}, {i}]"),
            Column::LookupFixedTable(i) => format!("t[{i}]"),
            Column::LookupAggregation => String::from("φ"),
        }
    }

    fn ocaml(&self, _cache: &mut HashMap<CacheId, Self>) -> String {
        // FIXME
        unimplemented!("Not used at the moment")
    }

    fn is_alpha(&self) -> bool {
        // FIXME
        unimplemented!("Not used at the moment")
    }
}

/// A datatype expressing a generalized column, but with potentially
/// more convenient interface than a bare column.
pub trait ColumnIndexer: core::fmt::Debug + Copy + Eq + Ord {
    /// Total number of columns in this index.
    const N_COL: usize;

    /// Flatten the column "alias" into the integer-like column.
    fn to_column(self) -> Column;
}

// Implementation to be compatible with folding if we use generic column
// constraints
impl FoldingColumnTrait for Column {
    fn is_witness(&self) -> bool {
        match self {
            // Witness
            Column::Relation(_) => true,
            Column::DynamicSelector(_) => true,
            // FIXME: check if we want to treat lookups differently
            Column::LookupPartialSum(_) => true,
            Column::LookupMultiplicity(_) => true,
            Column::LookupAggregation => true,
            // Not witness/public values
            Column::FixedSelector(_) => false,
            Column::LookupFixedTable(_) => false,
        }
    }
}