Skip to main content

rsplug_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        /// References to dependents items
56        pub dependents_iter: DagDependentsIterator<'a, D>,
57    }
58
59    /// Dag Iterator with mapping function
60    pub struct DagIterator<T, D: DagNode, F: FnMut(DagIteratorMapFuncArgs<D>) -> T> {
61        pub(super) inner: Vec<DagItem<D>>,
62        pub(super) map_func: F,
63    }
64
65    impl<T, D: DagNode, F: FnMut(DagIteratorMapFuncArgs<D>) -> T> Iterator for DagIterator<T, D, F> {
66        type Item = T;
67
68        fn next(&mut self) -> Option<Self::Item> {
69            let Self { inner, map_func } = self;
70
71            inner.pop().map(|item| {
72                let dependents_iter = DagDependentsIterator {
73                    inner,
74                    seen: vec![false; inner.len()],
75                    idxes: item.dependents_indexes,
76                };
77                map_func(DagIteratorMapFuncArgs {
78                    inner: item.inner,
79                    dependents_iter,
80                })
81            })
82        }
83    }
84}
85
86/// Dag Node Trait
87pub trait DagNode {
88    fn id(&self) -> &str;
89    fn depends(&self) -> impl IntoIterator<Item = &impl AsRef<str>>;
90}
91
92/// Dag Resolution Error
93#[derive(Debug, Error)]
94pub enum DagError {
95    #[error("duplicate node: {0}")]
96    DuplicateName(String),
97    #[error("unknown dependency: {dep} (referred by {by})")]
98    UnknownDependency { dep: String, by: String },
99    #[error("cycle detected; remaining: {0:?}")]
100    CycleDetected(Vec<String>),
101}
102
103pub mod tree {
104    use super::*;
105
106    /// Resolved DAG Tree
107    pub struct DagTree<D: DagNode> {
108        pub(super) inner: Vec<DagItem<D>>,
109    }
110
111    pub(super) struct DagItem<D: DagNode> {
112        pub(super) inner: D,
113        pub(super) dependents_indexes: Vec<usize>,
114    }
115}
116
117/// Extension to IntoIterator<D>: Allow DAG resolution to be called by the method
118pub trait TryDag<D: DagNode>: IntoIterator<Item = D> + Sized {
119    /// Consume self to resolve the DAG and return a topo-ordered DagTree
120    fn try_dag(self) -> Result<DagTree<D>, DagError> {
121        // 1) まず全ノードを集める(この段階では重複チェックしない)
122        let mut nodes: Vec<DagItem<D>> = self
123            .into_iter()
124            .map(|node| DagItem {
125                inner: node,
126                dependents_indexes: Vec::new(),
127            })
128            .collect();
129
130        let n = nodes.len();
131
132        let mut waiting = Vec::with_capacity(n);
133        let mut references: Vec<Vec<usize>> = vec![Vec::new(); n];
134        {
135            // 2) &str をキーにした id → index マップを作成(ここで重複検出)
136            let mut id_to_index: HashMap<&str, usize> = HashMap::with_capacity(n);
137            for (i, dag) in nodes.iter().enumerate() {
138                let id = dag.inner.id(); // &str
139                if id_to_index.insert(id, i).is_some() {
140                    // ここだけエラーメッセージ用に to_string()
141                    return Err(DagError::DuplicateName(id.to_string()));
142                }
143            }
144
145            // 3) 依存グラフ(Kahn法用)と dependents の一時格納
146            for (idx, item) in nodes.iter().enumerate() {
147                let deps: Vec<_> = item.inner.depends().into_iter().collect();
148                for dep in &deps {
149                    let dep = dep.as_ref();
150                    let &dep_idx =
151                        id_to_index
152                            .get(dep)
153                            .ok_or_else(|| DagError::UnknownDependency {
154                                dep: dep.to_string(),
155                                by: item.inner.id().to_string(),
156                            })?;
157                    references[dep_idx].push(idx);
158                }
159                waiting.push(deps.len());
160            }
161        }
162
163        // 4) Kahn法でトポロジカルソート
164        let mut q = Vec::new();
165        for (i, &deg) in waiting.iter().enumerate() {
166            if deg == 0 {
167                q.push(i);
168            }
169        }
170        let mut topo = VecDeque::with_capacity(n); // VecDequeなのは先入れ先出しで取り出すため
171        while let Some(i) = q.pop() {
172            topo.push_front(i);
173            for &to in &references[i] {
174                waiting[to] -= 1;
175                if waiting[to] == 0 {
176                    q.push(to);
177                }
178            }
179        }
180        drop(q);
181
182        if topo.len() != n {
183            // 残っているノードはサイクルの一部
184            return Err(DagError::CycleDetected(
185                waiting
186                    .into_iter()
187                    .enumerate()
188                    .filter_map(|(i, deg)| {
189                        if deg > 0 {
190                            Some(nodes[i].inner.id().to_string())
191                        } else {
192                            None
193                        }
194                    })
195                    .collect::<Vec<_>>(),
196            ));
197        } else {
198            drop(waiting);
199        }
200
201        // 7) 現時点での依存関係をセット
202        for (node, r) in nodes.iter_mut().zip(references) {
203            node.dependents_indexes = r
204                .into_iter()
205                .map(|e| {
206                    topo.iter()
207                        .enumerate()
208                        .find_map(|(i, t)| if *t == e { Some(i) } else { None })
209                        .unwrap()
210                })
211                .collect();
212        }
213
214        // 6) topo 順に取り出す(swap_removeは使わずtake)
215        let inner: Vec<DagItem<D>> = {
216            let mut nodes_opt: Vec<_> = nodes.into_iter().map(Some).collect();
217            topo.into_iter()
218                .map(|i| nodes_opt[i].take().unwrap())
219                .collect()
220        };
221
222        Ok(DagTree { inner })
223    }
224}
225
226// Automatic implementation for all IntoIterator<Item = D>.
227impl<D: DagNode, I: IntoIterator<Item = D>> TryDag<D> for I {}
228
229impl<D: DagNode> DagTree<D> {
230    /// Iterate with mapping
231    pub fn into_map_iter<T, F: FnMut(DagIteratorMapFuncArgs<D>) -> T>(
232        self,
233        map_func: F,
234    ) -> DagIterator<T, D, F> {
235        DagIterator {
236            inner: self.inner,
237            map_func,
238        }
239    }
240}
241
242impl<D: DagNode> IntoIterator for DagTree<D> {
243    type Item = D;
244
245    type IntoIter = DagIterator<D, D, fn(DagIteratorMapFuncArgs<D>) -> D>;
246
247    fn into_iter(self) -> Self::IntoIter {
248        fn unwrap<D: DagNode>(d: DagIteratorMapFuncArgs<D>) -> D {
249            d.inner
250        }
251        self.into_map_iter(unwrap)
252    }
253}