Merge pull request #19 from gizatechxyz/interaction-phase

Interaction phase

Cargo.toml

@@ -9,9 +9,11 @@ repository = "https://github.com/raphaelDkhn/Luminair"
 license = "MIT"
 
 [workspace.dependencies]
-stwo-prover = { git = "https://github.com/starkware-libs/stwo.git", rev = "36b17ac" }
+stwo-prover = {path = "/Users/raphaeldoukhan/Desktop/Cairo/stwo/crates/prover"}
+# stwo-prover = { git = "https://github.com/starkware-libs/stwo.git", rev = "36b17ac" }
 luminal = { git = "https://github.com/raphaelDkhn/luminal.git", rev = "9f13081" }
-numerair = { git = "https://github.com/gizatechxyz/NumerAir.git", rev = "3d9a9b9" }
+# numerair = { git = "https://github.com/gizatechxyz/NumerAir.git", rev = "3d9a9b9" }
+numerair = {path = "/Users/raphaeldoukhan/Desktop/numerair"}
 rayon = "1.10.0"
 tracing = "0.1"
 tracing-subscriber = "0.3"

crates/air/src/components/add/component.rs

@@ -1,7 +1,12 @@
+use num_traits::{One, Zero};
 use numerair::eval::EvalFixedPoint;
-use stwo_prover::constraint_framework::{EvalAtRow, FrameworkComponent, FrameworkEval};
+use stwo_prover::constraint_framework::{
+    EvalAtRow, FrameworkComponent, FrameworkEval, RelationEntry,
+};
 
-use crate::components::AddClaim;
+use crate::{components::AddClaim, pie::IOInfo};
+
+use super::table::AddElements;
 
 /// Implementation of `Component` and `ComponentProver` for the Add component.
 /// It targets the `SimdBackend` from the Stwo constraint framework, with the fallback
@@ -15,12 +20,18 @@ pub type AddComponent = FrameworkComponent<AddEval>;
 pub struct AddEval {
     /// The log size of the component's main trace height.
     log_size: u32,
+    /// The random elements used for the lookup protocol.
+    lookup_elements: AddElements,
+    /// Inputs/output information.
+    io_info: IOInfo,
 }
 
 impl AddEval {
-    pub const fn new(claim: &AddClaim) -> Self {
+    pub fn new(claim: &AddClaim, lookup_elements: AddElements) -> Self {
         Self {
             log_size: claim.log_size,
+            io_info: claim.io_info.clone(),
+            lookup_elements,
         }
     }
 }
@@ -55,9 +66,46 @@ impl FrameworkEval for AddEval {
         let lhs = eval.next_trace_mask();
         let rhs = eval.next_trace_mask();
         let out = eval.next_trace_mask();
+
+        eval.eval_fixed_add(lhs.clone(), rhs.clone(), out.clone());
+
+        let lhs_multiplicity = if self.io_info.inputs[0].is_initializer {
+            E::EF::zero()
+        } else {
+            -E::EF::one()
+        };
+
+        let rhs_multiplicity = if self.io_info.inputs[1].is_initializer {
+            E::EF::zero()
+        } else {
+            -E::EF::one()
+        };
+        let out_multiplicity = if self.io_info.output.is_final_output {
+            E::EF::zero()
+        } else {
+            E::EF::one()
+        };
+
+        eval.add_to_relation(RelationEntry::new(
+            &self.lookup_elements,
+            lhs_multiplicity,
+            &[lhs],
+        ));
+
+        eval.add_to_relation(RelationEntry::new(
+            &self.lookup_elements,
+            rhs_multiplicity,
+            &[rhs],
+        ));
 
-        eval.eval_fixed_add(lhs, rhs, out);
+        eval.add_to_relation(RelationEntry::new(
+            &self.lookup_elements,
+            out_multiplicity,
+            &[out],
+        ));
 
+        eval.finalize_logup();
+
         eval
     }
 }

crates/air/src/components/add/table.rs

@@ -1,21 +1,33 @@
-use num_traits::Zero;
+use num_traits::{One, Zero};
 use serde::{Deserialize, Serialize};
