Implement ComponentProver for RangeCheckUnitComponent.

stwo_cairo_prover/src/components/range_check_unit.rs → stwo_cairo_prover/src/components/range_check_unit/component.rs

@@ -1,6 +1,5 @@
 use itertools::{zip_eq, Itertools};
 use num_traits::Zero;
-
 use stwo_prover::core::air::accumulation::PointEvaluationAccumulator;
 use stwo_prover::core::air::mask::fixed_mask_points;
 use stwo_prover::core::air::Component;
@@ -14,15 +13,15 @@ use stwo_prover::core::fields::FieldExpOps;
 use stwo_prover::core::pcs::TreeVec;
 use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use stwo_prover::core::poly::BitReversedOrder;
-use stwo_prover::core::LookupValues;
-use stwo_prover::core::{ColumnVec, InteractionElements};
+use stwo_prover::core::prover::{BASE_TRACE, INTERACTION_TRACE};
+use stwo_prover::core::{ColumnVec, InteractionElements, LookupValues};
 use stwo_prover::trace_generation::registry::ComponentGenerationRegistry;
-use stwo_prover::trace_generation::ComponentGen;
-use stwo_prover::trace_generation::ComponentTraceGenerator;
+use stwo_prover::trace_generation::{ComponentGen, ComponentTraceGenerator};
 
 pub const RC_Z: &str = "RangeCheckUnit_Z";
 pub const RC_COMPONENT_ID: &str = "RC_UNIT";
 
+#[derive(Clone)]
 pub struct RangeCheckUnitComponent {
     pub log_n_instances: u32,
 }
@@ -42,8 +41,9 @@ impl RangeCheckUnitComponent {
         constraint_zero_domain: Coset,
     ) {
         let z = interaction_elements[RC_Z];
-        let value = SecureField::from_partial_evals(std::array::from_fn(|i| mask[1][i][0]));
-        let numerator = value * (z - mask[0][0][0]) - mask[0][1][0];
+        let value =
+            SecureField::from_partial_evals(std::array::from_fn(|i| mask[INTERACTION_TRACE][i][0]));
+        let numerator = value * (z - mask[BASE_TRACE][0][0]) - mask[BASE_TRACE][1][0];
         let denom = point_vanishing(constraint_zero_domain.at(0), point);
         evaluation_accumulator.accumulate(numerator / denom);
     }
