From 9dd03753bbe6463954be688c27e7554381a56281 Mon Sep 17 00:00:00 2001
From: Nam Ngo <ngochannam@gmail.com>
Date: Mon, 20 May 2024 11:54:59 +0700
Subject: [PATCH] markking inputs/outputs in json

---
 src/circuit.rs | 22 ++++++++++++++++++++++
 src/program.rs | 20 ++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/src/circuit.rs b/src/circuit.rs
index 09e7fc9..9a37b42 100644
--- a/src/circuit.rs
+++ b/src/circuit.rs
@@ -166,6 +166,8 @@ impl ArithmeticGate {
 #[derive(Default, Debug, Serialize, Deserialize)]
 pub struct ArithmeticCircuit {
     node_count: u32,
+    inputs: HashMap<u32, String>,
+    outputs: HashMap<u32, String>,
     signals: HashMap<u32, Signal>,
     nodes: HashMap<u32, Node>,
     gates: Vec<ArithmeticGate>,
@@ -176,12 +178,22 @@ impl ArithmeticCircuit {
     pub fn new() -> ArithmeticCircuit {
         ArithmeticCircuit {
             node_count: 0,
+            inputs: HashMap::new(),
+            outputs: HashMap::new(),
             signals: HashMap::new(),
             nodes: HashMap::new(),
             gates: Vec::new(),
         }
     }
 
+    pub fn add_inputs(&mut self, inputs: HashMap<u32, String>) {
+        self.inputs.extend(inputs);
+    }
+
+    pub fn add_outputs(&mut self, outputs: HashMap<u32, String>) {
+        self.outputs.extend(outputs);
+    }
+
     /// Adds a new signal to the circuit.
     pub fn add_signal(
         &mut self,
@@ -207,6 +219,16 @@ impl ArithmeticCircuit {
         Ok(())
     }
 
+    pub fn get_signals(&self, filter: String) -> HashMap<u32, String> {
+        let mut ret = HashMap::new();
+        for (signal_id, signal) in self.signals.iter() {
+            if signal.name.starts_with(filter.as_str()) {
+                ret.insert(*signal_id, signal.name.to_string());
+            }
+        }
+        ret
+    }
+
     /// Adds a new gate to the circuit.
     pub fn add_gate(
         &mut self,
diff --git a/src/program.rs b/src/program.rs
index 2378ff1..03e3d49 100644
--- a/src/program.rs
+++ b/src/program.rs
@@ -8,6 +8,7 @@ use crate::{
     process::{process_expression, process_statements},
     runtime::{DataAccess, DataType, Runtime, RuntimeError},
 };
+use circom_compiler::num_traits::sign;
 use circom_program_structure::ast::Expression;
 use std::io;
 use thiserror::Error;
@@ -51,6 +52,25 @@ pub fn build_circuit(input: &Input) -> Result<ArithmeticCircuit, ProgramError> {
             // Process the main component
             let statements = template_data.get_body_as_vec();
             process_statements(&mut circuit, &mut runtime, &program_archive, statements)?;
+
+            for (ikey, (ivs, ivh)) in template_data.get_inputs().iter() {
+                // println!("{ikey}:{ivs}");
+                let filter = format!("0.{}", ikey);
+                circuit.add_inputs(circuit.get_signals(filter));
+                // for ivhs in ivh.iter() {
+                //     println!("{ivhs}");
+                // }
+            }
+
+            for (okey, (ovs, ovh)) in template_data.get_outputs().iter() {
+                // println!("{okey}:{ovs}");
+                let filter = format!("0.{}", okey);
+                let signals = circuit.get_signals(filter);
+                circuit.add_outputs(signals);
+                // for ovhs in ovh.iter() {
+                //     println!("{ovhs}");
+                // }
+            }
         }
         _ => return Err(ProgramError::MainExpressionNotACall),
     }