codecraft 0.1.2

//! The symbols in a Rust file, read with `syn` without compiling anything, and how they tie together.
use std::collections::HashMap;
use std::path::Path;

use syn::spanned::Spanned;
use syn::{Attribute, Item, Type, UseTree, Visibility};

use super::graph::{Graph, Id, Kind, Node, Relation};

struct Impl {
    module: Id,
    type_name: String,
    trait_name: Option<String>,
    lines: usize,
}

struct Use {
    module: Id,
    path: Vec<String>,
}

/// Relations seen while reading, resolved by [`link`] once every file is in the graph.
#[derive(Default)]
pub struct Pending {
    impls: Vec<Impl>,
    uses: Vec<Use>,
}

impl Pending {
    pub fn extend(&mut self, other: Pending) {
        self.impls.extend(other.impls);
        self.uses.extend(other.uses);
    }
}

/// Reads the Rust source at `path` into the graph under its `file` node.
pub fn index_file(graph: &mut Graph, file: Id, path: &Path) -> Pending {
    let mut pending = Pending::default();
    let source = match std::fs::read_to_string(path) {
        Ok(source) => source,
        Err(error) => {
            log::warn!("{}: {error}", path.display());
            return pending;
        }
    };
    let ast = match syn::parse_file(&source) {
        Ok(ast) => ast,
        Err(error) => {
            log::warn!("{}: {error}", path.display());
            return pending;
        }
    };
    items(graph, &mut pending, file, &ast.items, path);
    pending
}

fn items(graph: &mut Graph, pending: &mut Pending, parent: Id, items: &[Item], file: &Path) {
    for item in items {
        let span = item.span();
        let line = span.start().line;
        let lines = span.end().line + 1 - line;
        let symbol = |kind: Kind, ident: &syn::Ident, vis: &Visibility| {
            Node::new(kind, ident.to_string())
                .at(file.to_path_buf(), line, lines)
                .visible(is_public(vis))
        };
        match item {
            // A `mod x;` is a file of its own and already in the tree; only inline modules are new.
            Item::Mod(module) => {
                if is_test(&module.attrs) {
                    continue;
                }
                if let Some((_, inner)) = &module.content {
                    let id = graph.add(parent, symbol(Kind::Module, &module.ident, &module.vis));
                    self::items(graph, pending, id, inner, file);
                }
            }
            Item::Struct(s) => {
                graph.add(parent, symbol(Kind::Struct, &s.ident, &s.vis));
            }
            Item::Enum(e) => {
                graph.add(parent, symbol(Kind::Enum, &e.ident, &e.vis));
            }
            Item::Union(u) => {
                graph.add(parent, symbol(Kind::Union, &u.ident, &u.vis));
            }
            Item::Trait(t) => {
                graph.add(parent, symbol(Kind::Trait, &t.ident, &t.vis));
            }
            Item::Fn(f) => {
                graph.add(parent, symbol(Kind::Fn, &f.sig.ident, &f.vis));
            }
            Item::Const(c) => {
                graph.add(parent, symbol(Kind::Const, &c.ident, &c.vis));
            }
            Item::Static(s) => {
                graph.add(parent, symbol(Kind::Static, &s.ident, &s.vis));
            }
            Item::Type(t) => {
                graph.add(parent, symbol(Kind::TypeAlias, &t.ident, &t.vis));
            }
            Item::Macro(m) => {
                if let Some(ident) = &m.ident {
                    let exported = m.attrs.iter().any(|a| a.path().is_ident("macro_export"));
                    let node = Node::new(Kind::Macro, ident.to_string())
                        .at(file.to_path_buf(), line, lines)
                        .visible(exported);
                    graph.add(parent, node);
                }
            }
            Item::Impl(imp) => {
                if let Some(type_name) = last_segment(&imp.self_ty) {
                    let trait_name = imp
                        .trait_
                        .as_ref()
                        .and_then(|(_, path, _)| path.segments.last())
                        .map(|segment| segment.ident.to_string());
                    pending.impls.push(Impl {
                        module: parent,
                        type_name,
                        trait_name,
                        lines,
                    });
                }
            }
            Item::Use(u) => {
                let mut paths = Vec::new();
                flatten(&u.tree, Vec::new(), &mut paths);
                pending.uses.extend(paths.into_iter().map(|path| Use {
                    module: parent,
                    path,
                }));
            }
            _ => {}
        }
    }
}

