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