Add BlakeGGate AIR component with full prove+verify

alon-f · claude · alon-f · commit 54289261bd69 · 2026-03-24T16:34:45.000+02:00
Integrates the auto-generated BlakeGGate into the circuit prover:
- BlakeGGate added to ComponentList, CircuitComponents, trace writing,
  interaction trace, statement, and prover
- Preprocessed column IDs renamed to match the AIR
  (blake_g_gate_input_addr_a/b/c/d/f0/f1, blake_g_gate_output_addr_a/b/c/d,
  blake_g_gate_multiplicity)
- Padding for blake_g gates in finalize_context
- Full prove+verify test (test_prove_and_stark_verify_blake_g_gate_context)
- Updated FIBONACCI_CIRCUIT_PREPROCESSED_ROOT for N_COMPONENTS=18

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/crates/circuit_air/src/components/mod.rs b/crates/circuit_air/src/components/mod.rs
@@ -42,6 +42,7 @@ define_component_list! {
     BlakeG,
     BlakeOutput,
     TripleXor32,
+    BlakeGGate,
     M31ToU32,
     TripleXor,
     VerifyBitwiseXor8,
@@ -62,6 +63,7 @@ pub struct CircuitComponents {
     pub blake_g: blake_g::Component,
     pub blake_output: blake_output::Component,
     pub triple_xor_32: triple_xor_32::Component,
+    pub blake_g_gate: blake_g_gate::Component,
     pub m_31_to_u_32: m_31_to_u_32::Component,
     pub triple_xor: triple_xor::Component,
     pub verify_bitwise_xor_8: verify_bitwise_xor_8::Component,
@@ -157,6 +159,16 @@ impl CircuitComponents {
             },
             interaction_claim.claimed_sums[ComponentList::TripleXor32 as usize],
         );
+        let blake_g_gate_component = blake_g_gate::Component::new(
+            tree_span_provider,
+            blake_g_gate::Eval {
+                claim: blake_g_gate::Claim {
+                    log_size: circuit_claim.log_sizes[ComponentList::BlakeGGate as usize],
+                },
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::BlakeGGate as usize],
+        );
         let m_31_to_u_32_component = m_31_to_u_32::Component::new(
             tree_span_provider,
             m_31_to_u_32::Eval {
@@ -242,6 +254,7 @@ impl CircuitComponents {
             blake_g: blake_g_component,
             blake_output: blake_output_component,
             triple_xor_32: triple_xor_32_component,
+            blake_g_gate: blake_g_gate_component,
             m_31_to_u_32: m_31_to_u_32_component,
             triple_xor: triple_xor_component,
             verify_bitwise_xor_8: verify_bitwise_xor_8_component,
@@ -264,6 +277,7 @@ impl CircuitComponents {
             Box::new(self.blake_g) as Box<dyn Component>,
             Box::new(self.blake_output) as Box<dyn Component>,
             Box::new(self.triple_xor_32) as Box<dyn Component>,
+            Box::new(self.blake_g_gate) as Box<dyn Component>,
             Box::new(self.m_31_to_u_32) as Box<dyn Component>,
             Box::new(self.triple_xor) as Box<dyn Component>,
             Box::new(self.verify_bitwise_xor_8) as Box<dyn Component>,
diff --git a/crates/circuit_air/src/statement.rs b/crates/circuit_air/src/statement.rs
@@ -1,5 +1,5 @@
 use crate::circuit_eval_components::{
-    blake_g, blake_gate, blake_output, blake_round, blake_round_sigma, m_31_to_u_32,
+    blake_g, blake_g_gate, blake_gate, blake_output, blake_round, blake_round_sigma, m_31_to_u_32,
     range_check_15, range_check_16, triple_xor, triple_xor_32, verify_bitwise_xor_4,
     verify_bitwise_xor_7, verify_bitwise_xor_8, verify_bitwise_xor_9, verify_bitwise_xor_12,
 };
@@ -152,6 +152,7 @@ pub fn all_circuit_components<Value: IValue>() -> Vec<Box<dyn CircuitEval<Value>
         Box::new(blake_g::Component {}),
         Box::new(blake_output::Component {}),
         Box::new(triple_xor_32::Component {}),
+        Box::new(blake_g_gate::Component {}),
         Box::new(m_31_to_u_32::Component {}),
         Box::new(triple_xor::Component {}),
         Box::new(verify_bitwise_xor_8::Component {}),
diff --git a/crates/circuit_common/src/finalize.rs b/crates/circuit_common/src/finalize.rs
@@ -1,6 +1,6 @@
 use crate::N_LANES;
 use circuits::blake::{HashValue, blake};
-use circuits::circuit::{M31ToU32Gate, TripleXorGate};
+use circuits::circuit::{BlakeGGate, M31ToU32Gate, TripleXorGate};
 use circuits::context::{Context, Var};
 use circuits::eval;
 use circuits::ivalue::{IValue, qm31_from_u32s};
@@ -54,6 +54,30 @@ fn pad_blake(context: &mut Context<impl IValue>) {
     }
 }
 
+fn pad_blake_g(context: &mut Context<impl IValue>) {
+    let n_rows = context.circuit.blake_g.len();
+    let padded = std::cmp::max(n_rows.next_power_of_two(), N_LANES);
+    let zero = context.zero();
+    for _ in n_rows..padded {
+        let out_a = context.new_var(IValue::from_qm31(0.into()));
+        let out_b = context.new_var(IValue::from_qm31(0.into()));
+        let out_c = context.new_var(IValue::from_qm31(0.into()));
+        let out_d = context.new_var(IValue::from_qm31(0.into()));
+        context.circuit.blake_g.push(BlakeGGate {
+            a: zero.idx,
+            b: zero.idx,
+            c: zero.idx,
+            d: zero.idx,
+            m0: zero.idx,
+            m1: zero.idx,
+            out_a: out_a.idx,
+            out_b: out_b.idx,
+            out_c: out_c.idx,
+            out_d: out_d.idx,
+        });
+    }
+}
+
 fn pad_m31_to_u32(context: &mut Context<impl IValue>) {
     let n_rows = context.circuit.m31_to_u32.len();
     let padded = std::cmp::max(n_rows.next_power_of_two(), N_LANES);
@@ -102,6 +126,7 @@ pub fn finalize_context(context: &mut Context<impl IValue>) {
     pad_eq(context);
     pad_qm31_ops(context);
     pad_blake(context);
+    pad_blake_g(context);
     pad_m31_to_u32(context);
     pad_triple_xor(context);
 }
diff --git a/crates/circuit_common/src/preprocessed.rs b/crates/circuit_common/src/preprocessed.rs
@@ -459,17 +459,17 @@ fn add_blake_g_to_preprocessed_trace(
     fill_blake_g_columns(&circuit.blake_g, multiplicities, &mut columns);
 
     let ids = [
-        "blake_g_a_address",
-        "blake_g_b_address",
-        "blake_g_c_address",
-        "blake_g_d_address",
-        "blake_g_m0_address",
-        "blake_g_m1_address",
-        "blake_g_out_a_address",
-        "blake_g_out_b_address",
-        "blake_g_out_c_address",
-        "blake_g_out_d_address",
-        "blake_g_mult",
+        "blake_g_gate_input_addr_a",
+        "blake_g_gate_input_addr_b",
+        "blake_g_gate_input_addr_c",
+        "blake_g_gate_input_addr_d",
+        "blake_g_gate_input_addr_f0",
+        "blake_g_gate_input_addr_f1",
+        "blake_g_gate_output_addr_a",
+        "blake_g_gate_output_addr_b",
+        "blake_g_gate_output_addr_c",
+        "blake_g_gate_output_addr_d",
+        "blake_g_gate_multiplicity",
     ];
     for (id, column) in zip_eq(ids, columns) {
         pp_trace.push_column(PreProcessedColumnId { id: id.to_owned() }, column);
diff --git a/crates/circuit_common/src/preprocessed_test.rs b/crates/circuit_common/src/preprocessed_test.rs
@@ -196,10 +196,12 @@ fn test_preprocess_decomposed_gates() {
     assert!(m31_mult.iter().all(|&m| m > 0));
 
     // Verify BlakeG columns exist and have correct addresses.
-    let bg_a = pp_trace.get_column(&PreProcessedColumnId { id: "blake_g_a_address".to_owned() });
-    let bg_out_a =
-        pp_trace.get_column(&PreProcessedColumnId { id: "blake_g_out_a_address".to_owned() });
-    let bg_mult = pp_trace.get_column(&PreProcessedColumnId { id: "blake_g_mult".to_owned() });
+    let bg_a = pp_trace
+        .get_column(&PreProcessedColumnId { id: "blake_g_gate_input_addr_a".to_owned() });
+    let bg_out_a = pp_trace
+        .get_column(&PreProcessedColumnId { id: "blake_g_gate_output_addr_a".to_owned() });
+    let bg_mult = pp_trace
+        .get_column(&PreProcessedColumnId { id: "blake_g_gate_multiplicity".to_owned() });
     assert_eq!(bg_a.len(), 16);
     assert_eq!(bg_a[0], 16); // First gate's a = wire 16.
     assert_eq!(bg_out_a[0], 32); // First gate's out_a = wire 32.
diff --git a/crates/circuit_prover/src/prover.rs b/crates/circuit_prover/src/prover.rs
@@ -64,6 +64,7 @@ pub fn to_component_provers(
         &components.blake_g as &dyn ComponentProver<SimdBackend>,
         &components.blake_output as &dyn ComponentProver<SimdBackend>,
         &components.triple_xor_32 as &dyn ComponentProver<SimdBackend>,
+        &components.blake_g_gate as &dyn ComponentProver<SimdBackend>,
         &components.m_31_to_u_32 as &dyn ComponentProver<SimdBackend>,
         &components.triple_xor as &dyn ComponentProver<SimdBackend>,
         &components.verify_bitwise_xor_8 as &dyn ComponentProver<SimdBackend>,
diff --git a/crates/circuit_prover/src/prover_test.rs b/crates/circuit_prover/src/prover_test.rs
@@ -6,7 +6,7 @@ use circuit_air::statement::{INTERACTION_POW_BITS, all_circuit_components};
 use circuit_air::verify::{CircuitConfig, verify_circuit};
 use circuit_common::finalize::finalize_context;
 use circuit_common::preprocessed::PreprocessedCircuit;
-use circuits::blake::{blake, m31_to_u32_gate, pack_u32, triple_xor_gate};
+use circuits::blake::{blake, blake_g_gate as blake_g_gate_fn, m31_to_u32_gate, pack_u32, triple_xor_gate};
 use circuits::context::Var;
 use circuits::eval;
 use circuits::ivalue::{IValue, qm31_from_u32s};
@@ -278,7 +278,7 @@ fn test_prove_and_stark_verify_fibonacci_context() {
 }
 
 const FIBONACCI_CIRCUIT_PREPROCESSED_ROOT: [u32; 8] =
-    [819694180, 1334134931, 1528443241, 180931882, 500090799, 733233906, 1222299448, 1633395937];
+    [623845414, 632928349, 1241137236, 579888351, 1255246133, 1278652435, 1440752206, 143572365];
 
 #[test]
 fn test_prove_and_circuit_verify_fibonacci_context() {
@@ -477,6 +477,90 @@ fn test_prove_and_stark_verify_triple_xor_context() {
     );
 }
 
+pub fn build_blake_g_gate_context() -> Context<QM31> {
+    use circuits::wrappers::U32Wrapper;
+    let mut context = Context::<QM31>::default();
+
+    // Create 32 blake_g gates (> N_LANES=16 so the trace is non-trivial).
+    for i in 0..32u32 {
+        let a = U32Wrapper::new_unsafe(guess(&mut context, pack_u32(i * 10 + 1)));
+        let b = U32Wrapper::new_unsafe(guess(&mut context, pack_u32(i * 10 + 2)));
+        let c = U32Wrapper::new_unsafe(guess(&mut context, pack_u32(i * 10 + 3)));
+        let d = U32Wrapper::new_unsafe(guess(&mut context, pack_u32(i * 10 + 4)));
+        let m0 = U32Wrapper::new_unsafe(guess(&mut context, pack_u32(i * 10 + 5)));
+        let m1 = U32Wrapper::new_unsafe(guess(&mut context, pack_u32(i * 10 + 6)));
+        let (_out_a, _out_b, _out_c, _out_d) =
+            blake_g_gate_fn(&mut context, a, b, c, d, m0, m1);
+    }
+    let last = guess(&mut context, QM31::zero());
+    output(&mut context, last);
+    context
+}
+
+#[test]
+fn test_prove_and_stark_verify_blake_g_gate_context() {
+    let mut context = build_blake_g_gate_context();
+    context.finalize_guessed_vars();
+    context.validate_circuit();
+
+    let preprocessed_circuit = PreprocessedCircuit::preprocess_circuit(&mut context);
+    let CircuitProof {
+        pcs_config,
+        claim,
+        interaction_pow_nonce,
+        interaction_claim,
+        components,
+        stark_proof,
+        channel_salt,
+    } = prove_circuit_assignment(
+        context.values(),
+        &preprocessed_circuit,
+        &BaseColumnPool::<SimdBackend>::new(),
+    );
+    assert!(stark_proof.is_ok(), "Got error: {}", stark_proof.err().unwrap());
+    let proof = stark_proof.unwrap();
+
+    let verifier_channel = &mut Blake2sM31Channel::default();
+    verifier_channel.mix_felts(&[channel_salt.into()]);
+    pcs_config.mix_into(verifier_channel);
+    let commitment_scheme =
+        &mut CommitmentSchemeVerifier::<Blake2sM31MerkleChannel>::new(pcs_config);
+
+    let sizes = TreeVec::concat_cols(components.iter().map(|c| c.trace_log_degree_bounds()));
+
+    commitment_scheme.commit(
+        proof.proof.commitments[0],
+        &preprocessed_circuit.preprocessed_trace.log_sizes(),
+        verifier_channel,
+    );
+    claim.mix_into(verifier_channel);
+    commitment_scheme.commit(proof.proof.commitments[1], &sizes[1], verifier_channel);
+    verifier_channel.verify_pow_nonce(INTERACTION_POW_BITS, interaction_pow_nonce);
+    verifier_channel.mix_u64(interaction_pow_nonce);
+    let interaction_elements = CircuitInteractionElements::draw(verifier_channel);
+    interaction_claim.mix_into(verifier_channel);
+    commitment_scheme.commit(proof.proof.commitments[2], &sizes[2], verifier_channel);
+    stwo::core::verifier::verify_ex(
+        &components.iter().map(|c| c.as_ref()).collect::<Vec<&dyn Component>>(),
+        verifier_channel,
+        commitment_scheme,
+        proof.proof,
+        true,
+    )
+    .unwrap();
+
+    assert_eq!(
+        lookup_sum(
+            &claim,
+            &interaction_claim,
+            &interaction_elements,
+            &preprocessed_circuit.params.output_addresses,
+            preprocessed_circuit.params.n_blake_gates
+        ),
+        QM31::zero()
+    );
+}
+
 #[test]
 fn test_finalize_context() {
     let mut context = build_fibonacci_context();
diff --git a/crates/circuit_prover/src/witness/trace.rs b/crates/circuit_prover/src/witness/trace.rs
@@ -10,6 +10,7 @@ use crate::witness::components::qm31_ops;
 use crate::witness::components::range_check_15;
 use crate::witness::components::range_check_16;
 use crate::witness::components::triple_xor_32;
+use crate::witness::components::blake_g_gate;
 use crate::witness::components::triple_xor;
 use crate::witness::components::m_31_to_u_32;
 use crate::witness::components::verify_bitwise_xor_4;
@@ -197,6 +198,19 @@ pub fn write_trace(
     trace_evals.extend(blake_output_trace.to_evals());
     trace_evals.extend(triple_xor_32_trace.to_evals());
 
+    // Write blake_g_gate component.
+    let (blake_g_gate_trace, blake_g_gate_claim, blake_g_gate_interaction_claim_gen) =
+        blake_g_gate::write_trace(
+            context_values,
+            preprocessed_trace_ref,
+            &verify_bitwise_xor_8_state,
+            &verify_bitwise_xor_12_state,
+            &verify_bitwise_xor_4_state,
+            &verify_bitwise_xor_9_state,
+            &verify_bitwise_xor_7_state,
+        );
+    trace_evals.extend(blake_g_gate_trace.to_evals());
+
     // Write m31_to_u32 component.
     let (m_31_to_u_32_trace, m_31_to_u_32_claim, m_31_to_u_32_interaction_claim_gen) =
         m_31_to_u_32::write_trace(context_values, preprocessed_trace_ref, &range_check_16_state);
@@ -295,6 +309,7 @@ pub fn write_trace(
                 blake_g_claim.log_size,
                 blake_output_claim.log_size,
                 triple_xor_32_claim.log_size,
+                blake_g_gate_claim.log_size,
                 m_31_to_u_32_claim.log_size,
                 triple_xor_claim.log_size,
                 circuit_air::components::verify_bitwise_xor_8::LOG_SIZE,
@@ -316,6 +331,7 @@ pub fn write_trace(
             blake_g: blake_g_interaction_claim_gen,
             blake_output: blake_output_interaction_claim_gen,
             triple_xor_32: triple_xor_32_interaction_claim_gen,
+            blake_g_gate: blake_g_gate_interaction_claim_gen,
             m_31_to_u_32: m_31_to_u_32_interaction_claim_gen,
             triple_xor: triple_xor_interaction_claim_gen,
             verify_bitwise_xor_8: verify_bitwise_xor_8_interaction_claim_gen,
@@ -338,6 +354,7 @@ pub struct CircuitInteractionClaimGenerator {
     pub blake_g: blake_g::InteractionClaimGenerator,
     pub blake_output: blake_output::InteractionClaimGenerator,
     pub triple_xor_32: triple_xor_32::InteractionClaimGenerator,
+    pub blake_g_gate: blake_g_gate::InteractionClaimGenerator,
     pub m_31_to_u_32: m_31_to_u_32::InteractionClaimGenerator,
     pub triple_xor: triple_xor::InteractionClaimGenerator,
     pub verify_bitwise_xor_8: verify_bitwise_xor_8::InteractionClaimGenerator,
@@ -409,6 +426,12 @@ pub fn write_interaction_trace(
             .write_interaction_trace(&interaction_elements.common_lookup_elements);
     tree_builder.extend_evals(triple_xor_32_trace);
 
+    // Write blake_g_gate interaction trace.
+    let (blake_g_gate_trace, blake_g_gate_interaction_claim) = circuit_interaction_claim_generator
+        .blake_g_gate
+        .write_interaction_trace(&interaction_elements.common_lookup_elements);
+    tree_builder.extend_evals(blake_g_gate_trace);
+
     // Write m31_to_u32 interaction trace.
     let (m_31_to_u_32_trace, m_31_to_u_32_interaction_claim) = circuit_interaction_claim_generator
         .m_31_to_u_32
@@ -478,6 +501,7 @@ pub fn write_interaction_trace(
             blake_g_interaction_claim.claimed_sum,
             blake_output_interaction_claim.claimed_sum,
             triple_xor_32_interaction_claim.claimed_sum,
+            blake_g_gate_interaction_claim.claimed_sum,
             m_31_to_u_32_interaction_claim.claimed_sum,
             triple_xor_interaction_claim.claimed_sum,
             verify_bitwise_xor_8_interaction_claim.claimed_sum,
