openvm_stark_sdk/config/
goldilocks_poseidon.rs

1use std::any::type_name;
2
3use openvm_stark_backend::{
4    config::StarkConfig,
5    interaction::fri_log_up::FriLogUpPhase,
6    p3_challenger::DuplexChallenger,
7    p3_commit::ExtensionMmcs,
8    p3_field::{extension::BinomialExtensionField, Field},
9};
10use p3_dft::Radix2DitParallel;
11use p3_fri::{FriConfig, TwoAdicFriPcs};
12use p3_goldilocks::{Goldilocks, MdsMatrixGoldilocks};
13use p3_merkle_tree::MerkleTreeMmcs;
14use p3_poseidon::Poseidon;
15use p3_symmetric::{CryptographicPermutation, PaddingFreeSponge, TruncatedPermutation};
16use rand::{rngs::StdRng, SeedableRng};
17
18use super::instrument::{HashStatistics, Instrumented, StarkHashStatistics};
19use crate::{
20    assert_sc_compatible_with_serde,
21    config::fri_params::SecurityParameters,
22    engine::{StarkEngine, StarkEngineWithHashInstrumentation},
23};
24
25const RATE: usize = 4;
26// permutation width
27const WIDTH: usize = 8; // rate + capacity
28const DIGEST_WIDTH: usize = 4;
29
30type Val = Goldilocks;
31type PackedVal = <Val as Field>::Packing;
32type Challenge = BinomialExtensionField<Val, 2>;
33type Perm = Poseidon<Val, MdsMatrixGoldilocks, WIDTH, 7>;
34type InstrPerm = Instrumented<Perm>;
35
36// Generic over P: CryptographicPermutation<[F; WIDTH]>
37type Hash<P> = PaddingFreeSponge<P, WIDTH, RATE, DIGEST_WIDTH>;
38type Compress<P> = TruncatedPermutation<P, 2, DIGEST_WIDTH, WIDTH>;
39type ValMmcs<P> =
40    MerkleTreeMmcs<PackedVal, <Val as Field>::Packing, Hash<P>, Compress<P>, DIGEST_WIDTH>;
41type ChallengeMmcs<P> = ExtensionMmcs<Val, Challenge, ValMmcs<P>>;
42pub type Challenger<P> = DuplexChallenger<Val, P, WIDTH, RATE>;
43type Dft = Radix2DitParallel<Val>;
44type Pcs<P> = TwoAdicFriPcs<Val, Dft, ValMmcs<P>, ChallengeMmcs<P>>;
45type RapPhase<P> = FriLogUpPhase<Val, Challenge, Challenger<P>>;
46
47pub type GoldilocksPermutationConfig<P> =
48    StarkConfig<Pcs<P>, RapPhase<P>, Challenge, Challenger<P>>;
49pub type GoldilocksPoseidonConfig = GoldilocksPermutationConfig<Perm>;
50pub type GoldilocksPoseidonEngine = GoldilocksPermutationEngine<Perm>;
51
52assert_sc_compatible_with_serde!(GoldilocksPoseidonConfig);
53
54pub struct GoldilocksPermutationEngine<P>
55where
56    P: CryptographicPermutation<[Val; WIDTH]>
57        + CryptographicPermutation<[PackedVal; WIDTH]>
58        + Clone,
59{
60    security_params: SecurityParameters,
61    pub config: GoldilocksPermutationConfig<P>,
62    pub perm: P,
63    pub max_constraint_degree: usize,
64}
65
66impl<P> StarkEngine<GoldilocksPermutationConfig<P>> for GoldilocksPermutationEngine<P>
67where
68    P: CryptographicPermutation<[Val; WIDTH]>
69        + CryptographicPermutation<[PackedVal; WIDTH]>
70        + Clone,
71{
72    fn config(&self) -> &GoldilocksPermutationConfig<P> {
73        &self.config
74    }
75
76    fn max_constraint_degree(&self) -> Option<usize> {
77        Some(self.max_constraint_degree)
78    }
79
80    fn new_challenger(&self) -> Challenger<P> {
81        Challenger::new(self.perm.clone())
82    }
83}
84
85impl<P> StarkEngineWithHashInstrumentation<GoldilocksPermutationConfig<Instrumented<P>>>
86    for GoldilocksPermutationEngine<Instrumented<P>>
87where
88    P: CryptographicPermutation<[Val; WIDTH]>
89        + CryptographicPermutation<[PackedVal; WIDTH]>
90        + Clone,
91{
92    fn clear_instruments(&mut self) {
93        self.perm.input_lens_by_type.lock().unwrap().clear();
94    }
95    fn stark_hash_statistics<T>(&self, custom: T) -> StarkHashStatistics<T> {
96        let counter = self.perm.input_lens_by_type.lock().unwrap();
97        let permutations = counter.iter().fold(0, |total, (name, lens)| {
98            if name == type_name::<[Val; WIDTH]>() {
99                let count: usize = lens.iter().sum();
100                println!("Permutation: {name}, Count: {count}");
101                total + count
102            } else {
103                panic!("Permutation type not yet supported: {}", name);
104            }
105        });
106
107        StarkHashStatistics {
108            name: type_name::<P>().to_string(),
109            stats: HashStatistics { permutations },
110            fri_params: self.security_params.fri_params,
111            custom,
112        }
113    }
114}
115
116/// `pcs_log_degree` is the upper bound on the log_2(PCS polynomial degree).
117pub fn default_engine() -> GoldilocksPoseidonEngine {
118    let perm = random_perm();
119    engine_from_perm(perm, SecurityParameters::standard_fast())
120}
121
122/// `pcs_log_degree` is the upper bound on the log_2(PCS polynomial degree).
123pub fn default_config(perm: &Perm) -> GoldilocksPoseidonConfig {
124    config_from_perm(perm, SecurityParameters::standard_fast())
125}
126
127pub fn engine_from_perm<P>(
128    perm: P,
129    security_params: SecurityParameters,
130) -> GoldilocksPermutationEngine<P>
131where
132    P: CryptographicPermutation<[Val; WIDTH]>
133        + CryptographicPermutation<[PackedVal; WIDTH]>
134        + Clone,
135{
136    let max_constraint_degree = security_params.fri_params.max_constraint_degree();
137    let config = config_from_perm(&perm, security_params.clone());
138    GoldilocksPermutationEngine {
139        config,
140        perm,
141        security_params,
142        max_constraint_degree,
143    }
144}
145
146pub fn config_from_perm<P>(
147    perm: &P,
148    security_params: SecurityParameters,
149) -> GoldilocksPermutationConfig<P>
150where
151    P: CryptographicPermutation<[Val; WIDTH]>
152        + CryptographicPermutation<[PackedVal; WIDTH]>
153        + Clone,
154{
155    let hash = Hash::new(perm.clone());
156    let compress = Compress::new(perm.clone());
157    let val_mmcs = ValMmcs::new(hash, compress);
158    let challenge_mmcs = ChallengeMmcs::new(val_mmcs.clone());
159    let dft = Dft::default();
160    let SecurityParameters {
161        fri_params,
162        log_up_params,
163    } = security_params;
164    let fri_config = FriConfig {
165        log_blowup: fri_params.log_blowup,
166        log_final_poly_len: fri_params.log_final_poly_len,
167        num_queries: fri_params.num_queries,
168        proof_of_work_bits: fri_params.proof_of_work_bits,
169        mmcs: challenge_mmcs,
170    };
171    let pcs = Pcs::new(dft, val_mmcs, fri_config);
172    let rap_phase = RapPhase::new(log_up_params);
173    GoldilocksPermutationConfig::new(pcs, rap_phase)
174}
175
176pub fn random_perm() -> Perm {
177    let seed = [42; 32];
178    let mut rng = StdRng::from_seed(seed);
179    Perm::new_from_rng(4, 22, MdsMatrixGoldilocks, &mut rng)
180}
181
182pub fn random_instrumented_perm() -> InstrPerm {
183    let perm = random_perm();
184    Instrumented::new(perm)
185}