techtree/techtree-manager/src/tree.rs

use std::collections::BTreeMap;

use petgraph::visit::EdgeRef as _;

const HASH_EPOCH: u32 = 0x00000001;

/// Element in the techtree
#[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Serialize)]
pub struct Element {
    /// Issue associated with this element
    pub issue_number: u64,
    /// Description of this element
    pub description: String,
    /// Type of this element
    pub ty: String,
    /// Completion status of this element.
    pub status: ElementStatus,
}

#[derive(Debug, PartialEq, Eq, Clone, Copy)]
pub enum ElementRole {
    Root,
    Ultimate,
    Intermediate,
    Disjoint,
}

impl Element {
    fn to_dot_node_attributes(
        &self,
        role: Option<ElementRole>,
        subtree_role: Option<SubtreeElementRole>,
    ) -> String {
        let Element {
            issue_number,
            description,
            ty,
            status,
        } = self;

        let mut attributes = Vec::new();

        let description = description
            .replace("<", "&lt;")
            .replace(">", "&gt;")
            .replace("&", "&amp;")
            .replace("'", "&apos;")
            .replace("\"", "&quot;");

        attributes.push(format!(
            r##"label = <{{{{#{issue_number} | {status}}}|<I>{ty}</I>|<B>{description}</B>}}>"##
        ));
        attributes.push(r#"shape = "record""#.to_owned());

        let (color, background) = match (subtree_role, status) {
            (Some(SubtreeElementRole::ElementOfInterest), _) => ("black", "white"),
            (Some(SubtreeElementRole::Dependant), _) => ("gray", "gray"),
            (_, ElementStatus::Missing) => (
                "#800",
                // Highlight root elements
                if role == Some(ElementRole::Root) {
                    "#ffddc1"
                } else {
                    "#fcc"
                },
            ),
            (_, ElementStatus::Assigned) => ("#a50", "#ffa"),
            (_, ElementStatus::Completed) => ("#080", "#afa"),
        };
        attributes.push(format!(r#"color = "{color}""#));
        attributes.push(format!(r#"fontcolor = "{color}""#));
        attributes.push(format!(r#"fillcolor = "{background}""#));
        attributes.push(format!(r#"style = "filled""#));

        attributes.join(", ")
    }

    fn to_mermaid_node(
        &self,
        index: ElementIndex,
        role: Option<SubtreeElementRole>,
        repo_url: &str,
        simple: bool,
    ) -> String {
        let Element {
            issue_number,
            description,
            ty,
            status,
        } = self;

        let label = if simple {
            format!(r##"#{issue_number} | {status}<br/><i>{ty}</i><br/><b>{description}</b>"##)
        } else {
            format!(
                r##"<a href='{repo_url}/issues/{issue_number}' target='_blank'>#{issue_number}</a> | {status}<br/><i>{ty}</i><br/><b>{description}</b>"##
            )
        };

        let class = match (role, status) {
            (Some(SubtreeElementRole::ElementOfInterest), _) => "eoi",
            (Some(SubtreeElementRole::Dependant), _) => "dependant",
            (_, ElementStatus::Missing) => "dep_missing",
            (_, ElementStatus::Assigned) => "dep_assigned",
            (_, ElementStatus::Completed) => "dep_completed",
        };

        format!(
            "    {index}:::{class}\n    {index}[\"{label}\"]",
            index = index.index(),
        )
    }
}

fn mermaid_classes() -> String {
    r##"
    classDef eoi fill:#fff, stroke:#000, color:#000;
    classDef dependant fill:#fff, stroke:#888, color:#888;
    classDef dep_missing fill:#fcc, stroke:#800, color:#000;
    classDef dep_assigned fill:#ffa, stroke:#a50, color:#000;
    classDef dep_completed fill:#afa, stroke:#080, color:#000;
"##
    .to_owned()
}

#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, serde::Serialize)]
pub enum ElementStatus {
    Missing,
    Assigned,
    Completed,
}

impl std::fmt::Display for ElementStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            ElementStatus::Missing => "MISSING",
            ElementStatus::Assigned => "ASSIGNED",
            ElementStatus::Completed => "COMPLETED",
        })
    }
}

pub type ElementIndex = petgraph::graph::NodeIndex;

#[derive(serde::Serialize)]
pub struct Tree {
    graph: petgraph::Graph<Element, ()>,
    #[serde(skip)]
    issue_map: BTreeMap<u64, ElementIndex>,
}

impl Tree {
    pub fn new() -> Self {
        Self {
            graph: petgraph::Graph::new(),
            issue_map: BTreeMap::new(),
        }
    }

    pub fn add_element(&mut self, element: Element) {
        let issue_number = element.issue_number;
        assert!(!self.issue_map.contains_key(&issue_number));
        let idx = self.graph.add_node(element);
        self.issue_map.insert(issue_number, idx);
    }

    pub fn add_dependency_by_issue_number(&mut self, dependant: u64, dependency: u64) {
        let a = self.find_element_by_issue_number(dependant).unwrap();
        let b = self.find_element_by_issue_number(dependency).unwrap();
        self.graph.add_edge(a, b, ());
    }

    pub fn find_element_by_issue_number(&self, issue_number: u64) -> Option<ElementIndex> {
        self.issue_map.get(&issue_number).copied()
    }

    pub fn subtree_for_issue<'a>(&'a self, issue_number: u64) -> Option<Subtree<'a>> {
        Some(Subtree::new_for_element(
            self,
            self.find_element_by_issue_number(issue_number)?,
        ))
    }

    pub fn iter_issues<'a>(&'a self) -> impl Iterator<Item = u64> + 'a {
        self.issue_map.keys().copied()
    }

    pub fn subtree_for_element<'a>(&'a self, element: ElementIndex) -> Subtree<'a> {
        Subtree::new_for_element(self, element)
    }

