Remove code from lib.rs / mod.rs files (#451)

Author: az-starkware

crates/circuit_air/src/blake2s_consts.rs

@@ -0,0 +1,9 @@
+pub const BLAKE2S_IV: [u32; 8] = [
+    0x6A09E667, 0xBB67AE85, 0x3C6EF372, 0xA54FF53A, 0x510E527F, 0x9B05688C, 0x1F83D9AB, 0x5BE0CD19,
+];
+
+pub fn blake2s_initial_state() -> [u32; 8] {
+    let mut h = BLAKE2S_IV;
+    h[0] ^= 0x01010020;
+    h
+}

crates/circuit_air/src/circuit_claim.rs

@@ -0,0 +1,137 @@
+use crate::blake2s_consts::blake2s_initial_state;
+use crate::circuit_components::N_COMPONENTS;
+use crate::relations::{BLAKE_STATE_RELATION_ID, CommonLookupElements, GATE_RELATION_ID};
+use circuits::ivalue::qm31_from_u32s;
+use circuits_stark_verifier::proof_from_stark_proof::{pack_component_log_sizes, pack_enable_bits};
+use itertools::zip_eq;
+use num_traits::Zero;
+use stwo::core::channel::Channel;
+use stwo::core::fields::FieldExpOps;
+use stwo::core::fields::m31::M31;
+use stwo::core::fields::qm31::QM31;
+use stwo_constraint_framework::Relation;
+
+pub type ComponentLogSize = u32;
+pub type ClaimedSum = QM31;
+
+#[derive(Debug, PartialEq)]
+pub struct CircuitClaim {
+    pub log_sizes: [ComponentLogSize; N_COMPONENTS],
+    pub output_values: Vec<QM31>,
+}
+impl CircuitClaim {
+    pub fn mix_into(&self, channel: &mut impl Channel) {
+        let Self { log_sizes, output_values } = self;
+
+        // mix the number of components.
+        let n_components = log_sizes.len();
+        channel.mix_felts(&[qm31_from_u32s(n_components as u32, 0, 0, 0)]);
+
+        // mix the enable bits into the channel.
+        channel.mix_felts(&pack_enable_bits(&[true; N_COMPONENTS]));
+        channel.mix_felts(&pack_component_log_sizes(log_sizes));
+        // mix the output values into the channel.
+        channel.mix_felts(output_values);
+    }
+}
+
+pub struct CircuitInteractionElements {
+    pub common_lookup_elements: CommonLookupElements,
+}
+impl CircuitInteractionElements {
+    pub fn draw(channel: &mut impl Channel) -> CircuitInteractionElements {
+        CircuitInteractionElements { common_lookup_elements: CommonLookupElements::draw(channel) }
+    }
+}
+
+#[derive(Debug, PartialEq)]
+pub struct CircuitInteractionClaim {
+    pub claimed_sums: [ClaimedSum; N_COMPONENTS],
+}
+impl CircuitInteractionClaim {
+    pub fn mix_into(&self, channel: &mut impl Channel) {
+        let Self { claimed_sums } = self;
+        channel.mix_felts(claimed_sums);
+    }
+}
+
+fn blake_iv_public_logup_sum(
+    n_blake_gates: usize,
+    common_lookup_elements: &CommonLookupElements,
+) -> QM31 {
+    // Each Blake gate uses the initial state once and creates one row in blake_output.
+    // Then blake_output is padded to a power of two, and each padding row uses the
+    // initial state once. In total we have n_blake_gates.next_power_of_two() uses.
+    let initial_state_uses = n_blake_gates.next_power_of_two();
+
+    let state_relation_id = M31::from(BLAKE_STATE_RELATION_ID);
+    let initial_state = blake2s_initial_state();
+    let initial_state_limbs = [
+        M31::from(initial_state[0] & 0xffff),
+        M31::from((initial_state[0] >> 16) & 0xffff),
+        M31::from(initial_state[1] & 0xffff),
+        M31::from((initial_state[1] >> 16) & 0xffff),
+        M31::from(initial_state[2] & 0xffff),
+        M31::from((initial_state[2] >> 16) & 0xffff),
+        M31::from(initial_state[3] & 0xffff),
+        M31::from((initial_state[3] >> 16) & 0xffff),
+        M31::from(initial_state[4] & 0xffff),
+        M31::from((initial_state[4] >> 16) & 0xffff),
+        M31::from(initial_state[5] & 0xffff),
+        M31::from((initial_state[5] >> 16) & 0xffff),
+        M31::from(initial_state[6] & 0xffff),
+        M31::from((initial_state[6] >> 16) & 0xffff),
+        M31::from(initial_state[7] & 0xffff),
+        M31::from((initial_state[7] >> 16) & 0xffff),
+    ];
+
+    let limbs = [
+        state_relation_id,
+        M31::from(0u32),
+        initial_state_limbs[0],
+        initial_state_limbs[1],
+        initial_state_limbs[2],
+        initial_state_limbs[3],
+        initial_state_limbs[4],
+        initial_state_limbs[5],
+        initial_state_limbs[6],
+        initial_state_limbs[7],
+        initial_state_limbs[8],
+        initial_state_limbs[9],
+        initial_state_limbs[10],
+        initial_state_limbs[11],
+        initial_state_limbs[12],
+        initial_state_limbs[13],
+        initial_state_limbs[14],
+        initial_state_limbs[15],
+    ];
+    let denom: QM31 = common_lookup_elements.combine(&limbs);
+    denom.inverse() * M31::from(initial_state_uses)
+}
+
+pub fn lookup_sum(
+    claim: &CircuitClaim,
+    interaction_claim: &CircuitInteractionClaim,
+    interaction_elements: &CircuitInteractionElements,
+    output_addresses: &[usize],
+    n_blake_gates: usize,
+) -> QM31 {
+    let CircuitInteractionClaim { claimed_sums } = interaction_claim;
+    let component_sum: QM31 = claimed_sums.iter().sum();
+
+    // Compute the public logup sum from output gates.
+    let mut output_sum = QM31::zero();
+    let gate_relation_id = M31::from(GATE_RELATION_ID);
+    for (addr, value) in zip_eq(output_addresses, &claim.output_values) {
+        let values =
+            [gate_relation_id, M31::from(*addr as u32), value.0.0, value.0.1, value.1.0, value.1.1];
+        let denom: QM31 = interaction_elements.common_lookup_elements.combine(&values);
+        output_sum += denom.inverse();
+    }
+
+    // Subtract the blake IV public logup sum (blake IV state is used but never yielded).
+    let blake_iv_sum =
+        blake_iv_public_logup_sum(n_blake_gates, &interaction_elements.common_lookup_elements);
+
+    component_sum + output_sum - blake_iv_sum
+}

crates/circuit_air/src/circuit_components.rs

@@ -0,0 +1,234 @@
+use crate::circuit_claim::{CircuitClaim, CircuitInteractionClaim, CircuitInteractionElements};
+use crate::components::{
+    blake_g, blake_gate, blake_output, blake_round, blake_round_sigma, eq, qm31_ops,
+    range_check_15, range_check_16, triple_xor_32, verify_bitwise_xor_4, verify_bitwise_xor_7,
+    verify_bitwise_xor_8, verify_bitwise_xor_9, verify_bitwise_xor_12,
+};
+use stwo::core::air::Component;
+use stwo_constraint_framework::TraceLocationAllocator;
+use stwo_constraint_framework::preprocessed_columns::PreProcessedColumnId;
+
+macro_rules! define_component_list {
+    ($($variant:ident),* $(,)?) => {
+        pub enum ComponentList {
+            $($variant),*
+        }
+        pub const N_COMPONENTS: usize = [$(stringify!($variant)),*].len();
+    };
+}
+
+define_component_list! {
+    Eq,
+    Qm31Ops,
+    BlakeGate,
+    BlakeRound,
+    BlakeRoundSigma,
+    BlakeG,
+    BlakeOutput,
+    TripleXor32,
+    VerifyBitwiseXor8,
+    VerifyBitwiseXor12,
+    VerifyBitwiseXor4,
+    VerifyBitwiseXor7,
+    VerifyBitwiseXor9,
+    RangeCheck15,
+    RangeCheck16,
+}
+
+pub struct CircuitComponents {
+    pub eq: eq::Component,
+    pub qm31_ops: qm31_ops::Component,
+    pub blake_gate: blake_gate::Component,
+    pub blake_round: blake_round::Component,
+    pub blake_round_sigma: blake_round_sigma::Component,
+    pub blake_g: blake_g::Component,
+    pub blake_output: blake_output::Component,
+    pub triple_xor_32: triple_xor_32::Component,
+    pub verify_bitwise_xor_8: verify_bitwise_xor_8::Component,
+    pub verify_bitwise_xor_12: verify_bitwise_xor_12::Component,
+    pub verify_bitwise_xor_4: verify_bitwise_xor_4::Component,
+    pub verify_bitwise_xor_7: verify_bitwise_xor_7::Component,
+    pub verify_bitwise_xor_9: verify_bitwise_xor_9::Component,
+    pub range_check_15: range_check_15::Component,
+    pub range_check_16: range_check_16::Component,
+}
+impl CircuitComponents {
+    pub fn new(
+        circuit_claim: &CircuitClaim,
+        interaction_elements: &CircuitInteractionElements,
+        interaction_claim: &CircuitInteractionClaim,
+        // Describes the structure of the preprocessed trace. Sensitive to order.
+        preprocessed_column_ids: &[PreProcessedColumnId],
+    ) -> Self {
+        let tree_span_provider =
+            &mut TraceLocationAllocator::new_with_preprocessed_columns(preprocessed_column_ids);
+
+        let eq_component = eq::Component::new(
+            tree_span_provider,
+            eq::Eval {
+                log_size: circuit_claim.log_sizes[ComponentList::Eq as usize],
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::Eq as usize],
+        );
+        let qm31_ops_component = qm31_ops::Component::new(
+            tree_span_provider,
+            qm31_ops::Eval {
+                log_size: circuit_claim.log_sizes[ComponentList::Qm31Ops as usize],
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::Qm31Ops as usize],
+        );
+        let blake_gate_component = blake_gate::Component::new(
+            tree_span_provider,
+            blake_gate::Eval {
+                claim: blake_gate::Claim {
+                    log_size: circuit_claim.log_sizes[ComponentList::BlakeGate as usize],
+                },
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::BlakeGate as usize],
+        );
+        let blake_round_component = blake_round::Component::new(
+            tree_span_provider,
+            blake_round::Eval {
+                claim: blake_round::Claim {
+                    log_size: circuit_claim.log_sizes[ComponentList::BlakeRound as usize],
+                },
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::BlakeRound as usize],
+        );
+        let blake_round_sigma_component = blake_round_sigma::Component::new(
+            tree_span_provider,
+            blake_round_sigma::Eval {
+                claim: blake_round_sigma::Claim {},
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::BlakeRoundSigma as usize],
+        );
+        let blake_g_component = blake_g::Component::new(
+            tree_span_provider,
+            blake_g::Eval {
+                claim: blake_g::Claim {
+                    log_size: circuit_claim.log_sizes[ComponentList::BlakeG as usize],
+                },
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::BlakeG as usize],
+        );
+        let blake_output_component = blake_output::Component::new(
+            tree_span_provider,
+            blake_output::Eval {
+                claim: blake_output::Claim {
+                    log_size: circuit_claim.log_sizes[ComponentList::BlakeOutput as usize],
+                },
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::BlakeOutput as usize],
+        );
+        let triple_xor_32_component = triple_xor_32::Component::new(
+            tree_span_provider,
+            triple_xor_32::Eval {
+                claim: triple_xor_32::Claim {
+                    log_size: circuit_claim.log_sizes[ComponentList::TripleXor32 as usize],
+                },
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::TripleXor32 as usize],
+        );
+        let verify_bitwise_xor_8_component = verify_bitwise_xor_8::Component::new(
+            tree_span_provider,
+            verify_bitwise_xor_8::Eval {
+                claim: verify_bitwise_xor_8::Claim {},
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::VerifyBitwiseXor8 as usize],
+        );
+        let verify_bitwise_xor_12_component = verify_bitwise_xor_12::Component::new(
+            tree_span_provider,
+            verify_bitwise_xor_12::Eval {
+                claim: verify_bitwise_xor_12::Claim {},
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::VerifyBitwiseXor12 as usize],
+        );
+        let verify_bitwise_xor_4_component = verify_bitwise_xor_4::Component::new(
+            tree_span_provider,
+            verify_bitwise_xor_4::Eval {
+                claim: verify_bitwise_xor_4::Claim {},
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::VerifyBitwiseXor4 as usize],
+        );
+        let verify_bitwise_xor_7_component = verify_bitwise_xor_7::Component::new(
+            tree_span_provider,
+            verify_bitwise_xor_7::Eval {
+                claim: verify_bitwise_xor_7::Claim {},
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::VerifyBitwiseXor7 as usize],
+        );
+        let verify_bitwise_xor_9_component = verify_bitwise_xor_9::Component::new(
+            tree_span_provider,
+            verify_bitwise_xor_9::Eval {
+                claim: verify_bitwise_xor_9::Claim {},
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::VerifyBitwiseXor9 as usize],
+        );
+        let range_check_15_component = range_check_15::Component::new(
+            tree_span_provider,
+            range_check_15::Eval {
+                claim: range_check_15::Claim {},
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::RangeCheck15 as usize],
+        );
+        let range_check_16_component = range_check_16::Component::new(
+            tree_span_provider,
+            range_check_16::Eval {
+                claim: range_check_16::Claim {},
+                common_lookup_elements: interaction_elements.common_lookup_elements.clone(),
+            },
+            interaction_claim.claimed_sums[ComponentList::RangeCheck16 as usize],
+        );
+        Self {
+            eq: eq_component,
+            qm31_ops: qm31_ops_component,
+            blake_gate: blake_gate_component,
+            blake_round: blake_round_component,
+            blake_round_sigma: blake_round_sigma_component,
+            blake_g: blake_g_component,
+            blake_output: blake_output_component,
+            triple_xor_32: triple_xor_32_component,
+            verify_bitwise_xor_8: verify_bitwise_xor_8_component,
+            verify_bitwise_xor_12: verify_bitwise_xor_12_component,
+            verify_bitwise_xor_4: verify_bitwise_xor_4_component,
+            verify_bitwise_xor_7: verify_bitwise_xor_7_component,
+            verify_bitwise_xor_9: verify_bitwise_xor_9_component,
+            range_check_15: range_check_15_component,
+            range_check_16: range_check_16_component,
+        }
+    }
+
+    pub fn components(self) -> Vec<Box<dyn Component>> {
+        vec![
+            Box::new(self.eq) as Box<dyn Component>,
+            Box::new(self.qm31_ops) as Box<dyn Component>,
+            Box::new(self.blake_gate) as Box<dyn Component>,
+            Box::new(self.blake_round) as Box<dyn Component>,
+            Box::new(self.blake_round_sigma) as Box<dyn Component>,
+            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.verify_bitwise_xor_8) as Box<dyn Component>,
+            Box::new(self.verify_bitwise_xor_12) as Box<dyn Component>,
+            Box::new(self.verify_bitwise_xor_4) as Box<dyn Component>,
+            Box::new(self.verify_bitwise_xor_7) as Box<dyn Component>,
+            Box::new(self.verify_bitwise_xor_9) as Box<dyn Component>,
+            Box::new(self.range_check_15) as Box<dyn Component>,
+            Box::new(self.range_check_16) as Box<dyn Component>,
+        ]
+    }
+}