Feat/batched opening proofs

book/src/how/r1cs_constraints.md

@@ -1,71 +1,79 @@
 # R1CS constraints
 
-Jolt usees R1CS constraints to enforce certain rules of the RISC-V fetch-decode-execute loop 
-and to ensure 
-consistency between the proofs for the different modules of Jolt ([instruction lookups](./instruction_lookups.md), [read-write memory](./read_write_memory.md), and [bytecode](./bytecode.md)). 
+Jolt uses R1CS constraints to enforce certain rules of the RISC-V fetch-decode-execute loop
+and to ensure
+consistency between the proofs for the different modules of Jolt ([instruction lookups](./instruction_lookups.md), [read-write memory](./read_write_memory.md), and [bytecode](./bytecode.md)).
 
 ## Uniformity
 
 Jolt's R1CS is uniform, which means
-that the constraint matrices for an entire program are just repeated copies of the constraint 
-matrices for a single CPU step. 
-Each step is conceptually simple and involves around 60 constraints and 80 variables. 
+that the constraint matrices for an entire program are just repeated copies of the constraint
+matrices for a single CPU step.
+Each step is conceptually simple and involves around 60 constraints and 80 variables.
 
 ## Input Variables and constraints
 
-The inputs required for the constraint system for a single CPU step are: 
+The inputs required for the constraint system for a single CPU step are:
 
 #### Pertaining to bytecode
-* Program counters (PCs): this is the only state passed between CPU steps. 
-* Bytecode read address: the address in the program code read at this step. 
-* The preprocessed ("5-tuple") representation of the instruction: (`bitflags`, `rs1`, `rs2`, `rd`, `imm`). 
+* Bytecode read address: the index in the program code read at this step.
+* The preprocessed representation of the instruction: (`elf_address`, `bitflags`, `rs1`, `rs2`, `rd`, `imm`).
 
 #### Pertaining to read-write memory
 * The (starting) RAM address read by the instruction: if the instruction is not a load/store, this is 0.
 * The bytes written to or read from memory.
 
 ####  Pertaining to instruction lookups
-* The chunks of the instruction's operands `x` and `y`. 
+* The chunks of the instruction's operands `x` and `y`.
 * The chunks of the lookup query. These are typically some combination of the operand chunks (e.g. the i-th chunk of the lookup query is often the concatenation of `x_i` and `y_i`).
