codecraft 0.1.2

//! The code graph: a tree of workspace, crates, modules and symbols, with the relations between them.
use std::path::PathBuf;

use petgraph::Direction;
use petgraph::graph::DiGraph;
use petgraph::visit::EdgeRef;

pub use petgraph::graph::NodeIndex as Id;

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum Kind {
    Workspace,
    Crate,
    Dir,
    File,
    Module,
    Struct,
    Enum,
    Union,
    Trait,
    Fn,
    Const,
    Static,
    TypeAlias,
    Macro,
}

impl Kind {
    pub const ALL: [Kind; 14] = [
        Kind::Workspace,
        Kind::Crate,
        Kind::Dir,
        Kind::File,
        Kind::Module,
        Kind::Struct,
        Kind::Enum,
        Kind::Union,
        Kind::Trait,
        Kind::Fn,
        Kind::Const,
        Kind::Static,
        Kind::TypeAlias,
        Kind::Macro,
    ];

    pub fn label(self) -> &'static str {
        match self {
            Kind::Workspace => "workspace",
            Kind::Crate => "crate",
            Kind::Dir => "folder",
            Kind::File => "file",
            Kind::Module => "mod",
            Kind::Struct => "struct",
            Kind::Enum => "enum",
            Kind::Union => "union",
            Kind::Trait => "trait",
            Kind::Fn => "fn",
            Kind::Const => "const",
            Kind::Static => "static",
            Kind::TypeAlias => "type",
            Kind::Macro => "macro",
        }
    }

    /// Whether nodes of this kind hold other nodes.
    pub fn is_container(self) -> bool {
        matches!(
            self,
            Kind::Workspace | Kind::Crate | Kind::Dir | Kind::File | Kind::Module
        )
    }

    /// Whether this is something on disk rather than something in a file.
    pub fn is_on_disk(self) -> bool {
        matches!(self, Kind::Workspace | Kind::Crate | Kind::Dir | Kind::File)
    }

    /// Whether an `impl` block can be for this kind.
    pub fn is_type(self) -> bool {
        matches!(
            self,
            Kind::Struct | Kind::Enum | Kind::Union | Kind::TypeAlias
        )
    }
}

#[derive(Clone, Debug)]
pub struct Node {
    pub kind: Kind,
    pub name: String,
    /// `dir/file.rs::module::Name` from the workspace, whose own path is empty.
    pub path: String,
    pub file: Option<PathBuf>,
    pub line: usize,
    /// Lines of source, with a type's impl blocks counted in.
    pub lines: usize,
    pub public: bool,
}

impl Node {
    pub fn new(kind: Kind, name: impl Into<String>) -> Self {
        let name = name.into();
        Self {
            kind,
            path: name.clone(),
            name,
            file: None,
            line: 0,
            lines: 0,
            public: true,
        }
    }

    pub fn at(mut self, file: PathBuf, line: usize, lines: usize) -> Self {
        self.file = Some(file);
        self.line = line;
        self.lines = lines;
        self
    }

    pub fn visible(mut self, public: bool) -> Self {
        self.public = public;
        self
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Relation {
    Contains,
    DependsOn,
    Implements,
    Uses,
}

impl Relation {
    pub fn label(self) -> &'static str {
        match self {
            Relation::Contains => "contains",
            Relation::DependsOn => "depends on",
            Relation::Implements => "implements",
            Relation::Uses => "uses",
        }
    }
}

#[derive(Clone, Debug)]
pub struct Graph {
    inner: DiGraph<Node, Relation>,
    root: Id,
}

impl Graph {
    pub fn new(workspace: impl Into<String>) -> Self {
        let mut inner = DiGraph::new();
        let mut node = Node::new(Kind::Workspace, workspace);
        node.path = String::new();
        let root = inner.add_node(node);
        Self { inner, root }
    }

    pub fn root(&self) -> Id {
        self.root
    }

    /// Adds `node` under `parent`, giving it its path: `/` between things on disk, `::` inside a file.
    pub fn add(&mut self, parent: Id, mut node: Node) -> Id {
        let parent_node = &self.inner[parent];
        if parent_node.kind != Kind::Workspace {
            let join = match parent_node.kind.is_on_disk() && node.kind.is_on_disk() {
                true => "/",
                false => "::",
            };
            node.path = format!("{}{join}{}", parent_node.path, node.name);
        }
        let id = self.inner.add_node(node);
        self.inner.add_edge(parent, id, Relation::Contains);
        id
    }

    /// Adds a relation once; a node never relates to itself.
    pub fn relate(&mut self, from: Id, to: Id, relation: Relation) {
        if from == to {
            return;
        }
        let already = self
            .inner
            .edges_connecting(from, to)
            .any(|edge| *edge.weight() == relation);
        if !already {
            self.inner.add_edge(from, to, relation);
        }
    }

    pub fn node(&self, id: Id) -> &Node {
        &self.inner[id]
    }

    pub fn node_mut(&mut self, id: Id) -> &mut Node {
        &mut self.inner[id]
    }

    pub fn ids(&self) -> impl Iterator<Item = Id> + '_ {
        self.inner.node_indices()
    }

    pub fn len(&self) -> usize {
        self.inner.node_count()
    }

    pub fn is_empty(&self) -> bool {
        self.inner.node_count() == 0
    }

    pub fn count(&self, kind: Kind) -> usize {
        self.inner
            .node_weights()
            .filter(|node| node.kind == kind)
            .count()
    }

    pub fn parent(&self, id: Id) -> Option<Id> {
        self.inner
            .edges_directed(id, Direction::Incoming)
            .find(|edge| *edge.weight() == Relation::Contains)
            .map(|edge| edge.source())
    }

