dag/
lib.rs

1use hibitset::BitSet;
2use std::{
3    collections::{BinaryHeap, HashMap, HashSet, VecDeque},
4    convert::TryInto,
5    fmt::Debug,
6    hash::Hash,
7};
8
9/// Topologically sort `index_set` and return a sorted `Vec`.
10/// For the nodes without order-before relationship, the ones with smaller
11/// `order_indicator` output will be ordered first.
12pub fn topological_sort<InIndex, OutIndex, F, OrderIndicator, FOrd, Set>(
13    index_set: Set, predecessor_edges: F, order_indicator: FOrd,
14) -> Vec<OutIndex>
15where
16    InIndex: Copy + Hash + Eq + PartialEq + Ord + TryInto<OutIndex>,
17    <InIndex as TryInto<OutIndex>>::Error: Debug,
18    F: Fn(InIndex) -> Vec<InIndex>,
19    OrderIndicator: Ord,
20    FOrd: Fn(InIndex) -> OrderIndicator,
21    Set: SetLike<InIndex> + Default + Clone + IntoIterator<Item = InIndex>,
22{
23    let mut num_next_edges = HashMap::new();
24
25    for me in index_set.clone() {
26        num_next_edges.entry(me).or_insert(0);
27        for prev in predecessor_edges(me) {
28            if index_set.contains(&prev) {
29                *num_next_edges.entry(prev).or_insert(0) += 1;
30            }
31        }
32    }
33
34    let mut candidates = BinaryHeap::new();
35    let mut reversed_indices = Vec::new();
36
37    for me in index_set.clone() {
38        if num_next_edges[&me] == 0 {
39            candidates.push((order_indicator(me), me));
40        }
41    }
42    while let Some((_, me)) = candidates.pop() {
43        reversed_indices.push(me.try_into().expect("index in range"));
44
45        for prev in predecessor_edges(me) {
46            if index_set.contains(&prev) {
47                num_next_edges.entry(prev).and_modify(|e| *e -= 1);
48                if num_next_edges[&prev] == 0 {
49                    candidates.push((order_indicator(prev), prev));
50                }
51            }
52        }
53    }
54    reversed_indices.reverse();
55    reversed_indices
56}
57
58/// Return the future set of the nodes in `index_set`.
59/// The future set (including itself) of a node whose `stop_condition` is `true`
60/// will not be included.
61pub fn get_future<'a, InIndex, OutIndex, F, FStop, Set, Iter>(
62    index_set: Iter, successor_edges: F, stop_condition: FStop,
63) -> Set
64where
65    InIndex: 'a + Copy + TryInto<OutIndex>,
66    <InIndex as TryInto<OutIndex>>::Error: Debug,
67    OutIndex: 'a + Copy + Hash + Eq + PartialEq + Ord,
68    F: Fn(InIndex) -> Vec<InIndex>,
69    FStop: Fn(InIndex) -> bool,
70    Set: SetLike<OutIndex> + Default,
71    Iter: IntoIterator<Item = InIndex>,
72{
73    let mut queue = VecDeque::new();
74    let mut visited = Set::default();
75    for i in index_set {
76        visited.insert(i.try_into().expect("index in range"));
77        queue.push_back(i);
78    }
79    while let Some(x) = queue.pop_front() {
80        for succ in successor_edges(x) {
81            if stop_condition(succ) {
82                continue;
83            }
84            let out_index = succ.try_into().expect("index in range");
85            if !visited.contains(&out_index) {
86                queue.push_back(succ);
87                visited.insert(out_index);
88            }
89        }
90    }
91    visited
92}
93
94pub trait Graph {
95    type NodeIndex: 'static + Copy + Hash + Eq + PartialEq + Ord;
96}
97
98// TODO: Decide if returning Iterator is better than returning `Vec`?
99pub trait DAG: Graph {
100    fn predecessor_edges(
101        &self, node_index: Self::NodeIndex,
102    ) -> Vec<Self::NodeIndex>;
103
104    fn topological_sort_with_order_indicator<OrderIndicator, FOrd, Set>(
105        &self, index_set: Set, order_indicator: FOrd,
106    ) -> Vec<Self::NodeIndex>
107    where
108        OrderIndicator: Ord,
109        FOrd: Fn(Self::NodeIndex) -> OrderIndicator,
110        Set: SetLike<Self::NodeIndex>
111            + Default
112            + Clone
113            + IntoIterator<Item = Self::NodeIndex>,
114    {
115        topological_sort(
116            index_set,
117            |i| self.predecessor_edges(i),
118            order_indicator,
119        )
120    }
121
122    fn topological_sort<Set>(&self, index_set: Set) -> Vec<Self::NodeIndex>
123    where Set: SetLike<Self::NodeIndex>
124            + Default
125            + Clone
126            + IntoIterator<Item = Self::NodeIndex> {
127        // Any topological order will work, so just return a constant for
128        // `order_indicator`.
129        self.topological_sort_with_order_indicator(index_set, |_| true)
130    }
131}
132
133pub trait RichDAG: DAG {
134    fn successor_edges(
135        &self, node_index: Self::NodeIndex,
136    ) -> Vec<Self::NodeIndex>;
137
138    fn get_future_with_stop_condition<FStop, Set, Iter>(
139        &self, index_set: Iter, stop_condition: FStop,
140    ) -> Set
141    where
142        FStop: Fn(Self::NodeIndex) -> bool,
143        Set: SetLike<Self::NodeIndex> + Default,
144        Iter: IntoIterator<Item = Self::NodeIndex>,
145    {
146        get_future(index_set, |i| self.successor_edges(i), stop_condition)
147    }
148
149    fn get_future<Set, Iter>(&self, index_set: Iter) -> Set
150    where
151        Set: SetLike<Self::NodeIndex> + Default,
152        Iter: IntoIterator<Item = Self::NodeIndex>,
153    {
154        self.get_future_with_stop_condition(index_set, |_| false)
155    }
156}
157
158pub trait TreeGraph: Graph {
159    fn parent(&self, node_index: Self::NodeIndex) -> Option<Self::NodeIndex>;
160    fn referees(&self, node_index: Self::NodeIndex) -> Vec<Self::NodeIndex>;
161}
162
163pub trait RichTreeGraph: TreeGraph {
164    fn children(&self, node_index: Self::NodeIndex) -> Vec<Self::NodeIndex>;
165    fn referrers(&self, node_index: Self::NodeIndex) -> Vec<Self::NodeIndex>;
166}
167
168impl<T: TreeGraph> DAG for T {
169    fn predecessor_edges(
170        &self, node_index: Self::NodeIndex,
171    ) -> Vec<Self::NodeIndex> {
172        let mut predecessor_edges = self.referees(node_index);
173        if let Some(p) = self.parent(node_index) {
174            predecessor_edges.push(p);
175        }
176        predecessor_edges
177    }
178}
179
180impl<T: RichTreeGraph + DAG> RichDAG for T {
181    fn successor_edges(
182        &self, node_index: Self::NodeIndex,
183    ) -> Vec<Self::NodeIndex> {
184        let mut successor_edges = self.children(node_index);
185        successor_edges.append(&mut self.referrers(node_index));
186        successor_edges
187    }
188}
189
190pub trait SetLike<T> {
191    fn insert(&mut self, i: T) -> bool;
192    fn contains(&self, i: &T) -> bool;
193}
194
195impl<T: Eq + Hash> SetLike<T> for HashSet<T> {
196    fn insert(&mut self, i: T) -> bool { self.insert(i) }
197
198    fn contains(&self, i: &T) -> bool { self.contains(i) }
199}
200
201impl SetLike<u32> for BitSet {
202    fn insert(&mut self, i: u32) -> bool { self.add(i) }
203
204    fn contains(&self, i: &u32) -> bool { self.contains(*i) }
205}
dag/lib.rs

dag/
lib.rs
