gen_trace supporting different operators

gizatechxyz · Jan 24, 2025 · 750ed2a · 750ed2a
1 parent 2b0b92f
commit 750ed2a
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 9 deletions.
diff --git a/crates/air/src/lib.rs b/crates/air/src/lib.rs
@@ -42,8 +42,15 @@ impl LuminairClaim {
     }
 
     pub fn log_sizes(&self) -> TreeVec<Vec<u32>> {
-        // Combine log sizes from all components, both adds and muls
-        let mut log_sizes = TreeVec::concat_cols(self.add.iter().map(|claim| claim.log_sizes()));
+        // Combine log sizes from all components
+        let mut log_sizes = TreeVec::concat_cols(
+            self.add
+                .iter()
+                .chain(
+                    [], // We will chain other ops here. E.g; `self.mul.iter()`
+                )
+                .map(|claim| claim.log_sizes()),
+        );
 
         // Overwrite preprocessed column claim
         log_sizes[PREPROCESSED_TRACE_IDX] = IS_FIRST_LOG_SIZES.to_vec();

diff --git a/crates/compiler/src/graph.rs b/crates/compiler/src/graph.rs
@@ -1,4 +1,7 @@
-use crate::{data::{OutputConverter, StwoData}, op::HasProcessTrace};
+use crate::{
+    data::{OutputConverter, StwoData},
+    op::HasProcessTrace,
+};
 use luminair_air::{
     components::add::trace::AddColumn,
     prover::{prover, verifier, LuminairProof},
@@ -30,9 +33,8 @@ impl LuminairGraph for Graph {
         let mut consumers = self.consumers_map.as_ref().unwrap().clone();
         let mut dim_stack = Vec::new();
 
-        // Store all add traces
+        // Initialize trace collectors for different operators
         let mut add_traces = Vec::new();
-        // Store all add claims
         let mut add_claims = Vec::new();
 
         for (node, src_ids) in self.linearized_graph.as_ref().unwrap() {
@@ -50,19 +52,28 @@ impl LuminairGraph for Graph {
 
             // Get operator and try to use process_trace if available
             let node_op = &mut *self.graph.node_weight_mut(*node).unwrap();
+
             let tensors =
                 if <Box<dyn Operator> as HasProcessTrace<AddColumn>>::has_process_trace(node_op) {
                     let (tensors, claim, trace) =
                         <Box<dyn Operator> as HasProcessTrace<AddColumn>>::call_process_trace(
                             node_op, srcs,
                         )
                         .unwrap();
-
                     add_traces.push(trace);
                     add_claims.push(claim);
-
                     tensors
-                } else {
+                }
+                // else if <Box<dyn Operator> as HasProcessTrace<MulColumn>>::has_process_trace(node_op) {
+                //     let (tensors, claim, trace) =
+                //         <Box<dyn Operator> as HasProcessTrace<MulColumn>>::call_process_trace(node_op, srcs)
+                //         .unwrap();
+                //     mul_traces.push(trace);
+                //     mul_claims.push(claim);
+                //     tensors
+                // }
+                else {
+                    // Handle other operators or fallback
                     node_op.process(srcs)
                 };
 
@@ -79,7 +90,7 @@ impl LuminairGraph for Graph {
         self.reset();
 
         let luminair_trace = LuminairTrace {
-            traces: add_traces,
+            traces: [add_traces].concat(),
             claims: LuminairClaim { add: add_claims },
         };