    /// In the order they were added.
    pub fn children(&self, id: Id) -> Vec<Id> {
        let mut children: Vec<Id> = self
            .inner
            .edges_directed(id, Direction::Outgoing)
            .filter(|edge| *edge.weight() == Relation::Contains)
            .map(|edge| edge.target())
            .collect();
        children.sort();
        children
    }

    pub fn child_named(&self, id: Id, name: &str) -> Option<Id> {
        self.children(id)
            .into_iter()
            .find(|&child| self.inner[child].name == name)
    }

    /// Nearest first, ending at the root.
    pub fn ancestors(&self, id: Id) -> Vec<Id> {
        let mut ancestors = Vec::new();
        let mut current = id;
        while let Some(parent) = self.parent(current) {
            ancestors.push(parent);
            current = parent;
        }
        ancestors
    }

    pub fn depth(&self, id: Id) -> usize {
        self.ancestors(id).len()
    }

    /// Every node under `id`, not counting itself.
    pub fn descendants(&self, id: Id) -> Vec<Id> {
        let mut out = Vec::new();
        let mut stack = self.children(id);
        while let Some(next) = stack.pop() {
            out.push(next);
            stack.extend(self.children(next));
        }
        out
    }

    /// Every crate, wherever it sits in the tree.
    pub fn crates(&self) -> Vec<Id> {
        self.ids()
            .filter(|&id| self.inner[id].kind == Kind::Crate)
            .collect()
    }

    pub fn crate_named(&self, name: &str) -> Option<Id> {
        self.crates()
            .into_iter()
            .find(|&id| self.inner[id].name == name)
    }

    pub fn crate_of(&self, id: Id) -> Option<Id> {
        std::iter::once(id)
            .chain(self.ancestors(id))
            .find(|&node| self.inner[node].kind == Kind::Crate)
    }

    pub fn dependencies(&self, krate: Id) -> Vec<Id> {
        self.related(krate, Relation::DependsOn)
    }

    /// What `id` points at with `relation`.
    pub fn related(&self, id: Id, relation: Relation) -> Vec<Id> {
        let mut out: Vec<Id> = self
            .inner
            .edges_directed(id, Direction::Outgoing)
            .filter(|edge| *edge.weight() == relation)
            .map(|edge| edge.target())
            .collect();
        out.sort();
        out
    }

    /// Every relation but containment, in both directions: `(other, relation, outgoing)`.
    pub fn relations(&self, id: Id) -> Vec<(Id, Relation, bool)> {
        let mut out = Vec::new();
        for edge in self.inner.edges_directed(id, Direction::Outgoing) {
            if *edge.weight() != Relation::Contains {
                out.push((edge.target(), *edge.weight(), true));
            }
        }
        for edge in self.inner.edges_directed(id, Direction::Incoming) {
            if *edge.weight() != Relation::Contains {
                out.push((edge.source(), *edge.weight(), false));
            }
        }
        out.sort_by_key(|&(other, relation, outgoing)| (!outgoing, relation.label(), other));
        out
    }

    pub fn find(&self, path: &str) -> Option<Id> {
        self.ids().find(|&id| self.inner[id].path == path)
    }
}

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

    fn sample() -> (Graph, Id, Id, Id) {
        let mut graph = Graph::new("ws");
        let krate = graph.add(graph.root(), Node::new(Kind::Crate, "app"));
        let module = graph.add(krate, Node::new(Kind::Module, "net"));
        let symbol = graph.add(module, Node::new(Kind::Struct, "Socket"));
        (graph, krate, module, symbol)
    }

    #[test]
    fn paths_follow_containment() {
        let (mut graph, krate, module, symbol) = sample();
        assert_eq!(graph.node(krate).path, "app");
        assert_eq!(graph.node(module).path, "app::net");
        assert_eq!(graph.node(symbol).path, "app::net::Socket");
        assert_eq!(graph.find("app::net::Socket"), Some(symbol));
        let dir = graph.add(krate, Node::new(Kind::Dir, "src"));
        let file = graph.add(dir, Node::new(Kind::File, "lib.rs"));
        let item = graph.add(file, Node::new(Kind::Fn, "main"));
        assert_eq!(graph.node(item).path, "app/src/lib.rs::main");
        assert_eq!(graph.crate_named("app"), Some(krate));
    }

    #[test]
    fn the_tree_can_be_walked_both_ways() {
        let (graph, krate, module, symbol) = sample();
        assert_eq!(graph.parent(symbol), Some(module));
        assert_eq!(graph.ancestors(symbol), vec![module, krate, graph.root()]);
        assert_eq!(graph.depth(symbol), 3);
        assert_eq!(graph.children(krate), vec![module]);
        assert_eq!(graph.child_named(module, "Socket"), Some(symbol));
        assert_eq!(graph.descendants(krate), vec![module, symbol]);
        assert_eq!(graph.crate_of(symbol), Some(krate));
        assert_eq!(graph.crates(), vec![krate]);
    }

    #[test]
    fn relations_are_deduplicated_and_never_reflexive() {
        let (mut graph, krate, module, symbol) = sample();
        graph.relate(module, symbol, Relation::Uses);
        graph.relate(module, symbol, Relation::Uses);
        graph.relate(symbol, symbol, Relation::Uses);
        assert_eq!(graph.related(module, Relation::Uses), vec![symbol]);
        assert_eq!(
            graph.relations(symbol),
            vec![(module, Relation::Uses, false)]
        );
        assert!(graph.relations(krate).is_empty());
        assert_eq!(graph.count(Kind::Struct), 1);
    }
}