1use crate::ast::visitor::{Visitor, walk_expr, walk_param, walk_program, walk_stmt};
2use crate::ast::*;
3use crate::span::Span;
4use std::collections::HashMap;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum SymbolKind {
8 Variable,
9 Function,
10 Class,
11 Interface,
12 Trait,
13 Enum,
14 EnumCase,
15 Parameter,
16}
17
18#[derive(Debug, Clone)]
19pub struct Symbol {
20 pub name: String,
21 pub kind: SymbolKind,
22 pub span: Span,
23}
24
25#[derive(Debug, Default)]
26pub struct Scope {
27 pub symbols: HashMap<String, Symbol>,
28 pub parent: Option<usize>,
29 pub children: Vec<usize>,
30}
31
32impl Scope {
33 pub fn new(parent: Option<usize>) -> Self {
34 Self {
35 symbols: HashMap::new(),
36 parent,
37 children: Vec::new(),
38 }
39 }
40
41 pub fn add(&mut self, name: String, kind: SymbolKind, span: Span) {
42 self.symbols
43 .insert(name.clone(), Symbol { name, kind, span });
44 }
45
46 pub fn get(&self, name: &str) -> Option<&Symbol> {
47 self.symbols.get(name)
48 }
49}
50
51#[derive(Debug)]
52pub struct SymbolTable {
53 pub scopes: Vec<Scope>,
54 pub current_scope_idx: usize,
55}
56
57impl Default for SymbolTable {
58 fn default() -> Self {
59 Self {
60 scopes: vec![Scope::new(None)], current_scope_idx: 0,
62 }
63 }
64}
65
66impl SymbolTable {
67 pub fn new() -> Self {
68 Self::default()
69 }
70
71 pub fn enter_scope(&mut self) {
72 let new_scope_idx = self.scopes.len();
73 let new_scope = Scope::new(Some(self.current_scope_idx));
74 self.scopes.push(new_scope);
75
76 self.scopes[self.current_scope_idx]
78 .children
79 .push(new_scope_idx);
80
81 self.current_scope_idx = new_scope_idx;
82 }
83
84 pub fn exit_scope(&mut self) {
85 if let Some(parent) = self.scopes[self.current_scope_idx].parent {
86 self.current_scope_idx = parent;
87 } else {
88 eprintln!("Warning: Attempted to exit root scope");
90 }
91 }
92
93 pub fn add_symbol(&mut self, name: String, kind: SymbolKind, span: Span) {
94 self.scopes[self.current_scope_idx].add(name, kind, span);
95 }
96
97 pub fn lookup(&self, name: &str) -> Option<&Symbol> {
98 let mut current = Some(self.current_scope_idx);
99 while let Some(idx) = current {
100 if let Some(sym) = self.scopes[idx].get(name) {
101 return Some(sym);
102 }
103 current = self.scopes[idx].parent;
104 }
105 None
106 }
107}
108
109pub struct SymbolVisitor<'src> {
110 pub table: SymbolTable,
111 pub source: &'src [u8],
112}
113
114impl<'src> SymbolVisitor<'src> {
115 pub fn new(source: &'src [u8]) -> Self {
116 Self {
117 table: SymbolTable::new(),
118 source,
119 }
120 }
121
122 fn get_text(&self, span: Span) -> String {
123 String::from_utf8_lossy(span.as_str(self.source)).to_string()
124 }
125}
126
127impl<'ast, 'src> Visitor<'ast> for SymbolVisitor<'src> {
128 fn visit_program(&mut self, program: &'ast Program<'ast>) {
129 walk_program(self, program);
131 }
132
133 fn visit_stmt(&mut self, stmt: StmtId<'ast>) {
134 match stmt {
135 Stmt::Function {
136 name,
137 params,
138 body,
139 span,
140 ..
141 } => {
142 let func_name = self.get_text(name.span);
143 self.table
144 .add_symbol(func_name, SymbolKind::Function, *span);
145
146 self.table.enter_scope();
147 for param in *params {
148 self.visit_param(param);
149 }
150 for s in *body {
151 self.visit_stmt(s);
152 }
153 self.table.exit_scope();
154 }
155 Stmt::Class {
156 name,
157 members,
158 span,
159 ..
160 } => {
161 let class_name = self.get_text(name.span);
162 self.table.add_symbol(class_name, SymbolKind::Class, *span);
163 self.table.enter_scope();
164 for member in *members {
165 self.visit_class_member(member);
166 }
167 self.table.exit_scope();
168 }
169 Stmt::Interface {
170 name,
171 members,
172 span,
173 ..
174 } => {
175 let interface_name = self.get_text(name.span);
176 self.table
177 .add_symbol(interface_name, SymbolKind::Interface, *span);
178 self.table.enter_scope();
179 for member in *members {
180 self.visit_class_member(member);
181 }
182 self.table.exit_scope();
183 }
184 Stmt::Trait {
185 name,
186 members,
187 span,
188 ..
189 } => {
190 let trait_name = self.get_text(name.span);
191 self.table.add_symbol(trait_name, SymbolKind::Trait, *span);
192 self.table.enter_scope();
193 for member in *members {
194 self.visit_class_member(member);
195 }
196 self.table.exit_scope();
197 }
198 Stmt::Enum {
199 name,
200 members,
201 span,
202 ..
203 } => {
204 let enum_name = self.get_text(name.span);
205 self.table.add_symbol(enum_name, SymbolKind::Enum, *span);
206 self.table.enter_scope();
207 for member in *members {
208 self.visit_class_member(member);
209 }
210 self.table.exit_scope();
211 }
212 _ => walk_stmt(self, stmt),
213 }
214 }
215
216 fn visit_param(&mut self, param: &'ast Param<'ast>) {
217 let name = self.get_text(param.name.span);
218 self.table
219 .add_symbol(name, SymbolKind::Parameter, param.span);
220 walk_param(self, param);
221 }
222
223 fn visit_expr(&mut self, expr: ExprId<'ast>) {
224 match expr {
225 Expr::Assign { var, .. } => {
226 if let Expr::Variable { name, span } = var {
227 let var_name = self.get_text(*name);
228 if self.table.scopes[self.table.current_scope_idx]
229 .get(&var_name)
230 .is_none()
231 {
232 self.table.add_symbol(var_name, SymbolKind::Variable, *span);
233 }
234 }
235 walk_expr(self, expr);
236 }
237 _ => walk_expr(self, expr),
238 }
239 }
240}
241
242impl<'src> SymbolVisitor<'src> {
243 fn visit_class_member<'ast>(&mut self, member: &'ast ClassMember<'ast>) {
244 match member {
245 ClassMember::Method {
246 name,
247 params,
248 body,
249 span,
250 ..
251 } => {
252 let method_name = self.get_text(name.span);
253 self.table
254 .add_symbol(method_name, SymbolKind::Function, *span);
255 self.table.enter_scope();
256 for param in *params {
257 self.visit_param(param);
258 }
259 for stmt in *body {
260 self.visit_stmt(stmt);
261 }
262 self.table.exit_scope();
263 }
264 ClassMember::Property { entries, .. } => {
265 for entry in *entries {
266 let prop_name = self.get_text(entry.name.span);
267 self.table
268 .add_symbol(prop_name, SymbolKind::Variable, entry.span);
269 }
270 }
271 ClassMember::Case { name, span, .. } => {
272 let case_name = self.get_text(name.span);
273 self.table
275 .add_symbol(case_name, SymbolKind::EnumCase, *span);
276 }
277 _ => {}
278 }
279 }
280}