@@ -73,7 +73,10 @@ impl Component for RangeCheckUnitComponent {
         &self,
         point: CirclePoint<SecureField>,
     ) -> TreeVec<ColumnVec<Vec<CirclePoint<SecureField>>>> {
-        TreeVec::new(vec![fixed_mask_points(&vec![vec![0_usize]], point)])
+        TreeVec::new(vec![
+            fixed_mask_points(&vec![vec![0_usize]; 2], point),
+            vec![vec![point]; SECURE_EXTENSION_DEGREE],
+        ])
     }
 
     fn evaluate_constraint_quotients_at_point(
@@ -180,28 +183,12 @@ impl ComponentTraceGenerator<CpuBackend> for RangeCheckUnitTraceGenerator {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::components::range_check_unit::tests::register_test_rc;
 
     #[test]
     fn test_rc_unit_trace() {
         let mut registry = ComponentGenerationRegistry::default();
-        registry.register(RC_COMPONENT_ID, RangeCheckUnitTraceGenerator::new(8));
-        let inputs = vec![
-            vec![BaseField::from_u32_unchecked(0); 3],
-            vec![BaseField::from_u32_unchecked(1); 1],
-            vec![BaseField::from_u32_unchecked(2); 2],
-            vec![BaseField::from_u32_unchecked(3); 5],
-            vec![BaseField::from_u32_unchecked(4); 10],
-            vec![BaseField::from_u32_unchecked(5); 1],
-            vec![BaseField::from_u32_unchecked(6); 0],
-            vec![BaseField::from_u32_unchecked(7); 1],
-        ]
-        .into_iter()
-        .flatten()
-        .collect_vec();
-        registry
-            .get_generator_mut::<RangeCheckUnitTraceGenerator>(RC_COMPONENT_ID)
-            .add_inputs(&inputs);
-
+        register_test_rc(&mut registry);
         let trace = RangeCheckUnitTraceGenerator::write_trace(RC_COMPONENT_ID, &mut registry);
         let random_value = SecureField::from_u32_unchecked(1, 2, 3, 117);
         let interaction_elements =

stwo_cairo_prover/src/components/range_check_unit/component_prover.rs

@@ -0,0 +1,66 @@
+use stwo_prover::core::air::accumulation::{ColumnAccumulator, DomainEvaluationAccumulator};
+use stwo_prover::core::air::{Component, ComponentProver, ComponentTrace};
+use stwo_prover::core::backend::CpuBackend;
+use stwo_prover::core::circle::Coset;
+use stwo_prover::core::fields::m31::BaseField;
+use stwo_prover::core::fields::qm31::SecureField;
+use stwo_prover::core::pcs::TreeVec;
+use stwo_prover::core::poly::circle::{CanonicCoset, CircleDomain, CircleEvaluation};
+use stwo_prover::core::poly::BitReversedOrder;
+use stwo_prover::core::prover::{BASE_TRACE, INTERACTION_TRACE};
+use stwo_prover::core::utils::point_vanish_denominator_inverses;
+use stwo_prover::core::{InteractionElements, LookupValues};
+
+use super::component::{RangeCheckUnitComponent, RC_Z};
+
+impl RangeCheckUnitComponent {
+    fn evaluate_lookup_boundary_constraints(
+        &self,
+        trace_evals: &TreeVec<Vec<&CircleEvaluation<CpuBackend, BaseField, BitReversedOrder>>>,
+        trace_eval_domain: CircleDomain,
+        zero_domain: Coset,
+        accum: &mut ColumnAccumulator<'_, CpuBackend>,
+        interaction_elements: &InteractionElements,
+    ) {
+        let denom_inverses =
+            point_vanish_denominator_inverses(trace_eval_domain, zero_domain.at(0));
+        let z = interaction_elements[RC_Z];
+        for (i, denom_inverse) in denom_inverses.iter().enumerate() {
+            let value = SecureField::from_m31_array(std::array::from_fn(|j| {
+                trace_evals[INTERACTION_TRACE][j][i]
+            }));
+            let numerator = accum.random_coeff_powers[0]
+                * (value * (z - trace_evals[BASE_TRACE][0][i]) - trace_evals[BASE_TRACE][1][i]);
+            accum.accumulate(i, numerator * *denom_inverse);
+        }
+    }
+}
+
+impl ComponentProver<CpuBackend> for RangeCheckUnitComponent {
+    fn evaluate_constraint_quotients_on_domain(
+        &self,
+        trace: &ComponentTrace<'_, CpuBackend>,
+        evaluation_accumulator: &mut DomainEvaluationAccumulator<CpuBackend>,
+        interaction_elements: &InteractionElements,
+        _lookup_values: &LookupValues,
+    ) {
+        let max_constraint_degree = self.max_constraint_log_degree_bound();
+        let trace_eval_domain = CanonicCoset::new(max_constraint_degree).circle_domain();
+        let trace_evals = &trace.evals;
+        let zero_domain = CanonicCoset::new(self.log_n_instances).coset;
+        let [mut accum] =
+            evaluation_accumulator.columns([(max_constraint_degree, self.n_constraints())]);
+
+        self.evaluate_lookup_boundary_constraints(
+            trace_evals,
+            trace_eval_domain,
+            zero_domain,
+            &mut accum,
+            interaction_elements,
+        );
+    }
+
+    fn lookup_values(&self, _trace: &ComponentTrace<'_, CpuBackend>) -> LookupValues {
+        LookupValues::default()
+    }
+}