    pub fn iter_ultimate_elements<'a>(&'a self) -> impl Iterator<Item = ElementIndex> + 'a {
        self.graph.node_indices().filter(|index| {
            // If there are no incoming edges, then this is an ultimate element!
            self.graph
                .neighbors_directed(*index, petgraph::Direction::Incoming)
                .next()
                .is_none()
        })
    }

    pub fn get_element_role(&self, element: ElementIndex) -> ElementRole {
        let has_dependencies = self
            .graph
            .neighbors_directed(element, petgraph::Direction::Outgoing)
            .next()
            .is_some();
        let has_dependants = self
            .graph
            .neighbors_directed(element, petgraph::Direction::Incoming)
            .next()
            .is_some();

        match (has_dependencies, has_dependants) {
            (false, true) => ElementRole::Root,
            (true, false) => ElementRole::Ultimate,
            (true, true) => ElementRole::Intermediate,
            (false, false) => ElementRole::Disjoint,
        }
    }

    pub fn to_dot(&self) -> String {
        let to_node_attributes = |_g, (id, element): (ElementIndex, &Element)| {
            element.to_dot_node_attributes(Some(self.get_element_role(id)), None)
        };
        let dot = petgraph::dot::Dot::with_attr_getters(
            &self.graph,
            &[
                petgraph::dot::Config::EdgeNoLabel,
                petgraph::dot::Config::NodeNoLabel,
                petgraph::dot::Config::RankDir(petgraph::dot::RankDir::BT),
                petgraph::dot::Config::GraphContentOnly,
            ],
            &|_g, _edge_id| "".to_string(),
            &to_node_attributes,
        );

        let ultimate_elements: Vec<_> = self
            .iter_ultimate_elements()
            .map(|idx| idx.index().to_string())
            .collect();
        let ultimate_element_list = ultimate_elements.join("; ");

        format!(
            r#"digraph {{
    ranksep=1.2
    {{ rank=same; {ultimate_element_list}; }}
{:?}
}}
"#,
            dot
        )
    }

    pub fn to_mermaid(&self, repo_url: &str, simple: bool) -> String {
        let mut mermaid = String::new();
        mermaid.push_str("flowchart BT\n");
        mermaid.push_str(&mermaid_classes());

        for index in self.graph.node_indices() {
            mermaid.push_str(&self.graph[index].to_mermaid_node(index, None, repo_url, simple));
            mermaid.push_str("\n");
        }

        for edge in self.graph.edge_references() {
            mermaid.push_str(&format!(
                "    {} --> {}\n",
                edge.source().index(),
                edge.target().index()
            ));
        }

        mermaid
    }
}

#[derive(Debug, Clone, Hash, PartialEq, Eq, serde::Serialize)]
pub struct SubtreeElement<'a> {
    element: &'a Element,
    role: SubtreeElementRole,
}

