openvm_stark_sdk/config/
goldilocks_poseidon.rs

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

openvm_stark_sdk/config/
goldilocks_poseidon.rs
