xsd_parser/code/
module.rs

1use std::collections::{BTreeMap, BTreeSet};
2use std::mem::replace;
3use std::str::FromStr;
4
5use proc_macro2::{Ident as Ident2, TokenStream};
6use quote::{format_ident, quote, ToTokens};
7
8use super::IdentPath;
9
10/// Manages the code that is generated for a specific module.
11///
12/// This type stores the generated code, the related using directives as well as
13/// any sub-module for a specific module generated by the code generator.
14#[derive(Default, Debug)]
15pub struct Module {
16    /// The actual code that is manages by this module.
17    pub code: TokenStream,
18
19    /// A set of using directives this module needs.
20    pub usings: BTreeSet<String>,
21
22    /// A map of sub-modules contained inside this module.
23    pub modules: BTreeMap<String, Module>,
24}
25
26impl Module {
27    /// Append the passed `code` to this modules code.
28    pub fn append(&mut self, code: TokenStream) -> &mut Self {
29        self.code.extend(code);
30
31        self
32    }
33
34    /// Prepend the passed `code` to this modules code.
35    pub fn prepend(&mut self, code: TokenStream) -> &mut Self {
36        let code = replace(&mut self.code, code);
37
38        self.append(code)
39    }
40
41    /// Add using directives to the set of this module.
42    pub fn usings<I>(&mut self, usings: I) -> &mut Self
43    where
44        I: IntoIterator,
45        I::Item: ToString,
46    {
47        for using in usings {
48            self.usings.insert(using.to_string());
49        }
50
51        self
52    }
53
54    /// Get a reference to a sub-module identified by the passed `ident`.
55    ///
56    /// If the module does not exist `None` is returned.
57    pub fn module<T>(&self, ident: T) -> Option<&Module>
58    where
59        T: AsRef<str>,
60    {
61        self.modules.get(ident.as_ref())
62    }
63
64    /// Get a mutable reference to a sub-module identified by the passed `ident`.
65    ///
66    /// If the module does not exist it will be created.
67    pub fn module_mut<T>(&mut self, ident: T) -> &mut Module
68    where
69        T: Into<String>,
70    {
71        self.modules.entry(ident.into()).or_default()
72    }
73}
74
75impl ToTokens for Module {
76    fn to_tokens(&self, tokens: &mut TokenStream) {
77        let Self {
78            code,
79            usings,
80            modules,
81        } = self;
82        let usings = render_usings(usings.iter());
83
84        tokens.extend(quote! {
85            #usings
86            #code
87        });
88
89        for (ident, module) in modules {
90            let name = format_ident!("{ident}");
91
92            tokens.extend(quote! {
93                pub mod #name {
94                    #module
95                }
96            });
97        }
98    }
99}
100
101fn render_usings<I>(usings: I) -> TokenStream
102where
103    I: IntoIterator,
104    I::Item: AsRef<str>,
105{
106    #[derive(Default)]
107    struct Module {
108        usings: BTreeSet<Ident2>,
109        sub_modules: BTreeMap<Ident2, Module>,
110    }
111
112    impl Module {
113        fn render(&self) -> TokenStream {
114            let count = self.usings.len() + self.sub_modules.len();
115
116            let usings = self.usings.iter().map(|ident| quote!(#ident));
117            let sub_modules = self.sub_modules.iter().map(|(ident, module)| {
118                let using = module.render();
119
120                quote!(#ident::#using)
121            });
122
123            let items = usings.chain(sub_modules);
124
125            if count > 1 {
126                quote!({ #( #items ),* })
127            } else {
128                quote!(#( #items )*)
129            }
130        }
131    }
132
133    let mut root = Module::default();
134
135    for using in usings {
136        let using = using.as_ref();
137        let Ok(ident) = IdentPath::from_str(using) else {
138            continue;
139        };
140
141        let (ident, path) = ident.into_parts();
142
143        let mut module = &mut root;
144        for part in path.into_iter().flat_map(|x| x.0) {
145            module = module.sub_modules.entry(part).or_default();
146        }
147
148        module.usings.insert(ident);
149    }
150
151    let mut ret = TokenStream::new();
152    for (ident, module) in &root.sub_modules {
153        let using = module.render();
154        ret.extend(quote!(use #ident::#using;));
155    }
156
157    ret
158}