Untitled

use std::collections::HashMap;
use syn::Ident;
use quote::quote;
use proc_macro2::TokenStream;

struct Module {
    name: Ident,
    content: Vec<TokenStream>,
    submodules: HashMap<String, Module>,
}

impl Module {
    fn new(name: &str) -> Self {
        Self {
            name: Ident::new(name, proc_macro2::Span::call_site()),
            content: Vec::new(),
            submodules: HashMap::new(),
        }
    }

    fn add_content(&mut self, content: TokenStream) {
        self.content.push(content);
    }

    fn get_or_create_submodule(&mut self, name: &str) -> &mut Module {
        if !self.submodules.contains_key(name) {
            self.submodules.insert(name.to_string(), Module::new(name));
        }
        self.submodules.get_mut(name).unwrap()
    }

    fn generate(&self) -> TokenStream {
        let name = &self.name;
        let content = &self.content;
        let submodules: Vec<TokenStream> = self.submodules.values().map(|m| m.generate()).collect();
        quote! {
            pub mod #name {
                #( #content )*
                #( #submodules )*
            }
        }
    }
}

struct ModuleComposer {
    root: Module,
}

impl ModuleComposer {
    fn new() -> Self {
        Self {
            root: Module::new("root"),
        }
    }

    fn add(&mut self, path: &str, content: TokenStream) {
        let mut current_module = &mut self.root;
        for name in path.split("::") {
            current_module = current_module.get_or_create_submodule(name);
        }
        current_module.add_content(content);
    }

    fn generate(&self) -> TokenStream {
        self.root.generate()
    }
}