Skip to main content

vyre_debug/
source_walker.rs

1use vyre_foundation::ir::{Node, Program};
2
3#[derive(Clone)]
4struct Scope {
5    is_loop_boundary: bool,
6    defs: std::collections::HashSet<String>,
7}
8
9pub fn walk_source_assigns<F>(program: &Program, mut callback: F)
10where
11    F: FnMut(&str, Vec<String>),
12{
13    let mut env = vec![Scope {
14        is_loop_boundary: false,
15        defs: std::collections::HashSet::new(),
16    }];
17    let mut current_loop_path = Vec::new();
18
19    fn walk<F>(
20        nodes: &[Node],
21        env: &mut Vec<Scope>,
22        current_loop_path: &mut Vec<String>,
23        callback: &mut F,
24    ) where
25        F: FnMut(&str, Vec<String>),
26    {
27        for node in nodes {
28            match node {
29                Node::Let { name, .. } => {
30                    if let Some(scope) = env.last_mut() {
31                        scope.defs.insert(name.as_ref().to_string());
32                    }
33                }
34                Node::Assign { name, .. } => {
35                    let name_str = name.as_ref().to_string();
36                    if current_loop_path.is_empty() {
37                        if let Some(scope) = env.last_mut() {
38                            scope.defs.insert(name_str);
39                        }
40                        continue;
41                    }
42
43                    let mut defined_outside_loop = false;
44                    let mut loop_boundaries_crossed = 0;
45
46                    for scope in env.iter().rev() {
47                        if scope.is_loop_boundary {
48                            loop_boundaries_crossed += 1;
49                        }
50
51                        if scope.defs.contains(&name_str) {
52                            if loop_boundaries_crossed > 0 {
53                                defined_outside_loop = true;
54                            }
55                            break;
56                        }
57                    }
58
59                    if defined_outside_loop {
60                        callback(&name_str, current_loop_path.clone());
61                    }
62
63                    if let Some(scope) = env.last_mut() {
64                        scope.defs.insert(name_str);
65                    }
66                }
67                Node::Block(stmts) => {
68                    env.push(Scope {
69                        is_loop_boundary: false,
70                        defs: std::collections::HashSet::new(),
71                    });
72                    walk(stmts, env, current_loop_path, callback);
73                    env.pop();
74                }
75                Node::If {
76                    then, otherwise, ..
77                } => {
78                    env.push(Scope {
79                        is_loop_boundary: false,
80                        defs: std::collections::HashSet::new(),
81                    });
82                    walk(then, env, current_loop_path, callback);
83                    env.pop();
84
85                    env.push(Scope {
86                        is_loop_boundary: false,
87                        defs: std::collections::HashSet::new(),
88                    });
89                    walk(otherwise, env, current_loop_path, callback);
90                    env.pop();
91                }
92                Node::Loop { var, body, .. } => {
93                    let loop_name_str = var.as_ref().to_string();
94                    current_loop_path.push(loop_name_str.clone());
95
96                    env.push(Scope {
97                        is_loop_boundary: true,
98                        defs: std::collections::HashSet::new(),
99                    });
100                    if let Some(scope) = env.last_mut() {
101                        scope.defs.insert(loop_name_str);
102                    }
103
104                    walk(body, env, current_loop_path, callback);
105
106                    env.pop();
107                    current_loop_path.pop();
108                }
109                Node::Region { body, .. } => {
110                    env.push(Scope {
111                        is_loop_boundary: false,
112                        defs: std::collections::HashSet::new(),
113                    });
114                    walk(body, env, current_loop_path, callback);
115                    env.pop();
116                }
117                _ => {}
118            }
119        }
120    }
121
122    walk(
123        &program.entry,
124        &mut env,
125        &mut current_loop_path,
126        &mut callback,
127    );
128}