fix: fixed codegen

Author: shreyas-londhe

compiler/src/nfa/codegen/circom.rs

@@ -14,7 +14,6 @@
 //! - Start/accept state validation
 
 use serde::Serialize;
-use std::collections::BTreeSet;
 
 use crate::nfa::NFAGraph;
 use crate::nfa::error::{NFAError, NFAResult};
@@ -35,10 +34,10 @@ pub struct CircomInputs {
     pub next_states: Vec<usize>,
     #[serde(skip_serializing_if = "Option::is_none")]
     #[serde(rename = "captureGroupIds")]
-    pub capture_group_ids: Option<Vec<usize>>,
+    pub capture_group_ids: Option<Vec<Vec<usize>>>,
     #[serde(skip_serializing_if = "Option::is_none")]
     #[serde(rename = "captureGroupStarts")]
-    pub capture_group_starts: Option<Vec<u8>>,
+    pub capture_group_starts: Option<Vec<Vec<u8>>>,
     #[serde(skip_serializing_if = "Option::is_none")]
     #[serde(rename = "captureGroupStartIndices")]
     pub capture_group_start_indices: Option<Vec<usize>>,
@@ -85,18 +84,12 @@ pub fn generate_circom_code(
 
     let (start_states, accept_states, transitions) = generate_circuit_data(nfa)?;
 
-    // Validate capture groups
-    let capture_group_set: BTreeSet<_> = transitions
-        .iter()
-        .filter_map(|(_, _, _, _, cap)| cap.map(|(id, _)| id))
-        .collect();
-
-    if !capture_group_set.is_empty() {
+    if nfa.num_capture_groups > 0 {
         if let Some(max_bytes) = max_substring_bytes.as_ref() {
-            if max_bytes.len() < capture_group_set.len() {
+            if max_bytes.len() != nfa.num_capture_groups {
                 return Err(NFAError::InvalidCapture(format!(
                     "Insufficient max_substring_bytes: need {} but got {}",
-                    capture_group_set.len(),
+                    nfa.num_capture_groups,
                     max_bytes.len()
                 )));
             }
@@ -114,8 +107,6 @@ pub fn generate_circom_code(
         }
     }
 
-    let has_capture_groups = !capture_group_set.is_empty();
-
     let mut code = String::new();
 
     code.push_str("pragma circom 2.1.5;\n\n");
@@ -143,9 +134,21 @@ pub fn generate_circom_code(
     code.push_str("    signal input nextStates[maxMatchBytes];\n");
 
     // Only add capture group signals if needed
-    if has_capture_groups {
-        code.push_str("    signal input captureGroupIds[maxMatchBytes];\n");
-        code.push_str("    signal input captureGroupStarts[maxMatchBytes];\n\n");
+    if nfa.num_capture_groups > 0 {
+        for i in 0..nfa.num_capture_groups {
+            code.push_str(
+                format!("    signal input captureGroup{}Id[maxMatchBytes];\n", i + 1).as_str(),
+            );
+        }
+        for i in 0..nfa.num_capture_groups {
+            code.push_str(
+                format!(
+                    "    signal input captureGroup{}Start[maxMatchBytes];\n",
+                    i + 1
+                )
+                .as_str(),
+            );
+        }
     }
 
     code.push_str("    signal output isValid;\n\n");
@@ -234,63 +237,103 @@ pub fn generate_circom_code(
     code.push_str("            isTransitionLinked[i] === isWithinPathLengthMinusOne[i];\n");
     code.push_str("        }\n\n");
 
-    if has_capture_groups {
+    if nfa.num_capture_groups > 0 {
+        // Prepare strings for input signal arrays, used in each transition's Circom call
+        let input_signal_cg_ids_list_str = (1..=nfa.num_capture_groups)
+            .map(|k| format!("captureGroup{}Ids[i]", k))
+            .collect::<Vec<_>>()
+            .join(", ");
+        let input_signal_cg_starts_list_str = (1..=nfa.num_capture_groups)
+            .map(|k| format!("captureGroup{}Starts[i]", k))
+            .collect::<Vec<_>>()
+            .join(", ");
+
         for (transition_idx, (curr_state, start, end, next_state, capture_info)) in
             transitions.iter().enumerate()
         {
-            let (capture_group_id, capture_group_start) = match capture_info {
-                Some(capture_info) => (capture_info.0, capture_info.1 as u8),
-                None => (0, 0),
+            // These vectors store the properties of *this specific transition*
+            // regarding which capture groups it affects and how.
+            let mut transition_prop_cg_ids = vec![0; nfa.num_capture_groups];
+            let mut transition_prop_cg_starts = vec![0; nfa.num_capture_groups];
+
+            let capture_details_for_comment = capture_info
+                .as_ref()
+                .map(|infos| {
+                    infos
+                        .iter()
+                        .map(|(id, is_start_bool)| {
+                            if *id > 0 && *id <= nfa.num_capture_groups {
+                                transition_prop_cg_ids[*id - 1] = *id;
+                                transition_prop_cg_starts[*id - 1] = *is_start_bool as u8;
+                            }
+                            format!("({}, {})", id, *is_start_bool as u8)
+                        })
+                        .collect::<Vec<String>>()
+                        .join(", ")
+                })
+                .unwrap_or_else(|| "".to_string());
+
+            let capture_comment_segment = if capture_details_for_comment.is_empty() {
+                "Capture Group: []".to_string()
+            } else {
+                format!("Capture Group:[ {}]", capture_details_for_comment)
             };
 
+            // String representation of this transition's capture group properties
+            let transition_prop_cg_ids_str = transition_prop_cg_ids
+                .iter()
+                .map(ToString::to_string)
+                .collect::<Vec<_>>()
+                .join(", ");
+            let transition_prop_cg_starts_str = transition_prop_cg_starts
+                .iter()
+                .map(ToString::to_string)
+                .collect::<Vec<_>>()
+                .join(", ");
+
             if start == end {
                 code.push_str(
                     format!(
-                        "        // Transition {}: {} -[{}]-> {} | Capture Group: ({}, {})\n",
-                        transition_idx,
-                        curr_state,
-                        start,
-                        next_state,
-                        capture_group_id,
-                        capture_group_start
+                        "        // Transition {}: {} -[{}]-> {} | {}\n",
+                        transition_idx, curr_state, start, next_state, capture_comment_segment
                     )
                     .as_str(),
                 );
                 code.push_str(
                     format!(
-                        "        isValidTransition[{}][i] <== CheckByteTransitionWithCapture()({}, {}, {}, {}, {}, currStates[i], nextStates[i], haystack[i], captureGroupIds[i], captureGroupStarts[i]);\n",
+                        "        isValidTransition[{}][i] <== CheckByteTransitionWithCapture({})({}, {}, {}, [{}], [{}], currStates[i], nextStates[i], haystack[i], [{}], [{}]);\n",
                         transition_idx,
+                        nfa.num_capture_groups,
                         curr_state,
                         next_state,
                         start,
-                        capture_group_id,
-                        capture_group_start
+                        transition_prop_cg_ids_str,
+                        transition_prop_cg_starts_str,
+                        input_signal_cg_ids_list_str, // Array of input signals
+                        input_signal_cg_starts_list_str // Array of input signals
                     ).as_str()
                 );
             } else {
                 code.push_str(
                     format!(
-                        "        // Transition {}: {} -[{}-{}]-> {} | Capture Group: ({}, {})\n",
-                        transition_idx,
-                        curr_state,
-                        start,
-                        end,
-                        next_state,
-                        capture_group_id,
-                        capture_group_start
+                        "        // Transition {}: {} -[{}-{}]-> {} | {}\n",
+                        transition_idx, curr_state, start, end, next_state, capture_comment_segment
                     )
                     .as_str(),
                 );
                 code.push_str(
                     format!(
-                        "        isValidTransition[{}][i] <== CheckByteRangeTransitionWithCapture()({}, {}, {}, {}, {}, {}, currStates[i], nextStates[i], haystack[i], captureGroupIds[i], captureGroupStarts[i]);\n",
+                        "        isValidTransition[{}][i] <== CheckByteRangeTransitionWithCapture({})({}, {}, {}, {}, [{}], [{}], currStates[i], nextStates[i], haystack[i], [{}], [{}]);\n",
                         transition_idx,
+                        nfa.num_capture_groups,
                         curr_state,
                         next_state,
                         start,
                         end,
-                        capture_group_id,
-                        capture_group_start
+                        transition_prop_cg_ids_str,
+                        transition_prop_cg_starts_str,
+                        input_signal_cg_ids_list_str, // Array of input signals
+                        input_signal_cg_starts_list_str // Array of input signals
                     ).as_str()
                 );
             }
@@ -375,15 +418,15 @@ pub fn generate_circom_code(
     code.push_str("    }\n\n");
     code.push_str("    isValid <== isValidRegex[maxMatchBytes-1];\n\n");
 
-    if has_capture_groups {
+    if nfa.num_capture_groups > 0 {
         code.push_str(
             format!(
                 "    signal input captureGroupStartIndices[{}];\n\n",
-                capture_group_set.len()
+                nfa.num_capture_groups
             )
             .as_str(),
         );
-        for capture_group_id in capture_group_set {
+        for capture_group_id in 1..=nfa.num_capture_groups {
             let max_substring_bytes =
                 if let Some(max_substring_bytes) = max_substring_bytes.as_ref() {
                     max_substring_bytes[capture_group_id - 1]
@@ -397,12 +440,14 @@ pub fn generate_circom_code(
             code.push_str(format!("    // Capture Group {}\n", capture_group_id).as_str());
             code.push_str(
                 format!(
-                    "    signal output capture{}[{}] <== CaptureSubstring(maxMatchBytes, {}, {})(captureGroupStartIndices[{}], haystack, captureGroupIds, captureGroupStarts);\n",
+                    "    signal output capture{}[{}] <== CaptureSubstring(maxMatchBytes, {}, {})(captureGroupStartIndices[{}], haystack, captureGroup{}Id, captureGroup{}Start);\n",
                     capture_group_id,
                     max_substring_bytes,
                     max_substring_bytes,
                     capture_group_id,
-                    capture_group_id - 1
+                    capture_group_id,
+                    capture_group_id,
+                    capture_group_id
                 ).as_str()
             );
         }

compiler/src/nfa/codegen/mod.rs

@@ -7,7 +7,7 @@ use circom::CircomInputs;
 use noir::NoirInputs;
 use regex_automata::meta::Regex;
 use serde::Serialize;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, BTreeSet};
 
 use crate::{
     ProverInputs, ProvingFramework,
@@ -25,9 +25,9 @@ pub struct CircuitInputs {
     curr_states: Vec<usize>,
     next_states: Vec<usize>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    capture_group_ids: Option<Vec<usize>>,
+    capture_group_ids: Option<Vec<Vec<usize>>>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    capture_group_starts: Option<Vec<u8>>,
+    capture_group_starts: Option<Vec<Vec<u8>>>,
     #[serde(skip_serializing_if = "Option::is_none")]
     capture_group_start_indices: Option<Vec<usize>>,
 }
@@ -37,7 +37,7 @@ pub fn generate_circuit_data(
 ) -> NFAResult<(
     Vec<usize>,
     Vec<usize>,
-    Vec<(usize, u8, u8, usize, Option<(usize, bool)>)>,
+    Vec<(usize, u8, u8, usize, Option<BTreeSet<(usize, bool)>>)>,
 )> {
     if nfa.start_states.is_empty() {
         return Err(NFAError::Verification("NFA has no start states".into()));
@@ -60,7 +60,8 @@ pub fn generate_circuit_data(
 
     // Group and convert to ranges - use BTreeMap for deterministic ordering
     let mut range_transitions = Vec::new();
-    let mut grouped: BTreeMap<(usize, usize, Option<(usize, bool)>), Vec<u8>> = BTreeMap::new();
+    let mut grouped: BTreeMap<(usize, usize, Option<BTreeSet<(usize, bool)>>), Vec<u8>> =
+        BTreeMap::new();
 
     for (src, byte, dst, capture) in transitions {
         if src >= nfa.nodes.len() || dst >= nfa.nodes.len() {
@@ -75,7 +76,12 @@ pub fn generate_circuit_data(
     // Convert to ranges
     for ((src, dst, capture), mut bytes) in grouped {
         if bytes.is_empty() {
-            continue;
+            // This case should ideally not be reached if the grouping logic is correct
+            // and transitions always have associated bytes.
+            return Err(NFAError::InvalidTransition(format!(
+                "Found an empty byte list for transition group (src: {}, dst: {}, capture: {:?}). This indicates an issue with NFA processing.",
+                src, dst, capture
+            )));
         }
 
         bytes.sort_unstable();
@@ -84,12 +90,12 @@ pub fn generate_circuit_data(
 
         for &byte in &bytes[1..] {
             if byte != prev + 1 {
-                range_transitions.push((src, start, prev, dst, capture));
+                range_transitions.push((src, start, prev, dst, capture.clone()));
                 start = byte;
             }
             prev = byte;
         }
-        range_transitions.push((src, start, prev, dst, capture));
+        range_transitions.push((src, start, prev, dst, capture.clone()));
     }
 
     Ok((start_states, accept_states, range_transitions))
@@ -144,14 +150,33 @@ pub fn generate_circuit_inputs(
     // Handle capture groups if they exist
     let (capture_group_ids, capture_group_starts, capture_group_start_indices) =
         if path.iter().any(|(_, _, _, c)| c.is_some()) {
-            let mut ids = path
-                .iter()
-                .map(|(_, _, _, c)| c.map(|(id, _)| id).unwrap_or(0))
-                .collect::<Vec<_>>();
-            let mut starts = path
-                .iter()
-                .map(|(_, _, _, c)| c.map(|(_, start)| start as u8).unwrap_or(0))
-                .collect::<Vec<_>>();
+            // Initialize structures:
+            // capture_group_ids[group_idx_0_based][step_idx]
+            let mut capture_group_ids: Vec<Vec<usize>> =
+                vec![vec![0; max_match_len]; nfa.num_capture_groups];
+            // capture_group_starts[group_idx_0_based][step_idx]
+            let mut capture_group_starts: Vec<Vec<u8>> =
+                vec![vec![0; max_match_len]; nfa.num_capture_groups];
+
+            // Populate these based on the actual path traversal
+            // path_len is the actual number of steps taken for the match
+            for step_idx in 0..path_len {
+                // path[step_idx].3 is the Option<BTreeSet<(usize group_id, bool is_start)>>
+                if let Some(capture_set) = &path[step_idx].3 {
+                    for (group_id, is_start) in capture_set.iter() {
+                        // group_id is 1-based from the regex engine
+                        if *group_id > 0 && *group_id <= nfa.num_capture_groups {
+                            let group_vector_idx = *group_id - 1; // Convert to 0-based for vector access
+
+                            // Record the group ID if active
+                            capture_group_ids[group_vector_idx][step_idx] = *group_id;
+                            // Record if it's a start or end/continuation
+                            capture_group_starts[group_vector_idx][step_idx] =
+                                if *is_start { 1 } else { 0 };
+                        }
+                    }
+                }
+            }
 
             // Use regex_automata to get capture start indices
             let re = Regex::new(&nfa.regex).map_err(|e| {
@@ -165,11 +190,11 @@ pub fn generate_circuit_inputs(
                 .map(|m| m.start)
                 .collect();
 
-            // Pad arrays
-            ids.resize(max_match_len, 0);
-            starts.resize(max_match_len, 0);
-
-            (Some(ids), Some(starts), Some(start_indices))
+            (
+                Some(capture_group_ids),
+                Some(capture_group_starts),
+                Some(start_indices),
+            )
         } else {
             (None, None, None)
         };