ocsf_codegen/
module.rs

1use std::collections::HashMap;
2use std::fs::create_dir_all;
3use std::path::PathBuf;
4
5use codegen::{Enum, Function, Scope, Variant};
6use itertools::{any, Itertools};
7use log::{debug, error, info, trace};
8
9use crate::enums::EnumData;
10use crate::{
11    collapsed_title_case, read_file_to_value, write_source_file, CustomScopeThings, DirPaths,
12};
13
14#[derive(Debug)]
15pub struct ModuleEnumWithU8 {
16    name: String,
17    pub variants: HashMap<u8, EnumData>,
18}
19
20impl ModuleEnumWithU8 {
21    pub fn new(paths: &DirPaths, name: String) -> Self {
22        let variants = Self::get_enum_defaults(paths).unwrap();
23
24        ModuleEnumWithU8 { name, variants }
25    }
26
27    pub fn get_enum_defaults(
28        paths: &DirPaths,
29    ) -> Result<HashMap<u8, EnumData>, Box<dyn std::error::Error>> {
30        let defaults = format!("{}enums/defaults.json", paths.schema_path);
31        trace!("Pulling defaults from {defaults}");
32        let defaults: crate::enums::EnumFile =
33            serde_json::from_value(read_file_to_value(&defaults)?)?;
34        let defaults = defaults.elements;
35
36        trace!("Defaults: {defaults:#?}");
37        Ok(defaults)
38    }
39
40    pub fn add_to_scope(&self, scope: &mut Scope) {
41        let mut new_enum = Enum::new(&self.name);
42        new_enum
43            .vis("pub")
44            .derive("Debug")
45            .derive("serde::Serialize")
46            .derive("serde::Deserialize");
47
48        let mut enum_to_u8 = Function::new("from");
49        enum_to_u8.arg("input", &self.name).ret("u8");
50        enum_to_u8.line("match input {");
51
52        let mut try_u8_to_enum = Function::new("try_from");
53        try_u8_to_enum
54            .arg("input", "u8")
55            .ret("Result<Self, String>");
56        try_u8_to_enum.line("let res = match input {");
57
58        self.variants.keys().sorted().for_each(|key| {
59            let value = self.variants.get(key).unwrap();
60            let variant_name = collapsed_title_case(&value.caption);
61            let mut variant = Variant::new(&variant_name);
62            if let Some(description) = &value.description {
63                variant.annotation(&format!("/// {}", description));
64            }
65            new_enum.push_variant(variant);
66
67            enum_to_u8.line(&format!("    {}::{} => {},", self.name, variant_name, &key));
68            try_u8_to_enum.line(&format!("    {} => {}::{},", &key, self.name, variant_name));
69        });
70
71        debug!("Adding enum called {} to scope...", &self.name);
72        enum_to_u8.line("}");
73        try_u8_to_enum.line("_ => return Err(\"invalid value\".to_string()),");
74        try_u8_to_enum.line("}; Ok(res)");
75
76        scope.push_enum(new_enum);
77
78        let enum_to_u8_impl = scope.new_impl("u8");
79        enum_to_u8_impl.impl_trait(&format!("From<{}>", &self.name));
80        enum_to_u8_impl.push_fn(enum_to_u8);
81
82        let try_u8_to_enum_impl = scope.new_impl(&self.name);
83        try_u8_to_enum_impl
84            .impl_trait("TryFrom<u8>")
85            .associate_type("Error", "String");
86        try_u8_to_enum_impl.push_fn(try_u8_to_enum);
87    }
88}
89
90#[derive(Debug)]
91pub struct ModuleStruct {
92    name: String,
93    pub scope: Scope,
94}
95
96impl ModuleStruct {
97    pub fn new(name: impl ToString) -> Self {
98        Self {
99            name: name.to_string(),
100            scope: Scope::new(),
101        }
102    }
103}
104
105#[derive(Debug)]
106pub struct Module {
107    name: String,
108    pub children: HashMap<String, Module>,
109    pub enums: Vec<ModuleEnumWithU8>,
110    pub structs: Vec<ModuleStruct>,
111    pub is_root: bool,
112    pub scope: codegen::Scope,
113    pub imports: Vec<String>,
114}
115
116impl Default for Module {
117    fn default() -> Self {
118        Module {
119            name: "".to_string(),
120            children: HashMap::new(),
121            enums: vec![],
122            structs: vec![],
123            is_root: false,
124            scope: codegen::Scope::new(),
125            imports: vec![],
126        }
127    }
128}
129
130impl Module {
131    pub fn new(name: String, is_root: bool) -> Self {
132        Self {
133            name,
134            is_root,
135            scope: codegen::Scope::new(),
136            ..Default::default()
137        }
138    }
139
140    /// add an empty child module
141    pub fn add_child(&mut self, name: String) {
142        self.children.insert(name.clone(), Module::new(name, false));
143    }
144
145    /// add an empty struct module
146    pub fn add_struct(&mut self, name: impl ToString) {
147        self.structs.push(ModuleStruct::new(name));
148    }
149
150    pub fn has_enum(&self, name: &str) -> bool {
151        if !self.is_root {
152            panic!("Only check for enums from the root module please!");
153        }
154        any(self.enums.iter(), |f| f.name == name)
155    }
156
157    pub fn has_struct(&self, name: &str) -> bool {
158        any(self.structs.iter(), |f| f.name == name)
159    }
160
161    pub fn has_child(&self, name: &str) -> bool {
162        self.children.contains_key(name)
163    }
164
165    /// add an enum to the module, only do it to the root module though!
166    pub fn add_enum(&mut self, name: String) {
167        if !self.is_root {
168            panic!("Only add enums to the root module please!");
169        }
170        self.enums.push(ModuleEnumWithU8 {
171            name,
172            variants: HashMap::new(),
173        })
174    }
175
176    /// write all the things!
177    pub fn write_module(
178        &mut self,
179        expected_paths: &mut Vec<String>,
180        parent_dirname: &PathBuf,
181    ) -> Result<(), Box<dyn std::error::Error>> {
182        let my_filename = parent_dirname.join(format!("{}.rs", self.name));
183        debug!("My filename: {:#?}", my_filename);
184
185        if !expected_paths.contains(&my_filename.to_str().unwrap().to_owned()) {
186            debug!(
187                "Adding expected path from filename {}",
188                my_filename.to_str().unwrap()
189            );
190            expected_paths.push(my_filename.to_str().unwrap().to_owned());
191        }
192
193        if !parent_dirname.exists() {
194            debug!("Creating directory {:?}", parent_dirname);
195            create_dir_all(parent_dirname)?;
196        }
197
198        for import in self.imports.iter() {
199            self.scope.raw(&format!("{};", import));
200        }
201
202        let child_keys: Vec<String> = self.children.keys().cloned().collect();
203
204        info!("Child modules: {child_keys:#?}");
205
206        child_keys.iter().sorted().for_each(|key| {
207            if key.is_empty() {
208                panic!("Empty module name?");
209            }
210            self.scope.raw(&format!("pub mod {};", key));
211        });
212
213        child_keys.iter().sorted().for_each(|key| {
214            let child_path = match self.is_root {
215                true => parent_dirname.clone(),
216                false => parent_dirname.clone().join(&self.name),
217            };
218
219            info!("writing child module '{}' to {:?}", key, child_path);
220            let child = self.children.get_mut(key).unwrap();
221            // if self.is_root {
222            if let Err(err) = child.write_module(expected_paths, &child_path) {
223                error!(
224                    "Failed to write child module '{}' to {:?}: {:?}",
225                    key, parent_dirname, err
226                );
227            };
228
229            if !expected_paths.contains(&child_path.to_str().unwrap().to_owned()) {
230                debug!("Adding expected path {}", child_path.to_str().unwrap());
231                expected_paths.push(child_path.to_str().unwrap().to_owned());
232            }
233        });
234
235        self.enums.iter().for_each(|object| {
236            debug!("adding enum to scope {:?}", object.name);
237            object.add_to_scope(&mut self.scope)
238        });
239
240        if !self.scope.to_string().contains("automatically generated") {
241            self.scope.add_generation_timestamp_comment();
242        }
243        write_source_file(my_filename.to_str().unwrap(), &self.scope.to_string())
244    }
245}
246
247#[cfg(test)]
248mod tests {
249
250    #[test]
251    fn test_has_child() {
252        use super::Module;
253
254        let mut root = Module::new("lib".into(), true);
255
256        assert!(!root.has_child("cheese"));
257
258        root.add_child("cheese".to_string());
259
260        assert!(root.has_child("cheese"));
261    }
262}