vyre_foundation/visit/
traits.rs

1use crate::error::Result;
2use crate::ir_inner::model::expr::{Expr, GeneratorRef, Ident};
3use crate::ir_inner::model::generated::Node;
4use crate::ir_inner::model::node::NodeExtension;
5use crate::visit::VisitOrder;
6use smallvec::SmallVec;
7use std::ops::ControlFlow;
8
9/// Anything that can be lowered to a target representation.
10///
11/// Backends implement this trait for their target. The IR does not know
12/// what targets exist — it only knows that calling `.lower(&mut ctx)`
13/// walks the structure through the visitor contract.
14///
15/// # Errors
16///
17/// Backends report structured errors through their own context type.
18pub trait Lowerable<Ctx: ?Sized> {
19    /// Visit this IR structure and emit into the backend-specific context.
20    fn lower(&self, ctx: &mut Ctx) -> Result<()>;
21}
22
23/// Anything that can be executed against a runtime environment.
24///
25/// The reference interpreter and each backend implement this trait. Two
26/// `Evaluatable` implementations that produce the same output for the
27/// same input + environment are certifiably equivalent under the
28/// conform contract.
29pub trait Evaluatable<Env: ?Sized> {
30    /// The value type the evaluator produces (typically `Value` for the
31    /// reference interpreter, a typed handle for GPU backends).
32    type Value;
33
34    /// Evaluate this IR structure against the environment.
35    fn evaluate(&self, env: &mut Env) -> Result<Self::Value>;
36}
37
38/// Visitor over [`Node`] trees.
39///
40/// Implementors must handle every core node variant explicitly. Like
41/// [`crate::visit::ExprVisitor`], this trait is abstract-by-default so
42/// adding a new node variant forces downstream code to make a conscious
43/// decision.
44///
45/// Traversal order is explicit:
46/// - [`visit_node_preorder`] visits the current node before nested nodes.
47/// - [`visit_node_postorder`] visits nested nodes before the current node.
48///
49/// `NodeVisitor` traverses node structure only. If a visitor also needs
50/// to recurse into node-owned expressions, it should pair this trait
51/// with [`crate::visit::ExprVisitor`] and call the expression entry
52/// points from the relevant node hooks.
53pub trait NodeVisitor {
54    /// Break payload returned when traversal short-circuits.
55    type Break;
56
57    /// Variable declaration.
58    fn visit_let(&mut self, node: &Node, name: &Ident, value: &Expr) -> ControlFlow<Self::Break>;
59    /// Variable assignment.
60    fn visit_assign(&mut self, node: &Node, name: &Ident, value: &Expr)
61        -> ControlFlow<Self::Break>;
62    /// Buffer store.
63    fn visit_store(
64        &mut self,
65        node: &Node,
66        buffer: &Ident,
67        index: &Expr,
68        value: &Expr,
69    ) -> ControlFlow<Self::Break>;
70    /// Conditional branch.
71    fn visit_if(
72        &mut self,
73        node: &Node,
74        cond: &Expr,
75        then_nodes: &[Node],
76        otherwise: &[Node],
77    ) -> ControlFlow<Self::Break>;
78    /// Counted loop.
79    fn visit_loop(
80        &mut self,
81        node: &Node,
82        var: &Ident,
83        from: &Expr,
84        to: &Expr,
85        body: &[Node],
86    ) -> ControlFlow<Self::Break>;
87    /// Indirect dispatch source.
88    fn visit_indirect_dispatch(
89        &mut self,
90        node: &Node,
91        count_buffer: &Ident,
92        count_offset: u64,
93    ) -> ControlFlow<Self::Break>;
94    /// Async load node.
95    fn visit_async_load(
96        &mut self,
97        node: &Node,
98        source: &Ident,
99        destination: &Ident,
100        offset: &Expr,
101        size: &Expr,
102        tag: &Ident,
103    ) -> ControlFlow<Self::Break>;
104    /// Async store node.
105    fn visit_async_store(
106        &mut self,
107        node: &Node,
108        source: &Ident,
109        destination: &Ident,
110        offset: &Expr,
111        size: &Expr,
112        tag: &Ident,
113    ) -> ControlFlow<Self::Break>;
114    /// Async wait node.
115    fn visit_async_wait(&mut self, node: &Node, tag: &Ident) -> ControlFlow<Self::Break>;
116    /// Trap node.
117    fn visit_trap(&mut self, node: &Node, address: &Expr, tag: &Ident) -> ControlFlow<Self::Break>;
118    /// Resume node.
119    fn visit_resume(&mut self, node: &Node, tag: &Ident) -> ControlFlow<Self::Break>;
120    /// Return node.
121    fn visit_return(&mut self, node: &Node) -> ControlFlow<Self::Break>;
122    /// Barrier node.
123    fn visit_barrier(&mut self, node: &Node) -> ControlFlow<Self::Break>;
124    /// Block node.
125    fn visit_block(&mut self, node: &Node, body: &[Node]) -> ControlFlow<Self::Break>;
126    /// Region wrapper node.
127    fn visit_region(
128        &mut self,
129        node: &Node,
130        generator: &Ident,
131        source_region: &Option<GeneratorRef>,
132        body: &[Node],
133    ) -> ControlFlow<Self::Break>;
134    /// Downstream opaque node extension.
135    fn visit_opaque_node(
136        &mut self,
137        node: &Node,
138        extension: &dyn NodeExtension,
139    ) -> ControlFlow<Self::Break>;
140
141    /// Recursively walk this node's nested node children using the requested order.
142    fn walk_children_default(&mut self, node: &Node, order: VisitOrder) -> ControlFlow<Self::Break>
143    where
144        Self: Sized,
145    {
146        walk_node_children_default(self, node, order)
147    }
148}
149
150/// Visit a node tree in pre-order.
151pub fn visit_node<V: NodeVisitor>(visitor: &mut V, node: &Node) -> ControlFlow<V::Break> {
152    visit_node_preorder(visitor, node)
153}
154
155/// Visit a node tree in pre-order without recursive stack growth.
156pub fn visit_node_preorder<V: NodeVisitor>(visitor: &mut V, node: &Node) -> ControlFlow<V::Break> {
157    let mut stack = SmallVec::<[&Node; 32]>::new();
158    stack.push(node);
159    while let Some(current) = stack.pop() {
160        dispatch_node(visitor, current)?;
161        match current {
162            Node::If {
163                then, otherwise, ..
164            } => {
165                for n in otherwise.iter().rev() {
166                    stack.push(n);
167                }
168                for n in then.iter().rev() {
169                    stack.push(n);
170                }
171            }
172            Node::Loop { body, .. } | Node::Block(body) => {
173                for n in body.iter().rev() {
174                    stack.push(n);
175                }
176            }
177            Node::Region { body, .. } => {
178                for n in body.iter().rev() {
179                    stack.push(n);
180                }
181            }
182            _ => {}
183        }
184    }
185    ControlFlow::Continue(())
186}
187
188/// Visit a node tree in post-order without recursive stack growth.
189pub fn visit_node_postorder<V: NodeVisitor>(visitor: &mut V, node: &Node) -> ControlFlow<V::Break> {
190    enum Task<'a> {
191        Visit(&'a Node),
192        Dispatch(&'a Node),
193    }
194    let mut stack = SmallVec::<[Task<'_>; 32]>::new();
195    stack.push(Task::Visit(node));
196    while let Some(task) = stack.pop() {
197        match task {
198            Task::Visit(n) => {
199                stack.push(Task::Dispatch(n));
200                match n {
201                    Node::If {
202                        then, otherwise, ..
203                    } => {
204                        for child in otherwise.iter().rev() {
205                            stack.push(Task::Visit(child));
206                        }
207                        for child in then.iter().rev() {
208                            stack.push(Task::Visit(child));
209                        }
210                    }
211                    Node::Loop { body, .. } | Node::Block(body) => {
212                        for child in body.iter().rev() {
213                            stack.push(Task::Visit(child));
214                        }
215                    }
216                    Node::Region { body, .. } => {
217                        for child in body.iter().rev() {
218                            stack.push(Task::Visit(child));
219                        }
220                    }
221                    _ => {}
222                }
223            }
224            Task::Dispatch(n) => {
225                dispatch_node(visitor, n)?;
226            }
227        }
228    }
229    ControlFlow::Continue(())
230}
231
232/// Walk only the nested node children of `node`, leaving the current node to the caller.
233pub fn walk_node_children_default<V: NodeVisitor>(
234    visitor: &mut V,
235    node: &Node,
236    order: VisitOrder,
237) -> ControlFlow<V::Break> {
238    match node {
239        Node::If {
240            then, otherwise, ..
241        } => {
242            for child in then {
243                visit_node_with_order(visitor, child, order)?;
244            }
245            for child in otherwise {
246                visit_node_with_order(visitor, child, order)?;
247            }
248        }
249        Node::Loop { body, .. } | Node::Block(body) => {
250            for child in body {
251                visit_node_with_order(visitor, child, order)?;
252            }
253        }
254        Node::Region { body, .. } => {
255            for child in body.iter() {
256                visit_node_with_order(visitor, child, order)?;
257            }
258        }
259        _ => {}
260    }
261    ControlFlow::Continue(())
262}
263
264fn visit_node_with_order<V: NodeVisitor>(
265    visitor: &mut V,
266    node: &Node,
267    order: VisitOrder,
268) -> ControlFlow<V::Break> {
269    match order {
270        VisitOrder::Preorder => visit_node_preorder(visitor, node),
271        VisitOrder::Postorder => visit_node_postorder(visitor, node),
272    }
273}
274
275pub(crate) fn dispatch_node<V: NodeVisitor>(visitor: &mut V, node: &Node) -> ControlFlow<V::Break> {
276    match node {
277        Node::Let { name, value } => visitor.visit_let(node, name, value),
278        Node::Assign { name, value } => visitor.visit_assign(node, name, value),
279        Node::Store {
280            buffer,
281            index,
282            value,
283        } => visitor.visit_store(node, buffer, index, value),
284        Node::If {
285            cond,
286            then,
287            otherwise,
288        } => visitor.visit_if(node, cond, then, otherwise),
289        Node::Loop {
290            var,
291            from,
292            to,
293            body,
294        } => visitor.visit_loop(node, var, from, to, body),
295        Node::IndirectDispatch {
296            count_buffer,
297            count_offset,
298        } => visitor.visit_indirect_dispatch(node, count_buffer, *count_offset),
299        Node::AsyncLoad {
300            source,
301            destination,
302            offset,
303            size,
304            tag,
305        } => visitor.visit_async_load(node, source, destination, offset, size, tag),
306        Node::AsyncStore {
307            source,
308            destination,
309            offset,
310            size,
311            tag,
312        } => visitor.visit_async_store(node, source, destination, offset, size, tag),
313        Node::AsyncWait { tag } => visitor.visit_async_wait(node, tag),
314        Node::Trap { address, tag } => visitor.visit_trap(node, address, tag),
315        Node::Resume { tag } => visitor.visit_resume(node, tag),
316        Node::Return => visitor.visit_return(node),
317        Node::Barrier { .. } => visitor.visit_barrier(node),
318        Node::Block(body) => visitor.visit_block(node, body),
319        Node::Region {
320            generator,
321            source_region,
322            body,
323        } => visitor.visit_region(node, generator, source_region, body),
324        Node::Opaque(extension) => visitor.visit_opaque_node(node, extension.as_ref()),
325    }
326}
vyre_foundation/visit/traits.rs

vyre_foundation/visit/
traits.rs