fn is_public(vis: &Visibility) -> bool {
    matches!(vis, Visibility::Public(_))
}

fn is_test(attrs: &[Attribute]) -> bool {
    attrs.iter().any(|attr| {
        attr.path().is_ident("cfg")
            && matches!(&attr.meta, syn::Meta::List(list) if list.tokens.to_string().contains("test"))
    })
}

fn last_segment(ty: &Type) -> Option<String> {
    match ty {
        Type::Path(path) => path.path.segments.last().map(|s| s.ident.to_string()),
        Type::Reference(reference) => last_segment(&reference.elem),
        _ => None,
    }
}

fn flatten(tree: &UseTree, mut prefix: Vec<String>, out: &mut Vec<Vec<String>>) {
    match tree {
        UseTree::Path(path) => {
            prefix.push(path.ident.to_string());
            flatten(&path.tree, prefix, out);
        }
        UseTree::Name(name) => {
            if name.ident != "self" {
                prefix.push(name.ident.to_string());
            }
            out.push(prefix);
        }
        UseTree::Rename(rename) => {
            if rename.ident != "self" {
                prefix.push(rename.ident.to_string());
            }
            out.push(prefix);
        }
        UseTree::Glob(_) => out.push(prefix),
        UseTree::Group(group) => {
            for item in &group.items {
                flatten(item, prefix.clone(), out);
            }
        }
    }
}

/// Turns the pending impls and uses into `Implements` and `Uses` relations.
pub fn link(graph: &mut Graph, pending: Pending) {
    let index = Index::of(graph);

    for imp in pending.impls {
        let Some(krate) = graph.crate_of(imp.module) else {
            continue;
        };
        let Some(type_id) = index.local(
            graph,
            krate,
            &imp.type_name,
            |kind| kind.is_type(),
            imp.module,
        ) else {
            continue;
        };
        graph.node_mut(type_id).lines += imp.lines;
        if let Some(name) = &imp.trait_name
            && let Some(trait_id) = index
                .local(graph, krate, name, |kind| kind == Kind::Trait, imp.module)
                .or_else(|| index.trait_anywhere(name))
        {
            graph.relate(type_id, trait_id, Relation::Implements);
        }
    }

    for u in pending.uses {
        let Some(target) = index.resolve(graph, u.module, &u.path) else {
            continue;
        };
        if target == u.module || graph.ancestors(u.module).contains(&target) {
            continue;
        }
        graph.relate(u.module, target, Relation::Uses);
    }
}

struct Index {
    /// Per crate, every node by name.
    local: HashMap<Id, HashMap<String, Vec<Id>>>,
    traits: HashMap<String, Vec<Id>>,
    crates: HashMap<String, Id>,
}

impl Index {
    fn of(graph: &Graph) -> Self {
        let mut local: HashMap<Id, HashMap<String, Vec<Id>>> = HashMap::new();
        let mut traits: HashMap<String, Vec<Id>> = HashMap::new();
        let mut crates = HashMap::new();
        for krate in graph.crates() {
            crates.insert(graph.node(krate).name.replace('-', "_"), krate);
            let names = local.entry(krate).or_default();
            for id in graph.descendants(krate) {
                let node = graph.node(id);
                names.entry(node.name.clone()).or_default().push(id);
                if node.kind == Kind::Trait {
                    traits.entry(node.name.clone()).or_default().push(id);
                }
            }
        }
        Self {
            local,
            traits,
            crates,
        }
    }

