Merge pull request #289 from a16z/feat/trace-saver

feat: Save program traces to file

book/src/usage/troubleshooting.md

@@ -24,3 +24,14 @@ Jolt restricts the size of the inputs and outputs to 4096 bytes. This value is c
 
 ## Guest Attempts to Compile Standard Library
 Sometimes after installing the toolchain the guest does still tries to compile with the standard library which will fail with a large number of errors that certain items such as `Result` are referenced and not available. This generally happens when one tries to run jolt before installing the toolchain. To address, try rerunning `jolt install-toolchain`, restarting your terminal, and delete both your rust target directory and any files under `/tmp` that begin with jolt.
+
+
+## Getting Help
+If none of the above help, please serialize your program and send it along with a detailed bug report.
+
+Serializing a call to the "fib" function in the Jolt guest:
+```rust
+// let (prove_fib, verify_fib) = guest::build_fib();
+let program_summary = guest::analyze_fib(10);
+program_summary.write_to_file("fib_10.txt".into()).expect("should write");
+```

common/src/rv_trace.rs

@@ -5,14 +5,14 @@ use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use serde::{Deserialize, Serialize};
 use strum_macros::FromRepr;
 
-#[derive(Debug, PartialEq, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct RVTraceRow {
     pub instruction: ELFInstruction,
     pub register_state: RegisterState,
     pub memory_state: Option<MemoryState>,
 }
 
