// 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 _;use smallvec::smallvec_inline;use smallvec::SmallVec; /// 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.////// Panics if the graph has cycle.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));    let cycle_fn = |_| panic!("graph has cycle");    topo_order_forward_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_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. If the graph has cycle, `cycle_fn()` is called with one/// of the nodes involved in the cycle.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,    cycle_fn: impl FnOnce(T) -> E,) -> 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 {            if !visiting.insert(id.clone()) {                return Err(cycle_fn(node));            }            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.////// Panics if the graph has cycle.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));    let cycle_fn = |_| panic!("graph has cycle");    topo_order_reverse_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_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. If the graph has cycle, `cycle_fn()` is called with one/// of the nodes involved in the cycle.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,    cycle_fn: impl FnOnce(T) -> E,) -> 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, cycle_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.////// Panics if the graph has cycle.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));    let cycle_fn = |_| panic!("graph has cycle");    topo_order_reverse_lazy_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_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. If the graph has cycle, `cycle_fn()` is/// called with one of the nodes involved in the cycle.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,    mut cycle_fn: impl FnMut(T) -> E,) -> 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, &mut cycle_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,        mut cycle_fn: impl FnMut(T) -> E,    ) -> 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));            if self.emitted.insert(id_fn(&node)) {                return Some(Ok(node));            } else {                return Some(Err(cycle_fn(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 = match topo_order_forward_ok(                    start_ids.iter().map(Ok),                    |id| *id,                    |id| neighbor_ids_map[id].iter().map(Ok),                    |id| cycle_fn(node_map.remove(id).unwrap()),                ) {                    Ok(ids) => ids,                    Err(err) => return Some(Err(err)),                };                self.result.reserve(sorted_ids.len());                for id in sorted_ids {                    let (id, node) = node_map.remove_entry(id).unwrap();                    if self.emitted.insert(id) {                        self.result.push(Ok(node));                    } else {                        self.result.push(Err(cycle_fn(node)));                    }                }                self.result.pop()            }            Err(err) => Some(Err(err)),        }    }} /// Splits DAG at the first single fork point, and builds a list of nodes/// reachable from the `start` where neighbors come after the node itself.////// This is a building block for lazy DAG iterators similar to/// [`topo_order_reverse_lazy()`]. The `start` list will be updated to include/// the next nodes to visit.////// If the split chunk of the graph has cycle, `cycle_fn()` is called with one/// of the nodes involved in the cycle.pub fn topo_order_reverse_chunked<T, ID, E, NI>(    start: &mut Vec<T>,    id_fn: impl Fn(&T) -> ID,    mut neighbors_fn: impl FnMut(&T) -> NI,    mut cycle_fn: impl FnMut(T) -> E,) -> Result<SmallVec<[T; 1]>, E>where    T: Ord,    ID: Hash + Eq + Clone,    NI: IntoIterator<Item = Result<T, E>>,{    // Fast path for linear DAG    if start.len() <= 1 {        let Some(node) = start.pop() else {            return Ok(SmallVec::new());        };        let neighbors_iter = neighbors_fn(&node).into_iter();        start.reserve(neighbors_iter.size_hint().0);        for neighbor in neighbors_iter {            start.push(neighbor?);        }        return Ok(smallvec_inline![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 = start.iter().map(&id_fn).collect_vec();    let (mut node_map, neighbor_ids_map, remainder) =        look_ahead_sub_graph(mem::take(start), &id_fn, &mut neighbors_fn)?;    *start = remainder;    let sorted_ids = topo_order_forward_ok(        start_ids.iter().map(Ok),        |id| *id,        |id| neighbor_ids_map[id].iter().map(Ok),        |id| cycle_fn(node_map.remove(id).unwrap()),    )?;    let sorted_nodes = sorted_ids        .iter()        .rev()        .map(|&id| node_map.remove(id).unwrap())        .collect();    Ok(sorted_nodes)} /// 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.////// Panics if the graph has cycle.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));    let cycle_fn = |_| panic!("graph has cycle");    topo_order_reverse_ord_ok(to_ok_iter(start), id_fn, neighbors_fn, cycle_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. If the graph has cycle, `cycle_fn()` is called with one/// of the nodes involved in the cycle.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,    cycle_fn: impl FnOnce(T) -> E,) -> 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 {

Lines 1–500 of 1470. Load more