circom_lsp_program_structure/program_library/
template_data.rs

1use super::ast;
2use super::ast::{FillMeta, Statement};
3use super::file_definition::FileID;
4use crate::file_definition::FileLocation;
5use std::collections::{HashMap, HashSet, BTreeMap};
6
7pub type TemplateInfo = HashMap<String, TemplateData>;
8pub type TagInfo = HashSet<String>;
9type SignalInfo = BTreeMap<String, (usize, TagInfo)>;
10type SignalDeclarationOrder = Vec<(String, usize)>;
11
12#[derive(Clone)]
13pub struct TemplateData {
14    file_id: FileID,
15    name: String,
16    body: Statement,
17    num_of_params: usize,
18    name_of_params: Vec<String>,
19    param_location: FileLocation,
20    input_signals: SignalInfo,
21    output_signals: SignalInfo,
22    is_parallel: bool,
23    is_custom_gate: bool,
24    /* Only used to know the order in which signals are declared.*/
25    input_declarations: SignalDeclarationOrder,
26    output_declarations: SignalDeclarationOrder,
27}
28
29impl TemplateData {
30    pub fn new(
31        name: String,
32        file_id: FileID,
33        mut body: Statement,
34        num_of_params: usize,
35        name_of_params: Vec<String>,
36        param_location: FileLocation,
37        elem_id: &mut usize,
38        is_parallel: bool,
39        is_custom_gate: bool,
40    ) -> TemplateData {
41        body.fill(file_id, elem_id);
42        let mut input_signals = SignalInfo::new();
43        let mut output_signals = SignalInfo::new();
44        let mut input_declarations =  SignalDeclarationOrder::new();
45        let mut output_declarations = SignalDeclarationOrder::new();
46        fill_inputs_and_outputs(&body, &mut input_signals, &mut output_signals, &mut input_declarations, &mut output_declarations);
47        TemplateData {
48            name,
49            file_id,
50            body,
51            num_of_params,
52            name_of_params,
53            param_location,
54            input_signals,
55            output_signals,
56            is_parallel,
57            is_custom_gate,
58            input_declarations,
59            output_declarations
60        }
61    }
62
63    pub fn copy(
64        name: String,
65        file_id: FileID,
66        body: Statement,
67        num_of_params: usize,
68        name_of_params: Vec<String>,
69        param_location: FileLocation,
70        input_signals: SignalInfo,
71        output_signals: SignalInfo,
72        is_parallel: bool,
73        is_custom_gate: bool,
74        input_declarations :SignalDeclarationOrder,
75        output_declarations : SignalDeclarationOrder
76    ) -> TemplateData {
77        TemplateData {
78            name,
79            file_id,
80            body,
81            num_of_params,
82            name_of_params,
83            param_location,
84            input_signals,
85            output_signals,
86            is_parallel,
87            is_custom_gate,
88            input_declarations,
89            output_declarations
90        }
91    }
92    pub fn get_file_id(&self) -> FileID {
93        self.file_id
94    }
95    pub fn get_body(&self) -> &Statement {
96        &self.body
97    }
98    pub fn get_body_as_vec(&self) -> &Vec<Statement> {
99        match &self.body {
100            Statement::Block { stmts, .. } => stmts,
101            _ => panic!("Function body should be a block"),
102        }
103    }
104    pub fn get_mut_body(&mut self) -> &mut Statement {
105        &mut self.body
106    }
107    pub fn get_mut_body_as_vec(&mut self) -> &mut Vec<Statement> {
108        match &mut self.body {
109            Statement::Block { stmts, .. } => stmts,
110            _ => panic!("Function body should be a block"),
111        }
112    }
113    pub fn get_num_of_params(&self) -> usize {
114        self.num_of_params
115    }
116    pub fn get_param_location(&self) -> FileLocation {
117        self.param_location.clone()
118    }
119    pub fn get_name_of_params(&self) -> &Vec<String> {
120        &self.name_of_params
121    }
122    pub fn get_input_info(&self, name: &str) -> Option<&(usize, TagInfo)> {
123        self.input_signals.get(name)
124    }
125    pub fn get_output_info(&self, name: &str) -> Option<&(usize, TagInfo)> {
126        self.output_signals.get(name)
127    }
128    pub fn get_inputs(&self) -> &SignalInfo {
129        &self.input_signals
130    }
131    pub fn get_outputs(&self) -> &SignalInfo {
132        &self.output_signals
133    }
134    pub fn get_declaration_inputs(&self) -> &SignalDeclarationOrder {
135        &&self.input_declarations
136    }
137    pub fn get_declaration_outputs(&self) -> &SignalDeclarationOrder {
138        &self.output_declarations
139    }
140    pub fn get_name(&self) -> &str {
141        &self.name
142    }
143    pub fn is_parallel(&self) -> bool {
144        self.is_parallel
145    }
146    pub fn is_custom_gate(&self) -> bool {
147        self.is_custom_gate
148    }
149}
150
151fn fill_inputs_and_outputs(
152    template_statement: &Statement,
153    input_signals: &mut SignalInfo,
154    output_signals: &mut SignalInfo,
155    input_declarations : &mut SignalDeclarationOrder,
156    output_declarations : &mut SignalDeclarationOrder
157) {
158    match template_statement {
159        Statement::IfThenElse { if_case, else_case, .. } => {
160            fill_inputs_and_outputs(if_case, input_signals, output_signals, input_declarations, output_declarations);
161            if let Option::Some(else_value) = else_case {
162                fill_inputs_and_outputs(else_value, input_signals, output_signals, input_declarations, output_declarations);
163            }
164        }
165        Statement::Block { stmts, .. } => {
166            for stmt in stmts.iter() {
167                fill_inputs_and_outputs(stmt, input_signals, output_signals, input_declarations, output_declarations);
168            }
169        }
170        Statement::While { stmt, .. } => {
171            fill_inputs_and_outputs(stmt, input_signals, output_signals, input_declarations, output_declarations);
172        }
173        Statement::InitializationBlock { initializations, .. } => {
174            for initialization in initializations.iter() {
175                fill_inputs_and_outputs(initialization, input_signals, output_signals, input_declarations, output_declarations);
176            }
177        }
178        Statement::Declaration { xtype, name, dimensions, .. } => {
179            if let ast::VariableType::Signal(stype, tag_list) = xtype {
180                let signal_name = name.clone();
181                let dim = dimensions.len();
182                let mut tag_info = HashSet::new();
183                for tag in tag_list{
184                    tag_info.insert(tag.clone());
185                }
186
187                match stype {
188                    ast::SignalType::Input => {
189                        input_signals.insert(signal_name.clone(), (dim, tag_info));
190                        input_declarations.push((signal_name,dim));
191                    }
192                    ast::SignalType::Output => {
193                        output_signals.insert(signal_name.clone(), (dim, tag_info));
194                        output_declarations.push((signal_name,dim));
195                    }
196                    _ => {} //no need to deal with intermediate signals
197                }
198            }
199        }
200        _ => {}
201    }
202}