#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, serde::Serialize)]
pub enum SubtreeElementRole {
    ElementOfInterest,
    Dependency,
    Dependant,
}

pub struct Subtree<'a> {
    original: &'a Tree,
    graph: petgraph::Graph<SubtreeElement<'a>, ()>,
    index_map: BTreeMap<ElementIndex, ElementIndex>,
}

impl<'a> Subtree<'a> {
    pub fn new_for_element(tree: &'a Tree, element: ElementIndex) -> Subtree<'a> {
        let mut this = Self {
            original: tree,
            graph: petgraph::Graph::new(),
            index_map: BTreeMap::new(),
        };
        let graph = &this.original.graph;

        this.add_node_if_missing(element, SubtreeElementRole::ElementOfInterest);

        let mut dfs = petgraph::visit::Dfs::new(&graph, element);
        while let Some(idx) = dfs.next(&graph) {
            this.add_node_if_missing(idx, SubtreeElementRole::Dependency);
        }

        for idx in this.index_map.keys().copied() {
            for neighbor in graph.neighbors_directed(idx, petgraph::Direction::Outgoing) {
                if !this.index_map.contains_key(&neighbor) {
                    log::warn!(
                        "Probably missed a dependency from #{from} to #{to} while generating subtree for #{el}",
                        el = graph[element].issue_number,
                        from = graph[idx].issue_number,
                        to = graph[neighbor].issue_number
                    );
                    continue;
                }

                this.graph
                    .add_edge(this.index_map[&idx], this.index_map[&neighbor], ());
            }
        }

        for neighbor in graph.neighbors_directed(element, petgraph::Direction::Incoming) {
            if this.index_map.contains_key(&neighbor) {
                log::warn!(
                    "Already found #{from} in dependencies of #{to}, but it should have been the other way around?!",
                    from = graph[neighbor].issue_number,
                    to = graph[element].issue_number
                );
            } else {
                this.add_node_if_missing(neighbor, SubtreeElementRole::Dependant);
            }

            this.graph
                .add_edge(this.index_map[&neighbor], this.index_map[&element], ());
        }

        this
    }

    fn add_node_if_missing(&mut self, index: ElementIndex, role: SubtreeElementRole) -> bool {
        if self.index_map.contains_key(&index) {
            debug_assert!(self.graph.node_weight(self.index_map[&index]).is_some());
            return false;
        }

        let subtree_element = SubtreeElement {
            element: &self.original.graph[index],
            role,
        };

        let new_idx = self.graph.add_node(subtree_element);
        self.index_map.insert(index, new_idx);

        return true;
    }

    pub fn to_dot(&self) -> String {
        let dot = petgraph::dot::Dot::with_attr_getters(
            &self.graph,
            &[
                petgraph::dot::Config::EdgeNoLabel,
                petgraph::dot::Config::NodeNoLabel,
                petgraph::dot::Config::RankDir(petgraph::dot::RankDir::BT),
            ],
            &|_g, _edge_id| "".to_string(),
            &|_g, (_, element)| {
                element
                    .element
                    .to_dot_node_attributes(None, Some(element.role))
            },
        );

        format!("{:?}", dot)
    }

    pub fn to_mermaid(&self, repo_url: &str, simple: bool) -> String {
        let mut mermaid = String::new();
        mermaid.push_str("flowchart BT\n");
        mermaid.push_str(&mermaid_classes());

        for index in self.graph.node_indices() {
            let node = &self.graph[index];
            mermaid.push_str(&node.element.to_mermaid_node(
                index,
                Some(node.role),
                repo_url,
                simple,
            ));
            mermaid.push_str("\n");
        }

        for edge in self.graph.edge_references() {
            mermaid.push_str(&format!(
                "    {} --> {}\n",
                edge.source().index(),
                edge.target().index()
            ));
        }

        mermaid
    }

    pub fn stable_hash(&self) -> String {
        let mut nodes: Vec<_> = self.graph.node_weights().collect();
        nodes.sort_by_key(|n| n.element.issue_number);
        let mut edges: Vec<_> = self
            .graph
            .edge_references()
            .map(|edge| {
                (
                    self.graph[edge.source()].element.issue_number,
                    self.graph[edge.target()].element.issue_number,
                )
            })
            .collect();
        edges.sort();

        let json_data = serde_json::ser::to_string(&(nodes, edges, HASH_EPOCH)).unwrap();
        sha256::digest(json_data)
    }
}