// Copyright 2020 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. //! General-purpose DAG algorithms. use std::collections::BinaryHeap;use std::collections::HashMap;use std::collections::HashSet;use std::convert::Infallible;use std::hash::Hash;use std::iter;use std::mem; use itertools::Itertools as _; /// Traverses nodes from `start` in depth-first order.pub fn dfs<T, ID, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> impl Iterator<Item = T>where    ID: Hash + Eq,    II: IntoIterator<Item = T>,    NI: IntoIterator<Item = T>,{    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));    dfs_ok(to_ok_iter(start), id_fn, neighbors_fn).map(Result::unwrap)} /// Traverses nodes from `start` in depth-first order.////// An `Err` is emitted as a node with no neighbors. Caller may decide to/// short-circuit on it.pub fn dfs_ok<T, ID, E, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> impl Iterator<Item = Result<T, E>>where    ID: Hash + Eq,    II: IntoIterator<Item = Result<T, E>>,    NI: IntoIterator<Item = Result<T, E>>,{    let mut work: Vec<Result<T, E>> = start.into_iter().collect();    let mut visited: HashSet<ID> = HashSet::new();    iter::from_fn(move || loop {        let c = match work.pop() {            Some(Ok(c)) => c,            r @ (Some(Err(_)) | None) => return r,        };        let id = id_fn(&c);        if visited.contains(&id) {            continue;        }        for p in neighbors_fn(&c) {            work.push(p);        }        visited.insert(id);        return Some(Ok(c));    })} /// Builds a list of nodes reachable from the `start` where neighbors come/// before the node itself.pub fn topo_order_forward<T, ID, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> Vec<T>where    ID: Hash + Eq + Clone,    II: IntoIterator<Item = T>,    NI: IntoIterator<Item = T>,{    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));    topo_order_forward_ok(to_ok_iter(start), id_fn, neighbors_fn).unwrap()} /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come/// before the node itself.////// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and/// returns the error.pub fn topo_order_forward_ok<T, ID, E, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> Result<Vec<T>, E>where    ID: Hash + Eq + Clone,    II: IntoIterator<Item = Result<T, E>>,    NI: IntoIterator<Item = Result<T, E>>,{    let mut stack: Vec<(T, bool)> = start.into_iter().map(|r| Ok((r?, false))).try_collect()?;    let mut visiting = HashSet::new();    let mut emitted = HashSet::new();    let mut result = vec![];    while let Some((node, neighbors_visited)) = stack.pop() {        let id = id_fn(&node);        if emitted.contains(&id) {            continue;        }        if !neighbors_visited {            assert!(visiting.insert(id.clone()), "graph has cycle");            let neighbors_iter = neighbors_fn(&node).into_iter();            stack.reserve(neighbors_iter.size_hint().0 + 1);            stack.push((node, true));            for neighbor in neighbors_iter {                stack.push((neighbor?, false));            }        } else {            visiting.remove(&id);            emitted.insert(id);            result.push(node);        }    }    Ok(result)} /// Builds a list of nodes reachable from the `start` where neighbors come after/// the node itself.pub fn topo_order_reverse<T, ID, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> Vec<T>where    ID: Hash + Eq + Clone,    II: IntoIterator<Item = T>,    NI: IntoIterator<Item = T>,{    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));    topo_order_reverse_ok(to_ok_iter(start), id_fn, neighbors_fn).unwrap()} /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come/// after the node itself.////// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and/// returns the error.pub fn topo_order_reverse_ok<T, ID, E, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    neighbors_fn: impl FnMut(&T) -> NI,) -> Result<Vec<T>, E>where    ID: Hash + Eq + Clone,    II: IntoIterator<Item = Result<T, E>>,    NI: IntoIterator<Item = Result<T, E>>,{    let mut result = topo_order_forward_ok(start, id_fn, neighbors_fn)?;    result.reverse();    Ok(result)} /// Like `topo_order_reverse()`, but can iterate linear DAG lazily.////// The DAG is supposed to be (mostly) topologically ordered by `T: Ord`./// For example, topological order of chronological data should respect/// timestamp (except a few outliers caused by clock skew.)////// Use `topo_order_reverse()` if the DAG is heavily branched. This can/// only process linear part lazily.pub fn topo_order_reverse_lazy<T, ID, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> impl Iterator<Item = T>where    T: Ord,    ID: Hash + Eq + Clone,    II: IntoIterator<Item = T>,    NI: IntoIterator<Item = T>,{    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));    topo_order_reverse_lazy_ok(to_ok_iter(start), id_fn, neighbors_fn).map(Result::unwrap)} /// Like `topo_order_reverse_ok()`, but can iterate linear DAG lazily.////// The returned iterator short-circuits at an `Err`. Pending non-linear nodes/// before the `Err` will be discarded.pub fn topo_order_reverse_lazy_ok<T, ID, E, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> impl Iterator<Item = Result<T, E>>where    T: Ord,    ID: Hash + Eq + Clone,    II: IntoIterator<Item = Result<T, E>>,    NI: IntoIterator<Item = Result<T, E>>,{    let mut inner = TopoOrderReverseLazyInner::empty();    inner.extend(start);    iter::from_fn(move || inner.next(&id_fn, &mut neighbors_fn))} #[derive(Clone, Debug)]struct TopoOrderReverseLazyInner<T, ID, E> {    start: Vec<T>,    result: Vec<Result<T, E>>,    emitted: HashSet<ID>,} impl<T: Ord, ID: Hash + Eq + Clone, E> TopoOrderReverseLazyInner<T, ID, E> {    fn empty() -> Self {        TopoOrderReverseLazyInner {            start: Vec::new(),            result: Vec::new(),            emitted: HashSet::new(),        }    }     fn extend(&mut self, iter: impl IntoIterator<Item = Result<T, E>>) {        let iter = iter.into_iter();        self.start.reserve(iter.size_hint().0);        for res in iter {            if let Ok(node) = res {                self.start.push(node);            } else {                // Emit the error and terminate                self.start.clear();                self.result.insert(0, res);                return;            }        }    }     fn next<NI: IntoIterator<Item = Result<T, E>>>(        &mut self,        id_fn: impl Fn(&T) -> ID,        mut neighbors_fn: impl FnMut(&T) -> NI,    ) -> Option<Result<T, E>> {        if let Some(res) = self.result.pop() {            return Some(res);        }         // Fast path for linear DAG        if self.start.len() <= 1 {            let node = self.start.pop()?;            self.extend(neighbors_fn(&node));            assert!(self.emitted.insert(id_fn(&node)), "graph has cycle");            return Some(Ok(node));        }         // Extract graph nodes based on T's order, and sort them by using ids        // (because we wouldn't want to clone T itself)        let start_ids = self.start.iter().map(&id_fn).collect_vec();        match look_ahead_sub_graph(mem::take(&mut self.start), &id_fn, &mut neighbors_fn) {            Ok((mut node_map, neighbor_ids_map, remainder)) => {                self.start = remainder;                let sorted_ids =                    topo_order_forward(&start_ids, |id| *id, |id| &neighbor_ids_map[id]);                self.result.reserve(sorted_ids.len());                for id in sorted_ids {                    let (id, node) = node_map.remove_entry(id).unwrap();                    assert!(self.emitted.insert(id), "graph has cycle");                    self.result.push(Ok(node));                }                self.result.pop()            }            Err(err) => Some(Err(err)),        }    }} /// Splits DAG at single fork point, and extracts branchy part as sub graph.////// ```text///  o | C///  | o B///  |/ <---- split here (A->B or A->C would create cycle)///  o A/// ```////// If a branch reached to root (empty neighbors), the graph can't be split/// anymore because the other branch may be connected to a descendant of/// the rooted branch.////// ```text///  o | C///  | o B///  |  <---- can't split here (there may be edge A->B)///  o A/// ```////// We assume the graph is (mostly) topologically ordered by `T: Ord`.#[expect(clippy::type_complexity)]fn look_ahead_sub_graph<T, ID, E, NI>(    start: Vec<T>,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> Result<(HashMap<ID, T>, HashMap<ID, Vec<ID>>, Vec<T>), E>where    T: Ord,    ID: Hash + Eq + Clone,    NI: IntoIterator<Item = Result<T, E>>,{    let mut queue: BinaryHeap<T> = start.into();    // Build separate node/neighbors maps since lifetime is different at caller    let mut node_map: HashMap<ID, T> = HashMap::new();    let mut neighbor_ids_map: HashMap<ID, Vec<ID>> = HashMap::new();    let mut has_reached_root = false;    while queue.len() > 1 || node_map.is_empty() || has_reached_root {        let Some(node) = queue.pop() else {            break;        };        let node_id = id_fn(&node);        if node_map.contains_key(&node_id) {            continue;        }         let mut neighbor_ids = Vec::new();        let mut neighbors_iter = neighbors_fn(&node).into_iter().peekable();        has_reached_root |= neighbors_iter.peek().is_none();        for neighbor in neighbors_iter {            let neighbor = neighbor?;            neighbor_ids.push(id_fn(&neighbor));            queue.push(neighbor);        }        node_map.insert(node_id.clone(), node);        neighbor_ids_map.insert(node_id, neighbor_ids);    }     assert!(queue.len() <= 1, "order of remainder shouldn't matter");    let remainder = queue.into_vec();     // Omit unvisited neighbors    if let Some(unvisited_id) = remainder.first().map(&id_fn) {        for neighbor_ids in neighbor_ids_map.values_mut() {            neighbor_ids.retain(|id| *id != unvisited_id);        }    }     Ok((node_map, neighbor_ids_map, remainder))} /// Builds a list of nodes reachable from the `start` where neighbors come after/// the node itself.////// Unlike `topo_order_reverse()`, nodes are sorted in reverse `T: Ord` order so/// long as they can respect the topological requirement.pub fn topo_order_reverse_ord<T, ID, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> Vec<T>where    T: Ord,    ID: Hash + Eq + Clone,    II: IntoIterator<Item = T>,    NI: IntoIterator<Item = T>,{    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));    topo_order_reverse_ord_ok(to_ok_iter(start), id_fn, neighbors_fn).unwrap()} /// Builds a list of `Ok` nodes reachable from the `start` where neighbors come/// after the node itself.////// Unlike `topo_order_reverse_ok()`, nodes are sorted in reverse `T: Ord` order/// so long as they can respect the topological requirement.////// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and/// returns the error.pub fn topo_order_reverse_ord_ok<T, ID, E, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> Result<Vec<T>, E>where    T: Ord,    ID: Hash + Eq + Clone,    II: IntoIterator<Item = Result<T, E>>,    NI: IntoIterator<Item = Result<T, E>>,{    struct InnerNode<T> {        node: Option<T>,        indegree: usize,    }     // DFS to accumulate incoming edges    let mut stack: Vec<T> = start.into_iter().try_collect()?;    let mut head_node_map: HashMap<ID, T> = HashMap::new();    let mut inner_node_map: HashMap<ID, InnerNode<T>> = HashMap::new();    let mut neighbor_ids_map: HashMap<ID, Vec<ID>> = HashMap::new();    while let Some(node) = stack.pop() {        let node_id = id_fn(&node);        if neighbor_ids_map.contains_key(&node_id) {            continue; // Already visited        }         let neighbors_iter = neighbors_fn(&node).into_iter();        let pos = stack.len();        stack.reserve(neighbors_iter.size_hint().0);        for neighbor in neighbors_iter {            stack.push(neighbor?);        }        let neighbor_ids = stack[pos..].iter().map(&id_fn).collect_vec();        if let Some(inner) = inner_node_map.get_mut(&node_id) {            inner.node = Some(node);        } else {            head_node_map.insert(node_id.clone(), node);        }         for neighbor_id in &neighbor_ids {            if let Some(inner) = inner_node_map.get_mut(neighbor_id) {                inner.indegree += 1;            } else {                let inner = InnerNode {                    node: head_node_map.remove(neighbor_id),                    indegree: 1,                };                inner_node_map.insert(neighbor_id.clone(), inner);            }        }        neighbor_ids_map.insert(node_id, neighbor_ids);    }     debug_assert!(head_node_map        .keys()        .all(|id| !inner_node_map.contains_key(id)));    debug_assert!(inner_node_map.values().all(|inner| inner.node.is_some()));    debug_assert!(inner_node_map.values().all(|inner| inner.indegree > 0));     // Using Kahn's algorithm    let mut queue: BinaryHeap<T> = head_node_map.into_values().collect();    let mut result = Vec::new();    while let Some(node) = queue.pop() {        let node_id = id_fn(&node);        result.push(node);        for neighbor_id in neighbor_ids_map.remove(&node_id).unwrap() {            let inner = inner_node_map.get_mut(&neighbor_id).unwrap();            inner.indegree -= 1;            if inner.indegree == 0 {                queue.push(inner.node.take().unwrap());                inner_node_map.remove(&neighbor_id);            }        }    }     assert!(inner_node_map.is_empty(), "graph has cycle");    Ok(result)} /// Find nodes in the start set that are not reachable from other nodes in the/// start set.pub fn heads<T, ID, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> HashSet<T>where    T: Hash + Eq + Clone,    ID: Hash + Eq,    II: IntoIterator<Item = T>,    NI: IntoIterator<Item = T>,{    let neighbors_fn = move |node: &T| to_ok_iter(neighbors_fn(node));    heads_ok(to_ok_iter(start), id_fn, neighbors_fn).unwrap()} /// Finds `Ok` nodes in the start set that are not reachable from other nodes in/// the start set.////// If `start` or `neighbors_fn()` yields an `Err`, this function terminates and/// returns the error.pub fn heads_ok<T, ID, E, II, NI>(    start: II,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,) -> Result<HashSet<T>, E>where    T: Hash + Eq + Clone,    ID: Hash + Eq,    II: IntoIterator<Item = Result<T, E>>,    NI: IntoIterator<Item = Result<T, E>>,{    let mut heads: HashSet<T> = start.into_iter().try_collect()?;    // Do a BFS until we have only one item left in the frontier. That frontier must    // have originated from one of the heads, and since there can't be cycles,    // it won't be able to eliminate any other heads.    let mut frontier: Vec<T> = heads.iter().cloned().collect();    let mut visited: HashSet<ID> = heads.iter().map(&id_fn).collect();    let mut root_reached = false;    while frontier.len() > 1 || (!frontier.is_empty() && root_reached) {        frontier = frontier

Lines 1–500 of 1241. Load more