1648 lines · 59.2 KB · Rust
RawDownload
// Copyright 2026 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. //! Generic implementation of the "closest common dominator" algorithm for//! directed graphs.//!//! Generic implementation of the Common Dominator algorithm for directed//! graphs, using the Cooper-Harvey-Kennedy iterative algorithm. Loosely//! speaking the algorithm finds the "choke point" for a set of nodes S in a//! directed graph (going from the "entry" node to nodes in S), closest to S.//!//! Dominance://!//! * A flow graph is a directed graph with a designated entry node.//! * A node z is said to dominate a node n if all paths from the entry node to//!   n must go through z. Every node dominates itself, and the entry node//!   dominates all nodes.//! * A node can have one or more dominators.//! * A node z strictly dominates n if z dominates n and z != n.//! * The immediate dominator of a node n is the dominator of n that doesn't//!   strictly dominate any other strict dominators of n. Informally it is the//!   "closest" choke point on all paths from the entry node to n.//! * Let S be a subset of the nodes in the graph. The intersection of the//!   dominators of each node in S is the set of common dominators of S.//! * The closest common dominator of S is the common dominator of S that//!   doesn't strictly dominate any other common dominator of S. Informally, it//!   is the choke point closest to S such that all paths from the entry node to//!   S must go through it.//!//! Dominator Tree://!//! For any flow graph G there is a corresponding dominator tree defined as//! follows://! * The nodes of the dominator tree are the same as the nodes of G//! * The root of the dominator tree is the entry node of G//! * In the dominator tree, the children of a node are the nodes it immediately//!   dominates//!//! The closest common dominator of S is the Lowest Common Ancestor (LCA)//! of S in the graph's dominator tree.//!//! This implementation constructs the Dominator Tree by first determining//! the Immediate Dominator for every node (using the standard iterative//! algorithm), and then calculating the LCA for the set S. See://!//! * <http://www.hipersoft.rice.edu/grads/publications/dom14.pdf>//! * <https://en.wikipedia.org/wiki/Dominator_(graph_theory)>//!//! The running time is O(V+E+|S|*V)in the worst case, the space complexity is//! O(V+E), where V is the number of nodes and E is the number of edges. In//! practice the algorithm is fast and efficient for typical use cases because//! the number of nodes that dominate any given node is typically small, and//! the dominator tree is typically shallow. use std::collections::HashMap;use std::collections::HashSet;use std::hash::Hash;use std::iter;use std::rc::Rc; use futures::future::try_join_all;use indexmap::IndexMap;use indexmap::IndexSet;use itertools::Itertools as _;use thiserror::Error; /// An immutable directed graph with nodes of type N and a minimal interface for/// iterating over nodes and their adjacent nodes.#[derive(Clone, Eq, PartialEq, Debug)]pub struct SimpleDirectedGraph<N>where    N: Clone + Eq + Hash,{    /// The adjacency map of the graph. Each key is a node, and the    /// corresponding value is the set of adjacent nodes (i.e., the children of    /// the key node). The adjacency map is in canonical form: for every    /// u->v edge, there is an entry in adj with key v (even if v has no    /// outgoing edges).    adj: IndexMap<N, IndexSet<N>>,} impl<N> SimpleDirectedGraph<N>where    N: Clone + Eq + Hash,{    /// Constructs a new SimpleDirectedGraph from a list of edges.    pub fn new<EI>(edges: EI) -> Self    where        EI: IntoIterator<Item = (N, N)>,    {        let mut adj: IndexMap<N, IndexSet<N>> = IndexMap::new();        for (parent, child) in edges {            adj.entry(parent).or_default().insert(child.clone());            adj.entry(child).or_default();        }        Self { adj }    }     /// Returns the nodes in this graph.    pub fn nodes(&self) -> impl Iterator<Item = &N> {        self.adj.keys()    }     /// Returns the nodes in this graph.    pub fn num_nodes(&self) -> usize {        self.adj.len()    }     /// Returns the edges in this graph.    pub fn edges(&self) -> impl Iterator<Item = (&N, &N)> {        self.adj            .iter()            .flat_map(|(parent, adj_set)| adj_set.iter().map(move |child| (parent, child)))    }     /// Returns the adjacent nodes for the given node, or None if the node is    /// not in the graph.    pub fn adjacent_nodes(&self, node: &N) -> Option<impl DoubleEndedIterator<Item = &N>> {        self.adj.get(node).map(|adj_set| adj_set.iter())    }     /// Returns true if this graph contains the given node.    pub fn contains_node(&self, node: &N) -> bool {        self.adj.contains_key(node)    }     /// Returns a postorder traversal of the nodes in this graph starting from    /// the given node.    pub fn get_postorder<'a>(&'a self, start_node: &'a N) -> Vec<&'a N> {        post_order(start_node, |&u| self.adjacent_nodes(u).unwrap()).collect_vec()    }} /// A FlowGraph is a directed graph with a designated start node.////// Any node in the graph can be the start node. There are no reachability/// requirements whatsoever: some nodes may be unreachable from the start node,/// the start node could have incoming edges, the graph could be disconnected,/// etc.#[derive(Clone, Eq, PartialEq, Debug)]pub struct FlowGraph<N>where    N: Clone + Eq + Hash,{    /// The graph.    pub graph: SimpleDirectedGraph<N>,    /// The start node.    pub start_node: N,} /// Calculates the dominators in a flow graph. Also has a method for finding the/// closest common dominator of a set of nodes.pub struct DominatorFinder<'a, N> {    /// Map from nodes to integers in [0, N-1] range, in postorder (the start    /// node has index N-1).    node_to_id: HashMap<&'a N, InternalId>,    /// The inverse of node_to_id.    id_to_node: Vec<&'a N>,    /// The immediate dominator for each node (by index). NOTE: the immediate    /// dominator of the start node is itself.    immediate_dominators: Vec<InternalId>,} /// Errors that can occur while finding dominators.#[derive(Debug, Error, PartialEq)]pub enum DominatorFinderError {    /// The flow graph is invalid.    #[error("The flow graph is invalid: some nodes are unreachable from the start node")]    UnreachableNodesInFlowGraph,    /// The target set is empty.    #[error("Target set must not be empty")]    EmptyTargetSet,    /// The target set is invalid.    #[error("Target set contains a node which is not in the flow graph")]    UnknownNodeInTargetSet,} /// The dominator algorithm assigns consecutive numeric IDs to nodes, for/// efficiency reasons. We use this type alias for clarity.type InternalId = usize; impl<'a, N> DominatorFinder<'a, N>where    N: Clone + Eq + Hash,{    /// Constructs a new DominatorFinder. Returns an error if the flow graph is    /// invalid: e.g. if some node is unreachable from the start node.    pub fn calculate(flow_graph: &'a FlowGraph<N>) -> Result<Self, DominatorFinderError> {        // Get postorder traversal of the graph starting from the start node.        let postorder = flow_graph.graph.get_postorder(&flow_graph.start_node);        if postorder.len() != flow_graph.graph.num_nodes() {            return Err(DominatorFinderError::UnreachableNodesInFlowGraph);        }         // Map generic types to integer IDs        let mut node_to_id = HashMap::new();        let mut id_to_node = Vec::new();        for (index, &node) in postorder.iter().enumerate() {            id_to_node.push(node);            node_to_id.insert(node, index);        }         // Build graph using internal IDs.        let num_nodes = node_to_id.len();        let mut rev_adj = vec![vec![]; num_nodes];        for (u, v) in flow_graph.graph.edges() {            rev_adj[node_to_id[v]].push(node_to_id[u]);        }         // Find the immediate dominators for each node using the Cooper-Harvey-Kennedy        // iterative algorithm.        let immediate_dominators = Self::calculate_immediate_dominators(&rev_adj);         Ok(Self {            node_to_id,            id_to_node,            immediate_dominators,        })    }     /// Returns a map from each node to its immediate dominator. NOTE: the    /// immediate dominator of the start node is itself.    pub fn get_immediate_dominators(&self) -> HashMap<N, N> {        self.immediate_dominators            .iter()            .enumerate()            .map(|(index, &idom)| {                (                    self.id_to_node[index].clone(),                    self.id_to_node[idom].clone(),                )            })            .collect()    }     /// Finds the closest common dominator for the given flow graph and set of    /// nodes S (target_set).    pub fn find_closest_common_dominator<NI>(        &self,        target_set: NI,    ) -> Result<N, DominatorFinderError>    where        NI: IntoIterator<Item = N>,    {        // Convert generic target_set to internal IDs        let target_ids: Vec<InternalId> = target_set            .into_iter()            .map(|node| match self.node_to_id.get(&node) {                Some(id) => Ok(*id),                None => Err(DominatorFinderError::UnknownNodeInTargetSet),            })            .try_collect()?;        if target_ids.is_empty() {            return Err(DominatorFinderError::EmptyTargetSet);        }         // The closest common dominator of a set of nodes is the lowest common ancestor        // of those nodes in the dominator tree.        let closest_common_dominator =            Self::find_lowest_common_ancestor(&target_ids, &self.immediate_dominators);         // Map the internal ID back to generic type N.        Ok(self.id_to_node[closest_common_dominator].clone())    }     // Applies the Cooper-Harvey-Kennedy iterative algorithm to find the immediate    // dominators for each node in the graph.    // See http://www.hipersoft.rice.edu/grads/publications/dom14.pdf for details on how this function works.    fn calculate_immediate_dominators(rev_adj: &[Vec<InternalId>]) -> Vec<InternalId> {        // Step 1: Compute Dominators on Reverse Graph        let num_nodes = rev_adj.len();        let start_node_id = num_nodes - 1;         // We hold the immediate dominator for each node in the following vector, in        // index position (the kth entry is the immediate dominator of the node with ID        // k). We initialize the immediate dominator of every node to usize::MAX to        // represent that those nodes are not processed yet. Once a node is        // processed, its immediate dominator is guaranteed to be a valid node        // index.        let mut immediate_dominators: Vec<InternalId> = vec![usize::MAX; num_nodes];        // NOTE: technically speaking the immediate dominator is NOT defined for the        // start node, but it is convenient for the algorithm to set it to itself; this        // is consistent with the literature and specifically with        // Cooper-Harvey-Kennedy.        immediate_dominators[start_node_id] = start_node_id;         loop {            // Each iteration of the loop processes all nodes in reverse postorder, trying            // to improve the immediate dominator for each node. The loop continues until we            // have an iteration where no immediate dominator is changed. Note that the            // entries in immediate_dominators are only guaranteed to be correct when the            // loop terminates.            let mut changed = false;             // Iterate in reverse postorder, skipping the start node.            for u in (0..start_node_id).rev() {                let mut new_idom = usize::MAX;                // Process predecessors (nodes that flow INTO u).                let preds = &rev_adj[u];                for &p in preds {                    if immediate_dominators[p] == usize::MAX {                        // Skip predecessors that have not been processed yet.                        continue;                    }                    if new_idom == usize::MAX {                        // This is the first predecessor of u that has been processed so far. We use                        // it as the starting point for finding the new "improved" immediate                        // dominator for u.                        new_idom = p;                    } else {                        // "Intersect" the current new_idom with p's idom.                        new_idom = Self::intersect(new_idom, p, &immediate_dominators);                    }                }                if new_idom == usize::MAX {                    // None of the predecessors of u have been processed yet. That's fine, we will                    // try again of the next iteration of the outer loop.                    continue;                }                if immediate_dominators[u] != new_idom {                    // We "improved" the immediate dominator for u!                    immediate_dominators[u] = new_idom;                    changed = true;                }            }             if !changed {                // We reached the fixed point. We are done.                break;            }        }         // At this point we know the immediate dominator of every node, but we keep the        // Option wrapper so that we can use the intersect function during        // find_lowest_common_ancestor.        immediate_dominators    }     // See http://www.hipersoft.rice.edu/grads/publications/dom14.pdf for details on how this function works.    fn intersect(        mut b1: InternalId,        mut b2: InternalId,        immediate_dominators: &[InternalId],    ) -> InternalId {        while b1 != b2 {            while b1 < b2 {                b1 = immediate_dominators[b1];            }            while b2 < b1 {                b2 = immediate_dominators[b2];            }        }        b1    }     // See http://www.hipersoft.rice.edu/grads/publications/dom14.pdf for details on how this function works.    fn find_lowest_common_ancestor(        targets: &[InternalId],        immediate_dominators: &[InternalId],    ) -> InternalId {        targets            .iter()            .copied()            .reduce(|a, b| Self::intersect(a, b, immediate_dominators))            .expect("targets must not be empty")    }} /// Errors that can occur while finding a dominator value (i.e. a dominator in a/// value flow graph).#[derive(Debug, Error, PartialEq)]pub enum FindDominatorValueError<E> {    /// An error occurred while computing the value of a node.    #[error(transparent)]    ValueFnError(E),    /// An error occurred while finding the dominator.    #[error(transparent)]    DominatorFinderError(#[from] DominatorFinderError),} /// Helper struct for constructing a value flow graph. It caches the results/// of applying value_fn to nodes, and also keeps track of the mapping from/// values to nodes and nodes to values.pub struct ValueCache<N, V, VF> {    /// The function that emits values.    value_fn: VF,    /// Maps nodes to their corresponding values.    node_values: HashMap<N, Rc<V>>,    /// Maps values to the nodes that have that value.    value_to_nodes: HashMap<Rc<V>, Vec<N>>,} impl<N, V, VF, E> ValueCache<N, V, VF>where    N: Hash + Eq + Clone,    V: Hash + Eq,    VF: AsyncFn(&N) -> Result<V, E>,{    /// Creates a new ValueCache that uses the given function to get values.    pub fn new(value_fn: VF) -> Self {        Self {            value_fn,            node_values: HashMap::new(),            value_to_nodes: HashMap::new(),        }    }     /// Returns the value for the given node, computing it if it is not already    /// cached.    pub async fn get_value(&mut self, node: &N) -> Result<Rc<V>, E> {        let values = self.get_values([node]).await?;        assert_eq!(values.len(), 1);        Ok(values[0].clone())    }     /// Returns the nodes that have the given value.    pub fn get_nodes_for_value(&self, value: &Rc<V>) -> Option<&Vec<N>> {        self.value_to_nodes.get(value)    }     /// Returns the value of the given nodes, computing them (if they are not    /// already cached) asynchronously and concurrently. The result will have    /// the same cardinality as `nodes`, and may contain duplicate values    /// (e.g. if multiple nodes have the same value, or if    /// there are repeated nodes). The order of the result is not specified    /// (any order is possible).    async fn get_values<'a, NI>(&mut self, nodes: NI) -> Result<Vec<Rc<V>>, E>    where        N: 'a,        NI: IntoIterator<Item = &'a N>,    {        let mut values = vec![];        let mut futures = vec![];         // 1. Filter out nodes already in the map and create futures for new nodes.        for node in nodes {            match self.node_values.get(node) {                Some(value) => {                    values.push(value.clone());                }                None => {                    // This node is not cached, we will compute its value.                    futures.push(async {                        let value = (self.value_fn)(node).await?;                        Ok((node.clone(), Rc::new(value)))                    });                }            }        }        // 2. Run all new futures concurrently        let new_results: Vec<(N, Rc<V>)> = try_join_all(futures).await?;        // 3. Insert the new entries into the maps.        for (node, value) in new_results {            self.node_values.insert(node.clone(), value.clone());            self.value_to_nodes                .entry(value.clone())                .or_default()                .push(node);            values.push(value);        }        Ok(values)    }} impl<N> FlowGraph<N>where    N: Clone + Eq + Hash,{    /// Constructs a new FlowGraph.    pub fn new(graph: SimpleDirectedGraph<N>, start_node: N) -> Self {        Self { graph, start_node }    }     /// Creates a flow graph of values from a flow graph of nodes.    ///    /// More precisely, let G be a FlowGraph of nodes with start node S. The    /// value flow graph G' is a FlowGraph derived from G. Let v(g) be the    /// result of applying value_fn to g. The nodes of G' are the set of values    /// v(g), for all g in G. For each edge g1->g2 in G, there is a    /// corresponding edge v(g1)->v(g2) in G'. The start node in G' is v(S).    ///    /// Returns an error if any value_fn invocation fails.    pub fn create_value_flow_graph<'a, V>(&self, node_values: &'a HashMap<N, V>) -> FlowGraph<&'a V>    where        V: Eq + Hash,    {        let mut edges = vec![];        let start_value = node_values.get(&self.start_node).expect("cached");        for (parent, children) in &self.graph.adj {

Lines 1–500 of 1648.