o1vm/interpreters/keccak/
witness.rs

1//! This file contains the witness for the Keccak hash function for the zkVM project.
2//! It assigns the witness values to the corresponding columns of KeccakWitness in the environment.
3//!
4//! The actual witness generation code makes use of the code which is already present in Kimchi,
5//! to avoid code duplication and reduce error-proneness.
6//!
7//! For a pseudo code implementation of Keccap-f, see
8//! <https://keccak.team/keccak_specs_summary.html>
9use std::collections::HashMap;
10
11use crate::{
12    interpreters::keccak::{
13        column::KeccakWitness,
14        helpers::{ArithHelpers, BoolHelpers, LogupHelpers},
15        interpreter::{Interpreter, KeccakInterpreter},
16        Constraint, Error, KeccakColumn,
17    },
18    lookups::{
19        FixedLookupTables, Lookup, LookupTable,
20        LookupTableIDs::{self, *},
21    },
22};
23
24use ark_ff::Field;
25use kimchi::o1_utils::Two;
26use kimchi_msm::LookupTableID;
27
28/// This struct contains all that needs to be kept track of during the execution of the Keccak step interpreter
29// TODO: the fixed tables information should be inferred from the general environment
30#[derive(Clone, Debug)]
31pub struct Env<F> {
32    /// The full state of the Keccak gate (witness)
33    pub witness: KeccakWitness<F>,
34    /// The fixed tables used in the Keccak gate
35    pub tables: HashMap<LookupTableIDs, LookupTable<F>>,
36    /// The multiplicities of each lookup entry. Should not be cleared between steps.
37    pub multiplicities: HashMap<LookupTableIDs, Vec<u32>>,
38    /// If any, an error that occurred during the execution of the constraints, to help with debugging
39    pub(crate) errors: Vec<Error>,
40}
41
42impl<F: Field> Default for Env<F> {
43    fn default() -> Self {
44        Self {
45            witness: KeccakWitness::default(),
46            tables: {
47                let mut t = HashMap::new();
48                t.insert(PadLookup, LookupTable::table_pad());
49                t.insert(RoundConstantsLookup, LookupTable::table_round_constants());
50                t.insert(AtMost4Lookup, LookupTable::table_at_most_4());
51                t.insert(ByteLookup, LookupTable::table_byte());
52                t.insert(RangeCheck16Lookup, LookupTable::table_range_check_16());
53                t.insert(SparseLookup, LookupTable::table_sparse());
54                t.insert(ResetLookup, LookupTable::table_reset());
55                t
56            },
57            multiplicities: {
58                let mut m = HashMap::new();
59                m.insert(PadLookup, vec![0; PadLookup.length()]);
60                m.insert(RoundConstantsLookup, vec![0; RoundConstantsLookup.length()]);
61                m.insert(AtMost4Lookup, vec![0; AtMost4Lookup.length()]);
62                m.insert(ByteLookup, vec![0; ByteLookup.length()]);
63                m.insert(RangeCheck16Lookup, vec![0; RangeCheck16Lookup.length()]);
64                m.insert(SparseLookup, vec![0; SparseLookup.length()]);
65                m.insert(ResetLookup, vec![0; ResetLookup.length()]);
66                m
67            },
68            errors: vec![],
69        }
70    }
71}
72
73impl<F: Field> ArithHelpers<F> for Env<F> {
74    fn two_pow(x: u64) -> <Env<F> as Interpreter<F>>::Variable {
75        Self::constant_field(F::two_pow(x))
76    }
77}
78
79impl<F: Field> BoolHelpers<F> for Env<F> {}
80
81impl<F: Field> LogupHelpers<F> for Env<F> {}
82
83impl<F: Field> Interpreter<F> for Env<F> {
84    type Variable = F;
85
86    fn constant(x: u64) -> Self::Variable {
87        Self::constant_field(F::from(x))
88    }
89
90    fn constant_field(x: F) -> Self::Variable {
91        x
92    }
93
94    fn variable(&self, column: KeccakColumn) -> Self::Variable {
95        self.witness[column]
96    }
97
98    /// Checks the constraint `tag` by checking that the input `x` is zero
99    fn constrain(&mut self, tag: Constraint, if_true: Self::Variable, x: Self::Variable) {
100        if if_true == Self::Variable::one() && x != F::zero() {
101            self.errors.push(Error::Constraint(tag));
102        }
103    }
104
105    fn add_lookup(&mut self, if_true: Self::Variable, lookup: Lookup<Self::Variable>) {
106        // Keep track of multiplicities for fixed lookups
107        if if_true == Self::Variable::one() && lookup.table_id.is_fixed() {
108            // Only when reading. We ignore the other values.
109            if lookup.magnitude == Self::one() {
110                // Check that the lookup value is in the table
111                if let Some(idx) = LookupTable::is_in_table(
112                    self.tables.get_mut(&lookup.table_id).unwrap(),
113                    lookup.value,
114                ) {
115                    self.multiplicities.get_mut(&lookup.table_id).unwrap()[idx] += 1;
116                } else {
117                    self.errors.push(Error::Lookup(lookup.table_id));
118                }
119            }
120        }
121    }
122}
123
124impl<F: Field> KeccakInterpreter<F> for Env<F> {}