feat(fr): BN254 Fr native-field coprocessor on the modular stack

Cargo.toml

@@ -59,6 +59,7 @@ members = [
   "jolt-inlines/secp256k1",
   "jolt-inlines/grumpkin",
   "jolt-inlines/p256",
+  "jolt-inlines/bn254-fr",
   "examples/btreemap/host",
   "examples/btreemap/guest",
   "examples/collatz",
@@ -83,6 +84,9 @@ members = [
   "examples/alloc/guest",
   "examples/stdlib",
   "examples/stdlib/guest",
+  "examples/bn254-eqpoly",
+  "examples/bn254-eqpoly/inline-guest",
+  "examples/bn254-eqpoly/native-guest",
   "examples/muldiv",
   "examples/muldiv/guest",
   "examples/overflow",
@@ -414,3 +418,4 @@ jolt-inlines-bigint = { path = "./jolt-inlines/bigint", default-features = false
 jolt-inlines-secp256k1 = { path = "./jolt-inlines/secp256k1", default-features = false }
 jolt-inlines-grumpkin = { path = "./jolt-inlines/grumpkin", default-features = false }
 jolt-inlines-p256 = { path = "./jolt-inlines/p256", default-features = false }
+jolt-inlines-bn254-fr = { path = "./jolt-inlines/bn254-fr", default-features = false }

crates/bolt/src/protocols/jolt/artifacts.rs

@@ -740,6 +740,7 @@ pub struct JoltProverWitnessInputs<'a, CommitmentInputs> {{
     pub instruction_ra_virtual_d: usize,
     pub stage7_openings: &'a [stage7::Stage7OpeningInputValue<{field_type}>],
     pub evaluation_openings: Option<&'a [stage7::Stage7OpeningInputValue<{field_type}>]>,
+    pub field_reg_replay: Option<&'a jolt_witness::field_reg::FieldRegReplay>,
 }}
 
 pub fn prove_jolt_with_witness_inputs<CommitmentInputs, T>(
@@ -768,10 +769,13 @@ where
     let stage3 = stage3_prover_inputs(inputs.stage3_openings, inputs.stage3_cycles);
     drop(_stage3_input_span);
     let _stage45_witness_span = tracing::info_span!("bolt.prove.inputs.stage45_witness").entered();
-    let stage45_witness = stage4::stage4_5_sparse_trace_witness_from_accesses(
+    let mut stage45_witness = stage4::stage4_5_sparse_trace_witness_from_accesses(
         inputs.register_accesses,
         inputs.ram.accesses,
     );
+    if let Some(replay) = inputs.field_reg_replay {{
+        stage45_witness = stage45_witness.with_field_reg_replay(replay.clone());
+    }}
     drop(_stage45_witness_span);
     let _stage4_input_span = tracing::info_span!("bolt.prove.inputs.stage4").entered();
     let stage4 = stage4_prover_inputs(
@@ -1551,6 +1555,12 @@ pub fn stage6_verifier_data_from_witness_entries(
                     right_is_rs2: entry.right_is_rs2,
                     right_is_imm: entry.right_is_imm,
                     is_noop: entry.is_noop,
+                    frd: entry.frd,
+                    frs1: entry.frs1,
+                    frs2: entry.frs2,
+                    reads_frs1: entry.reads_frs1,
+                    reads_frs2: entry.reads_frs2,
+                    writes_frd: entry.writes_frd,
                 }})
                 .collect(),
             entry_bytecode_index,

crates/bolt/src/protocols/jolt/emit/rust/commitment.rs

@@ -1724,9 +1724,11 @@ fn rust_input_field(source: &str) -> Result<&'static str, EmitError> {
     match source {
         "trace.rd_inc" => Ok("rd_inc"),
         "trace.ram_inc" => Ok("ram_inc"),
+        "trace.field_reg_inc" => Ok("field_reg_inc"),
         "trace.instruction_keys" => Ok("instruction_keys"),
         "trace.ram_addresses" => Ok("ram_addresses"),
         "trace.bytecode_indices" => Ok("bytecode_indices"),
+        "trace.field_reg_indices" => Ok("field_reg_indices"),
         "advice.untrusted" => Ok("untrusted_advice"),
         "advice.trusted" => Ok("trusted_advice"),
         _ => Err(EmitError::new(format!(

crates/bolt/src/protocols/jolt/emit/rust/stage3.rs

@@ -624,6 +624,7 @@ impl Stage3CpuProgram {
                 "jolt.stage3.spartan_shift" => "jolt_stage3_spartan_shift",
                 "jolt.stage3.instruction_input" => "jolt_stage3_instruction_input",
                 "jolt.stage3.registers_claim_reduction" => "jolt_stage3_registers_claim_reduction",
+                "jolt.stage3.field_reg_claim_reduction" => "jolt_stage3_field_reg_claim_reduction",
                 "jolt.stage3.batched" => "jolt_stage3_batched",
                 _ => {
                     return Err(EmitError::new(format!(
@@ -1946,6 +1947,9 @@ fn expected_batched_output_claim(
             "jolt.stage3.registers_claim_reduction" => {
                 expected_registers(store, evals, local_point)?
             }
+            "jolt.stage3.field_reg_claim_reduction" => {
+                expected_field_regs(store, evals, local_point)?
+            }
             _ => {
                 return Err(VerifyStage3Error::UnsupportedRelation {
                     relation: instance.relation,
@@ -2032,6 +2036,24 @@ fn expected_registers(
                 * eval_by_name(evals, "stage3.registers_claim_reduction.eval.Rs2Value")?))
 }
 
+fn expected_field_regs(
+    store: &super::common::ValueStore<Fr>,
+    evals: &[Stage3NamedEval<Fr>],
+    local_point: &[Fr],
+) -> Result<Fr, VerifyStage3Error> {
+    let opening_point = reverse_slice(local_point);
+    let eq_eval = EqPolynomial::<Fr>::mle(
+        &opening_point,
+        super::common::store_point(store, "stage3.input.stage1.FieldRdWriteValue")?,
+    );
+    Ok(eq_eval
+        * (eval_by_name(evals, "stage3.field_reg_claim_reduction.eval.FieldRdWriteValue")?
+            + super::common::store_scalar(store, "stage3.field_reg.gamma")?
+                * eval_by_name(evals, "stage3.field_reg_claim_reduction.eval.FieldRs1Value")?
+            + super::common::store_scalar(store, "stage3.field_reg.gamma2")?
+                * eval_by_name(evals, "stage3.field_reg_claim_reduction.eval.FieldRs2Value")?))
+}
+
 "#
     }
 }

crates/bolt/src/protocols/jolt/emit/rust/stage4.rs

@@ -642,6 +642,7 @@ impl Stage4CpuProgram {
             }
             let expected_abi = match kernel.relation.as_str() {
                 "jolt.stage4.registers_read_write" => "jolt_stage4_registers_read_write",
+                "jolt.stage4.field_reg_rw" => "jolt_stage4_field_reg_rw",
                 "jolt.stage4.ram_val_check" => "jolt_stage4_ram_val_check",
                 "jolt.stage4.batched" => "jolt_stage4_batched",
                 _ => {
@@ -2087,6 +2088,9 @@ fn observe_stage4_sumcheck_output<F: Field>(
                 "stage4_registers_rw" => {
                     point = normalize_stage4_registers_rw_point(program, output.driver, &point)?;
                 }
+                "stage4_field_reg_rw" => {
+                    point = normalize_stage4_field_reg_rw_point(program, output.driver, &point)?;
+                }
                 _ => {
                     return Err(VerifyStage4Error::InvalidProof {
                         driver: output.driver,
@@ -2154,6 +2158,9 @@ fn expected_batched_output_claim(
             "jolt.stage4.registers_read_write" => {
                 expected_registers_read_write(store, evals, local_point)?
             }
+            "jolt.stage4.field_reg_rw" => {
+                expected_field_reg_rw(store, evals, local_point)?
+            }
             "jolt.stage4.ram_val_check" => {
                 expected_ram_val_check(store, evals, local_point)?
             }
@@ -2190,6 +2197,29 @@ fn expected_registers_read_write(
             + gamma * (rs1_ra * registers_val + gamma * rs2_ra * registers_val)))
 }
 
+fn expected_field_reg_rw(
+    store: &super::common::ValueStore<Fr>,
+    evals: &[Stage4NamedEval<Fr>],
+    local_point: &[Fr],
+) -> Result<Fr, VerifyStage4Error> {
+    let trace_point = super::common::store_point(store, "stage4.input.stage3.field_reg.FieldRdWriteValue")?;
+    let r_cycle = normalize_stage4_registers_rw_cycle_point(
+        local_point,
+        trace_point.len(),
+        "stage4.field_reg_rw.instance",
+    )?;
+    let eq_eval = EqPolynomial::<Fr>::mle(&r_cycle, trace_point);
+    let field_reg_val = eval_by_name(evals, "stage4.field_reg_rw.eval.FieldRegVal")?;
+    let frs1_ra = eval_by_name(evals, "stage4.field_reg_rw.eval.FrRs1Ra")?;
+    let frs2_ra = eval_by_name(evals, "stage4.field_reg_rw.eval.FrRs2Ra")?;
+    let frd_wa = eval_by_name(evals, "stage4.field_reg_rw.eval.FrdWa")?;
+    let frd_inc = eval_by_name(evals, "stage4.field_reg_rw.eval.FrdInc")?;
+    let gamma = super::common::store_scalar(store, "stage4.field_reg_rw.gamma")?;
+    Ok(eq_eval
+        * (frd_wa * (field_reg_val + frd_inc)
+            + gamma * (frs1_ra * field_reg_val + gamma * frs2_ra * field_reg_val)))
+}
+
 fn expected_ram_val_check(
     store: &super::common::ValueStore<Fr>,
     evals: &[Stage4NamedEval<Fr>],
@@ -2272,6 +2302,37 @@ fn normalize_stage4_registers_rw_cycle_point<F: Field>(
     Ok(cycle.iter().rev().copied().collect())
 }
 
+fn normalize_stage4_field_reg_rw_point<F: Field>(
+    program: &'static Stage4VerifierProgramPlan,
+    driver: &'static str,
+    point: &[F],
+) -> Result<Vec<F>, VerifyStage4Error> {
+    let driver_plan = find_plan(program.drivers, driver).ok_or(VerifyStage4Error::MissingProof {
+        driver,
+    })?;
+    if driver_plan.round_schedule.is_empty() {
+        return Err(VerifyStage4Error::InvalidProof {
+            driver,
+            reason: "stage4 field_reg point normalization requires non-empty schedule",
+        });
+    }
+    let cycle_rounds = driver_plan.round_schedule[0];
+    if point.len() < cycle_rounds {
+        return Err(VerifyStage4Error::InvalidInputLength {
+            input: "stage4.field_reg_rw.instance",
+            expected: cycle_rounds,
+            actual: point.len(),
+        });
+    }
+    let (cycle, address) = point.split_at(cycle_rounds);
+    Ok(address
+        .iter()
+        .rev()
+        .copied()
+        .chain(cycle.iter().rev().copied())
+        .collect())
+}
+
 "#
             }
         }

crates/bolt/src/protocols/jolt/emit/rust/stage5.rs

@@ -644,6 +644,7 @@ impl Stage5CpuProgram {
                 "jolt.stage5.instruction_read_raf" => "jolt_stage5_instruction_read_raf",
                 "jolt.stage5.ram_ra_claim_reduction" => "jolt_stage5_ram_ra_claim_reduction",
                 "jolt.stage5.registers_val_evaluation" => "jolt_stage5_registers_val_evaluation",
+                "jolt.stage5.field_reg_val_evaluation" => "jolt_stage5_field_reg_val_evaluation",
                 "jolt.stage5.batched" => "jolt_stage5_batched",
                 _ => {
                     return Err(EmitError::new(format!(
@@ -2162,6 +2163,9 @@ fn expected_batched_output_claim(
             "jolt.stage5.registers_val_evaluation" => {
                 expected_registers_val_evaluation(store, evals, local_point)?
             }
+            "jolt.stage5.field_reg_val_evaluation" => {
+                expected_field_reg_val_evaluation(store, evals, local_point)?
+            }
             _ => return Err(VerifyStage5Error::UnsupportedRelation { relation }),
         };
         expected += *coefficient * value;
@@ -2274,6 +2278,24 @@ fn expected_registers_val_evaluation(
     Ok(rd_inc * rd_wa * lt_eval)
 }
 
+fn expected_field_reg_val_evaluation(
+    store: &super::common::ValueStore<Fr>,
+    evals: &[Stage5NamedEval<Fr>],
+    local_point: &[Fr],
+) -> Result<Fr, VerifyStage5Error> {
+    let field_reg_val_point = super::common::store_point(store, "stage5.input.stage4.field_reg.FieldRegVal")?;
+    let r_cycle = suffix_point(
+        field_reg_val_point,
+        local_point.len(),
+        "stage5.input.stage4.field_reg.FieldRegVal",
+    )?;
+    let r_reduced = reverse_slice(local_point);
+    let lt_eval = lt_polynomial_eval(&r_reduced, r_cycle);
+    let frd_inc = eval_by_name(evals, "stage5.field_reg_val_evaluation.eval.FrdInc")?;
+    let frd_wa = eval_by_name(evals, "stage5.field_reg_val_evaluation.eval.FrdWa")?;
+    Ok(frd_inc * frd_wa * lt_eval)
+}
+
 "#
             }
         }

crates/bolt/src/protocols/jolt/emit/rust/stage6.rs

@@ -1243,7 +1243,7 @@ pub type Stage6VerifierProgramPlan = Stage6CpuProgramPlan;
 pub struct Stage6BytecodeEntry {
     pub address: Fr,
     pub imm: Fr,
-    pub circuit_flags: [bool; 14],
+    pub circuit_flags: [bool; 23],
     pub rd: Option<usize>,
     pub rs1: Option<usize>,
     pub rs2: Option<usize>,
@@ -1255,12 +1255,18 @@ pub struct Stage6BytecodeEntry {
     pub right_is_rs2: bool,
     pub right_is_imm: bool,
     pub is_noop: bool,
+    pub frd: Option<usize>,
+    pub frs1: Option<usize>,
+    pub frs2: Option<usize>,
+    pub reads_frs1: bool,
+    pub reads_frs2: bool,
+    pub writes_frd: bool,
 }
 
 impl Stage67BytecodeEntry for Stage6BytecodeEntry {
     fn address(&self) -> Fr { self.address }
     fn imm(&self) -> Fr { self.imm }
-    fn circuit_flags(&self) -> &[bool; 14] { &self.circuit_flags }
+    fn circuit_flags(&self) -> &[bool; 23] { &self.circuit_flags }
     fn rd(&self) -> Option<usize> { self.rd }
     fn rs1(&self) -> Option<usize> { self.rs1 }
     fn rs2(&self) -> Option<usize> { self.rs2 }
@@ -1272,6 +1278,12 @@ impl Stage67BytecodeEntry for Stage6BytecodeEntry {
     fn right_is_rs2(&self) -> bool { self.right_is_rs2 }
     fn right_is_imm(&self) -> bool { self.right_is_imm }
     fn is_noop(&self) -> bool { self.is_noop }
+    fn frd(&self) -> Option<usize> { self.frd }
+    fn frs1(&self) -> Option<usize> { self.frs1 }
+    fn frs2(&self) -> Option<usize> { self.frs2 }
+    fn reads_frs1(&self) -> bool { self.reads_frs1 }
+    fn reads_frs2(&self) -> bool { self.reads_frs2 }
+    fn writes_frd(&self) -> bool { self.writes_frd }
 }
 
 
@@ -1326,20 +1338,26 @@ const STAGE6_BYTECODE_SYMBOLS: Stage67BytecodeSymbols = Stage67BytecodeSymbols {
         "stage6.bytecode_read_raf.stage3_gamma",
         "stage6.bytecode_read_raf.stage4_gamma",
         "stage6.bytecode_read_raf.stage5_gamma",
+        "stage6.bytecode_read_raf.stage_fr_gamma",
     ],
     stage_cycle_points: [
         "stage6.input.stage1.Imm",
         "stage6.input.stage2.OpFlagJump",
         "stage6.input.stage3.spartan_shift.UnexpandedPC",
         "stage6.input.stage4.Rs1Ra",
         "stage6.input.stage5.registers_val_evaluation.RdWa",
+        "stage6.input.stage4.field_reg_rw.FrRs1Ra",
     ],
     stage4_register_point: "stage6.input.stage4.Rs1Ra",
     stage5_register_point: "stage6.input.stage5.registers_val_evaluation.RdWa",
+    stage_fr_register_point: "stage6.input.stage4.field_reg_rw.FrRs1Ra",
     entry_rd: "stage6.bytecode.entry.rd",
     entry_rs1: "stage6.bytecode.entry.rs1",
     entry_rs2: "stage6.bytecode.entry.rs2",
     entry_lookup_table: "stage6.bytecode.entry.lookup_table",
+    entry_frd: "stage6.bytecode.entry.frd",
+    entry_frs1: "stage6.bytecode.entry.frs1",
+    entry_frs2: "stage6.bytecode.entry.frs2",
 };
 
 #[derive(Debug)]