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

pickles_o1vm/
main.rs
