circom_lsp_program_structure/utils/
environment.rs

1use std::collections::HashMap;
2use std::hash::Hash;
3use std::marker::PhantomData;
4
5pub trait VarInfo {}
6pub trait SignalInfo {}
7pub trait ComponentInfo {}
8
9#[derive(Clone)]
10pub struct OnlyVars;
11impl VarInfo for OnlyVars {}
12#[derive(Clone)]
13pub struct OnlySignals;
14impl SignalInfo for OnlySignals {}
15#[derive(Clone)]
16pub struct OnlyComponents;
17impl ComponentInfo for OnlyComponents {}
18#[derive(Clone)]
19pub struct FullEnvironment;
20impl VarInfo for FullEnvironment {}
21impl SignalInfo for FullEnvironment {}
22impl ComponentInfo for FullEnvironment {}
23
24pub type VarEnvironment<VC> = RawEnvironment<OnlyVars, (), (), VC>;
25pub type SignalEnvironment<SC> = RawEnvironment<OnlySignals, (), SC, ()>;
26pub type ComponentEnvironment<CC> = RawEnvironment<OnlyComponents, CC, (), ()>;
27pub type CircomEnvironment<CC, SC, VC> = RawEnvironment<FullEnvironment, CC, SC, VC>;
28
29pub enum CircomEnvironmentError {
30    NonExistentSymbol,
31}
32
33#[derive(Clone)]
34pub struct RawEnvironment<T, CC, SC, VC> {
35    components: HashMap<String, CC>,
36    inputs: HashMap<String, SC>,
37    outputs: HashMap<String, SC>,
38    intermediates: HashMap<String, SC>,
39    variables: Vec<VariableBlock<VC>>,
40    behaviour: PhantomData<T>,
41}
42impl<T, CC, SC, VC> Default for RawEnvironment<T, CC, SC, VC> {
43    fn default() -> Self {
44        let variables = vec![VariableBlock::new()];
45        RawEnvironment {
46            components: HashMap::new(),
47            inputs: HashMap::new(),
48            outputs: HashMap::new(),
49            intermediates: HashMap::new(),
50            variables,
51            behaviour: PhantomData,
52        }
53    }
54}
55impl<T, CC, SC, VC> RawEnvironment<T, CC, SC, VC>
56where
57    T: VarInfo + SignalInfo + ComponentInfo,
58{
59    pub fn has_symbol(&self, symbol: &str) -> bool {
60        self.has_signal(symbol) || self.has_component(symbol) || self.has_variable(symbol)
61    }
62}
63impl<T, CC, SC, VC> RawEnvironment<T, CC, SC, VC> {
64    pub fn merge(
65        left: RawEnvironment<T, CC, SC, VC>,
66        right: RawEnvironment<T, CC, SC, VC>,
67        using: fn(VC, VC) -> VC,
68    ) -> RawEnvironment<T, CC, SC, VC> {
69        let mut components = left.components;
70        let mut inputs = left.inputs;
71        let mut outputs = left.outputs;
72        let mut intermediates = left.intermediates;
73        components.extend(right.components);
74        inputs.extend(right.inputs);
75        outputs.extend(right.outputs);
76        intermediates.extend(right.intermediates);
77        let mut variables_left = left.variables;
78        let mut variables_right = right.variables;
79        let mut variables = Vec::new();
80        while !variables_left.is_empty() && !variables_right.is_empty() {
81            let left_block = variables_left.pop().unwrap();
82            let right_block = variables_right.pop().unwrap();
83            let merged_blocks = VariableBlock::merge(left_block, right_block, using);
84            variables.push(merged_blocks);
85        }
86        variables.reverse();
87        RawEnvironment {
88            components,
89            inputs,
90            intermediates,
91            outputs,
92            variables,
93            behaviour: PhantomData,
94        }
95    }
96}
97impl<T, CC, SC, VC> RawEnvironment<T, CC, SC, VC>
98where
99    T: VarInfo,
100{
101    fn block_with_variable_symbol(&self, symbol: &str) -> Option<&VariableBlock<VC>> {
102        let variables = &self.variables;
103        let mut act = variables.len();
104        while act > 0 {
105            if VariableBlock::contains_variable(&variables[act - 1], symbol) {
106                return Option::Some(&variables[act - 1]);
107            }
108            act -= 1;
109        }
110        Option::None
111    }
112    fn mut_block_with_variable_symbol(&mut self, symbol: &str) -> Option<&mut VariableBlock<VC>> {
113        let variables = &mut self.variables;
114        let mut act = variables.len();
115        while act > 0 {
116            if VariableBlock::contains_variable(&variables[act - 1], symbol) {
117                return Option::Some(&mut variables[act - 1]);
118            }
119            act -= 1;
120        }
121        Option::None
122    }
123    pub fn new() -> RawEnvironment<T, CC, SC, VC> {
124        RawEnvironment::default()
125    }
126    pub fn add_variable_block(&mut self) {
127        self.variables.push(VariableBlock::new());
128    }
129    pub fn remove_variable_block(&mut self) {
130        assert!(!self.variables.is_empty());
131        self.variables.pop();
132    }
133    pub fn add_variable(&mut self, variable_name: &str, content: VC) {
134        assert!(!self.variables.is_empty());
135        let last_block = self.variables.last_mut().unwrap();
136        last_block.add_variable(variable_name, content);
137    }
138    pub fn has_variable(&self, symbol: &str) -> bool {
139        self.block_with_variable_symbol(symbol).is_some()
140    }
141
142    pub fn get_variable(&self, symbol: &str) -> Option<&VC> {
143        let possible_block = self.block_with_variable_symbol(symbol);
144        if let Option::Some(block) = possible_block {
145            Option::Some(block.get_variable(symbol))
146        } else {
147            Option::None
148        }
149    }
150    pub fn get_mut_variable(&mut self, symbol: &str) -> Option<&mut VC> {
151        let possible_block = self.mut_block_with_variable_symbol(symbol);
152        if let Option::Some(block) = possible_block {
153            Option::Some(block.get_mut_variable(symbol))
154        } else {
155            Option::None
156        }
157    }
158    pub fn get_variable_res(&self, symbol: &str) -> Result<&VC, CircomEnvironmentError> {
159        let possible_block = self.block_with_variable_symbol(symbol);
160        if let Option::Some(block) = possible_block {
161            Result::Ok(block.get_variable(symbol))
162        } else {
163            Result::Err(CircomEnvironmentError::NonExistentSymbol)
164        }
165    }
166    pub fn remove_variable(&mut self, symbol: &str) {
167        let possible_block = self.mut_block_with_variable_symbol(symbol);
168        if let Option::Some(block) = possible_block {
169            block.remove_variable(symbol)
170        }
171    }
172    pub fn get_variable_or_break(&self, symbol: &str, file: &str, line: u32) -> &VC {
173        assert!(self.has_variable(symbol), "Method call in file {} line {}", file, line);
174        if let Result::Ok(v) = self.get_variable_res(symbol) {
175            v
176        } else {
177            unreachable!();
178        }
179    }
180    pub fn get_mut_variable_mut(
181        &mut self,
182        symbol: &str,
183    ) -> Result<&mut VC, CircomEnvironmentError> {
184        let possible_block = self.mut_block_with_variable_symbol(symbol);
185        if let Option::Some(block) = possible_block {
186            Result::Ok(block.get_mut_variable(symbol))
187        } else {
188            Result::Err(CircomEnvironmentError::NonExistentSymbol)
189        }
190    }
191    pub fn get_mut_variable_or_break(&mut self, symbol: &str, file: &str, line: u32) -> &mut VC {
192        assert!(self.has_variable(symbol), "Method call in file {} line {}", file, line);
193        if let Result::Ok(v) = self.get_mut_variable_mut(symbol) {
194            v
195        } else {
196            unreachable!();
197        }
198    }
199}
200
201impl<T, CC, SC, VC> RawEnvironment<T, CC, SC, VC>
202where
203    T: ComponentInfo,
204{
205    pub fn add_component(&mut self, component_name: &str, content: CC) {
206        self.components.insert(component_name.to_string(), content);
207    }
208    pub fn remove_component(&mut self, component_name: &str) {
209        self.components.remove(component_name);
210    }
211    pub fn has_component(&self, symbol: &str) -> bool {
212        self.components.contains_key(symbol)
213    }
214    pub fn get_component(&self, symbol: &str) -> Option<&CC> {
215        self.components.get(symbol)
216    }
217    pub fn get_mut_component(&mut self, symbol: &str) -> Option<&mut CC> {
218        self.components.get_mut(symbol)
219    }
220    pub fn get_component_res(&self, symbol: &str) -> Result<&CC, CircomEnvironmentError> {
221        self.components.get(symbol).ok_or_else(|| CircomEnvironmentError::NonExistentSymbol)
222    }
223    pub fn get_component_or_break(&self, symbol: &str, file: &str, line: u32) -> &CC {
224        assert!(self.has_component(symbol), "Method call in file {} line {}", file, line);
225        self.components.get(symbol).unwrap()
226    }
227    pub fn get_mut_component_res(
228        &mut self,
229        symbol: &str,
230    ) -> Result<&mut CC, CircomEnvironmentError> {
231        self.components.get_mut(symbol).ok_or_else(|| CircomEnvironmentError::NonExistentSymbol)
232    }
233    pub fn get_mut_component_or_break(&mut self, symbol: &str, file: &str, line: u32) -> &mut CC {
234        assert!(self.has_component(symbol), "Method call in file {} line {}", file, line);
235        self.components.get_mut(symbol).unwrap()
236    }
237    pub fn get_components_ref(&self)-> &HashMap<String, CC>{
238        &self.components
239    }
240}
241
242impl<T, CC, SC, VC> RawEnvironment<T, CC, SC, VC>
243where
244    T: SignalInfo,
245{
246    pub fn add_input(&mut self, input_name: &str, content: SC) {
247        self.inputs.insert(input_name.to_string(), content);
248    }
249    pub fn remove_input(&mut self, input_name: &str) {
250        self.inputs.remove(input_name);
251    }
252    pub fn add_output(&mut self, output_name: &str, content: SC) {
253        self.outputs.insert(output_name.to_string(), content);
254    }
255    pub fn remove_output(&mut self, output_name: &str) {
256        self.outputs.remove(output_name);
257    }
258    pub fn add_intermediate(&mut self, intermediate_name: &str, content: SC) {
259        self.intermediates.insert(intermediate_name.to_string(), content);
260    }
261    pub fn remove_intermediate(&mut self, intermediate_name: &str) {
262        self.intermediates.remove(intermediate_name);
263    }
264    pub fn has_input(&self, symbol: &str) -> bool {
265        self.inputs.contains_key(symbol)
266    }
267    pub fn has_output(&self, symbol: &str) -> bool {
268        self.outputs.contains_key(symbol)
269    }
270    pub fn has_intermediate(&self, symbol: &str) -> bool {
271        self.intermediates.contains_key(symbol)
272    }
273    pub fn has_signal(&self, symbol: &str) -> bool {
274        self.has_input(symbol) || self.has_output(symbol) || self.has_intermediate(symbol)
275    }
276    pub fn get_input(&self, symbol: &str) -> Option<&SC> {
277        self.inputs.get(symbol)
278    }
279    pub fn get_mut_input(&mut self, symbol: &str) -> Option<&mut SC> {
280        self.inputs.get_mut(symbol)
281    }
282    pub fn get_input_res(&self, symbol: &str) -> Result<&SC, CircomEnvironmentError> {
283        self.inputs.get(symbol).ok_or_else(|| CircomEnvironmentError::NonExistentSymbol)
284    }
285    pub fn get_input_or_break(&self, symbol: &str, file: &str, line: u32) -> &SC {
286        assert!(self.has_input(symbol), "Method call in file {} line {}", file, line);
287        self.inputs.get(symbol).unwrap()
288    }
289    pub fn get_mut_input_res(&mut self, symbol: &str) -> Result<&mut SC, CircomEnvironmentError> {
290        self.inputs.get_mut(symbol).ok_or_else(|| CircomEnvironmentError::NonExistentSymbol)
291    }
292    pub fn get_mut_input_or_break(&mut self, symbol: &str, file: &str, line: u32) -> &mut SC {
293        assert!(self.has_input(symbol), "Method call in file {} line {}", file, line);
294        self.inputs.get_mut(symbol).unwrap()
295    }
296
297    pub fn get_output(&self, symbol: &str) -> Option<&SC> {
298        self.outputs.get(symbol)
299    }
300    pub fn get_mut_output(&mut self, symbol: &str) -> Option<&mut SC> {
301        self.outputs.get_mut(symbol)
302    }
303    pub fn get_output_res(&self, symbol: &str) -> Result<&SC, CircomEnvironmentError> {
304        self.outputs.get(symbol).ok_or_else(|| CircomEnvironmentError::NonExistentSymbol)
305    }
306    pub fn get_output_or_break(&self, symbol: &str, file: &str, line: u32) -> &SC {
307        assert!(self.has_output(symbol), "Method call in file {} line {}", file, line);
308        self.outputs.get(symbol).unwrap()
309    }
310    pub fn get_mut_output_res(&mut self, symbol: &str) -> Result<&mut SC, CircomEnvironmentError> {
311        self.outputs.get_mut(symbol).ok_or_else(|| CircomEnvironmentError::NonExistentSymbol)
312    }
313    pub fn get_mut_output_or_break(&mut self, symbol: &str, file: &str, line: u32) -> &mut SC {
314        assert!(self.has_output(symbol), "Method call in file {} line {}", file, line);
315        self.outputs.get_mut(symbol).unwrap()
316    }
317
318    pub fn get_intermediate(&self, symbol: &str) -> Option<&SC> {
319        self.intermediates.get(symbol)
320    }
321    pub fn get_mut_intermediate(&mut self, symbol: &str) -> Option<&mut SC> {
322        self.intermediates.get_mut(symbol)
323    }
324    pub fn get_intermediate_res(&self, symbol: &str) -> Result<&SC, CircomEnvironmentError> {
325        self.intermediates.get(symbol).ok_or_else(|| CircomEnvironmentError::NonExistentSymbol)
326    }
327    pub fn get_intermediate_or_break(&self, symbol: &str, file: &str, line: u32) -> &SC {
328        assert!(self.has_intermediate(symbol), "Method call in file {} line {}", file, line);
329        self.intermediates.get(symbol).unwrap()
330    }
331    pub fn get_mut_intermediate_res(
332        &mut self,
333        symbol: &str,
334    ) -> Result<&mut SC, CircomEnvironmentError> {
335        self.intermediates.get_mut(symbol).ok_or_else(|| CircomEnvironmentError::NonExistentSymbol)
336    }
337    pub fn get_mut_intermediate_or_break(
338        &mut self,
339        symbol: &str,
340        file: &str,
341        line: u32,
342    ) -> &mut SC {
343        assert!(self.has_intermediate(symbol), "Method call in file {} line {}", file, line);
344        self.intermediates.get_mut(symbol).unwrap()
345    }
346
347    pub fn get_signal(&self, symbol: &str) -> Option<&SC> {
348        if self.has_input(symbol) {
349            self.get_input(symbol)
350        } else if self.has_output(symbol) {
351            self.get_output(symbol)
352        } else if self.has_intermediate(symbol) {
353            self.get_intermediate(symbol)
354        } else {
355            Option::None
356        }
357    }
358    pub fn get_mut_signal(&mut self, symbol: &str) -> Option<&mut SC> {
359        if self.has_input(symbol) {
360            self.get_mut_input(symbol)
361        } else if self.has_output(symbol) {
362            self.get_mut_output(symbol)
363        } else if self.has_intermediate(symbol) {
364            self.get_mut_intermediate(symbol)
365        } else {
366            Option::None
367        }
368    }
369    pub fn get_signal_res(&self, symbol: &str) -> Result<&SC, CircomEnvironmentError> {
370        if self.has_input(symbol) {
371            self.get_input_res(symbol)
372        } else if self.has_output(symbol) {
373            self.get_output_res(symbol)
374        } else if self.has_intermediate(symbol) {
375            self.get_intermediate_res(symbol)
376        } else {
377            Result::Err(CircomEnvironmentError::NonExistentSymbol)
378        }
379    }
380    pub fn get_signal_or_break(&self, symbol: &str, file: &str, line: u32) -> &SC {
381        assert!(self.has_signal(symbol), "Method call in file {} line {}", file, line);
382        if let Result::Ok(v) = self.get_signal_res(symbol) {
383            v
384        } else {
385            unreachable!();
386        }
387    }
388    pub fn get_mut_signal_res(&mut self, symbol: &str) -> Result<&mut SC, CircomEnvironmentError> {
389        if self.has_input(symbol) {
390            self.get_mut_input_res(symbol)
391        } else if self.has_output(symbol) {
392            self.get_mut_output_res(symbol)
393        } else if self.has_intermediate(symbol) {
394            self.get_mut_intermediate_res(symbol)
395        } else {
396            Result::Err(CircomEnvironmentError::NonExistentSymbol)
397        }
398    }
399    pub fn get_mut_signal_or_break(&mut self, symbol: &str, file: &str, line: u32) -> &mut SC {
400        assert!(self.has_signal(symbol), "Method call in file {} line {}", file, line);
401        if let Result::Ok(v) = self.get_mut_signal_res(symbol) {
402            v
403        } else {
404            unreachable!();
405        }
406    }
407}
408
409#[derive(Clone)]
410struct VariableBlock<VC> {
411    variables: HashMap<String, VC>,
412}
413impl<VC> Default for VariableBlock<VC> {
414    fn default() -> Self {
415        VariableBlock { variables: HashMap::new() }
416    }
417}
418impl<VC> VariableBlock<VC> {
419    pub fn new() -> VariableBlock<VC> {
420        VariableBlock::default()
421    }
422    pub fn add_variable(&mut self, symbol: &str, content: VC) {
423        self.variables.insert(symbol.to_string(), content);
424    }
425    pub fn remove_variable(&mut self, symbol: &str) {
426        self.variables.remove(symbol);
427    }
428    pub fn contains_variable(&self, symbol: &str) -> bool {
429        self.variables.contains_key(symbol)
430    }
431    pub fn get_variable(&self, symbol: &str) -> &VC {
432        assert!(self.contains_variable(symbol));
433        self.variables.get(symbol).unwrap()
434    }
435    pub fn get_mut_variable(&mut self, symbol: &str) -> &mut VC {
436        assert!(self.contains_variable(symbol));
437        self.variables.get_mut(symbol).unwrap()
438    }
439    pub fn merge(
440        left: VariableBlock<VC>,
441        right: VariableBlock<VC>,
442        using: fn(VC, VC) -> VC,
443    ) -> VariableBlock<VC> {
444        let left_block = left.variables;
445        let right_block = right.variables;
446        let result_block = hashmap_union(left_block, right_block, using);
447        VariableBlock { variables: result_block }
448    }
449}
450
451fn hashmap_union<K, V>(
452    l: HashMap<K, V>,
453    mut r: HashMap<K, V>,
454    merge_function: fn(V, V) -> V,
455) -> HashMap<K, V>
456where
457    K: Hash + Eq,
458{
459    let mut result = HashMap::new();
460    for (k, v) in l {
461        if let Option::Some(r_v) = r.remove(&k) {
462            result.insert(k, merge_function(v, r_v));
463        } else {
464            result.insert(k, v);
465        }
466    }
467    for (k, v) in r {
468        result.entry(k).or_insert(v);
469    }
470    result
471}