flams_utils/
treelike.rs

1use std::collections::VecDeque;
2use std::fmt::{Display, Write};
3
4pub trait TreeChild<T: TreeLike> {
5    fn children<'a>(&self) -> Option<T::RefIter<'a>>
6    where
7        Self: 'a;
8}
9
10pub trait TreeLike: Sized {
11    type Child<'a>: TreeChild<Self>
12    where
13        Self: 'a;
14    type RefIter<'a>: Iterator<Item = Self::Child<'a>> + TreeChildIter<'a, Self>
15    where
16        Self: 'a;
17    fn children(&self) -> Option<Self::RefIter<'_>>;
18    #[inline]
19    fn dfs(&self) -> Option<DFSIter<Self>> {
20        self.children().map(TreeChildIter::dfs)
21    }
22    #[inline]
23    fn bfs(&self) -> Option<BFSIter<Self>> {
24        self.children().map(TreeChildIter::bfs)
25    }
26
27    #[cfg(feature = "rayon")]
28    fn bfs_par<'a>(&'a self) -> Option<spliter::ParSpliter<BFSIter<'a, Self>>>
29    where
30        Self::Child<'a>: Send,
31        Self::RefIter<'a>: Send,
32    {
33        self.children().map(TreeChildIter::par_iter)
34    }
35
36    #[inline]
37    #[allow(clippy::missing_errors_doc)]
38    fn dfs_with_close<
39        'a,
40        R,
41        SG,
42        SL,
43        Open: FnMut(&mut SG, Self::Child<'a>) -> Result<DFSContinuation<SL>, R>,
44        Close: FnMut(&mut SG, SL) -> Result<(), R>,
45    >(
46        &'a self,
47        state: &mut SG,
48        open: Open,
49        close: Close,
50    ) -> Result<(), R> {
51        self.children().map_or(Ok(()), |d| {
52            TreeChildIter::<'a, Self>::dfs_with_close(d, state, open, close)
53        })
54    }
55
56    #[inline]
57    #[allow(clippy::missing_errors_doc)]
58    fn display_nested<'a>(
59        &'a self,
60        f: &mut std::fmt::Formatter<'_>,
61        open: impl Fn(
62            &Self::Child<'a>,
63            &mut Indentor,
64            &mut std::fmt::Formatter<'_>,
65        ) -> Result<DFSContinuation<()>, std::fmt::Error>,
66        close: impl Fn(&Self::Child<'a>, &mut std::fmt::Formatter<'_>) -> std::fmt::Result,
67        indent: Option<Indentor>,
68    ) -> std::fmt::Result {
69        self.children().map_or(Ok(()), |d| {
70            TreeChildIter::<'a, Self>::display_nested(d, f, open, close, indent)
71        })
72    }
73
74    #[inline]
75    #[allow(clippy::missing_errors_doc)]
76    fn display_children<'a, I: Into<Self::RefIter<'a>>>(
77        i: I,
78        f: &mut std::fmt::Formatter<'_>,
79        open: impl Fn(
80            &Self::Child<'a>,
81            &mut Indentor,
82            &mut std::fmt::Formatter<'_>,
83        ) -> Result<DFSContinuation<()>, std::fmt::Error>,
84        close: impl Fn(&Self::Child<'a>, &mut std::fmt::Formatter<'_>) -> std::fmt::Result,
85        indent: Option<Indentor>,
86    ) -> std::fmt::Result
87    where
88        Self: 'a,
89    {
90        i.into().display_nested(f, open, close, indent)
91    }
92}
93
94#[derive(Clone)]
95pub struct Indentor<'a>(&'a str, i32, std::cell::Cell<bool>);
96impl Default for Indentor<'_> {
97    fn default() -> Self {
98        Self(Self::DEFAULT, 0, std::cell::Cell::new(false))
99    }
100}
101impl Indentor<'_> {
102    const DEFAULT: &'static str = "  ";
103    pub fn scoped<R>(&mut self, f: impl FnOnce(&mut Self) -> R) -> R {
104        self.1 += 1;
105        let r = f(self);
106        self.1 -= 1;
107        r
108    }
109    #[must_use]
110    pub const fn new(indent_str: &str, indent: i32) -> Indentor<'_> {
111        Indentor(indent_str, indent, std::cell::Cell::new(false))
112    }
113    #[must_use]
114    pub const fn with(indent: i32) -> Indentor<'static> {
115        Indentor(Self::DEFAULT, indent, std::cell::Cell::new(false))
116    }
117    pub fn skip_next(&self) {
118        self.2.set(true);
119    }
120}
121impl Display for Indentor<'_> {
122    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123        if self.2.replace(false) || self.1 == 0 {
124            return Ok(());
125        }
126        f.write_char('\n')?;
127        for _ in 0..self.1 {
128            write!(f, "{}", self.0)?;
129        }
130        Ok(())
131    }
132}
133
134pub enum DFSContinuation<SL> {
135    Continue,
136    OnClose(SL),
137    SkipChildren,
138    SkipNext,
139    SkipNextAndClose(SL),
140}
141
142pub trait TreeChildIter<'a, T: TreeLike<RefIter<'a> = Self> + 'a>: Iterator
143where
144    Self: Sized + 'a,
145{
146    #[inline]
147    fn dfs(self) -> DFSIter<'a, T> {
148        DFSIter {
149            stack: Vec::new(),
150            current: self,
151        }
152    }
153    #[inline]
154    fn bfs(self) -> BFSIter<'a, T> {
155        BFSIter {
156            stack: VecDeque::new(),
157            current: self,
158        }
159    }
160    #[cfg(feature = "rayon")]
161    #[inline]
162    fn par_iter(self) -> spliter::ParSpliter<BFSIter<'a, T>>
163    where
164        Self: Send,
165        T::Child<'a>: Send,
166    {
167        use rayon::iter::IntoParallelIterator;
168        use spliter::ParallelSpliterator;
169        self.bfs().par_split().into_par_iter()
170    }
171
172    #[allow(clippy::missing_errors_doc)]
173    #[allow(clippy::unnecessary_map_or)]
174    fn dfs_with_close<
175        R,
176        SG,
177        SL,
178        Open: FnMut(&mut SG, T::Child<'a>) -> Result<DFSContinuation<SL>, R>,
179        Close: FnMut(&mut SG, SL) -> Result<(), R>,
180    >(
181        self,
182        state: &mut SG,
183        mut open: Open,
184        mut close: Close,
185    ) -> Result<(), R> {
186        fn i_next<'a, SL, T: TreeLike + 'a>(
187            current: &mut (Option<SL>, T::RefIter<'a>),
188            stack: &mut Vec<(Option<SL>, T::RefIter<'a>)>,
189        ) -> Option<(bool, T::Child<'a>)> {
190            current.1.next().map(|c| {
191                (
192                    c.children().map_or(false, |children| {
193                        stack.push(std::mem::replace(current, (None, children)));
194                        true
195                    }),
196                    c,
197                )
198            })
199        }
200
201        let mut stack = Vec::new();
202        let mut current = (None, self);
203        loop {
204            let r = if let Some(r) = i_next::<SL, T>(&mut current, &mut stack) {
205                Some(r)
206            } else {
207                loop {
208                    if let Some(next) = stack.pop() {
209                        let (old, _) = std::mem::replace(&mut current, next);
210                        if let Some(s) = old {
211                            close(state, s)?;
212                        }
213                        if let Some(c) = i_next::<SL, T>(&mut current, &mut stack) {
214                            break Some(c);
215                        }
216                    } else {
217                        break None;
218                    }
219                }
220            };
221            if let Some((has_children, r)) = r {
222                match open(state, r)? {
223                    DFSContinuation::Continue => (),
224                    DFSContinuation::OnClose(s) if has_children => {
225                        current.0 = Some(s);
226                    }
227                    DFSContinuation::OnClose(s) => close(state, s)?,
228                    DFSContinuation::SkipChildren => {
229                        if has_children {
230                            current = stack.pop().unwrap_or_else(|| unreachable!());
231                        }
232                    }
233                    DFSContinuation::SkipNext => {
234                        current.1.next();
235                    }
236                    DFSContinuation::SkipNextAndClose(s) => {
237                        if has_children {
238                            current.1.next();
239                            current.0 = Some(s);
240                        } else {
241                            close(state, s)?;
242                        }
243                    }
244                }
245            } else {
246                break;
247            }
248        }
249        Ok(())
250    }
251
252    #[allow(clippy::cast_sign_loss)]
253    #[allow(clippy::missing_errors_doc)]
254    fn display_nested(
255        self,
256        f: &mut std::fmt::Formatter<'_>,
257        open: impl Fn(
258            &T::Child<'a>,
259            &mut Indentor,
260            &mut std::fmt::Formatter<'_>,
261        ) -> Result<DFSContinuation<()>, std::fmt::Error>,
262        close: impl Fn(&T::Child<'a>, &mut std::fmt::Formatter<'_>) -> std::fmt::Result,
263        ind: Option<Indentor>,
264    ) -> std::fmt::Result {
265        let mut ind = ind.unwrap_or_default();
266        self.dfs_with_close(
267            &mut (f, &mut ind),
268            |(f, ind), c| {
269                ind.fmt(f)?;
270                Ok(match open(&c, ind, f)? {
271                    DFSContinuation::OnClose(()) => {
272                        ind.1 += 1;
273                        DFSContinuation::OnClose(c)
274                    }
275                    DFSContinuation::Continue => DFSContinuation::Continue,
276                    DFSContinuation::SkipChildren => DFSContinuation::SkipChildren,
277                    DFSContinuation::SkipNext => DFSContinuation::SkipNext,
278                    DFSContinuation::SkipNextAndClose(()) => {
279                        ind.1 += 1;
280                        DFSContinuation::SkipNextAndClose(c)
281                    }
282                })
283            },
284            |(f, ind), c| {
285                ind.1 -= 1;
286                if ind.1 == 0 {
287                    f.write_char('\n')
288                } else {
289                    ind.fmt(f)
290                }?;
291                close(&c, f)
292            },
293        )
294    }
295}
296impl<'a, T: TreeLike<RefIter<'a> = I> + 'a, I: Iterator<Item = T::Child<'a>> + 'a>
297    TreeChildIter<'a, T> for I
298{
299}
300
301pub struct DFSIter<'a, T: TreeLike + 'a> {
302    stack: Vec<T::RefIter<'a>>,
303    current: T::RefIter<'a>,
304}
305impl<'a, T: TreeLike> DFSIter<'a, T> {
306    fn i_next(&mut self) -> Option<T::Child<'a>> {
307        if let Some(c) = self.current.next() {
308            if let Some(children) = c.children() {
309                self.stack
310                    .push(std::mem::replace(&mut self.current, children));
311            }
312            Some(c)
313        } else {
314            None
315        }
316    }
317}
318impl<'a, T: TreeLike + 'a> Iterator for DFSIter<'a, T> {
319    type Item = T::Child<'a>;
320    fn next(&mut self) -> Option<Self::Item> {
321        self.i_next().or_else(|| {
322            while let Some(next) = self.stack.pop() {
323                self.current = next;
324                if let Some(c) = self.i_next() {
325                    return Some(c);
326                }
327            }
328            None
329        })
330    }
331}
332
333pub struct BFSIter<'a, T: TreeLike + 'a> {
334    stack: VecDeque<T::RefIter<'a>>,
335    current: T::RefIter<'a>,
336}
337impl<'a, T: TreeLike + 'a> Iterator for BFSIter<'a, T> {
338    type Item = T::Child<'a>;
339    fn next(&mut self) -> Option<Self::Item> {
340        if let Some(c) = self.current.next() {
341            if let Some(children) = c.children() {
342                self.stack.push_back(children);
343            }
344            Some(c)
345        } else {
346            while let Some(next) = self.stack.pop_front() {
347                self.current = next;
348                if let Some(c) = self.current.next() {
349                    return Some(c);
350                }
351            }
352            None
353        }
354    }
355}
356#[cfg(feature = "rayon")]
357impl<'a, T: TreeLike + 'a> spliter::Spliterator for BFSIter<'a, T> {
358    fn split(&mut self) -> Option<Self> {
359        if self.stack.len() < 2 {
360            return None;
361        }
362        let split = self.stack.len() / 2;
363        let mut new_stack = self.stack.split_off(split);
364        let new_curr = new_stack.pop_front().unwrap_or_else(|| unreachable!());
365        Some(BFSIter {
366            stack: new_stack,
367            current: new_curr,
368        })
369    }
370}
371
372impl TreeChild<Self> for std::fs::DirEntry {
373    fn children<'a>(&self) -> Option<<Self as TreeLike>::RefIter<'a>> {
374        <Self as TreeLike>::children(self)
375    }
376}
377
378impl TreeLike for std::fs::DirEntry {
379    type Child<'a> = Self;
380    type RefIter<'a> =
381        std::iter::FilterMap<std::fs::ReadDir, fn(std::io::Result<Self>) -> Option<Self>>;
382    #[allow(clippy::option_if_let_else)]
383    fn children(&self) -> Option<Self::RefIter<'_>> {
384        if let Ok(p) = std::fs::read_dir(self.path()) {
385            Some(p.filter_map(Result::ok))
386        } else {
387            None
388        }
389    }
390    /*fn children<'a>(c:&Self::Item<'a>) -> Option<Self::RefIter<'a>> {
391        if let Ok(p) = std::fs::read_dir(c.path()) {
392            Some(p.filter_map(Result::ok))
393        } else { None }
394    }*/
395}
396
397#[cfg(test)]
398mod tests {
399    use super::*;
400    use std::fmt::{Display, Formatter, Write};
401
402    fn foo() -> TopFoo {
403        TopFoo {
404            children: vec![
405                FooChild::Leaf("a"),
406                FooChild::Foo1(Foo1 {
407                    name: "b",
408                    children: vec![
409                        FooChild::Leaf("c"),
410                        FooChild::Foo2(Foo2 {
411                            name: "d",
412                            children: vec![FooChild::Leaf("e"), FooChild::Leaf("f")],
413                        }),
414                    ],
415                }),
416                FooChild::Foo2(Foo2 {
417                    name: "g",
418                    children: vec![FooChild::Leaf("h"), FooChild::Leaf("i")],
419                }),
420            ],
421        }
422    }
423
424    #[test]
425    fn dfs() {
426        let foo = foo();
427        let dfs = foo
428            .dfs()
429            .unwrap_or_else(|| unreachable!())
430            .map(FooChild::name)
431            .collect::<Vec<_>>();
432        assert_eq!(dfs, vec!["a", "b", "c", "d", "e", "f", "g", "h", "i"]);
433    }
434    #[test]
435    fn bfs() {
436        let foo = foo();
437        let bfs = foo
438            .bfs()
439            .unwrap_or_else(|| unreachable!())
440            .map(FooChild::name)
441            .collect::<Vec<_>>();
442        assert_eq!(bfs, vec!["a", "b", "g", "c", "d", "h", "i", "e", "f"]);
443    }
444    #[test]
445    fn display() {
446        let foo = foo();
447        let expect = "<TopFoo>
448  <Leaf a/>
449  <Foo1>
450    <Leaf c/>
451    <Foo2>
452      <Leaf e/>
453      <Leaf f/>
454    </Foo2>
455  </Foo1>
456  <Foo2>
457    <Leaf h/>
458    <Leaf i/>
459  </Foo2>
460</TopFoo>";
461        let mut s = String::new();
462        write!(s, "{foo}").unwrap();
463        assert_eq!(s, expect);
464    }
465    struct TopFoo {
466        children: Vec<FooChild>,
467    }
468    impl TreeLike for TopFoo {
469        type Child<'a> = &'a FooChild;
470        type RefIter<'a> = std::slice::Iter<'a, FooChild>;
471        fn children(&self) -> Option<Self::RefIter<'_>> {
472            Some(self.children.iter())
473        }
474    }
475    impl Display for TopFoo {
476        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
477            write!(f, "<TopFoo>")?;
478            self.display_nested(
479                f,
480                |c, _, f| match *c {
481                    FooChild::Leaf(s) => {
482                        write!(f, "<Leaf {s}/>")?;
483                        Ok(DFSContinuation::Continue)
484                    }
485                    FooChild::Foo1(_) => {
486                        write!(f, "<Foo1>")?;
487                        Ok(DFSContinuation::OnClose(()))
488                    }
489                    FooChild::Foo2(_) => {
490                        write!(f, "<Foo2>")?;
491                        Ok(DFSContinuation::OnClose(()))
492                    }
493                },
494                |c, f| match *c {
495                    FooChild::Leaf(_) => unreachable!(),
496                    FooChild::Foo1(_) => write!(f, "</Foo1>"),
497                    FooChild::Foo2(_) => write!(f, "</Foo2>"),
498                },
499                Some(Indentor::with(1)),
500            )?;
501            write!(f, "\n</TopFoo>")
502        }
503    }
504    struct Foo1 {
505        name: &'static str,
506        children: Vec<FooChild>,
507    }
508    struct Foo2 {
509        name: &'static str,
510        children: Vec<FooChild>,
511    }
512    enum FooChild {
513        Leaf(&'static str),
514        Foo1(Foo1),
515        Foo2(Foo2),
516    }
517    impl FooChild {
518        const fn name(&self) -> &'static str {
519            match self {
520                Self::Leaf(s) => s,
521                Self::Foo1(f) => f.name,
522                Self::Foo2(f) => f.name,
523            }
524        }
525    }
526    impl TreeChild<TopFoo> for &FooChild {
527        fn children<'b>(&self) -> Option<<TopFoo as TreeLike>::RefIter<'b>>
528        where
529            Self: 'b,
530        {
531            match self {
532                FooChild::Leaf(_) => None,
533                FooChild::Foo1(f) => Some(f.children.iter()),
534                FooChild::Foo2(f) => Some(f.children.iter()),
535            }
536        }
537    }
538}
flams_utils/treelike.rs

flams_utils/
treelike.rs
