ryo_source/generator/

tree.rs

//! Module tree structure for source generation.
//!
//! `ModuleTree` represents a hierarchical module structure that can be
//! converted into Rust source code. It captures:
//! - Module name and visibility
//! - Use statements
//! - Inner attributes (module-level doc comments)
//! - Items (structs, functions, etc.)
//! - Child modules (nested `mod {}` blocks)

use crate::pure::{PureAttribute, PureItem, PureUse, PureVis};

/// A hierarchical module tree for source generation.
///
/// This structure represents a complete module hierarchy that can be
/// transformed into source code by a `SourceGenerator`.
///
/// # Example
///
/// ```ignore
/// // Build a tree representing:
/// // mod utils {
/// //     use std::io;
/// //     pub fn helper() {}
/// // }
///
/// let tree = ModuleTree::new("utils")
///     .with_vis(PureVis::Public)
///     .with_use(PureUse { ... })
///     .with_item(PureItem::Fn(helper_fn));
/// ```
#[derive(Debug, Clone, Default)]
pub struct ModuleTree {
    /// Module name (empty string for crate root).
    pub name: String,
    /// Module visibility.
    pub vis: PureVis,
    /// Use statements for this module.
    pub uses: Vec<PureUse>,
    /// Inner attributes (e.g., `//! doc`, `#![allow(...)]`).
    pub inner_attrs: Vec<PureAttribute>,
    /// Items in this module (structs, functions, etc.).
    pub items: Vec<PureItem>,
    /// Child modules.
    pub children: Vec<ModuleTree>,
}

