1574 lines · 45.6 KB · Rust
RawDownload
// Copyright 2021-2023 The Jujutsu Authors//// Licensed under the Apache License, Version 2.0 (the "License");// you may not use this file except in compliance with the License.// You may obtain a copy of the License at//// https://www.apache.org/licenses/LICENSE-2.0//// Unless required by applicable law or agreed to in writing, software// distributed under the License is distributed on an "AS IS" BASIS,// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.// See the License for the specific language governing permissions and// limitations under the License. #![allow(missing_docs)] use std::collections::HashMap;use std::collections::HashSet;use std::collections::VecDeque;use std::hash::Hash; use crate::revset::RevsetEvaluationError; /// Node and edges pair of type `N`.pub type GraphNode<N> = (N, Vec<GraphEdge<N>>); #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]pub struct GraphEdge<N> {    pub target: N,    pub edge_type: GraphEdgeType,} impl<N> GraphEdge<N> {    pub fn missing(target: N) -> Self {        Self {            target,            edge_type: GraphEdgeType::Missing,        }    }     pub fn direct(target: N) -> Self {        Self {            target,            edge_type: GraphEdgeType::Direct,        }    }     pub fn indirect(target: N) -> Self {        Self {            target,            edge_type: GraphEdgeType::Indirect,        }    }     pub fn map<M>(self, f: impl FnOnce(N) -> M) -> GraphEdge<M> {        GraphEdge {            target: f(self.target),            edge_type: self.edge_type,        }    }} #[derive(Debug, PartialEq, Eq, Clone, Copy, Hash)]pub enum GraphEdgeType {    Missing,    Direct,    Indirect,} fn reachable_targets<N>(edges: &[GraphEdge<N>]) -> impl DoubleEndedIterator<Item = &N> {    edges        .iter()        .filter(|edge| edge.edge_type != GraphEdgeType::Missing)        .map(|edge| &edge.target)} pub struct ReverseGraphIterator<N> {    items: Vec<GraphNode<N>>,} impl<N> ReverseGraphIterator<N>where    N: Hash + Eq + Clone,{    pub fn new(        input: impl Iterator<Item = Result<GraphNode<N>, RevsetEvaluationError>>,    ) -> Result<Self, RevsetEvaluationError> {        let mut entries = vec![];        let mut reverse_edges: HashMap<N, Vec<GraphEdge<N>>> = HashMap::new();        for item in input {            let (node, edges) = item?;            for GraphEdge { target, edge_type } in edges {                reverse_edges.entry(target).or_default().push(GraphEdge {                    target: node.clone(),                    edge_type,                });            }            entries.push(node);        }         let mut items = vec![];        for node in entries.into_iter() {            let edges = reverse_edges.get(&node).cloned().unwrap_or_default();            items.push((node, edges));        }        Ok(Self { items })    }} impl<N> Iterator for ReverseGraphIterator<N> {    type Item = Result<GraphNode<N>, RevsetEvaluationError>;     fn next(&mut self) -> Option<Self::Item> {        self.items.pop().map(Ok)    }} /// Graph iterator adapter to group topological branches.////// Basic idea is DFS from the heads. At fork point, the other descendant/// branches will be visited. At merge point, the second (or the last) ancestor/// branch will be visited first. This is practically [the same as Git][Git].////// The branch containing the first commit in the input iterator will be emitted/// first. It is often the working-copy ancestor branch. The other head branches/// won't be enqueued eagerly, and will be emitted as late as possible.////// [Git]: https://github.blog/2022-08-30-gits-database-internals-ii-commit-history-queries/#topological-sorting#[derive(Clone, Debug)]pub struct TopoGroupedGraphIterator<N, I> {    input_iter: I,    /// Graph nodes read from the input iterator but not yet emitted.    nodes: HashMap<N, TopoGroupedGraphNode<N>>,    /// Stack of graph nodes to be emitted.    emittable_ids: Vec<N>,    /// List of new head nodes found while processing unpopulated nodes.    new_head_ids: VecDeque<N>,    /// Set of nodes which may be ancestors of `new_head_ids`.    blocked_ids: HashSet<N>,} #[derive(Clone, Debug)]struct TopoGroupedGraphNode<N> {    /// Graph nodes which must be emitted before.    child_ids: HashSet<N>,    /// Graph edges to parent nodes. `None` until this node is populated.    edges: Option<Vec<GraphEdge<N>>>,} impl<N> Default for TopoGroupedGraphNode<N> {    fn default() -> Self {        Self {            child_ids: Default::default(),            edges: None,        }    }} impl<N, E, I> TopoGroupedGraphIterator<N, I>where    N: Clone + Hash + Eq,    I: Iterator<Item = Result<GraphNode<N>, E>>,{    /// Wraps the given iterator to group topological branches. The input    /// iterator must be topologically ordered.    pub fn new(input_iter: I) -> Self {        TopoGroupedGraphIterator {            input_iter,            nodes: HashMap::new(),            emittable_ids: Vec::new(),            new_head_ids: VecDeque::new(),            blocked_ids: HashSet::new(),        }    }     fn populate_one(&mut self) -> Result<Option<()>, E> {        let (current_id, edges) = match self.input_iter.next() {            Some(Ok(data)) => data,            Some(Err(err)) => {                return Err(err);            }            None => {                return Ok(None);            }        };         // Set up reverse reference        for parent_id in reachable_targets(&edges) {            let parent_node = self.nodes.entry(parent_id.clone()).or_default();            parent_node.child_ids.insert(current_id.clone());        }         // Populate the current node        if let Some(current_node) = self.nodes.get_mut(&current_id) {            assert!(current_node.edges.is_none());            current_node.edges = Some(edges);        } else {            let current_node = TopoGroupedGraphNode {                edges: Some(edges),                ..Default::default()            };            self.nodes.insert(current_id.clone(), current_node);            // Push to non-emitting list so the new head wouldn't be interleaved            self.new_head_ids.push_back(current_id);        }         Ok(Some(()))    }     /// Enqueues the first new head which will unblock the waiting ancestors.    ///    /// This does not move multiple head nodes to the queue at once because    /// heads may be connected to the fork points in arbitrary order.    fn flush_new_head(&mut self) {        assert!(!self.new_head_ids.is_empty());        if self.blocked_ids.is_empty() || self.new_head_ids.len() <= 1 {            // Fast path: orphaned or no choice            let new_head_id = self.new_head_ids.pop_front().unwrap();            self.emittable_ids.push(new_head_id);            self.blocked_ids.clear();            return;        }         // Mark descendant nodes reachable from the blocking nodes        let mut to_visit: Vec<&N> = self            .blocked_ids            .iter()            .map(|id| {                // Borrow from self.nodes so self.blocked_ids can be mutated later                let (id, _) = self.nodes.get_key_value(id).unwrap();                id            })            .collect();        let mut visited: HashSet<&N> = to_visit.iter().copied().collect();        while let Some(id) = to_visit.pop() {            let node = self.nodes.get(id).unwrap();            to_visit.extend(node.child_ids.iter().filter(|id| visited.insert(id)));        }         // Pick the first reachable head        let index = self            .new_head_ids            .iter()            .position(|id| visited.contains(id))            .expect("blocking head should exist");        let new_head_id = self.new_head_ids.remove(index).unwrap();         // Unmark ancestors of the selected head so they won't contribute to future        // new-head resolution within the newly-unblocked sub graph. The sub graph        // can have many fork points, and the corresponding heads should be picked in        // the fork-point order, not in the head appearance order.        to_visit.push(&new_head_id);        visited.remove(&new_head_id);        while let Some(id) = to_visit.pop() {            let node = self.nodes.get(id).unwrap();            if let Some(edges) = &node.edges {                to_visit.extend(reachable_targets(edges).filter(|id| visited.remove(id)));            }        }        self.blocked_ids.retain(|id| visited.contains(id));        self.emittable_ids.push(new_head_id);    }     fn next_node(&mut self) -> Result<Option<GraphNode<N>>, E> {        // Based on Kahn's algorithm        loop {            if let Some(current_id) = self.emittable_ids.last() {                let Some(current_node) = self.nodes.get_mut(current_id) else {                    // Queued twice because new children populated and emitted                    self.emittable_ids.pop().unwrap();                    continue;                };                if !current_node.child_ids.is_empty() {                    // New children populated after emitting the other                    let current_id = self.emittable_ids.pop().unwrap();                    self.blocked_ids.insert(current_id);                    continue;                }                let Some(edges) = current_node.edges.take() else {                    // Not yet populated                    self.populate_one()?.expect("parent node should exist");                    continue;                };                // The second (or the last) parent will be visited first                let current_id = self.emittable_ids.pop().unwrap();                self.nodes.remove(&current_id).unwrap();                for parent_id in reachable_targets(&edges) {                    let parent_node = self.nodes.get_mut(parent_id).unwrap();                    parent_node.child_ids.remove(&current_id);                    if parent_node.child_ids.is_empty() {                        let reusable_id = self.blocked_ids.take(parent_id);                        let parent_id = reusable_id.unwrap_or_else(|| parent_id.clone());                        self.emittable_ids.push(parent_id);                    } else {                        self.blocked_ids.insert(parent_id.clone());                    }                }                return Ok(Some((current_id, edges)));            } else if !self.new_head_ids.is_empty() {                self.flush_new_head();            } else {                // Populate the first or orphan head                if self.populate_one()?.is_none() {                    return Ok(None);                }            }        }    }} impl<N, E, I> Iterator for TopoGroupedGraphIterator<N, I>where    N: Clone + Hash + Eq,    I: Iterator<Item = Result<GraphNode<N>, E>>,{    type Item = Result<GraphNode<N>, E>;     fn next(&mut self) -> Option<Self::Item> {        match self.next_node() {            Ok(Some(node)) => Some(Ok(node)),            Ok(None) => {                assert!(self.nodes.is_empty(), "all nodes should have been emitted");                None            }            Err(err) => Some(Err(err)),        }    }} #[cfg(test)]mod tests {    use std::convert::Infallible;     use itertools::Itertools as _;    use renderdag::Ancestor;    use renderdag::GraphRowRenderer;    use renderdag::Renderer as _;     use super::*;     fn missing(c: char) -> GraphEdge<char> {        GraphEdge::missing(c)    }     fn direct(c: char) -> GraphEdge<char> {        GraphEdge::direct(c)    }     fn indirect(c: char) -> GraphEdge<char> {        GraphEdge::indirect(c)    }     fn format_edge(edge: &GraphEdge<char>) -> String {        let c = edge.target;        match edge.edge_type {            GraphEdgeType::Missing => format!("missing({c})"),            GraphEdgeType::Direct => format!("direct({c})"),            GraphEdgeType::Indirect => format!("indirect({c})"),        }    }     fn format_graph(graph_iter: impl IntoIterator<Item = GraphNode<char>>) -> String {        let mut renderer = GraphRowRenderer::new()            .output()            .with_min_row_height(2)            .build_box_drawing();        graph_iter            .into_iter()            .map(|(id, edges)| {                let glyph = id.to_string();                let message = edges.iter().map(format_edge).join(", ");                let parents = edges                    .into_iter()                    .map(|edge| match edge.edge_type {                        GraphEdgeType::Missing => Ancestor::Anonymous,                        GraphEdgeType::Direct => Ancestor::Parent(edge.target),                        GraphEdgeType::Indirect => Ancestor::Ancestor(edge.target),                    })                    .collect();                renderer.next_row(id, parents, glyph, message)            })            .collect()    }     #[test]    fn test_format_graph() {        let graph = vec![            ('D', vec![direct('C'), indirect('B')]),            ('C', vec![direct('A')]),            ('B', vec![missing('X')]),            ('A', vec![]),        ];        insta::assert_snapshot!(format_graph(graph), @r###"        D    direct(C), indirect(B)        ├─╮        C ╷  direct(A)        │ ╷        │ B  missing(X)        │ │        │ ~        │        A         "###);    }     fn topo_grouped<I, E>(graph_iter: I) -> TopoGroupedGraphIterator<char, I::IntoIter>    where        I: IntoIterator<Item = Result<GraphNode<char>, E>>,    {        TopoGroupedGraphIterator::new(graph_iter.into_iter())    }     fn infallible(input: GraphNode<char>) -> Result<GraphNode<char>, Infallible> {        Ok(input)    }     #[test]    fn test_topo_grouped_multiple_roots() {        let graph = [            ('C', vec![missing('Y')]),            ('B', vec![missing('X')]),            ('A', vec![]),        ];        insta::assert_snapshot!(format_graph(graph.iter().cloned()), @r###"        C  missing(Y)        │        ~         B  missing(X)        │        ~         A        "###);        insta::assert_snapshot!(format_graph(topo_grouped(graph.iter().cloned().map(infallible)).map(Result::unwrap)), @r###"        C  missing(Y)        │        ~         B  missing(X)        │        ~         A        "###);         // All nodes can be lazily emitted.        let mut iter = topo_grouped(graph.iter().cloned().map(infallible).peekable());        assert_eq!(iter.next().unwrap().unwrap().0, 'C');        assert_eq!(iter.input_iter.peek().unwrap().as_ref().unwrap().0, 'B');        assert_eq!(iter.next().unwrap().unwrap().0, 'B');        assert_eq!(iter.input_iter.peek().unwrap().as_ref().unwrap().0, 'A');    }     #[test]    fn test_topo_grouped_trivial_fork() {        let graph = [            ('E', vec![direct('B')]),            ('D', vec![direct('A')]),            ('C', vec![direct('B')]),            ('B', vec![direct('A')]),            ('A', vec![]),        ];        insta::assert_snapshot!(format_graph(graph.iter().cloned()), @r###"        E  direct(B)        │        │ D  direct(A)        │ │        │ │ C  direct(B)        ├───╯        B │  direct(A)        ├─╯        A         "###);        // D-A is found earlier than B-A, but B is emitted first because it belongs to        // the emitting branch.        insta::assert_snapshot!(format_graph(topo_grouped(graph.iter().cloned().map(infallible)).map(Result::unwrap)), @r###"        E  direct(B)        │        │ C  direct(B)        ├─╯        B  direct(A)        │        │ D  direct(A)        ├─╯        A         "###);         // E can be lazy, then D and C will be queued.        let mut iter = topo_grouped(graph.iter().cloned().map(infallible).peekable());        assert_eq!(iter.next().unwrap().unwrap().0, 'E');        assert_eq!(iter.input_iter.peek().unwrap().as_ref().unwrap().0, 'D');        assert_eq!(iter.next().unwrap().unwrap().0, 'C');        assert_eq!(iter.input_iter.peek().unwrap().as_ref().unwrap().0, 'B');        assert_eq!(iter.next().unwrap().unwrap().0, 'B');        assert_eq!(iter.input_iter.peek().unwrap().as_ref().unwrap().0, 'A');    }

Lines 1–500 of 1574.