lex_ast/
ids.rs

1//! Node IDs (§5.2). A NodeId encodes the path from a stage root to a node
2//! as `n_0[.<i>]*`, where each `<i>` is the position in the parent's
3//! children array.
4
5use crate::canonical::*;
6use std::collections::HashMap;
7
8#[derive(Debug, Clone, PartialEq, Eq, Hash)]
9pub struct NodeId(pub String);
10
11impl NodeId {
12    pub fn root() -> Self { NodeId("n_0".into()) }
13    pub fn child(&self, i: usize) -> Self { NodeId(format!("{}.{}", self.0, i)) }
14    pub fn as_str(&self) -> &str { &self.0 }
15}
16
17/// Walk a stage and emit all NodeIds with their underlying nodes (referenced).
18/// Order is depth-first, child-index ordered.
19pub fn collect_ids(stage: &Stage) -> Vec<(NodeId, NodeRef<'_>)> {
20    let mut out = Vec::new();
21    let root = NodeId::root();
22    out.push((root.clone(), NodeRef::Stage(stage)));
23    walk_stage(stage, &root, &mut out);
24    out
25}
26
27/// Map every CExpr in a stage to its NodeId, keyed by the CExpr's address.
28/// Caller borrows the same `stage` instance the bytecode compiler walks.
29pub fn expr_ids(stage: &Stage) -> HashMap<*const CExpr, NodeId> {
30    let mut out = HashMap::new();
31    for (id, n) in collect_ids(stage) {
32        if let NodeRef::CExpr(p) = n {
33            out.insert(p as *const CExpr, id);
34        }
35    }
36    out
37}
38
39#[derive(Debug)]
40pub enum NodeRef<'a> {
41    Stage(&'a Stage),
42    CExpr(&'a CExpr),
43    Pattern(&'a Pattern),
44    TypeExpr(&'a TypeExpr),
45}
46
47fn walk_stage<'a>(s: &'a Stage, parent: &NodeId, out: &mut Vec<(NodeId, NodeRef<'a>)>) {
48    match s {
49        Stage::FnDecl(fd) => {
50            // children: params (0..n_params), return_type (n), body (n+1)
51            for (i, p) in fd.params.iter().enumerate() {
52                let id = parent.child(i);
53                out.push((id.clone(), NodeRef::TypeExpr(&p.ty)));
54                walk_type(&p.ty, &id, out);
55            }
56            let rid = parent.child(fd.params.len());
57            out.push((rid.clone(), NodeRef::TypeExpr(&fd.return_type)));
58            walk_type(&fd.return_type, &rid, out);
59            let bid = parent.child(fd.params.len() + 1);
60            out.push((bid.clone(), NodeRef::CExpr(&fd.body)));
61            walk_expr(&fd.body, &bid, out);
62        }
63        Stage::TypeDecl(td) => {
64            let id = parent.child(0);
65            out.push((id.clone(), NodeRef::TypeExpr(&td.definition)));
66            walk_type(&td.definition, &id, out);
67        }
68        Stage::Import(_) => {}
69    }
70}
71
72fn walk_expr<'a>(e: &'a CExpr, parent: &NodeId, out: &mut Vec<(NodeId, NodeRef<'a>)>) {
73    let mut idx = 0;
74    let emit_expr = |child: &'a CExpr, idx: &mut usize, out: &mut Vec<(NodeId, NodeRef<'a>)>| {
75        let id = parent.child(*idx);
76        out.push((id.clone(), NodeRef::CExpr(child)));
77        walk_expr(child, &id, out);
78        *idx += 1;
79    };
80    let emit_pat = |p: &'a Pattern, idx: &mut usize, out: &mut Vec<(NodeId, NodeRef<'a>)>| {
81        let id = parent.child(*idx);
82        out.push((id.clone(), NodeRef::Pattern(p)));
83        walk_pat(p, &id, out);
84        *idx += 1;
85    };
86    match e {
87        CExpr::Literal { .. } | CExpr::Var { .. } => {}
88        CExpr::Call { callee, args } => {
89            emit_expr(callee, &mut idx, out);
90            for a in args { emit_expr(a, &mut idx, out); }
91        }
92        CExpr::Let { value, body, .. } => {
93            emit_expr(value, &mut idx, out);
94            emit_expr(body, &mut idx, out);
95        }
96        CExpr::Match { scrutinee, arms } => {
97            emit_expr(scrutinee, &mut idx, out);
98            for arm in arms {
99                emit_pat(&arm.pattern, &mut idx, out);
100                emit_expr(&arm.body, &mut idx, out);
101            }
102        }
103        CExpr::Block { statements, result } => {
104            for s in statements { emit_expr(s, &mut idx, out); }
105            emit_expr(result, &mut idx, out);
106        }
107        CExpr::Constructor { args, .. } => {
108            for a in args { emit_expr(a, &mut idx, out); }
109        }
110        CExpr::RecordLit { fields } => {
111            for f in fields { emit_expr(&f.value, &mut idx, out); }
112        }
113        CExpr::TupleLit { items } | CExpr::ListLit { items } => {
114            for it in items { emit_expr(it, &mut idx, out); }
115        }
116        CExpr::FieldAccess { value, .. } => {
117            emit_expr(value, &mut idx, out);
118        }
119        CExpr::Lambda { body, .. } => {
120            emit_expr(body, &mut idx, out);
121        }
122        CExpr::BinOp { lhs, rhs, .. } => {
123            emit_expr(lhs, &mut idx, out);
124            emit_expr(rhs, &mut idx, out);
125        }
126        CExpr::UnaryOp { expr, .. } => {
127            emit_expr(expr, &mut idx, out);
128        }
129        CExpr::Return { value } => {
130            emit_expr(value, &mut idx, out);
131        }
132    }
133}
134
135fn walk_pat<'a>(p: &'a Pattern, parent: &NodeId, out: &mut Vec<(NodeId, NodeRef<'a>)>) {
136    let mut idx = 0;
137    match p {
138        Pattern::PLiteral { .. } | Pattern::PVar { .. } | Pattern::PWild => {}
139        Pattern::PConstructor { args, .. } => {
140            for a in args {
141                let id = parent.child(idx);
142                out.push((id.clone(), NodeRef::Pattern(a)));
143                walk_pat(a, &id, out);
144                idx += 1;
145            }
146        }
147        Pattern::PRecord { fields } => {
148            for f in fields {
149                let id = parent.child(idx);
150                out.push((id.clone(), NodeRef::Pattern(&f.pattern)));
151                walk_pat(&f.pattern, &id, out);
152                idx += 1;
153            }
154        }
155        Pattern::PTuple { items } => {
156            for it in items {
157                let id = parent.child(idx);
158                out.push((id.clone(), NodeRef::Pattern(it)));
159                walk_pat(it, &id, out);
160                idx += 1;
161            }
162        }
163    }
164}
165
166fn walk_type<'a>(t: &'a TypeExpr, parent: &NodeId, out: &mut Vec<(NodeId, NodeRef<'a>)>) {
167    let mut idx = 0;
168    match t {
169        TypeExpr::Named { args, .. } => {
170            for a in args {
171                let id = parent.child(idx);
172                out.push((id.clone(), NodeRef::TypeExpr(a)));
173                walk_type(a, &id, out);
174                idx += 1;
175            }
176        }
177        TypeExpr::Record { fields } => {
178            for f in fields {
179                let id = parent.child(idx);
180                out.push((id.clone(), NodeRef::TypeExpr(&f.ty)));
181                walk_type(&f.ty, &id, out);
182                idx += 1;
183            }
184        }
185        TypeExpr::RecordWithSpreads { fields, .. } => {
186            for f in fields {
187                let id = parent.child(idx);
188                out.push((id.clone(), NodeRef::TypeExpr(&f.ty)));
189                walk_type(&f.ty, &id, out);
190                idx += 1;
191            }
192        }
193        TypeExpr::Tuple { items } => {
194            for it in items {
195                let id = parent.child(idx);
196                out.push((id.clone(), NodeRef::TypeExpr(it)));
197                walk_type(it, &id, out);
198                idx += 1;
199            }
200        }
201        TypeExpr::Function { params, ret, .. } => {
202            for p in params {
203                let id = parent.child(idx);
204                out.push((id.clone(), NodeRef::TypeExpr(p)));
205                walk_type(p, &id, out);
206                idx += 1;
207            }
208            let id = parent.child(idx);
209            out.push((id.clone(), NodeRef::TypeExpr(ret)));
210            walk_type(ret, &id, out);
211        }
212        TypeExpr::Union { variants } => {
213            for v in variants {
214                if let Some(p) = &v.payload {
215                    let id = parent.child(idx);
216                    out.push((id.clone(), NodeRef::TypeExpr(p)));
217                    walk_type(p, &id, out);
218                }
219                idx += 1;
220            }
221        }
222        TypeExpr::Refined { base, predicate, .. } => {
223            // The base type and predicate are children for NodeId
224            // attribution. The binding name is metadata, not a node.
225            let id = parent.child(idx);
226            out.push((id.clone(), NodeRef::TypeExpr(base)));
227            walk_type(base, &id, out);
228            idx += 1;
229            // Walk into the predicate so its sub-expressions get NodeIds
230            // too — same pattern as walking a function body.
231            let pid = parent.child(idx);
232            out.push((pid.clone(), NodeRef::CExpr(predicate)));
233            walk_expr(predicate, &pid, out);
234        }
235    }
236}
lex_ast/ids.rs

lex_ast/
ids.rs