-use stwo_prover::core::{
-    backend::{
-        simd::{m31::PackedBaseField, SimdBackend},
-        Col, Column,
+use stwo_prover::{
+    constraint_framework::{logup::LogupTraceGenerator, Relation},
+    core::{
+        backend::{
+            simd::{
+                m31::{PackedBaseField, LOG_N_LANES},
+                qm31::PackedSecureField,
+                SimdBackend,
+            },
+            Col, Column,
+        },
+        fields::m31::BaseField,
+        poly::circle::{CanonicCoset, CircleEvaluation},
     },
-    fields::m31::BaseField,
-    poly::circle::{CanonicCoset, CircleEvaluation},
+    relation,
 };
 
-use crate::components::{AddClaim, TraceColumn, TraceEval};
+use crate::{
+    components::{AddClaim, InteractionClaim, TraceColumn, TraceError, TraceEval},
+    pie::IOInfo,
+};
 
 /// Generate trace for element-wise addition of two vectors.
 pub fn trace_evaluation(
     log_size: u32,
     lhs: &[PackedBaseField],
     rhs: &[PackedBaseField],
+    io_info: &IOInfo,
 ) -> (TraceEval, AddClaim, Vec<PackedBaseField>) {
     // Calculate trace size
     let trace_size = 1 << log_size;
@@ -25,7 +37,7 @@ pub fn trace_evaluation(
     let domain = CanonicCoset::new(log_size).circle_domain();
 
     // Initialize result vector
-    let mut result = Vec::with_capacity(3);
+    let mut main_trace = Vec::with_capacity(3);
 
     // Prepare output data
     let mut output = Vec::with_capacity(size);
@@ -36,30 +48,28 @@ pub fn trace_evaluation(
 
         for i in 0..trace_size {
             if i < size {
-                // Get values with broadcasting
                 let lhs_val = lhs[i % lhs.len()];
                 let rhs_val = rhs[i % rhs.len()];
                 let out_val = lhs_val + rhs_val;
 
-                // Set appropriate value based on column index
                 match column_idx {
                     0 => column.set(i, lhs_val.to_array()[0]),
                     1 => column.set(i, rhs_val.to_array()[0]),
-                    2 => column.set(i, out_val.to_array()[0]),
+                    2 => {
+                        column.set(i, out_val.to_array()[0]);
+                        output.push(out_val);
+                    }
                     _ => unreachable!(),
                 }
-
-                output.push(out_val)
             } else {
-                // Pad with zeros
                 column.set(i, BaseField::zero());
             }
         }
 
-        result.push(CircleEvaluation::new(domain, column));
+        main_trace.push(CircleEvaluation::new(domain, column));
     }
 
-    (result, AddClaim::new(log_size), output)
+    (main_trace, AddClaim::new(log_size, io_info.clone()), output)
 }
 
 /// Enum representing the column indices in the Add trace.
@@ -88,6 +98,72 @@ impl AddColumn {
 
 impl TraceColumn for AddColumn {
     fn count() -> (usize, usize) {
-        (3, 1)
+        (3, 3)
+    }
+}
+
+relation!(AddElements, 3);
+
+/// Creates the interaction trace from the main trace evaluation
+/// and the interaction elements for the Add component.
+pub fn interaction_trace_evaluation(
+    main_trace_eval: &TraceEval,
+    lookup_elements: &AddElements,
+    io_info: &IOInfo,
+) -> Result<(TraceEval, InteractionClaim), TraceError> {
+    if main_trace_eval.is_empty() {
+        return Err(TraceError::EmptyTrace);
+    }
+
+    let log_size = main_trace_eval[0].domain.log_size();
+    let mut logup_gen = LogupTraceGenerator::new(log_size);
+
+    // Create trace for LHS
+    let lhs_col = &main_trace_eval[AddColumn::Lhs.index()].data;
+    let mut col_lhs = logup_gen.new_col();
+    for row in 0..1 << (log_size - LOG_N_LANES) {
+        let lhs = lhs_col[row];
+        let multiplicity = if io_info.inputs[0].is_initializer {
+            PackedSecureField::zero()
+        } else {
+            -PackedSecureField::one()
+        };
+
+        col_lhs.write_frac(row, multiplicity, lookup_elements.combine(&[lhs]));
     }
+    col_lhs.finalize_col();
+
+    // Create trace for RHS
+    let rhs_col = &main_trace_eval[AddColumn::Rhs.index()].data;
+    let mut col_rhs = logup_gen.new_col();
+    for row in 0..1 << (log_size - LOG_N_LANES) {
+        let rhs = rhs_col[row];
+        let multiplicity = if io_info.inputs[1].is_initializer {
+            PackedSecureField::zero()
+        } else {
+            -PackedSecureField::one()
+        };
+
+        col_rhs.write_frac(row, multiplicity, lookup_elements.combine(&[rhs]));
+    }
+    col_rhs.finalize_col();
+
+    // Create trace for output
+    let out_col = &main_trace_eval[AddColumn::Out.index()].data;
+    let mut col_out = logup_gen.new_col();
+    for row in 0..1 << (log_size - LOG_N_LANES) {
+        let out = out_col[row];
+        let multiplicity = if io_info.output.is_final_output {
+            PackedSecureField::zero()
+        } else {
+            PackedSecureField::one()
+        };
+
+        col_out.write_frac(row, multiplicity, lookup_elements.combine(&[out]));
+    }
+    col_out.finalize_col();
+
+    let (trace, claimed_sum) = logup_gen.finalize_last();
+
+    Ok((trace, InteractionClaim { claimed_sum }))
 }