impl ModuleTree {
    /// Create a new module tree with the given name.
    ///
    /// For crate root, use an empty string or "crate".
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            vis: PureVis::Private,
            uses: Vec::new(),
            inner_attrs: Vec::new(),
            items: Vec::new(),
            children: Vec::new(),
        }
    }

    /// Create a crate root module tree.
    pub fn crate_root() -> Self {
        Self::new("")
    }

    /// Set visibility.
    pub fn with_vis(mut self, vis: PureVis) -> Self {
        self.vis = vis;
        self
    }

    /// Add a use statement.
    pub fn with_use(mut self, use_stmt: PureUse) -> Self {
        self.uses.push(use_stmt);
        self
    }

    /// Add multiple use statements.
    pub fn with_uses(mut self, uses: impl IntoIterator<Item = PureUse>) -> Self {
        self.uses.extend(uses);
        self
    }

    /// Add an inner attribute.
    pub fn with_inner_attr(mut self, attr: PureAttribute) -> Self {
        self.inner_attrs.push(attr);
        self
    }

    /// Add an item.
    pub fn with_item(mut self, item: PureItem) -> Self {
        self.items.push(item);
        self
    }

    /// Add multiple items.
    pub fn with_items(mut self, items: impl IntoIterator<Item = PureItem>) -> Self {
        self.items.extend(items);
        self
    }

    /// Add a child module.
    pub fn with_child(mut self, child: ModuleTree) -> Self {
        self.children.push(child);
        self
    }

    /// Add multiple child modules.
    pub fn with_children(mut self, children: impl IntoIterator<Item = ModuleTree>) -> Self {
        self.children.extend(children);
        self
    }

    /// Check if this is the crate root.
    pub fn is_crate_root(&self) -> bool {
        self.name.is_empty() || self.name == "crate"
    }

    /// Get total item count (including children).
    pub fn total_items(&self) -> usize {
        self.items.len() + self.children.iter().map(|c| c.total_items()).sum::<usize>()
    }

    /// Get total module count (including self and children).
    pub fn total_modules(&self) -> usize {
        1 + self
            .children
            .iter()
            .map(|c| c.total_modules())
            .sum::<usize>()
    }

    /// Find a child module by name.
    pub fn find_child(&self, name: &str) -> Option<&ModuleTree> {
        self.children.iter().find(|c| c.name == name)
    }

    /// Find a child module by name (mutable).
    pub fn find_child_mut(&mut self, name: &str) -> Option<&mut ModuleTree> {
        self.children.iter_mut().find(|c| c.name == name)
    }

    /// Get or create a child module by name.
    ///
    /// If a child with the given name exists, returns a mutable reference to it.
    /// Otherwise, creates a new child module and returns a mutable reference.
    pub fn get_or_create_child(&mut self, name: &str) -> &mut ModuleTree {
        if !self.children.iter().any(|c| c.name == name) {
            self.children.push(ModuleTree::new(name));
        }
        self.find_child_mut(name)
            .expect("child with matching name was pushed above when absent")
    }

    /// Navigate to a nested module by path.
    ///
    /// Path is a slice of module names, e.g., `["foo", "bar"]` for `foo::bar`.
    pub fn navigate(&self, path: &[&str]) -> Option<&ModuleTree> {
        if path.is_empty() {
            return Some(self);
        }
        self.find_child(path[0])?.navigate(&path[1..])
    }

    /// Navigate to a nested module by path (mutable).
    pub fn navigate_mut(&mut self, path: &[&str]) -> Option<&mut ModuleTree> {
        if path.is_empty() {
            return Some(self);
        }
        self.find_child_mut(path[0])?.navigate_mut(&path[1..])
    }

    /// Get or create nested modules by path.
    ///
    /// Creates any missing intermediate modules.
    pub fn get_or_create_path(&mut self, path: &[&str]) -> &mut ModuleTree {
        if path.is_empty() {
            return self;
        }
        let child = self.get_or_create_child(path[0]);
        child.get_or_create_path(&path[1..])
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pure::{PureBlock, PureFields, PureFn, PureGenerics, PureStruct, PureVis};

    fn make_struct(name: &str) -> PureItem {
        PureItem::Struct(PureStruct {
            attrs: vec![],
            vis: PureVis::Public,
            name: name.to_string(),
            generics: PureGenerics::default(),
            fields: PureFields::Unit,
        })
    }

    fn make_fn(name: &str) -> PureItem {
        PureItem::Fn(PureFn {
            attrs: vec![],
            vis: PureVis::Public,
            is_async: false,
            is_async_inferred: false,
            is_const: false,
            is_unsafe: false,
            abi: None,
            name: name.to_string(),
            generics: PureGenerics::default(),
            params: vec![],
            ret: None,
            body: PureBlock::default(),
        })
    }

    #[test]
    fn test_module_tree_basic() {
        let tree = ModuleTree::new("utils")
            .with_vis(PureVis::Public)
            .with_item(make_fn("helper"));

        assert_eq!(tree.name, "utils");
        assert_eq!(tree.vis, PureVis::Public);
        assert_eq!(tree.items.len(), 1);
        assert_eq!(tree.total_items(), 1);
        assert_eq!(tree.total_modules(), 1);
    }

    #[test]
    fn test_module_tree_nested() {
        let tree = ModuleTree::crate_root()
            .with_item(make_struct("Config"))
            .with_child(
                ModuleTree::new("models")
                    .with_item(make_struct("User"))
                    .with_child(ModuleTree::new("nested").with_item(make_struct("Deep"))),
            );

        assert!(tree.is_crate_root());
        assert_eq!(tree.total_items(), 3);
        assert_eq!(tree.total_modules(), 3);

        let models = tree.find_child("models").unwrap();
        assert_eq!(models.items.len(), 1);
        assert_eq!(models.children.len(), 1);
    }

    #[test]
    fn test_navigate() {
        let tree = ModuleTree::crate_root().with_child(
            ModuleTree::new("foo").with_child(ModuleTree::new("bar").with_item(make_struct("Baz"))),
        );

        let bar = tree.navigate(&["foo", "bar"]).unwrap();
        assert_eq!(bar.name, "bar");
        assert_eq!(bar.items.len(), 1);

        assert!(tree.navigate(&["nonexistent"]).is_none());
    }

    #[test]
    fn test_get_or_create_path() {
        let mut tree = ModuleTree::crate_root();

        let nested = tree.get_or_create_path(&["a", "b", "c"]);
        nested.items.push(make_struct("Deep"));

        assert!(tree.navigate(&["a", "b", "c"]).is_some());
        assert_eq!(tree.total_modules(), 4); // root + a + b + c
    }
}