    /// A node called `name` in `krate`, preferring one in `near`.
    fn local(
        &self,
        graph: &Graph,
        krate: Id,
        name: &str,
        kind: impl Fn(Kind) -> bool,
        near: Id,
    ) -> Option<Id> {
        let candidates = self.local.get(&krate)?.get(name)?;
        candidates
            .iter()
            .copied()
            .filter(|&id| kind(graph.node(id).kind))
            .min_by_key(|&id| (graph.parent(id) != Some(near), id))
    }

    /// A trait by name across every crate, when there is exactly one.
    fn trait_anywhere(&self, name: &str) -> Option<Id> {
        match self.traits.get(name).map(Vec::as_slice) {
            Some([only]) => Some(*only),
            _ => None,
        }
    }

    /// The deepest node a `use` path reaches, starting from `module`.
    fn resolve(&self, graph: &Graph, module: Id, path: &[String]) -> Option<Id> {
        let (mut current, rest) = match path.first()?.as_str() {
            "crate" => (crate_root(graph, graph.crate_of(module)?), &path[1..]),
            "self" => (module, &path[1..]),
            "super" => {
                let mut up = module;
                let mut rest = path;
                while rest.first().is_some_and(|s| s == "super") {
                    up = super_of(graph, up)?;
                    rest = &rest[1..];
                }
                (up, rest)
            }
            first => match self.crates.get(first) {
                Some(&krate) => (crate_root(graph, krate), &path[1..]),
                None => (descend(graph, module, first)?, &path[1..]),
            },
        };
        for segment in rest {
            match descend(graph, current, segment) {
                Some(child) => current = child,
                None => break,
            }
        }
        Some(current)
    }
}

/// The file a crate's paths start from: `src/lib.rs` or `src/main.rs`, else the crate itself.
fn crate_root(graph: &Graph, krate: Id) -> Id {
    graph
        .child_named(krate, "src")
        .and_then(|src| module_file(graph, src))
        .unwrap_or(krate)
}

/// The file that speaks for a directory: `mod.rs`, `lib.rs` or `main.rs`.
fn module_file(graph: &Graph, dir: Id) -> Option<Id> {
    ["mod.rs", "lib.rs", "main.rs"]
        .into_iter()
        .find_map(|name| graph.child_named(dir, name))
}

/// One path segment on from `node`: an item inside it, a sibling file, or a directory's file.
fn descend(graph: &Graph, node: Id, segment: &str) -> Option<Id> {
    let file = format!("{segment}.rs");
    match graph.node(node).kind {
        Kind::File | Kind::Module => graph.child_named(node, segment).or_else(|| {
            let dir = graph.parent(node)?;
            graph
                .child_named(dir, &file)
                .or_else(|| graph.child_named(dir, segment))
        }),
        _ => graph
            .child_named(node, segment)
            .or_else(|| graph.child_named(node, &file))
            .or_else(|| descend(graph, module_file(graph, node)?, segment)),
    }
}

