xsd_parser/models/code/

module.rs

use std::collections::BTreeMap;
use std::fs::{create_dir_all, write};
use std::io::Error as IoError;
use std::mem::replace;
use std::path::Path;
use std::str::FromStr;

use proc_macro2::{Ident as Ident2, TokenStream};
use quote::{format_ident, quote, ToTokens};

use super::IdentPath;

/// Manages the code that is generated for a specific module.
///
/// This type stores the generated code, the related using directives as well as
/// any sub-module for a specific module generated by the code generator.
#[derive(Default, Debug)]
pub struct Module {
    /// The actual code that is manages by this module.
    pub code: TokenStream,

    /// A set of using directives this module needs.
    pub usings: BTreeMap<String, bool>,

    /// A map of sub-modules contained inside this module.
    pub modules: BTreeMap<String, Module>,
}

/// Defines how the sub modules of the [`Module`] will be generated when calling
/// [`Module::to_code()`].
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SubModules {
    /// No code is generated for any sub-module.
    None,

    /// The sub-modules are references as separate files.
    ///
    /// This will generate something like this:
    ///
    /// ```rust,ignore
    /// pub mod sub_module;
    /// pub mod another_module;
    ///
    /// // Actual module code
    /// ```
    Files,

    /// The sub-modules are generated as inline modules.
    ///
    /// This will generate something like this:
    ///
    /// ```rust,ignore
    /// // Actual module code
    ///
    /// pub mod sub_module {
    ///     // Sub-Module Code
    /// }
    ///
    /// pub mod another_module {
    ///     // Sub-Module Code
    /// }
    /// ```
    Inline,
}

impl Module {
    /// Append the passed `code` to this module code.
    pub fn append(&mut self, code: TokenStream) -> &mut Self {
        self.code.extend(code);

        self
    }

    /// Prepend the passed `code` to this module code.
    pub fn prepend(&mut self, code: TokenStream) -> &mut Self {
        let code = replace(&mut self.code, code);

        self.append(code)
    }

    /// Add using directives to the set of this module.
    pub fn usings<I>(&mut self, anonymous: bool, usings: I) -> &mut Self
    where
        I: IntoIterator,
        I::Item: ToString,
    {
        for using in usings {
            *self.usings.entry(using.to_string()).or_insert(anonymous) &= anonymous;
        }

        self
    }

    /// Get a reference to a sub-module identified by the passed `ident`.
    ///
    /// If the module does not exist `None` is returned.
    pub fn module<T>(&self, ident: T) -> Option<&Module>
    where
        T: AsRef<str>,
    {
        self.modules.get(ident.as_ref())
    }

    /// Get a mutable reference to a sub-module identified by the passed `ident`.
    ///
    /// If the module does not exist it will be created.
    pub fn module_mut<T>(&mut self, ident: T) -> &mut Module
    where
        T: Into<String>,
    {
        self.modules.entry(ident.into()).or_default()
    }

    /// Write the module to the passed `token` stream.
    ///
    /// This method writes the current module to the passed `token` stream. The
    /// `sub_modules` parameter defines how the sub-modules of this module will
    /// be rendered.
    pub fn to_code(&self, tokens: &mut TokenStream, sub_modules: SubModules) {
        let Self {
            code,
            usings,
            modules,
        } = self;

        if sub_modules == SubModules::Files {
            let modules = modules.iter().map(|(ident, _module)| {
                let name = format_ident!("{ident}");

                quote!(pub mod #name;)
            });

            tokens.extend(quote! {
                #( #modules )*
            });
        }

        let usings = render_usings(usings.iter());

        tokens.extend(quote! {
            #usings
            #code
        });

        if sub_modules == SubModules::Inline {
            for (ident, module) in modules {
                let name = format_ident!("{ident}");

                tokens.extend(quote! {
                    pub mod #name {
                        #module
                    }
                });
            }
        }
    }

    /// Write the code of the current module to the passed `directory`.
    ///
    /// This will split up the current module into it's sub modules and
    /// writes them to the passed `directory`. Each sub-module is written to
    /// it's own module file.
    ///
    /// # Errors
    /// Returns an [`IoError`] if writing a module to a file failed.
    pub fn write_to_files<P>(&self, directory: P) -> Result<(), IoError>
    where
        P: AsRef<Path>,
    {
        let directory = directory.as_ref();

        self.write_to_files_with(|module: &Module, path: &Path| {
            let filename = if module.modules.is_empty() {
                directory.join(path).with_extension("rs")
            } else {
                directory.join(path).join("mod.rs")
            };
            create_dir_all(filename.parent().unwrap())?;

            let mut code = TokenStream::new();
            module.to_code(&mut code, SubModules::Files);

            let code = code.to_string();
            write(filename, code)?;

            Ok(())
        })
    }

    /// Write the code of the different sub-modules by using the passed
    /// function `f`.
    ///
    /// This will split up the current module into it's sub modules and
    /// passes the generated code for each module to the passed function `f`.
    ///
    /// # Errors
    /// Forwards the error raised by `f`.
    pub fn write_to_files_with<F, E>(&self, mut f: F) -> Result<(), E>
    where
        F: FnMut(&Module, &Path) -> Result<(), E>,
    {
        self.write_to_files_with_impl("", &mut f)
    }

    fn write_to_files_with_impl<P, F, E>(&self, path: P, f: &mut F) -> Result<(), E>
    where
        P: AsRef<Path>,
        F: FnMut(&Module, &Path) -> Result<(), E>,
    {
        let path = path.as_ref();

        f(self, path)?;

        for (ident, module) in &self.modules {
            let path = path.join(ident);

            module.write_to_files_with_impl(path, f)?;
        }

        Ok(())
    }
}

impl ToTokens for Module {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.to_code(tokens, SubModules::Inline);
    }
}

fn render_usings<'x, I, X>(usings: I) -> TokenStream
where
    I: IntoIterator<Item = (X, &'x bool)>,
    X: AsRef<str>,
{
    #[derive(Default)]
    struct Module {
        usings: BTreeMap<Ident2, bool>,
        sub_modules: BTreeMap<Ident2, Module>,
    }

    impl Module {
        fn render(&self) -> TokenStream {
            let count = self.usings.len() + self.sub_modules.len();

            let usings = self.usings.iter().map(|(ident, anonymous)| {
                if *anonymous {
                    quote!(#ident as _)
                } else {
                    quote!(#ident)
                }
            });
            let sub_modules = self.sub_modules.iter().map(|(ident, module)| {
                let using = module.render();

                quote!(#ident::#using)
            });

            let items = usings.chain(sub_modules);

            if count > 1 {
                quote!({ #( #items ),* })
            } else {
                quote!(#( #items )*)
            }
        }
    }

    let mut root = Module::default();

    for (using, anonymous) in usings {
        let using = using.as_ref();
        let Ok(ident) = IdentPath::from_str(using) else {
            continue;
        };

        let (ident, path, _) = ident.into_parts();

        let mut module = &mut root;
        for part in path.into_iter().flat_map(|x| x.0) {
            module = module.sub_modules.entry(part).or_default();
        }

        *module.usings.entry(ident).or_insert(*anonymous) &= *anonymous;
    }

    let mut ret = TokenStream::new();
    for (ident, module) in &root.sub_modules {
        let using = module.render();
        ret.extend(quote!(use #ident::#using;));
    }

    ret
}