circomspect_program_structure/program_library/
template_library.rs

1use std::collections::HashMap;
2
3use crate::ast::Definition;
4use crate::file_definition::{FileID, FileLibrary};
5use crate::function_data::{FunctionData, FunctionInfo};
6use crate::template_data::{TemplateData, TemplateInfo};
7
8type Contents = HashMap<FileID, Vec<Definition>>;
9
10pub struct TemplateLibrary {
11    pub functions: FunctionInfo,
12    pub templates: TemplateInfo,
13    pub file_library: FileLibrary,
14}
15
16impl TemplateLibrary {
17    pub fn new(library_contents: Contents, file_library: FileLibrary) -> TemplateLibrary {
18        let mut functions = HashMap::new();
19        let mut templates = HashMap::new();
20
21        let mut elem_id = 0;
22        for (file_id, file_contents) in library_contents {
23            for definition in file_contents {
24                match definition {
25                    Definition::Function { name, args, arg_location, body, .. } => {
26                        functions.insert(
27                            name.clone(),
28                            FunctionData::new(
29                                name,
30                                file_id,
31                                body,
32                                args.len(),
33                                args,
34                                arg_location,
35                                &mut elem_id,
36                            ),
37                        );
38                    }
39                    Definition::Template {
40                        name,
41                        args,
42                        arg_location,
43                        body,
44                        parallel,
45                        is_custom_gate,
46                        ..
47                    } => {
48                        templates.insert(
49                            name.clone(),
50                            TemplateData::new(
51                                name,
52                                file_id,
53                                body,
54                                args.len(),
55                                args,
56                                arg_location,
57                                &mut elem_id,
58                                parallel,
59                                is_custom_gate,
60                            ),
61                        );
62                    }
63                }
64            }
65        }
66        TemplateLibrary { functions, templates, file_library }
67    }
68    // Template methods.
69    pub fn contains_template(&self, template_name: &str) -> bool {
70        self.templates.contains_key(template_name)
71    }
72
73    pub fn get_template(&self, template_name: &str) -> &TemplateData {
74        assert!(self.contains_template(template_name));
75        self.templates.get(template_name).unwrap()
76    }
77
78    pub fn get_template_mut(&mut self, template_name: &str) -> &mut TemplateData {
79        assert!(self.contains_template(template_name));
80        self.templates.get_mut(template_name).unwrap()
81    }
82
83    pub fn get_templates(&self) -> &TemplateInfo {
84        &self.templates
85    }
86
87    pub fn get_templates_mut(&mut self) -> &mut TemplateInfo {
88        &mut self.templates
89    }
90
91    // Function methods.
92    pub fn contains_function(&self, function_name: &str) -> bool {
93        self.functions.contains_key(function_name)
94    }
95
96    pub fn get_function(&self, function_name: &str) -> &FunctionData {
97        assert!(self.contains_function(function_name));
98        self.functions.get(function_name).unwrap()
99    }
100
101    pub fn get_function_mut(&mut self, function_name: &str) -> &mut FunctionData {
102        assert!(self.contains_function(function_name));
103        self.functions.get_mut(function_name).unwrap()
104    }
105
106    pub fn get_functions(&self) -> &FunctionInfo {
107        &self.functions
108    }
109
110    pub fn get_functions_mut(&mut self) -> &mut FunctionInfo {
111        &mut self.functions
112    }
113
114    pub fn get_file_library(&self) -> &FileLibrary {
115        &self.file_library
116    }
117}