/// The module above `node`: the file speaking for the directory above a file, or a parent item.
fn super_of(graph: &Graph, node: Id) -> Option<Id> {
    let n = graph.node(node);
    match n.kind {
        Kind::File => {
            let dir = graph.parent(node)?;
            let above = match n.name.as_str() {
                "mod.rs" | "lib.rs" | "main.rs" => graph.parent(dir)?,
                _ => dir,
            };
            Some(module_file(graph, above).unwrap_or(above))
        }
        _ => graph.parent(node),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;

    struct Tree(PathBuf);

    impl Tree {
        fn new(name: &str, files: &[(&str, &str)]) -> Self {
            let dir = std::env::temp_dir().join(format!("codecraft-{name}-{}", std::process::id()));
            let _ = std::fs::remove_dir_all(&dir);
            for (path, source) in files {
                let file = dir.join(path);
                std::fs::create_dir_all(file.parent().unwrap()).unwrap();
                std::fs::write(file, source).unwrap();
            }
            Self(dir)
        }
    }

    impl Drop for Tree {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.0);
        }
    }

    /// The tree as [`crate::code::index`] would build it, without cargo.
    fn indexed(tree: &Tree) -> Graph {
        fn walk(graph: &mut Graph, pending: &mut Pending, parent: Id, dir: &Path) {
            let mut entries: Vec<_> = std::fs::read_dir(dir).unwrap().flatten().collect();
            entries.sort_by_key(|e| e.file_name());
            for entry in entries {
                let path = entry.path();
                let name = entry.file_name().to_string_lossy().into_owned();
                if path.is_dir() {
                    let id = graph.add(parent, Node::new(Kind::Dir, name));
                    walk(graph, pending, id, &path);
                } else {
                    let node = Node::new(Kind::File, name).at(path.clone(), 1, 1);
                    let id = graph.add(parent, node);
                    if path.extension().is_some_and(|e| e == "rs") {
                        pending.extend(index_file(graph, id, &path));
                    }
                }
            }
        }
        let mut graph = Graph::new("ws");
        let krate = graph.add(graph.root(), Node::new(Kind::Crate, "app"));
        let mut pending = Pending::default();
        walk(&mut graph, &mut pending, krate, &tree.0);
        link(&mut graph, pending);
        graph
    }

    #[test]
    fn a_file_holds_its_items_and_inline_modules() {
        let tree = Tree::new(
            "files",
            &[
                (
                    "src/lib.rs",
                    "pub mod flat;\nmod inline { pub struct In; }\n#[cfg(test)]\nmod tests { fn t() {} }\npub struct Top;\n",
                ),
                ("src/flat.rs", "pub struct Flat;\n"),
            ],
        );
        let graph = indexed(&tree);
        assert!(graph.find("app/src/lib.rs::Top").is_some());
        assert!(graph.find("app/src/lib.rs::inline::In").is_some());
        assert!(graph.find("app/src/flat.rs::Flat").is_some());
        assert!(
            graph.find("app/src/lib.rs::tests").is_none(),
            "test modules are skipped"
        );
        assert!(
            graph.find("app/src/lib.rs::flat").is_none(),
            "a `mod x;` is the file, not a node"
        );
    }

    #[test]
    fn impls_and_uses_become_relations_across_files() {
        let tree = Tree::new(
            "relations",
            &[
                (
                    "src/lib.rs",
                    "pub mod shapes;\npub mod draw;\npub trait Draw { fn draw(&self); }\n",
                ),
                (
                    "src/shapes/mod.rs",
                    "use crate::Draw;\npub struct Circle;\nimpl Draw for Circle { fn draw(&self) {} }\nimpl Circle {\n    pub fn new() -> Self { Circle }\n}\nimpl Default for Circle { fn default() -> Self { Circle } }\n",
                ),
                (
                    "src/draw.rs",
                    "use super::shapes::{self, Circle};\nuse std::fmt::Display;\npub fn all() -> Circle { Circle }\n",
                ),
            ],
        );
        let graph = indexed(&tree);
        let circle = graph.find("app/src/shapes/mod.rs::Circle").unwrap();
        let draw_trait = graph.find("app/src/lib.rs::Draw").unwrap();
        assert_eq!(
            graph.related(circle, Relation::Implements),
            vec![draw_trait]
        );
        assert_eq!(
            graph.node(circle).lines,
            1 + 1 + 3 + 1,
            "impl blocks count towards the type"
        );
        let shapes = graph.find("app/src/shapes/mod.rs").unwrap();
        let draw = graph.find("app/src/draw.rs").unwrap();
        assert_eq!(graph.related(shapes, Relation::Uses), vec![draw_trait]);
        let mut used = graph.related(draw, Relation::Uses);
        used.sort();
        let mut wanted = vec![graph.find("app/src/shapes").unwrap(), circle];
        wanted.sort();
        assert_eq!(used, wanted);
    }
}