Skip to main content

dag/
lib.rs

1#[cfg(feature = "hashbrown")]
2use hashbrown::HashMap;
3#[cfg(not(feature = "hashbrown"))]
4use std::collections::HashMap;
5
6use std::collections::VecDeque;
7use thiserror::Error;
8
9use {
10    iterator::{DagIterator, DagIteratorMapFuncArgs},
11    tree::{DagItem, DagTree},
12};
13
14pub mod iterator {
15    use super::*;
16
17    pub struct DagDependentsIterator<'a, D: DagNode> {
18        inner: &'a Vec<DagItem<D>>,
19        seen: Vec<bool>,
20        idxes: Vec<usize>,
21    }
22
23    impl<'a, D: DagNode> Iterator for DagDependentsIterator<'a, D> {
24        type Item = Vec<&'a D>;
25        fn next(&mut self) -> Option<Self::Item> {
26            let next_idxes = self
27                .idxes
28                .iter()
29                .flat_map(|idx_before| {
30                    let idx_next = &self.inner[*idx_before].dependents_indexes;
31                    let mut res = Vec::new();
32                    for idx_next in idx_next {
33                        let seen = &mut self.seen[*idx_next];
34                        if !*seen {
35                            *seen = true;
36                            res.push(*idx_next);
37                        }
38                    }
39                    res
40                })
41                .collect();
42            let idxes = std::mem::replace(&mut self.idxes, next_idxes);
43            if idxes.is_empty() {
44                None
45            } else {
46                Some(idxes.into_iter().map(|i| &self.inner[i].inner).collect())
47            }
48        }
49    }
50
51    /// Arguments of the function which is used to map DagIterator
52    pub struct DagIteratorMapFuncArgs<'a, D: DagNode> {
53        /// Item itself
54        pub inner: D,
55        /// Index of the node within the original input order
56        pub index: usize,
57        /// Longest dependency chain depth (0 if no dependencies)
58        pub depth: usize,
59        /// References to dependents items
60        pub dependents_iter: DagDependentsIterator<'a, D>,
61    }
62
63    /// Dag Iterator with mapping function
64    pub struct DagIterator<T, D: DagNode, F: FnMut(DagIteratorMapFuncArgs<D>) -> T> {
65        pub(super) inner: Vec<DagItem<D>>,
66        pub(super) map_func: F,
67    }
68
69    impl<T, D: DagNode, F: FnMut(DagIteratorMapFuncArgs<D>) -> T> Iterator for DagIterator<T, D, F> {
70        type Item = T;
71
72        fn next(&mut self) -> Option<Self::Item> {
73            let Self { inner, map_func } = self;
74
75            inner.pop().map(|item| {
76                let dependents_iter = DagDependentsIterator {
77                    inner,
78                    seen: vec![false; inner.len()],
79                    idxes: item.dependents_indexes,
80                };
81                map_func(DagIteratorMapFuncArgs {
82                    inner: item.inner,
83                    index: item.original_index,
84                    depth: item.depth,
85                    dependents_iter,
86                })
87            })
88        }
89    }
90}
91
92/// Dag Node Trait
93pub trait DagNode {
94    /// 名前。`None` のノードは重複チェック・被依存・cycle の対象外となり、
95    /// ID を持たない「末端」ノードとして扱われる。
96    fn id(&self) -> Option<&str>;
97    fn depends(&self) -> impl IntoIterator<Item = &impl AsRef<str>>;
98}
99
100/// Dag Resolution Error
101#[derive(Debug, Error)]
102pub enum DagError {
103    #[error("duplicate node: {0}")]
104    DuplicateName(String),
105    #[error(
106        "unknown dependency: {dep} (referred by {by})",
107        by = by.as_deref().unwrap_or("<unnamed>")
108    )]
109    UnknownDependency { dep: String, by: Option<String> },
110    #[error("cycle detected; remaining: {0:?}")]
111    CycleDetected(Vec<String>),
112}
113
114pub mod tree {
115    use super::*;
116
117    /// Resolved DAG Tree
118    pub struct DagTree<D: DagNode> {
119        pub(super) inner: Vec<DagItem<D>>,
120    }
121
122    pub(super) struct DagItem<D: DagNode> {
123        pub(super) inner: D,
124        /// Index within the original input order
125        pub(super) original_index: usize,
126        /// Longest dependency chain depth (0 if no dependencies)
127        pub(super) depth: usize,
128        pub(super) dependents_indexes: Vec<usize>,
129    }
130}
131
132/// Extension to IntoIterator<D>: Allow DAG resolution to be called by the method
133pub trait TryDag<D: DagNode>: IntoIterator<Item = D> + Sized {
134    /// Consume self to resolve the DAG and return a topo-ordered DagTree
135    fn try_dag(self) -> Result<DagTree<D>, DagError> {
136        // 1) まず全ノードを集める(この段階では重複チェックしない)
137        let mut nodes: Vec<DagItem<D>> = self
138            .into_iter()
139            .enumerate()
140            .map(|(original_index, node)| DagItem {
141                inner: node,
142                original_index,
143                depth: 0,
144                dependents_indexes: Vec::new(),
145            })
146            .collect();
147
148        let n = nodes.len();
149
150        let mut waiting = Vec::with_capacity(n);
151        let mut references: Vec<Vec<usize>> = vec![Vec::new(); n];
152        {
153            // 2) &str をキーにした id → index マップを作成(ここで重複検出)。
154            //    名前なしノード(id() == None)は登録せず、被依存にもならない。
155            let mut id_to_index: HashMap<&str, usize> = HashMap::with_capacity(n);
156            for (i, dag) in nodes.iter().enumerate() {
157                if let Some(id) = dag.inner.id()
158                    && id_to_index.insert(id, i).is_some()
159                {
160                    // ここだけエラーメッセージ用に to_string()
161                    return Err(DagError::DuplicateName(id.to_string()));
162                }
163            }
164
165            // 3) 依存グラフ(Kahn法用)と dependents の一時格納
166            for (idx, item) in nodes.iter().enumerate() {
167                let deps: Vec<_> = item.inner.depends().into_iter().collect();
168                for dep in &deps {
169                    let dep = dep.as_ref();
170                    let &dep_idx =
171                        id_to_index
172                            .get(dep)
173                            .ok_or_else(|| DagError::UnknownDependency {
174                                dep: dep.to_string(),
175                                by: item.inner.id().map(str::to_string),
176                            })?;
177                    references[dep_idx].push(idx);
178                }
179                waiting.push(deps.len());
180            }
181        }
182
183        // 4) Kahn法でトポロジカルソート(ついでに依存の最深 depth を計算)。
184        //    i を処理する時点で nodes[i].depth は確定済み(依存先は全て先に処理される)ので、
185        //    1パスで正しい最深が入る。
186        let mut q = Vec::new();
187        for (i, &deg) in waiting.iter().enumerate() {
188            if deg == 0 {
189                q.push(i);
190            }
191        }
192        let mut topo = VecDeque::with_capacity(n); // VecDequeなのは先入れ先出しで取り出すため
193        while let Some(i) = q.pop() {
194            topo.push_front(i);
195            let depth_i = nodes[i].depth;
196            for &to in &references[i] {
197                // depth[to] = max(depth[to], depth[i] + 1)
198                let next = depth_i + 1;
199                if next > nodes[to].depth {
200                    nodes[to].depth = next;
201                }
202                waiting[to] -= 1;
203                if waiting[to] == 0 {
204                    q.push(to);
205                }
206            }
207        }
208        drop(q);
209
210        if topo.len() != n {
211            // 残っているノードはサイクルの一部。
212            // 名前なしノードは被依存になれないためサイクルに含まれず、
213            // id() == Some のノードのみ文字列化する。
214            return Err(DagError::CycleDetected(
215                waiting
216                    .into_iter()
217                    .enumerate()
218                    .filter_map(|(i, deg)| {
219                        if deg > 0 {
220                            nodes[i].inner.id().map(str::to_string)
221                        } else {
222                            None
223                        }
224                    })
225                    .collect::<Vec<_>>(),
226            ));
227        } else {
228            drop(waiting);
229        }
230
231        // 7) 現時点での依存関係をセット
232        let mut topo_positions = vec![0usize; n];
233        for (topo_index, &node_index) in topo.iter().enumerate() {
234            topo_positions[node_index] = topo_index;
235        }
236        for (node, r) in nodes.iter_mut().zip(references) {
237            node.dependents_indexes = r.into_iter().map(|e| topo_positions[e]).collect();
238        }
239
240        // 6) topo 順に取り出す(swap_removeは使わずtake)
241        let inner: Vec<DagItem<D>> = {
242            let mut nodes_opt: Vec<_> = nodes.into_iter().map(Some).collect();
243            topo.into_iter()
244                .map(|i| nodes_opt[i].take().unwrap())
245                .collect()
246        };
247
248        Ok(DagTree { inner })
249    }
250}
251
252// Automatic implementation for all IntoIterator<Item = D>.
253impl<D: DagNode, I: IntoIterator<Item = D>> TryDag<D> for I {}
254
255impl<D: DagNode> DagTree<D> {
256    /// Iterate with mapping
257    pub fn into_map_iter<T, F: FnMut(DagIteratorMapFuncArgs<D>) -> T>(
258        self,
259        map_func: F,
260    ) -> DagIterator<T, D, F> {
261        DagIterator {
262            inner: self.inner,
263            map_func,
264        }
265    }
266}
267
268impl<D: DagNode> IntoIterator for DagTree<D> {
269    type Item = D;
270
271    type IntoIter = DagIterator<D, D, fn(DagIteratorMapFuncArgs<D>) -> D>;
272
273    fn into_iter(self) -> Self::IntoIter {
274        fn unwrap<D: DagNode>(d: DagIteratorMapFuncArgs<D>) -> D {
275            d.inner
276        }
277        self.into_map_iter(unwrap)
278    }
279}