pickles_o1vm/
main.rs

1use ark_ff::{UniformRand, Zero};
2use clap::Parser;
3use kimchi::circuits::domains::EvaluationDomains;
4use log::debug;
5use mina_curves::pasta::{Fp, Vesta, VestaParameters};
6use mina_poseidon::{
7    constants::PlonkSpongeConstantsKimchi,
8    sponge::{DefaultFqSponge, DefaultFrSponge},
9};
10use o1vm::{
11    cannon::{self, Start, State},
12    cli, elf_loader,
13    interpreters::mips::{
14        column::N_MIPS_REL_COLS,
15        constraints as mips_constraints,
16        witness::{self as mips_witness},
17        Instruction,
18    },
19    pickles::{proof::ProofInputs, prover, verifier},
20    preimage_oracle::{NullPreImageOracle, PreImageOracle, PreImageOracleT},
21    test_preimage_read, E,
22};
23use poly_commitment::{ipa::SRS, precomputed_srs::TestSRS, SRS as _};
24use rand::rngs::ThreadRng;
25use std::{fs::File, io::BufReader, path::Path, process::ExitCode, time::Instant};
26
27pub fn cannon_main(args: cli::cannon::RunArgs) {
28    let mut rng = rand::thread_rng();
29
30    let configuration: cannon::VmConfiguration = args.vm_cfg.into();
31
32    let file =
33        File::open(&configuration.input_state_file).expect("Error opening input state file ");
34
35    let reader = BufReader::new(file);
36    // Read the JSON contents of the file as an instance of `State`.
37    let state: State = serde_json::from_reader(reader).expect("Error reading input state file");
38
39    let meta = &configuration.metadata_file.as_ref().map(|f| {
40        let meta_file =
41            File::open(f).unwrap_or_else(|_| panic!("Could not open metadata file {}", f));
42        serde_json::from_reader(BufReader::new(meta_file))
43            .unwrap_or_else(|_| panic!("Error deserializing metadata file {}", f))
44    });
45
46    // Initialize some data used for statistical computations
47    let start = Start::create(state.step as usize);
48
49    let (srs, domain_fp) = match &args.srs_cache {
50        Some(cache) => {
51            debug!("Loading SRS from cache {}", cache);
52            let file_path = Path::new(cache);
53            let file = File::open(file_path).expect("Error opening SRS cache file");
54            let srs: SRS<Vesta> = {
55                // By convention, proof systems serializes a TestSRS with filename 'test_<CURVE_NAME>.srs'.
56                // The benefit of using this is you don't waste time verifying the SRS.
57                if file_path
58                    .file_name()
59                    .unwrap()
60                    .to_str()
61                    .unwrap()
62                    .starts_with("test_")
63                {
64                    let test_srs: TestSRS<Vesta> = rmp_serde::from_read(&file).unwrap();
65                    From::from(test_srs)
66                } else {
67                    rmp_serde::from_read(&file).unwrap()
68                }
69            };
70            debug!("SRS loaded successfully from cache");
71            let domain_fp = EvaluationDomains::<Fp>::create(srs.size()).unwrap();
72            (srs, domain_fp)
73        }
74        None => {
75            debug!("No SRS cache provided. Creating SRS from scratch with domain size 2^16");
76            let domain_size = 1 << 16;
77            let srs = SRS::create(domain_size);
78            let domain_fp = EvaluationDomains::<Fp>::create(srs.size()).unwrap();
79            srs.get_lagrange_basis(domain_fp.d1);
80            debug!("SRS created successfully");
81            (srs, domain_fp)
82        }
83    };
84
85    // Initialize the environments
86    let mut mips_wit_env = match configuration.host.clone() {
87        Some(host) => {
88            let mut po = PreImageOracle::create(host);
89            let _child = po.start();
90            mips_witness::Env::<Fp, Box<dyn PreImageOracleT>>::create(
91                cannon::PAGE_SIZE as usize,
92                state,
93                Box::new(po),
94            )
95        }
96        None => {
97            debug!("No preimage oracle provided 🤞");
98            // warning: the null preimage oracle has no data and will crash the program if used
99            mips_witness::Env::<Fp, Box<dyn PreImageOracleT>>::create(
100                cannon::PAGE_SIZE as usize,
101                state,
102                Box::new(NullPreImageOracle),
103            )
104        }
105    };
106
107    let constraints = mips_constraints::get_all_constraints::<Fp>();
108    let domain_size = domain_fp.d1.size as usize;
109
110    let mut curr_proof_inputs: ProofInputs<Vesta> = ProofInputs::new(domain_size);
111    while !mips_wit_env.halt {
112        let _instr: Instruction = mips_wit_env.step(&configuration, meta, &start);
113        for (scratch, scratch_chunk) in mips_wit_env
114            .scratch_state
115            .iter()
116            .zip(curr_proof_inputs.evaluations.scratch.iter_mut())
117        {
118            scratch_chunk.push(*scratch);
119        }
120        for (scratch, scratch_chunk) in mips_wit_env
121            .scratch_state_inverse
122            .iter()
123            .zip(curr_proof_inputs.evaluations.scratch_inverse.iter_mut())
124        {
125            scratch_chunk.push(*scratch);
126        }
127        // Lookup state
128        {
129            let proof_inputs_length = curr_proof_inputs.evaluations.lookup_state.len();
130            let environment_length = mips_wit_env.lookup_state.len();
131            let lookup_state_size = std::cmp::max(proof_inputs_length, environment_length);
132            for idx in 0..lookup_state_size {
133                if idx >= environment_length {
134                    // We pad with 0s for dummy lookups missing from the environment.
135                    curr_proof_inputs.evaluations.lookup_state[idx].push(Fp::zero());
136                } else if idx >= proof_inputs_length {
137                    // We create a new column filled with 0s in the proof inputs.
138                    let mut new_vec =
139                        vec![Fp::zero(); curr_proof_inputs.evaluations.instruction_counter.len()];
140                    new_vec.push(Fp::from(mips_wit_env.lookup_state[idx]));
141                    curr_proof_inputs.evaluations.lookup_state.push(new_vec);
142                } else {
143                    // Push the value to the column.
144                    curr_proof_inputs.evaluations.lookup_state[idx]
145                        .push(Fp::from(mips_wit_env.lookup_state[idx]));
146                }
147            }
148        }
149        curr_proof_inputs
150            .evaluations
151            .instruction_counter
152            .push(Fp::from(mips_wit_env.instruction_counter));
153        // FIXME: Might be another value
154        curr_proof_inputs.evaluations.error.push(Fp::rand(&mut rng));
155
156        curr_proof_inputs
157            .evaluations
158            .selector
159            .push(Fp::from((mips_wit_env.selector - N_MIPS_REL_COLS) as u64));
160
161        if curr_proof_inputs.evaluations.instruction_counter.len() == domain_size {
162            prove_and_verify(domain_fp, &srs, &constraints, curr_proof_inputs, &mut rng);
163
164            curr_proof_inputs = ProofInputs::new(domain_size);
165        }
166    }
167
168    if curr_proof_inputs.evaluations.instruction_counter.len() < domain_size {
169        debug!("Padding witness for proof generation");
170        pad(&mips_wit_env, &mut curr_proof_inputs, &mut rng);
171        prove_and_verify(domain_fp, &srs, &constraints, curr_proof_inputs, &mut rng);
172    }
173}
174
175fn prove_and_verify(
176    domain_fp: EvaluationDomains<Fp>,
177    srs: &SRS<Vesta>,
178    constraints: &[E<Fp>],
179    curr_proof_inputs: ProofInputs<Vesta>,
180    rng: &mut ThreadRng,
181) {
182    let start_iteration = Instant::now();
183    let proof = prover::prove::<
184        Vesta,
185        DefaultFqSponge<VestaParameters, PlonkSpongeConstantsKimchi>,
186        DefaultFrSponge<Fp, PlonkSpongeConstantsKimchi>,
187        _,
188    >(domain_fp, srs, curr_proof_inputs, constraints, rng)
189    .unwrap();
190    debug!(
191        "Proof generated in {elapsed} μs",
192        elapsed = start_iteration.elapsed().as_micros()
193    );
194    let start_iteration = Instant::now();
195    let verif = verifier::verify::<
196        Vesta,
197        DefaultFqSponge<VestaParameters, PlonkSpongeConstantsKimchi>,
198        DefaultFrSponge<Fp, PlonkSpongeConstantsKimchi>,
199    >(domain_fp, srs, constraints, &proof);
200    debug!(
201        "Verification done in {elapsed} μs",
202        elapsed = start_iteration.elapsed().as_micros()
203    );
204    assert!(verif);
205}
206
207fn pad(
208    witness_env: &mips_witness::Env<Fp, Box<dyn PreImageOracleT>>,
209    curr_proof_inputs: &mut ProofInputs<Vesta>,
210    rng: &mut ThreadRng,
211) {
212    let zero = Fp::zero();
213    // FIXME: Find a better way to get instruction selectors that doesn't
214    // reveal internals.
215    let noop_selector: Fp = {
216        let noop: usize = Instruction::NoOp.into();
217        Fp::from((noop - N_MIPS_REL_COLS) as u64)
218    };
219    curr_proof_inputs
220        .evaluations
221        .scratch
222        .iter_mut()
223        .for_each(|x| x.resize(x.capacity(), zero));
224    curr_proof_inputs
225        .evaluations
226        .scratch_inverse
227        .iter_mut()
228        .for_each(|x| x.resize(x.capacity(), zero));
229    curr_proof_inputs.evaluations.instruction_counter.resize(
230        curr_proof_inputs.evaluations.instruction_counter.capacity(),
231        Fp::from(witness_env.instruction_counter),
232    );
233    curr_proof_inputs
234        .evaluations
235        .error
236        .resize_with(curr_proof_inputs.evaluations.error.capacity(), || {
237            Fp::rand(rng)
238        });
239    curr_proof_inputs.evaluations.selector.resize(
240        curr_proof_inputs.evaluations.selector.capacity(),
241        noop_selector,
242    );
243}
244
245fn gen_state_json(arg: cli::cannon::GenStateJsonArgs) -> Result<(), String> {
246    let path = Path::new(&arg.input);
247    let state = elf_loader::parse_elf(elf_loader::Architecture::Mips, path)?;
248    let file = File::create(&arg.output).expect("Error creating output state file");
249    serde_json::to_writer_pretty(file, &state).expect("Error writing output state file");
250    Ok(())
251}
252
253pub fn main() -> ExitCode {
254    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info")).init();
255    let args = cli::Commands::parse();
256    match args {
257        cli::Commands::Cannon(args) => match args {
258            cli::cannon::Cannon::Run(args) => {
259                cannon_main(args);
260            }
261            cli::cannon::Cannon::TestPreimageRead(args) => {
262                test_preimage_read::main(args);
263            }
264            cli::cannon::Cannon::GenStateJson(args) => {
265                gen_state_json(args).expect("Error generating state.json");
266            }
267        },
268    }
269    ExitCode::SUCCESS
270}
pickles_o1vm/main.rs

pickles_o1vm/
main.rs
