huff_neo_utils/
scope.rs

1//! Scope management for macro compilation.
2//!
3//! Tracks macro call stack context to ensure labels and arguments resolve correctly
4//! during bytecode generation.
5//!
6//! # Core Concepts
7//!
8//! ## ScopeId
9//! Uniquely identifies where a label is defined or where a jump should resolve to.
10//! Prevents label collisions between different macro invocations.
11//!
12//! ## ScopeManager
13//! Tracks the current macro call stack. Provides scope identifiers for label operations
14//! and access to the call stack for argument resolution.
15
16use crate::prelude::{MacroDefinition, MacroInvocation};
17use std::collections::BTreeMap;
18
19/// A single entry in the macro call stack.
20///
21/// Contains the macro definition, bytecode offset, and arguments passed at invocation.
22#[derive(Debug, Clone)]
23pub struct ScopeFrame<'a> {
24    /// The macro definition being executed
25    pub macro_def: &'a MacroDefinition,
26    /// The bytecode offset where this macro was invoked
27    pub offset: usize,
28    /// The invocation details (arguments passed), None for root macro
29    pub invocation: Option<MacroInvocation>,
30}
31
32/// Uniquely identifies a scope for label storage and resolution.
33///
34/// Combines a path of macro names with call depth to distinguish between different
35/// invocations of the same macro, preventing label name collisions.
36#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
37pub struct ScopeId {
38    /// Path of macro names with invocation offsets, e.g., ["MAIN", "M1_5", "M2_12"]
39    pub path: Vec<String>,
40    /// Depth in the macro call stack (0 for MAIN)
41    pub depth: usize,
42}
43
44impl ScopeId {
45    /// Creates a new scope identifier with the given path and depth.
46    pub fn new(path: Vec<String>, depth: usize) -> Self {
47        Self { path, depth }
48    }
49
50    /// Returns the scope identifier for the calling context, or None if already at root.
51    pub fn parent(&self) -> Option<ScopeId> {
52        if self.depth == 0 { None } else { Some(ScopeId { path: self.path[..self.path.len() - 1].to_vec(), depth: self.depth - 1 }) }
53    }
54}
55
56/// Tracks the macro call stack during compilation.
57///
58/// Provides scope identifiers for label operations and exposes the call stack
59/// for argument resolution.
60pub struct ScopeManager<'a> {
61    /// Stack of scope frames from root to current
62    pub frames: Vec<ScopeFrame<'a>>,
63    /// Macro sizes from the previous compile pass, for `__codesize` values that cannot be measured
64    /// in the code they are part of
65    pub codesize_hints: BTreeMap<String, usize>,
66    /// Macro sizes generated in this pass, the hints of the next pass
67    pub codesizes: BTreeMap<String, usize>,
68    /// Whether a hint was used, so the sizes must be compiled again until they are stable
69    pub codesize_hints_used: bool,
70    /// Frame indices where a `__codesize` measurement of another macro starts
71    measurements: Vec<usize>,
72}
73
74impl<'a> ScopeManager<'a> {
75    /// Creates a new empty scope manager.
76    pub fn new() -> Self {
77        Self {
78            frames: Vec::new(),
79            codesize_hints: BTreeMap::new(),
80            codesizes: BTreeMap::new(),
81            codesize_hints_used: false,
82            measurements: Vec::new(),
83        }
84    }
85
86    /// Starts measuring a macro for `__codesize`, in a frame pushed after this call
87    pub fn start_measurement(&mut self) {
88        self.measurements.push(self.frames.len());
89    }
90
91    /// Ends the innermost `__codesize` measurement
92    pub fn end_measurement(&mut self) {
93        self.measurements.pop();
94    }
95
96    /// The size from the previous pass for a `__codesize` of the enclosing macro `macro_name`.
97    ///
98    /// It is only needed when the reference sits in a measurement of another macro that does not
99    /// contain `macro_name`, which therefore cannot be measured there. Using it records that the
100    /// sizes must be compiled again until they are stable.
101    pub fn codesize_hint(&mut self, macro_name: &str) -> Option<usize> {
102        let enclosing = self.frames.iter().rposition(|frame| frame.macro_def.name == macro_name)?;
103        if !self.measurements.iter().any(|&start| start > enclosing) {
104            return None;
105        }
106        self.codesize_hints_used = true;
107        self.codesize_hints.get(macro_name).copied()
108    }
109
110    /// Returns the scope identifier for the currently executing macro.
111    ///
112    /// Used when defining labels or creating jumps within the current macro body.
113    pub fn current_scope(&self) -> ScopeId {
114        if self.frames.is_empty() {
115            return ScopeId::new(vec![], 0);
116        }
117        self.build_scope_id(self.frames.len() - 1)
118    }
119
120    /// Returns the scope identifier for the calling macro.
121    ///
122    /// Used when resolving label arguments that should be looked up in the caller's scope.
123    /// Returns the root scope if called from the top level.
124    pub fn parent_scope(&self) -> ScopeId {
125        if self.frames.len() <= 1 {
126            return ScopeId::new(vec![], 0);
127        }
128        self.build_scope_id(self.frames.len() - 2)
129    }
130
131    /// Enters a new macro context.
132    ///
133    /// Adds the macro to the call stack at the given bytecode offset.
134    /// Use for root macros (MAIN/CONSTRUCTOR) without invocation details.
135    pub fn push_macro(&mut self, macro_def: &'a MacroDefinition, offset: usize) {
136        self.frames.push(ScopeFrame { macro_def, offset, invocation: None });
137    }
138
139    /// Enters a new macro context with invocation arguments.
140    ///
141    /// Like `push_macro`, but includes the invocation details for argument resolution.
142    pub fn push_macro_with_invocation(&mut self, macro_def: &'a MacroDefinition, offset: usize, invocation: MacroInvocation) {
143        self.frames.push(ScopeFrame { macro_def, offset, invocation: Some(invocation) });
144    }
145
146    /// Exits the current macro context.
147    ///
148    /// Removes the top frame from the call stack and returns the macro definition.
149    pub fn pop_macro(&mut self) -> Option<&'a MacroDefinition> {
150        self.frames.pop().map(|frame| frame.macro_def)
151    }
152
153    /// Returns the macro definition at the top of the call stack, or None if empty.
154    pub fn current_macro(&self) -> Option<&'a MacroDefinition> {
155        self.frames.last().map(|frame| frame.macro_def)
156    }
157
158    /// Returns all macro definitions from root to current as a vector.
159    pub fn macro_stack(&self) -> Vec<&'a MacroDefinition> {
160        self.frames.iter().map(|frame| frame.macro_def).collect()
161    }
162
163    /// Returns the macro definition at the specified call stack depth, or None if out of bounds.
164    pub fn macro_at_depth(&self, depth: usize) -> Option<&'a MacroDefinition> {
165        self.frames.get(depth).map(|frame| frame.macro_def)
166    }
167
168    /// Returns the number of frames in the call stack.
169    pub fn depth(&self) -> usize {
170        self.frames.len()
171    }
172
173    /// Returns an iterator over (offset, invocation) pairs for all frames.
174    ///
175    /// Invocation is None for root macros.
176    pub fn invocation_stack(&self) -> impl Iterator<Item = (usize, Option<&MacroInvocation>)> + '_ {
177        self.frames.iter().map(|frame| (frame.offset, frame.invocation.as_ref()))
178    }
179
180    /// Returns (offset, invocation) pairs for all frames with invocation details.
181    ///
182    /// Root macros without invocations are excluded from the result.
183    pub fn invocation_vec(&self) -> Vec<(usize, &MacroInvocation)> {
184        self.frames.iter().filter_map(|frame| frame.invocation.as_ref().map(|inv| (frame.offset, inv))).collect()
185    }
186
187    /// Returns (offset, invocation) pairs up to the specified depth.
188    ///
189    /// Used to get the invocation context when resolving arguments for a macro at `target_depth`.
190    pub fn invocation_vec_up_to_depth(&self, target_depth: usize) -> Vec<(usize, &MacroInvocation)> {
191        self.frames.iter().take(target_depth).filter_map(|frame| frame.invocation.as_ref().map(|inv| (frame.offset, inv))).collect()
192    }
193
194    /// Returns the invocation details for the current macro, or None if unavailable.
195    pub fn current_invocation(&self) -> Option<&MacroInvocation> {
196        self.frames.last().and_then(|frame| frame.invocation.as_ref())
197    }
198
199    /// Returns the bytecode offset and invocation details for the current macro.
200    pub fn current_invocation_with_offset(&self) -> Option<(usize, Option<&MacroInvocation>)> {
201        self.frames.last().map(|frame| (frame.offset, frame.invocation.as_ref()))
202    }
203
204    /// Builds a scope ID at a specific depth in the call stack.
205    ///
206    /// # Arguments
207    /// * `target_depth` - The depth to build the scope ID for (0 = root)
208    ///
209    /// # Panics
210    /// Panics if `target_depth` is greater than the current stack depth.
211    fn build_scope_id(&self, target_depth: usize) -> ScopeId {
212        assert!(target_depth < self.frames.len(), "target_depth {} exceeds stack depth {}", target_depth, self.frames.len());
213
214        let path = if target_depth == 0 {
215            // Root scope: just the macro name, no offset
216            vec![self.frames[0].macro_def.name.clone()]
217        } else {
218            // Build path up to and including target_depth
219            let mut path = Vec::with_capacity(target_depth + 1);
220
221            // First element has no offset
222            path.push(self.frames[0].macro_def.name.clone());
223
224            // Subsequent elements include invocation offset for uniqueness
225            for i in 1..=target_depth {
226                let frame = &self.frames[i];
227                path.push(format!("{}_{}", frame.macro_def.name, frame.offset));
228            }
229
230            path
231        };
232
233        ScopeId::new(path, target_depth)
234    }
235}
236
237impl<'a> Default for ScopeManager<'a> {
238    fn default() -> Self {
239        Self::new()
240    }
241}
242
243#[cfg(test)]
244mod tests {
245    use super::*;
246    use crate::prelude::AstSpan;
247
248    fn make_test_macro(name: &str) -> MacroDefinition {
249        MacroDefinition {
250            name: name.to_string(),
251            decorator: None,
252            parameters: Vec::new(),
253            statements: Vec::new(),
254            takes: 0,
255            returns: 0,
256            outlined: false,
257            test: false,
258            span: AstSpan::default(),
259        }
260    }
261
262    #[test]
263    fn test_empty_scope_manager() {
264        let mgr = ScopeManager::new();
265        assert_eq!(mgr.depth(), 0);
266        assert_eq!(mgr.current_scope(), ScopeId::new(vec![], 0));
267        assert_eq!(mgr.parent_scope(), ScopeId::new(vec![], 0));
268    }
269
270    #[test]
271    fn test_single_macro() {
272        let main = make_test_macro("MAIN");
273        let mut mgr = ScopeManager::new();
274        mgr.push_macro(&main, 0);
275
276        assert_eq!(mgr.depth(), 1);
277        assert_eq!(mgr.current_scope(), ScopeId::new(vec!["MAIN".to_string()], 0));
278        assert_eq!(mgr.parent_scope(), ScopeId::new(vec![], 0));
279    }
280
281    #[test]
282    fn test_two_level_nesting() {
283        let main = make_test_macro("MAIN");
284        let m1 = make_test_macro("M1");
285
286        let mut mgr = ScopeManager::new();
287        mgr.push_macro(&main, 0);
288        mgr.push_macro(&m1, 5);
289
290        assert_eq!(mgr.depth(), 2);
291        assert_eq!(mgr.current_scope(), ScopeId::new(vec!["MAIN".to_string(), "M1_5".to_string()], 1));
292        assert_eq!(mgr.parent_scope(), ScopeId::new(vec!["MAIN".to_string()], 0));
293    }
294
295    #[test]
296    fn test_three_level_nesting() {
297        let main = make_test_macro("MAIN");
298        let m1 = make_test_macro("M1");
299        let m2 = make_test_macro("M2");
300
301        let mut mgr = ScopeManager::new();
302        mgr.push_macro(&main, 0);
303        mgr.push_macro(&m1, 5);
304        mgr.push_macro(&m2, 12);
305
306        assert_eq!(mgr.depth(), 3);
307        assert_eq!(mgr.current_scope(), ScopeId::new(vec!["MAIN".to_string(), "M1_5".to_string(), "M2_12".to_string()], 2));
308        assert_eq!(mgr.parent_scope(), ScopeId::new(vec!["MAIN".to_string(), "M1_5".to_string()], 1));
309    }
310
311    #[test]
312    fn test_push_pop() {
313        let main = make_test_macro("MAIN");
314        let m1 = make_test_macro("M1");
315
316        let mut mgr = ScopeManager::new();
317        mgr.push_macro(&main, 0);
318        mgr.push_macro(&m1, 5);
319
320        assert_eq!(mgr.depth(), 2);
321
322        mgr.pop_macro();
323        assert_eq!(mgr.depth(), 1);
324        assert_eq!(mgr.current_scope(), ScopeId::new(vec!["MAIN".to_string()], 0));
325
326        mgr.pop_macro();
327        assert_eq!(mgr.depth(), 0);
328        assert_eq!(mgr.current_scope(), ScopeId::new(vec![], 0));
329    }
330
331    #[test]
332    fn test_scope_id_parent() {
333        let scope = ScopeId::new(vec!["MAIN".to_string(), "M1_5".to_string(), "M2_12".to_string()], 2);
334
335        let parent = scope.parent().unwrap();
336        assert_eq!(parent.path, vec!["MAIN".to_string(), "M1_5".to_string()]);
337        assert_eq!(parent.depth, 1);
338
339        let grandparent = parent.parent().unwrap();
340        assert_eq!(grandparent.path, vec!["MAIN".to_string()]);
341        assert_eq!(grandparent.depth, 0);
342
343        assert!(grandparent.parent().is_none());
344    }
345
346    #[test]
347    fn test_parent_scope() {
348        let main = make_test_macro("MAIN");
349        let m1 = make_test_macro("M1");
350
351        let mut mgr = ScopeManager::new();
352        mgr.push_macro(&main, 0);
353        mgr.push_macro(&m1, 5);
354
355        // parent_scope should return MAIN when in M1
356        assert_eq!(mgr.parent_scope(), ScopeId::new(vec!["MAIN".to_string()], 0));
357    }
358}
huff_neo_utils/scope.rs

huff_neo_utils/
scope.rs