-* The lookup output. 
-
-### Circuit and instruction flags: 
-* There are nine circuit flags used to guide the constraints and are dependent only on the opcode of the instruction. These are thus stored as part of the preprocessed bytecode in Jolt. 
-    1. `operand_x_flag`: 0 if the first operand is the value in `rs1` or the `PC`. 
-    2. `operand_y_flag`: 0 if the second operand is the value in `rs2` or the `imm`. 
-    3. `is_load_instr`
-    4. `is_store_instr`
-    5. `is_jump_instr`
-    6. `is_branch_instr`
-    7. `if_update_rd_with_lookup_output`: 1 if the lookup output is to be stored in `rd` at the end of the step. 
-    8. `sign_imm_flag`: used in load/store and branch instructions where the instruction is added as constraints. 
-    9. `is_concat`: indicates whether the instruction performs a concat-type lookup. 
-* Instruction flags: these are the unary bits used to indicate instruction is executed at a given step. There are as many per step as the number of unique instruction lookup tables in Jolt, which is 19. 
-
-#### Constraint system 
-
-The constraints for a CPU step are detailed in the `get_jolt_matrices()` function in [`constraints.rs`](https://github.com/a16z/jolt/blob/main/jolt-core/src/r1cs/jolt_constraints.rs). 
-
-### Reusing commitments 
-
-As with most SNARK backends, Spartan requires computing a commitment to the inputs 
-to the constraint system. 
-A catch (and an optimization feature) in Jolt is that most of the inputs 
-are also used as inputs to proofs in the other modules. For example, 
-the address and values pertaining to the bytecode are used in the bytecode memory-checking proof, 
-and the lookup chunks, output and flags are used in the instruction lookup proof. 
-For Jolt to be sound, it must be ensured that the same inputs are fed to all relevant proofs. 
-We do this by re-using the commitments themselves. 
-This can be seen in the `format_commitments()` function in the `r1cs/snark` module. 
-Spartan is adapted to take pre-committed witness variables. 
-
-## Exploiting uniformity 
-
-The uniformity of the constraint system allows us to heavily optimize both the prover and verifier. 
-The main changes involved in making this happen are: 
-- Spartan is modified to only take in the constraint matrices a single step, and the total number of steps. Using this, the prover and verifier can efficiently calculate the multilinear extensions of the full R1CS matrices. 
-- The commitment format of the witness values is changed to reflect uniformity. All versions of a variable corresponding to each time step is committed together. This affects nearly all variables committed to in Jolt. 
-- The inputs and witnesses are provided to the constraint system as segments. 
-- Additional constraints are used to enforce consistency of the state transferred between CPU steps. 
+* The lookup output.
+
+### Circuit and instruction flags:
+* There are twelve circuit flags (`opflags` in the Jolt paper) used in Jolt's R1CS constraints.
+They are enumatered in `CircuitFlags` and computed in `to_circuit_flags` (see [`rv_trace.rs`](https://github.com/a16z/jolt/blob/main/common/src/rv_trace.rs))
+Circuit flags depend only on the instruction as it appears in the bytecode, so they are computed as part of
+the preprocessed bytecode in Jolt.
+    1. `LeftOperandIsPC`: 1 if the first lookup operand is the program counter; 0 otherwise (first lookup operand is RS1 value).
+    1. `RightOperandIsImm`: 1 if the second lookup operand is `imm`; 0 otherwise (second lookup operand is RS2 value).
+    1. `Load`: 1 if the instruction is a load (i.e. `LB`, `LH`, etc.)
+    1. `Store`: 1 if the instruction is a store (i.e. `SB`, `SH`, etc.)
+    1. `Jump`: 1 if the instruction is a jump (i.e. `JAL`, `JALR`)
+    1. `Branch`: 1 if the instruction is a branch (i.e. `BEQ`, `BNE`, etc.)
+    1. `WriteLookupOutputToRD`: 1 if the lookup output is to be stored in `rd` at the end of the step.
+    1. `ImmSignBit`: Used in load/store and branch instructions where the immediate value used as an offset
+    1. `ConcatLookupQueryChunks`: Indicates whether the instruction performs a concat-type lookup.
+    1. `Virtual`: 1 if the instruction is "virtual", as defined in Section 6.1 of the Jolt paper.
+    1. `Assert`: 1 if the instruction is an assert, as defined in Section 6.1.1 of the Jolt paper.
+    1. `DoNotUpdatePC`: Used in virtual sequences; the program counter should be the same for the full seqeuence.
+* Instruction flags: these are the unary bits used to indicate instruction is executed at a given step.
+There are as many per step as the number of unique instruction lookup tables in Jolt.
+
+#### Constraint system
+
+The constraints for a CPU step are detailed in the `uniform_constraints` and `non_uniform_constraints` functions in [`constraints.rs`](https://github.com/a16z/jolt/blob/main/jolt-core/src/r1cs/constraints.rs).
+
+### Reusing commitments
+
+As with most SNARK backends, Spartan requires computing a commitment to the inputs
+to the constraint system.
+A catch (and an optimization feature) in Jolt is that most of the inputs
+are also used as inputs to proofs in the other modules. For example,
+the address and values pertaining to the bytecode are used in the bytecode memory-checking proof,
+and the lookup chunks, output and flags are used in the instruction lookup proof.
+For Jolt to be sound, it must be ensured that the same inputs are fed to all relevant proofs.
+We do this by re-using the commitments themselves.
+Spartan is adapted to take pre-committed witness variables.
+
+## Exploiting uniformity
+
+The uniformity of the constraint system allows us to heavily optimize both the prover and verifier.
+The main changes involved in making this happen are:
+- Spartan is modified to only take in the constraint matrices a single step, and the total number of steps.
+Using this, the prover and verifier can efficiently calculate the multilinear extensions of the full R1CS matrices.
+- The commitment format of the witness values is changed to reflect uniformity.
+All versions of a variable corresponding to each time step is committed together.
+This affects nearly all variables committed to in Jolt.
+- The inputs and witnesses are provided to the constraint system as segments.
+- Additional constraints are used to enforce consistency of the state transferred between CPU steps.
 
 These changes and their impact on the code are visible in [`spartan.rs`](https://github.com/a16z/jolt/blob/main/jolt-core/src/r1cs/spartan.rs).

common/Cargo.toml

@@ -7,5 +7,6 @@ edition = "2021"
 ark-serialize = { version = "0.4.2", features = ["derive"] }
 serde = { version = "1.0.193", features = ["derive"] }
 serde_json = "1.0.108"
-strum_macros = "0.25.3"
+strum_macros = "0.26.4"
+strum = "0.26.3"
 syn = { version = "1.0", features = ["full"] }

common/src/rv_trace.rs

@@ -3,7 +3,8 @@ use std::str::FromStr;
 use crate::constants::{MEMORY_OPS_PER_INSTRUCTION, RAM_START_ADDRESS, REGISTER_COUNT};
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 use serde::{Deserialize, Serialize};
-use strum_macros::FromRepr;
+use strum::EnumCount;
+use strum_macros::{EnumCount as EnumCountMacro, EnumIter, FromRepr};
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 pub struct RVTraceRow {
@@ -230,33 +231,52 @@ pub struct ELFInstruction {
     pub virtual_sequence_remaining: Option<usize>,
 }
 
-pub const NUM_CIRCUIT_FLAGS: usize = 12;
+/// Boolean flags used in Jolt's R1CS constraints (`opflags` in the Jolt paper).
+/// Note that the flags below deviate slightly from those described in Appendix A.1
+/// of the Jolt paper.
+#[derive(
+    Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Hash, Ord, EnumCountMacro, EnumIter, Default,
+)]
+pub enum CircuitFlags {
+    #[default] // Need a default so that we can derive EnumIter on `JoltR1CSInputs`
+    /// 1 if the first lookup operand is the program counter; 0 otherwise (first lookup operand is RS1 value).
+    LeftOperandIsPC,
+    /// 1 if the second lookup operand is `imm`; 0 otherwise (second lookup operand is RS2 value).
+    RightOperandIsImm,
+    /// 1 if the instruction is a load (i.e. `LB`, `LH`, etc.)
+    Load,
+    /// 1 if the instruction is a store (i.e. `SB`, `SH`, etc.)
+    Store,
+    /// 1 if the instruction is a jump (i.e. `JAL`, `JALR`)
+    Jump,
+    /// 1 if the instruction is a branch (i.e. `BEQ`, `BNE`, etc.)
+    Branch,
+    /// 1 if the lookup output is to be stored in `rd` at the end of the step.
+    WriteLookupOutputToRD,
+    /// Used in load/store and branch instructions where the immediate value used as an offset
+    ImmSignBit,
+    /// Indicates whether the instruction performs a concat-type lookup.
+    ConcatLookupQueryChunks,
+    /// 1 if the instruction is "virtual", as defined in Section 6.1 of the Jolt paper.
+    Virtual,
+    /// 1 if the instruction is an assert, as defined in Section 6.1.1 of the Jolt paper.
+    Assert,
+    /// Used in virtual sequences; the program counter should be the same for the full seqeuence.
+    DoNotUpdatePC,
+}
+pub const NUM_CIRCUIT_FLAGS: usize = CircuitFlags::COUNT;
 
 impl ELFInstruction {
     #[rustfmt::skip]
     pub fn to_circuit_flags(&self) -> [bool; NUM_CIRCUIT_FLAGS] {
-        // Jolt Appendix A.1
-        // 0: first_operand == rs1 (1 if PC)
-        // 1: second_operand == rs2 (1 if imm)
-        // 2: Load instruction
-        // 3: Store instruction
-        // 4: Jump instruction
-        // 5: Branch instruction
-        // 6: Instruction writes lookup output to rd
-        // 7: Sign-bit of imm
-        // 8: Is concat
-        // 9: Virtual instruction
-        // 10: Assert instruction
-        // 11: Don't update PC
-
         let mut flags = [false; NUM_CIRCUIT_FLAGS];
 
-        flags[0] = matches!(
+        flags[CircuitFlags::LeftOperandIsPC as usize] = matches!(
             self.opcode,
             RV32IM::JAL | RV32IM::LUI | RV32IM::AUIPC,
         );
 
-        flags[1] = matches!(
+        flags[CircuitFlags::RightOperandIsImm as usize] = matches!(
             self.opcode,
             RV32IM::ADDI
             | RV32IM::XORI
@@ -272,28 +292,28 @@ impl ELFInstruction {
             | RV32IM::JALR,
         );
 
-        flags[2] = matches!(
+        flags[CircuitFlags::Load as usize] = matches!(
             self.opcode,
             RV32IM::LB | RV32IM::LH | RV32IM::LW | RV32IM::LBU | RV32IM::LHU,
         );
 
-        flags[3] = matches!(
+        flags[CircuitFlags::Store as usize] = matches!(
             self.opcode,
             RV32IM::SB | RV32IM::SH | RV32IM::SW,
         );
 
-        flags[4] = matches!(
+        flags[CircuitFlags::Jump as usize] = matches!(
             self.opcode,
             RV32IM::JAL | RV32IM::JALR,
         );
 
-        flags[5] = matches!(
+        flags[CircuitFlags::Branch as usize] = matches!(
             self.opcode,
             RV32IM::BEQ | RV32IM::BNE | RV32IM::BLT | RV32IM::BGE | RV32IM::BLTU | RV32IM::BGEU,
         );
 
         // loads, stores, branches, jumps, and asserts do not store the lookup output to rd (they may update rd in other ways)
-        flags[6] = !matches!(
+        flags[CircuitFlags::WriteLookupOutputToRD as usize] = !matches!(
             self.opcode,
             RV32IM::SB
             | RV32IM::SH
@@ -315,9 +335,9 @@ impl ELFInstruction {
         );
 
         let mask = 1u32 << 31;
-        flags[7] = matches!(self.imm, Some(imm) if imm & mask == mask);
+        flags[CircuitFlags::ImmSignBit as usize] = matches!(self.imm, Some(imm) if imm & mask == mask);
 
-        flags[8] = matches!(
+        flags[CircuitFlags::ConcatLookupQueryChunks as usize] = matches!(
             self.opcode,
             RV32IM::XOR
             | RV32IM::XORI
@@ -348,9 +368,9 @@ impl ELFInstruction {
             | RV32IM::VIRTUAL_ASSERT_VALID_DIV0,
         );
 
-        flags[9] = self.virtual_sequence_remaining.is_some();
+        flags[CircuitFlags::Virtual as usize] = self.virtual_sequence_remaining.is_some();
 
-        flags[10] = matches!(self.opcode,
+        flags[CircuitFlags::Assert as usize] = matches!(self.opcode,
             RV32IM::VIRTUAL_ASSERT_EQ                        |
             RV32IM::VIRTUAL_ASSERT_LTE                       |
             RV32IM::VIRTUAL_ASSERT_VALID_SIGNED_REMAINDER    |
@@ -361,7 +381,7 @@ impl ELFInstruction {
         // All instructions in virtual sequence are mapped from the same
         // ELF address. Thus if an instruction is virtual (and not the last one
         // in its sequence), then we should *not* update the PC.
-        flags[11] = match self.virtual_sequence_remaining {
+        flags[CircuitFlags::DoNotUpdatePC as usize] = match self.virtual_sequence_remaining {
             Some(i) => i != 0,
             None => false
         };

jolt-core/Cargo.toml

@@ -27,7 +27,7 @@ ark-serialize = { version = "0.4.2", default-features = false, features = [
     "derive",
 ] }
 ark-std = { version = "0.4.0" }
-binius-field = { git = "https://gitlab.com/UlvetannaOSS/binius", package = "binius_field"}
+binius-field = { git = "https://gitlab.com/UlvetannaOSS/binius", package = "binius_field" }
 clap = { version = "4.3.10", features = ["derive"] }
 enum_dispatch = "0.3.12"
 fixedbitset = "0.5.0"
@@ -46,8 +46,8 @@ rgb = "0.8.37"
 serde = { version = "1.0.*", default-features = false }
 sha3 = "0.10.8"
 smallvec = "1.13.1"
-strum = "0.25.0"
-strum_macros = "0.25.2"
+strum = "0.26.3"
+strum_macros = "0.26.4"
 textplots = "0.8.4"
 thiserror = "1.0.58"
 tracing = "0.1.37"
@@ -101,9 +101,9 @@ default = [
 host = ["dep:reqwest", "dep:tokio"]
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
-memory-stats = "1.0.0" 
+memory-stats = "1.0.0"
 tokio = { version = "1.38.0", optional = true, features = ["rt-multi-thread"] }
 
 
 [target.'cfg(target_arch = "wasm32")'.dependencies]
-getrandom = { version = "0.2", features = ["js"] }
+getrandom = { version = "0.2", features = ["js"] }

jolt-core/src/benches/bench.rs

@@ -107,17 +107,13 @@ where
 
     let task = move || {
         let (bytecode, memory_init) = program.decode();
-        let (io_device, trace, circuit_flags) = program.trace();
+        let (io_device, trace) = program.trace();
 
-        let preprocessing: crate::jolt::vm::JoltPreprocessing<F, PCS> =
+        let preprocessing: crate::jolt::vm::JoltPreprocessing<C, F, PCS> =
             RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 22);
 
-        let (jolt_proof, jolt_commitments) = <RV32IJoltVM as Jolt<_, PCS, C, M>>::prove(
-            io_device,
-            trace,
-            circuit_flags,
-            preprocessing.clone(),
-        );
+        let (jolt_proof, jolt_commitments, _) =
+            <RV32IJoltVM as Jolt<_, PCS, C, M>>::prove(io_device, trace, preprocessing.clone());
 
         println!("Proof sizing:");
         serialize_and_print_size("jolt_commitments", &jolt_commitments);
@@ -133,7 +129,8 @@ where
             &jolt_proof.instruction_lookups,
         );
 
-        let verification_result = RV32IJoltVM::verify(preprocessing, jolt_proof, jolt_commitments);
+        let verification_result =
+            RV32IJoltVM::verify(preprocessing, jolt_proof, jolt_commitments, None);
         assert!(
             verification_result.is_ok(),
             "Verification failed with error: {:?}",
@@ -157,22 +154,19 @@ where
     let mut tasks = Vec::new();
     let mut program = host::Program::new("sha2-chain-guest");
     program.set_input(&[5u8; 32]);
-    program.set_input(&1024u32);
+    program.set_input(&1000u32);
 
     let task = move || {
         let (bytecode, memory_init) = program.decode();
-        let (io_device, trace, circuit_flags) = program.trace();
+        let (io_device, trace) = program.trace();
 
-        let preprocessing: crate::jolt::vm::JoltPreprocessing<F, PCS> =
+        let preprocessing: crate::jolt::vm::JoltPreprocessing<C, F, PCS> =
             RV32IJoltVM::preprocess(bytecode.clone(), memory_init, 1 << 20, 1 << 20, 1 << 22);
 
-        let (jolt_proof, jolt_commitments) = <RV32IJoltVM as Jolt<_, PCS, C, M>>::prove(
-            io_device,
-            trace,
-            circuit_flags,
-            preprocessing.clone(),
-        );
-        let verification_result = RV32IJoltVM::verify(preprocessing, jolt_proof, jolt_commitments);
+        let (jolt_proof, jolt_commitments, _) =
+            <RV32IJoltVM as Jolt<_, PCS, C, M>>::prove(io_device, trace, preprocessing.clone());
+        let verification_result =
+            RV32IJoltVM::verify(preprocessing, jolt_proof, jolt_commitments, None);
         assert!(
             verification_result.is_ok(),
             "Verification failed with error: {:?}",

jolt-core/src/field/binius.rs

@@ -273,3 +273,9 @@ impl<F: BiniusSpecific> ark_serialize::Valid for BiniusField<F> {
         todo!()
     }
 }
+
+impl<F: BiniusSpecific> std::fmt::Display for BiniusField<F> {
+    fn fmt(&self, _f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        todo!()
+    }
+}