tasm_lib/
lib.rs

1// Recursion limit for macro_rules expansions, used for
2// triton_asm!
3#![recursion_limit = "2048"]
4
5// This is needed for `#[derive(TasmObject)]` macro to work consistently across crates.
6// Specifically:
7// From inside the `tasm-lib` crate, we need to refer to `tasm-lib` by `crate`.
8// However, from outside the `tasm-lib` crate, we need to refer to it by `tasm_lib`.
9// The re-export below allows using identifier `tasm_lib` even from inside `tasm-lib`.
10//
11// See also:
12// https://github.com/bkchr/proc-macro-crate/issues/2#issuecomment-572914520
13extern crate self as tasm_lib;
14
15use std::collections::HashMap;
16use std::io::Write;
17use std::time::SystemTime;
18
19use itertools::Itertools;
20use memory::dyn_malloc;
21use num_traits::Zero;
22use triton_vm::isa::op_stack::NUM_OP_STACK_REGISTERS;
23use triton_vm::prelude::*;
24
25pub mod arithmetic;
26pub mod array;
27pub mod data_type;
28pub mod exported_snippets;
29pub mod hashing;
30pub mod io;
31pub mod library;
32pub mod linker;
33pub mod list;
34pub mod memory;
35pub mod mmr;
36pub mod neptune;
37pub mod prelude;
38pub mod rust_shadowing_helper_functions;
39pub mod snippet_bencher;
40pub mod structure;
41pub mod test_helpers;
42pub mod traits;
43pub mod verifier;
44
45// re-exports for types exposed in our public API.
46pub use triton_vm;
47use triton_vm::isa::instruction::AnInstruction;
48use triton_vm::prelude::TableId;
49pub use triton_vm::twenty_first;
50
51use crate::test_helpers::prepend_program_with_stack_setup;
52
53pub(crate) const U32_TO_USIZE_ERR: &str =
54    "internal error: type `usize` should have at least 32 bits";
55pub(crate) const USIZE_TO_U64_ERR: &str =
56    "internal error: type `usize` should have at most 64 bits";
57
58#[derive(Clone, Debug, Default)]
59pub struct InitVmState {
60    pub stack: Vec<BFieldElement>,
61    pub public_input: Vec<BFieldElement>,
62    pub nondeterminism: NonDeterminism,
63    pub sponge: Option<Tip5>,
64}
65
66impl InitVmState {
67    pub fn with_stack(stack: Vec<BFieldElement>) -> Self {
68        InitVmState {
69            stack,
70            public_input: vec![],
71            nondeterminism: NonDeterminism::default(),
72            sponge: None,
73        }
74    }
75
76    pub fn with_stack_and_memory(
77        stack: Vec<BFieldElement>,
78        memory: HashMap<BFieldElement, BFieldElement>,
79    ) -> Self {
80        InitVmState {
81            stack,
82            public_input: vec![],
83            nondeterminism: NonDeterminism::default().with_ram(memory),
84            sponge: None,
85        }
86    }
87}
88
89#[derive(Clone, Debug)]
90pub struct RustShadowOutputState {
91    pub public_output: Vec<BFieldElement>,
92    pub stack: Vec<BFieldElement>,
93    pub ram: HashMap<BFieldElement, BFieldElement>,
94    pub sponge: Option<Tip5>,
95}
96
97pub fn empty_stack() -> Vec<BFieldElement> {
98    vec![BFieldElement::zero(); NUM_OP_STACK_REGISTERS]
99}
100
101pub fn push_encodable<T: BFieldCodec>(stack: &mut Vec<BFieldElement>, value: &T) {
102    stack.extend(value.encode().into_iter().rev());
103}
104
105/// Pops an element of the specified, generic type from the stack.
106///
107/// ### Panics
108///
109/// Panics if
110/// - the generic type has [dynamic length](BFieldCodec::static_length)
111/// - the stack does not contain enough elements
112/// - the top of the stack does not correspond to a [`BFieldCodec`] encoded
113///   element of type `T`
114pub(crate) fn pop_encodable<T: BFieldCodec>(stack: &mut Vec<BFieldElement>) -> T {
115    let len = T::static_length().unwrap();
116    let limbs = (0..len).map(|_| stack.pop().unwrap()).collect_vec();
117    *T::decode(&limbs).unwrap()
118}
119
120/// Execute a Triton-VM program and test correct behavior indicators.
121/// Modify stack and memory. Panic if anything goes wrong.
122pub(crate) fn execute_test(
123    code: &[LabelledInstruction],
124    stack: &mut Vec<BFieldElement>,
125    expected_stack_diff: isize,
126    std_in: Vec<BFieldElement>,
127    nondeterminism: NonDeterminism,
128    maybe_sponge: Option<Tip5>,
129) -> VMState {
130    let init_stack = stack.to_owned();
131    let public_input = PublicInput::new(std_in.clone());
132    let program = Program::new(code);
133
134    let mut vm_state = VMState::new(
135        program.clone(),
136        public_input.clone(),
137        nondeterminism.clone(),
138    );
139    vm_state.op_stack.stack.clone_from(&init_stack);
140    vm_state.sponge = maybe_sponge;
141
142    maybe_write_debuggable_vm_state_to_disk(&vm_state);
143
144    if let Err(err) = vm_state.run() {
145        panic!("{err}\n\nFinal state was: {vm_state}")
146    }
147    let terminal_state = vm_state;
148
149    if !terminal_state.jump_stack.is_empty() {
150        panic!("Jump stack must be unchanged after code execution");
151    }
152
153    let final_stack_height = terminal_state.op_stack.stack.len() as isize;
154    let initial_stack_height = init_stack.len();
155    assert_eq!(
156        expected_stack_diff,
157        final_stack_height - initial_stack_height as isize,
158        "Code must grow/shrink stack with expected number of elements.\n
159        init height: {initial_stack_height}\n
160        end height:  {final_stack_height}\n
161        expected difference: {expected_stack_diff}\n\n
162        initial stack: {}\n
163        final stack:   {}",
164        init_stack.iter().skip(NUM_OP_STACK_REGISTERS).join(","),
165        terminal_state
166            .op_stack
167            .stack
168            .iter()
169            .skip(NUM_OP_STACK_REGISTERS)
170            .join(","),
171    );
172
173    // If this environment variable is set, all programs, including the code to prepare the state,
174    // will be proven and then verified.
175    // Notice that this is only done after the successful execution of the program above, so all
176    // produced proofs here should be valid.
177    // If you run this, make sure `opt-level` is set to 3.
178    if std::env::var("DYING_TO_PROVE").is_ok() {
179        prove_and_verify(program, &std_in, &nondeterminism, Some(init_stack));
180    }
181
182    stack.clone_from(&terminal_state.op_stack.stack);
183    terminal_state
184}
185
186/// If the environment variable TASMLIB_TRITON_TUI is set, write the initial VM state
187/// to file `vm_state.json`.
188///
189/// This file can be used to debug the program using the [Triton TUI]:
190/// ```sh
191/// triton-tui --initial-state vm_state.json
192/// ```
193///
194/// [Triton TUI]: https://crates.io/crates/triton-tui
195pub fn maybe_write_debuggable_vm_state_to_disk(vm_state: &VMState) {
196    let Ok(_) = std::env::var("TASMLIB_TRITON_TUI") else {
197        return;
198    };
199
200    let mut state_file = std::fs::File::create("vm_state.json").unwrap();
201    let state = serde_json::to_string(&vm_state).unwrap();
202    write!(state_file, "{state}").unwrap();
203}
204
205/// Prepare state and run Triton VM
206pub(crate) fn execute_with_terminal_state(
207    program: Program,
208    std_in: &[BFieldElement],
209    stack: &[BFieldElement],
210    nondeterminism: &NonDeterminism,
211    maybe_sponge: Option<Tip5>,
212) -> Result<VMState, InstructionError> {
213    let public_input = PublicInput::new(std_in.into());
214    let mut vm_state = VMState::new(program, public_input, nondeterminism.to_owned());
215    stack.clone_into(&mut vm_state.op_stack.stack);
216    vm_state.sponge = maybe_sponge;
217
218    maybe_write_debuggable_vm_state_to_disk(&vm_state);
219    match vm_state.run() {
220        Ok(()) => {
221            println!("Triton VM execution successful.");
222            Ok(vm_state)
223        }
224        Err(err) => {
225            if let Some(ref sponge) = vm_state.sponge {
226                println!("tasm final sponge:");
227                println!("{}", sponge.state.iter().join(", "));
228            }
229            println!("Triton VM execution failed. Final state:\n{vm_state}");
230            Err(err)
231        }
232    }
233}
234
235/// Run prover on the program, with stack-initialization converted to code.
236///
237/// Run the prover on the program. If `init_stack` is provided, the prover is run on a program
238/// with the code to setup the stack prepended, since the prover will always fail if the stack
239/// is not initialized to the minimal height. The first `NUM_OP_STACK_REGISTERS` of `init_stack`
240/// are ignored.
241/// If you run this, make sure `opt-level` is set to 3.
242pub fn prove_and_verify(
243    program: Program,
244    std_in: &[BFieldElement],
245    nondeterminism: &NonDeterminism,
246    init_stack: Option<Vec<BFieldElement>>,
247) {
248    let labelled_instructions = program.labelled_instructions();
249    let timing_report_label = match labelled_instructions.first().unwrap() {
250        LabelledInstruction::Instruction(AnInstruction::Call(func)) => func,
251        LabelledInstruction::Label(label) => label,
252        _ => "Some program",
253    };
254
255    // Construct the program that initializes the stack to the expected start value.
256    // If this is not done, a stack underflow will occur for most programs
257    let program = match &init_stack {
258        Some(init_stack) => prepend_program_with_stack_setup(init_stack, &program),
259        None => program,
260    };
261
262    let claim = Claim::about_program(&program).with_input(std_in.to_owned());
263    let (aet, public_output) = VM::trace_execution(
264        program.clone(),
265        PublicInput::new(std_in.to_owned()),
266        nondeterminism.clone(),
267    )
268    .unwrap();
269    let claim = claim.with_output(public_output);
270
271    let stark = Stark::default();
272    let tick = SystemTime::now();
273    triton_vm::profiler::start(timing_report_label);
274    let proof = stark.prove(&claim, &aet).unwrap();
275    let profile = triton_vm::profiler::finish();
276    let measured_time = tick.elapsed().expect("Don't mess with time");
277
278    let padded_height = proof.padded_height().unwrap();
279    let fri = stark.fri(padded_height).unwrap();
280    let report = profile
281        .with_cycle_count(aet.processor_trace.nrows())
282        .with_padded_height(padded_height)
283        .with_fri_domain_len(fri.domain.length);
284    println!("{report}");
285
286    println!("Done proving. Elapsed time: {measured_time:?}");
287    println!(
288        "Proof was generated from:
289        table lengths:
290          processor table: {}
291          hash table: {}
292          u32 table: {}
293          op-stack table: {}
294          RAM table: {}
295          Program table: {}
296          Cascade table: {}
297          Lookup table: {}",
298        aet.height_of_table(TableId::Processor),
299        aet.height_of_table(TableId::Hash),
300        aet.height_of_table(TableId::U32),
301        aet.height_of_table(TableId::OpStack),
302        aet.height_of_table(TableId::Ram),
303        aet.height_of_table(TableId::Program),
304        aet.height_of_table(TableId::Cascade),
305        aet.height_of_table(TableId::Lookup),
306    );
307
308    assert!(
309        triton_vm::verify(stark, &claim, &proof),
310        "Generated proof must verify for program:\n{program}",
311    );
312}
313
314/// A thin wrapper around [`VM::profile`].
315pub fn generate_full_profile(
316    name: &str,
317    program: Program,
318    public_input: &PublicInput,
319    nondeterminism: &NonDeterminism,
320) -> String {
321    let (_output, profile) =
322        VM::profile(program, public_input.clone(), nondeterminism.clone()).unwrap();
323    format!("{name}:\n{profile}")
324}
325
326/// Glob-import this module to reduce the number of imports in a test module.
327///
328/// Feel free to add anything you frequently `use` in a test module.
329#[cfg(test)]
330pub mod test_prelude {
331    pub use std::collections::HashMap;
332
333    pub use itertools::Itertools;
334    pub use proptest::prelude::*;
335    pub use proptest_arbitrary_interop::arb;
336    pub use rand::Rng;
337    pub use rand::RngCore;
338    pub use rand::SeedableRng;
339    pub use rand::prelude::IndexedMutRandom;
340    pub use rand::prelude::IndexedRandom;
341    pub use rand::prelude::IteratorRandom;
342    pub use rand::rngs::StdRng;
343    pub use test_strategy::proptest;
344
345    pub use crate::InitVmState;
346    pub use crate::memory::encode_to_memory;
347    pub(crate) use crate::pop_encodable;
348    pub use crate::push_encodable;
349    pub use crate::snippet_bencher::BenchmarkCase;
350    pub use crate::test_helpers::test_assertion_failure;
351    pub use crate::test_helpers::test_rust_equivalence_given_complete_state;
352    pub use crate::traits::accessor::Accessor;
353    pub use crate::traits::accessor::AccessorInitialState;
354    pub use crate::traits::accessor::ShadowedAccessor;
355    pub use crate::traits::algorithm::Algorithm;
356    pub use crate::traits::algorithm::AlgorithmInitialState;
357    pub use crate::traits::algorithm::ShadowedAlgorithm;
358    pub use crate::traits::closure::Closure;
359    pub use crate::traits::closure::ShadowedClosure;
360    pub use crate::traits::function::Function;
361    pub use crate::traits::function::FunctionInitialState;
362    pub use crate::traits::function::ShadowedFunction;
363    pub use crate::traits::mem_preserver::MemPreserver;
364    pub use crate::traits::mem_preserver::MemPreserverInitialState;
365    pub use crate::traits::mem_preserver::ShadowedMemPreserver;
366    pub use crate::traits::procedure::Procedure;
367    pub use crate::traits::procedure::ProcedureInitialState;
368    pub use crate::traits::procedure::ShadowedProcedure;
369    pub use crate::traits::read_only_algorithm::ReadOnlyAlgorithm;
370    pub use crate::traits::read_only_algorithm::ReadOnlyAlgorithmInitialState;
371    pub use crate::traits::read_only_algorithm::ShadowedReadOnlyAlgorithm;
372    pub use crate::traits::rust_shadow::RustShadow;
373}
tasm_lib/lib.rs

tasm_lib/
lib.rs
