rssn-advanced 0.1.0

This is rssn-advanced: The next generation symbolic core of rssn.
Documentation 
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

//! Local AST projection for stack-local computation.
//!
//! # Architecture
//!
//! The AST projection maps a subgraph of the global DAG into a flat,
//! stack-allocated tree using relative pointers. This gives algorithms
//! a familiar tree interface while metadata remains in the DAG arena.
//!
//! - `projection` — The `AstProjection` type and its buffer management.
//! - `pointer` — `RelPtr<i32>` / `RelPtr<i64>` relative pointer types.
//! - `convert` — DAG ↔ AST conversion routines.

pub mod convert;
pub mod pointer;
pub mod projection;

pub use projection::{AstChildList, AstNode, AstProjection, PostOrderVisitor};

/// Visitor for traversing an [`AstProjection`] tree node-by-node.
///
/// Each method corresponds to a `SymbolKind` variant. Default implementations
/// are no-ops so callers only implement the kinds they care about.
pub trait AstVisitor {
    /// Called for each variable node.
    fn visit_variable(
        &mut self,
        _idx: usize,
        _sym_id: crate::dag::symbol::SymbolId,
        _dag_id: crate::dag::node::DagNodeId,
    ) {
    }
    /// Called for each constant node.
    fn visit_constant(&mut self, _idx: usize, _val: f64, _dag_id: crate::dag::node::DagNodeId) {}
    /// Called for each operator node.
    fn visit_operator(
        &mut self,
        _idx: usize,
        _op: crate::dag::symbol::OpKind,
        _children: &AstChildList,
        _dag_id: crate::dag::node::DagNodeId,
    ) {
    }
    /// Called for each function-call node.
    fn visit_function(
        &mut self,
        _idx: usize,
        _fn_id: crate::dag::symbol::FnId,
        _children: &AstChildList,
        _dag_id: crate::dag::node::DagNodeId,
    ) {
    }
}

/// Iterates the flat `AstProjection` buffer in storage order, calling the
/// appropriate `visit_*` method on `visitor` for each node.
pub fn walk_projection<V: AstVisitor>(proj: &AstProjection, visitor: &mut V) {
    for (idx, node) in proj.nodes.iter().enumerate() {
        match node.kind {
            crate::dag::symbol::SymbolKind::Variable(sym_id) => {
                visitor.visit_variable(idx, sym_id, node.dag_id);
            }
            crate::dag::symbol::SymbolKind::Constant(val) => {
                visitor.visit_constant(idx, val, node.dag_id);
            }
            crate::dag::symbol::SymbolKind::Operator(op) => {
                visitor.visit_operator(idx, op, &node.children, node.dag_id);
            }
            crate::dag::symbol::SymbolKind::Function(fn_id) => {
                visitor.visit_function(idx, fn_id, &node.children, node.dag_id);
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::convert::dag_to_ast;
    use crate::dag::builder::DagBuilder;

    #[test]
    fn walk_counts_node_kinds() {
        let mut b = DagBuilder::new();
        let x = b.variable("x");
        let c = b.constant(2.0);
        let expr = b.mul(x, c);
        let proj = dag_to_ast(b.arena(), expr);

        struct Counter {
            vars: usize,
            consts: usize,
            ops: usize,
        }
        impl AstVisitor for Counter {
            fn visit_variable(
                &mut self,
                _idx: usize,
                _sym_id: crate::dag::symbol::SymbolId,
                _dag_id: crate::dag::node::DagNodeId,
            ) {
                self.vars += 1;
            }
            fn visit_constant(
                &mut self,
                _idx: usize,
                _val: f64,
                _dag_id: crate::dag::node::DagNodeId,
            ) {
                self.consts += 1;
            }
            fn visit_operator(
                &mut self,
                _idx: usize,
                _op: crate::dag::symbol::OpKind,
                _children: &AstChildList,
                _dag_id: crate::dag::node::DagNodeId,
            ) {
                self.ops += 1;
            }
        }
        let mut counter = Counter {
            vars: 0,
            consts: 0,
            ops: 0,
        };
        walk_projection(&proj, &mut counter);
        assert_eq!(counter.vars, 1);
        assert_eq!(counter.consts, 1);
        assert_eq!(counter.ops, 1);
    }
}