-#[derive(Debug, PartialEq, Clone)]
+#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
 pub enum MemoryOp {
     Read(u64, u64),  // (address, value)
     Write(u64, u64), // (address, new_value)
@@ -375,7 +375,7 @@ pub struct RegisterState {
     pub rd_post_val: Option<u64>,
 }
 
-#[derive(Debug, PartialEq, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub enum MemoryState {
     Read {
         address: u64,

examples/fibonacci/src/main.rs

@@ -1,5 +1,9 @@
 pub fn main() {
     let (prove_fib, verify_fib) = guest::build_fib();
+    let program_summary = guest::analyze_fib(10);
+    program_summary
+        .write_to_file("fib_10.txt".into())
+        .expect("should write");
 
     let (output, proof) = prove_fib(50);
     let is_valid = verify_fib(proof);

jolt-core/Cargo.toml

@@ -59,6 +59,7 @@ tracing-texray = "0.2.0"
 
 common = { path = "../common" }
 tracer = { path = "../tracer" }
+bincode = "1.3.3"
 
 [build-dependencies]
 common = { path = "../common" }

jolt-core/src/host/analyze.rs

@@ -0,0 +1,54 @@
+use std::{collections::HashMap, fs::File, io, path::PathBuf};
+
+use ark_ff::PrimeField;
+use serde::{Deserialize, Serialize};
+use tracer::{ELFInstruction, JoltDevice, RVTraceRow, RV32IM};
+
+use crate::jolt::vm::{bytecode::BytecodeRow, rv32i_vm::RV32I};
+
+use common::{constants::MEMORY_OPS_PER_INSTRUCTION, rv_trace::MemoryOp};
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct ProgramSummary {
+    pub raw_trace: Vec<RVTraceRow>,
+
+    pub bytecode: Vec<ELFInstruction>,
+    pub memory_init: Vec<(u64, u8)>,
+
+    pub io_device: JoltDevice,
+    pub bytecode_trace: Vec<BytecodeRow>,
+    pub instruction_trace: Vec<Option<RV32I>>,
+    pub memory_trace: Vec<[MemoryOp; MEMORY_OPS_PER_INSTRUCTION]>,
+    pub circuit_flags: Vec<bool>,
+}
+
+impl ProgramSummary {
+    pub fn trace_len(&self) -> usize {
+        self.memory_trace.len()
+    }
+
+    pub fn analyze<F: PrimeField>(&self) -> Vec<(RV32IM, usize)> {
+        let mut counts = HashMap::<RV32IM, usize>::new();
+        for row in self.raw_trace.iter() {
+            let op = row.instruction.opcode;
+            if let Some(count) = counts.get(&op) {
+                counts.insert(op, count + 1);
+            } else {
+                counts.insert(op, 1);
+            }
+        }
+
+        let mut counts: Vec<_> = counts.into_iter().collect();
+        counts.sort_by_key(|v| v.1);
+        counts.reverse();
+
+        counts
+    }
+
+    pub fn write_to_file(self, path: PathBuf) -> Result<(), Box<dyn std::error::Error>> {
+        let mut file = File::create(path)?;
+        let data = bincode::serialize(&self)?;
+        io::Write::write_all(&mut file, &data)?;
+        Ok(())
+    }
+}

jolt-core/src/host/mod.rs

@@ -1,6 +1,5 @@
 use core::{str::FromStr, u8};
 use std::{
-    collections::HashMap,
     fs::{self, File},
     io::{self, Write},
     path::PathBuf,
@@ -14,7 +13,7 @@ use serde::Serialize;
 
 use common::{
     constants::MEMORY_OPS_PER_INSTRUCTION,
-    rv_trace::{JoltDevice, MemoryOp, NUM_CIRCUIT_FLAGS, RV32IM},
+    rv_trace::{JoltDevice, MemoryOp, NUM_CIRCUIT_FLAGS},
 };
 use tracer::ELFInstruction;
 
@@ -23,6 +22,10 @@ use crate::{
     utils::thread::unsafe_allocate_zero_vec,
 };
 
+use self::analyze::ProgramSummary;
+
+pub mod analyze;
+
 const DEFAULT_MEMORY_SIZE: usize = 10 * 1024 * 1024;
 const DEFAULT_STACK_SIZE: usize = 4096;
 
@@ -173,27 +176,31 @@ impl Program {
         )
     }
 
-    pub fn trace_analyze(mut self) -> (usize, Vec<(RV32IM, usize)>) {
+    pub fn trace_analyze<F: PrimeField>(mut self) -> ProgramSummary {
         self.build();
-        let elf = self.elf.unwrap();
-        let (rows, _) = tracer::trace(&elf, self.input);
-        let trace_len = rows.len();
-
-        let mut counts = HashMap::<RV32IM, usize>::new();
-        for row in rows {
-            let op = row.instruction.opcode;
-            if let Some(count) = counts.get(&op) {
-                counts.insert(op, count + 1);
-            } else {
-                counts.insert(op, 1);
-            }
-        }
+        let elf = self.elf.as_ref().unwrap();
+        let (raw_trace, _) = tracer::trace(&elf, self.input.clone());
+
+        let (bytecode, memory_init) = self.decode();
+        let (io_device, bytecode_trace, instruction_trace, memory_trace, circuit_flags) =
+            self.trace();
+        let circuit_flags: Vec<bool> = circuit_flags
+            .into_iter()
+            .map(|flag: F| flag.is_one())
+            .collect();
 
-        let mut counts: Vec<_> = counts.into_iter().collect();
-        counts.sort_by_key(|v| v.1);
-        counts.reverse();
+        let program_summary = ProgramSummary {
+            raw_trace,
+            bytecode,
+            memory_init,
+            io_device,
+            bytecode_trace,
+            instruction_trace,
+            memory_trace,
+            circuit_flags,
+        };
 
-        (trace_len, counts)
+        program_summary
     }
 
     fn save_linker(&self) {

jolt-core/src/jolt/instruction/add.rs

@@ -1,6 +1,7 @@
 use ark_ff::PrimeField;
 use ark_std::log2;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{JoltInstruction, SubtableIndices};
 use crate::jolt::subtable::{
@@ -10,7 +11,7 @@ use crate::utils::instruction_utils::{
     add_and_chunk_operands, assert_valid_parameters, concatenate_lookups,
 };
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct ADDInstruction<const WORD_SIZE: usize>(pub u64, pub u64);
 
 impl<const WORD_SIZE: usize> JoltInstruction for ADDInstruction<WORD_SIZE> {

jolt-core/src/jolt/instruction/and.rs

@@ -1,12 +1,13 @@
 use ark_ff::PrimeField;
 use ark_std::log2;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{JoltInstruction, SubtableIndices};
 use crate::jolt::subtable::{and::AndSubtable, LassoSubtable};
 use crate::utils::instruction_utils::{chunk_and_concatenate_operands, concatenate_lookups};
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct ANDInstruction(pub u64, pub u64);
 
 impl JoltInstruction for ANDInstruction {

jolt-core/src/jolt/instruction/beq.rs

@@ -1,5 +1,6 @@
 use ark_ff::PrimeField;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::JoltInstruction;
 use crate::{
@@ -10,7 +11,7 @@ use crate::{
     utils::instruction_utils::chunk_and_concatenate_operands,
 };
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct BEQInstruction(pub u64, pub u64);
 
 impl JoltInstruction for BEQInstruction {

jolt-core/src/jolt/instruction/bge.rs

@@ -1,5 +1,6 @@
 use ark_ff::PrimeField;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{slt::SLTInstruction, JoltInstruction, SubtableIndices};
 use crate::{
@@ -10,7 +11,7 @@ use crate::{
     utils::instruction_utils::chunk_and_concatenate_operands,
 };
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct BGEInstruction(pub u64, pub u64);
 
 impl JoltInstruction for BGEInstruction {

jolt-core/src/jolt/instruction/bgeu.rs

@@ -1,13 +1,14 @@
 use ark_ff::PrimeField;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{sltu::SLTUInstruction, JoltInstruction, SubtableIndices};
 use crate::{
     jolt::subtable::{eq::EqSubtable, ltu::LtuSubtable, LassoSubtable},
     utils::instruction_utils::chunk_and_concatenate_operands,
 };
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct BGEUInstruction(pub u64, pub u64);
 
 impl JoltInstruction for BGEUInstruction {

jolt-core/src/jolt/instruction/bne.rs

@@ -1,5 +1,6 @@
 use ark_ff::PrimeField;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::JoltInstruction;
 use crate::{
@@ -10,7 +11,7 @@ use crate::{
     utils::instruction_utils::chunk_and_concatenate_operands,
 };
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct BNEInstruction(pub u64, pub u64);
 
 impl JoltInstruction for BNEInstruction {

jolt-core/src/jolt/instruction/lb.rs

@@ -1,13 +1,14 @@
 use ark_ff::PrimeField;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{JoltInstruction, SubtableIndices};
 use crate::jolt::subtable::{
     sign_extend::SignExtendSubtable, truncate_overflow::TruncateOverflowSubtable, LassoSubtable,
 };
 use crate::utils::instruction_utils::chunk_operand_usize;
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct LBInstruction(pub u64);
 
 impl JoltInstruction for LBInstruction {

jolt-core/src/jolt/instruction/lh.rs

@@ -1,13 +1,14 @@
 use ark_ff::PrimeField;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{JoltInstruction, SubtableIndices};
 use crate::jolt::subtable::{
     identity::IdentitySubtable, sign_extend::SignExtendSubtable, LassoSubtable,
 };
 use crate::utils::instruction_utils::chunk_operand_usize;
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct LHInstruction(pub u64);
 
 impl JoltInstruction for LHInstruction {

jolt-core/src/jolt/instruction/mod.rs

@@ -2,6 +2,7 @@ use ark_ff::PrimeField;
 use enum_dispatch::enum_dispatch;
 use fixedbitset::*;
 use rand::prelude::StdRng;
+use serde::Serialize;
 use std::marker::Sync;
 use std::ops::Range;
 use strum::{EnumCount, IntoEnumIterator};
@@ -12,7 +13,7 @@ use common::rv_trace::ELFInstruction;
 use std::fmt::Debug;
 
 #[enum_dispatch]
-pub trait JoltInstruction: Clone + Debug + Send + Sync {
+pub trait JoltInstruction: Clone + Debug + Send + Sync + Serialize {
     fn operands(&self) -> (u64, u64);
     /// Combines `vals` according to the instruction's "collation" polynomial `g`.
     /// If `vals` are subtable entries (as opposed to MLE evaluations), this function returns the

jolt-core/src/jolt/instruction/or.rs

@@ -1,12 +1,13 @@
 use ark_ff::PrimeField;
 use ark_std::log2;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{JoltInstruction, SubtableIndices};
 use crate::jolt::subtable::{or::OrSubtable, LassoSubtable};
 use crate::utils::instruction_utils::{chunk_and_concatenate_operands, concatenate_lookups};
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct ORInstruction(pub u64, pub u64);
 
 impl JoltInstruction for ORInstruction {

jolt-core/src/jolt/instruction/sb.rs

@@ -1,11 +1,12 @@
 use ark_ff::PrimeField;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{JoltInstruction, SubtableIndices};
 use crate::jolt::subtable::{truncate_overflow::TruncateOverflowSubtable, LassoSubtable};
 use crate::utils::instruction_utils::chunk_operand_usize;
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct SBInstruction(pub u64);
 
 impl JoltInstruction for SBInstruction {

jolt-core/src/jolt/instruction/sh.rs

@@ -1,12 +1,13 @@
 use ark_ff::PrimeField;
 
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{JoltInstruction, SubtableIndices};
 use crate::jolt::subtable::{identity::IdentitySubtable, LassoSubtable};
 use crate::utils::instruction_utils::chunk_operand_usize;
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct SHInstruction(pub u64);
 
 impl JoltInstruction for SHInstruction {

jolt-core/src/jolt/instruction/sll.rs

@@ -1,14 +1,15 @@
 use ark_ff::PrimeField;
 use ark_std::log2;
 use rand::prelude::StdRng;
+use serde::{Deserialize, Serialize};
 
 use super::{JoltInstruction, SubtableIndices};
 use crate::jolt::subtable::{sll::SllSubtable, LassoSubtable};
 use crate::utils::instruction_utils::{
     assert_valid_parameters, chunk_and_concatenate_for_shift, concatenate_lookups,
 };
 
-#[derive(Copy, Clone, Default, Debug)]
+#[derive(Copy, Clone, Default, Debug, Serialize, Deserialize)]
 pub struct SLLInstruction<const WORD_SIZE: usize>(pub u64, pub u64);
 
 impl<const WORD_SIZE: usize> JoltInstruction for SLLInstruction<WORD_SIZE> {