nickel_lang_core/
traverse.rs

1//! Traversal of trees of objects.
2
3use crate::bytecode::ast::alloc::{Allocable, AstAlloc};
4
5#[derive(Copy, Clone)]
6pub enum TraverseOrder {
7    TopDown,
8    BottomUp,
9}
10
11/// Flow control for tree traverals.
12pub enum TraverseControl<S, U> {
13    /// Normal control flow: continue recursing into the children.
14    ///
15    /// Pass the state &S to all children.
16    ContinueWithScope(S),
17    /// Normal control flow: continue recursing into the children.
18    ///
19    /// The state that was passed to the parent will be re-used for the children.
20    Continue,
21
22    /// Skip this branch of the tree.
23    SkipBranch,
24
25    /// Finish traversing immediately (and return a value).
26    Return(U),
27}
28
29impl<S, U> From<Option<U>> for TraverseControl<S, U> {
30    fn from(value: Option<U>) -> Self {
31        match value {
32            Some(u) => TraverseControl::Return(u),
33            None => TraverseControl::Continue,
34        }
35    }
36}
37
38pub trait Traverse<T>: Sized {
39    /// Apply a transformation on a object containing syntactic elements of type `T` (terms, types,
40    /// etc.) by mapping a faillible function `f` on each such node as prescribed by the order.
41    ///
42    /// `f` may return a generic error `E` and use the state `S` which is passed around.
43    fn traverse<F, E>(self, f: &mut F, order: TraverseOrder) -> Result<Self, E>
44    where
45        F: FnMut(T) -> Result<T, E>;
46
47    /// Recurse through the tree of objects top-down (a.k.a. pre-order), applying `f` to
48    /// each object.
49    ///
50    /// Through its return value, `f` can short-circuit one branch of the traversal or
51    /// the entire traversal.
52    ///
53    /// This traversal can make use of "scoped" state. The `scope` argument is passed to
54    /// each callback, and the callback can optionally override that scope just for its
55    /// own subtree in the traversal. For example, when traversing a tree of terms you can
56    /// maintain an environment. Most of the time the environment should get passed around
57    /// unchanged, but a let binder should override the environment of its subtree. It
58    /// does this by returning a `TraverseControl::ContinueWithScope` that contains the
59    /// new environment.
60    fn traverse_ref<S, U>(
61        &self,
62        f: &mut dyn FnMut(&T, &S) -> TraverseControl<S, U>,
63        scope: &S,
64    ) -> Option<U>;
65
66    fn find_map<S>(&self, mut pred: impl FnMut(&T) -> Option<S>) -> Option<S>
67    where
68        T: Clone,
69    {
70        self.traverse_ref(
71            &mut |t, _state: &()| {
72                if let Some(s) = pred(t) {
73                    TraverseControl::Return(s)
74                } else {
75                    TraverseControl::Continue
76                }
77            },
78            &(),
79        )
80    }
81}
82
83/// Similar to [Traverse], but takes an additional AST allocator for AST components that require
84/// such an allocator in order to build the result.
85pub trait TraverseAlloc<'ast, T>: Sized {
86    /// Same as [Traverse::traverse], but takes an additional AST allocator.
87    fn traverse<F, E>(
88        self,
89        alloc: &'ast AstAlloc,
90        f: &mut F,
91        order: TraverseOrder,
92    ) -> Result<Self, E>
93    where
94        F: FnMut(T) -> Result<T, E>;
95
96    /// Same as [Traverse::traverse_ref], but takes an additional AST allocator.
97    ///
98    /// There is as small difference though: this function guarantees that the lifetime of the
99    /// references is bound to the lifetime of the AST allocator, which the signature in
100    /// [Traverse::traverse_ref] does not. This is useful e.g. in the LSP to extract references and
101    /// store them in separate data structure. We can guarantee that those reference won't be
102    /// dangling as long as the allocator is around.
103    fn traverse_ref<S, U>(
104        &'ast self,
105        f: &mut dyn FnMut(&'ast T, &S) -> TraverseControl<S, U>,
106        scope: &S,
107    ) -> Option<U>;
108
109    fn find_map<S>(&'ast self, mut pred: impl FnMut(&'ast T) -> Option<S>) -> Option<S>
110    where
111        T: Clone + 'ast,
112    {
113        self.traverse_ref(
114            &mut |t, _state: &()| {
115                if let Some(s) = pred(t) {
116                    TraverseControl::Return(s)
117                } else {
118                    TraverseControl::Continue
119                }
120            },
121            &(),
122        )
123    }
124}
125
126/// Takes an iterator whose item type implements [TraverseAlloc], traverse each element, and
127/// collect the result as a slice allocated via `alloc`.
128pub fn traverse_alloc_many<'ast, T, U, I, F, E>(
129    alloc: &'ast AstAlloc,
130    it: I,
131    f: &mut F,
132    order: TraverseOrder,
133) -> Result<&'ast [U], E>
134where
135    U: TraverseAlloc<'ast, T> + Sized + Allocable,
136    I: IntoIterator<Item = U>,
137    F: FnMut(T) -> Result<T, E>,
138{
139    let collected: Result<Vec<_>, E> = it
140        .into_iter()
141        .map(|elt| elt.traverse(alloc, f, order))
142        .collect();
143
144    Ok(alloc.alloc_many(collected?))
145}
nickel_lang_core/traverse.rs

nickel_lang_core/
traverse.rs
