compiler/
symbol.rs

1use dynamic::{ConstIntOp, Dynamic, Type};
2use parser::Stmt;
3use smol_str::SmolStr;
4use std::{rc::Rc, sync::Arc};
5
6use super::Capture;
7
8#[derive(Debug, Clone, Default)]
9pub enum Symbol {
10    #[default]
11    Null,
12    Const {
13        value: Dynamic,
14        ty: Type,
15        is_pub: bool,
16    },
17    Static {
18        value: Option<Dynamic>,
19        ty: Type,
20        is_pub: bool,
21    },
22    Struct(Type, bool),
23    Fn {
24        ty: Type,
25        args: Vec<SmolStr>,
26        generic_params: Vec<Type>,
27        cap: Capture,
28        body: Arc<Stmt>,
29        is_pub: bool,
30    },
31    Native(Type),
32}
33
34impl Symbol {
35    pub fn native(tys: Vec<Type>, ret: Type) -> Self {
36        Self::Native(Type::Fn { tys, ret: Rc::new(ret) })
37    }
38
39    pub fn is_pub(&self) -> bool {
40        match self {
41            Self::Const { value: _, ty: _, is_pub } => *is_pub,
42            Self::Static { value: _, ty: _, is_pub } => *is_pub,
43            Self::Struct(_, is_pub) => *is_pub,
44            Self::Fn { ty: _, args: _, generic_params: _, cap: _, body: _, is_pub } => *is_pub,
45            _ => true,
46        }
47    }
48
49    pub fn is_fn(&self) -> bool {
50        match self {
51            Self::Fn { ty: _, args: _, generic_params: _, cap: _, body: _, is_pub: _ } => true,
52            Self::Native(_) => true,
53            _ => false,
54        }
55    }
56}
57
58use anyhow::{Result, anyhow};
59use indexmap::IndexMap;
60use std::{cell::RefCell, collections::HashMap};
61
62pub fn eval_const_int_type(ty: &Type) -> Option<i64> {
63    match ty {
64        Type::ConstInt(value) => Some(*value),
65        Type::ConstBinary { op, left, right } => {
66            let left = eval_const_int_type(left)?;
67            let right = eval_const_int_type(right)?;
68            // 用 checked 算术:溢出时返回 None(放弃常量折叠,回退到 ArrayParam),
69            // 而非静默回绕成错误长度(如 i64::MAX+1 -> i64::MIN 可能落在合法 u32 范围)。
70            match op {
71                ConstIntOp::Add => left.checked_add(right),
72                ConstIntOp::Sub => left.checked_sub(right),
73                ConstIntOp::Mul => left.checked_mul(right),
74                ConstIntOp::Div => (right != 0).then(|| left / right),
75                ConstIntOp::Mod => (right != 0).then(|| left % right),
76            }
77        }
78        _ => None,
79    }
80}
81
82pub fn substitute_type(ty: &Type, params: &[Type], args: &[Type]) -> Type {
83    match ty {
84        Type::Ident { name, params: nested } if nested.is_empty() => {
85            params.iter().position(|param| matches!(param, Type::Ident { name: param_name, params } if params.is_empty() && param_name == name)).map(|idx| args[idx].clone()).unwrap_or_else(|| ty.clone())
86        }
87        Type::Ident { name, params: nested } => Type::Ident { name: name.clone(), params: nested.iter().map(|param| substitute_type(param, params, args)).collect() },
88        Type::Struct { params: struct_params, fields } => Type::Struct {
89            params: struct_params.iter().map(|param| substitute_type(param, params, args)).collect(),
90            fields: fields.iter().map(|(name, field_ty)| (name.clone(), substitute_type(field_ty, params, args))).collect(),
91        },
92        Type::List(elem) => Type::List(Rc::new(substitute_type(elem, params, args))),
93        Type::Vec(elem, len) => Type::Vec(Rc::new(substitute_type(elem, params, args)), *len),
94        Type::Array(elem, len) => Type::Array(Rc::new(substitute_type(elem, params, args)), *len),
95        Type::ArrayParam(elem, len) => Type::ArrayParam(Rc::new(substitute_type(elem, params, args)), Rc::new(substitute_type(len, params, args))),
96        Type::ConstBinary { op, left, right } => {
97            let left = substitute_type(left, params, args);
98            let right = substitute_type(right, params, args);
99            let ty = Type::ConstBinary { op: *op, left: Rc::new(left), right: Rc::new(right) };
100            eval_const_int_type(&ty).map(Type::ConstInt).unwrap_or(ty)
101        }
102        Type::Fn { tys, ret } => Type::Fn { tys: tys.iter().map(|ty| substitute_type(ty, params, args)).collect(), ret: Rc::new(substitute_type(ret, params, args)) },
103        Type::Symbol { id, params: nested } => Type::Symbol { id: *id, params: nested.iter().map(|param| substitute_type(param, params, args)).collect() },
104        Type::Tuple(items) => Type::Tuple(items.iter().map(|item| substitute_type(item, params, args)).collect()),
105        _ => ty.clone(),
106    }
107}
108
109#[derive(Clone, Default)]
110pub struct SymbolTable {
111    pub symbols: IndexMap<SmolStr, Symbol>,
112    // 双 IndexMap:按模块名 + 模块内符号名 O(1) 查找;按插入顺序遍历。
113    modules: IndexMap<SmolStr, IndexMap<SmolStr, u32>>,
114    pub roots: Vec<SmolStr>,
115    lookup_cache: RefCell<HashMap<SmolStr, (u64, u32)>>,
116    lookup_epoch: u64,
117}
118
119impl SymbolTable {
120    fn invalidate_lookup_cache(&mut self) {
121        self.lookup_epoch = self.lookup_epoch.wrapping_add(1);
122        self.lookup_cache.get_mut().clear();
123    }
124
125    fn cached_id(&self, name: &str) -> Option<u32> {
126        let key = SmolStr::new(name);
127        self.lookup_cache.borrow().get(&key).and_then(|(epoch, id)| (*epoch == self.lookup_epoch).then_some(*id))
128    }
129
130    fn cache_id(&self, name: &str, id: u32) {
131        self.lookup_cache.borrow_mut().insert(SmolStr::new(name), (self.lookup_epoch, id));
132    }
133
134    pub fn add_to_module(&mut self, module: &str, name: SmolStr, s: Symbol) -> Result<u32> {
135        self.invalidate_lookup_cache();
136        let full_name: SmolStr = format!("{}::{}", module, name).into();
137        let id = self.symbols.insert_full(full_name, s).0 as u32;
138        let module_symbols = self.modules.get_mut(module).ok_or_else(|| anyhow!("模块 {} 不存在", module))?;
139        module_symbols.insert(name, id);
140        Ok(id)
141    }
142    pub fn get_symbol(&self, idx: u32) -> Result<(&SmolStr, &Symbol)> {
143        self.symbols.get_index(idx as usize).ok_or(anyhow!("未发现符号 {}", idx))
144    }
145
146    pub fn get_symbol_mut(&mut self, idx: u32) -> Option<(&SmolStr, &mut Symbol)> {
147        self.symbols.get_index_mut(idx as usize)
148    }
149
150    pub fn symbol(&self, name: &str) -> Vec<(SmolStr, u32)> {
151        self.modules.get(name).map(|m| m.iter().map(|(name, id)| (name.clone(), *id)).collect()).unwrap_or(Vec::new())
152    }
153
154    pub fn disassemble(&self, name: &str) -> Result<String> {
155        let id = self.get_id(name)?;
156        let (name, s) = self.get_symbol(id)?;
157        if let Symbol::Fn { ty, args, generic_params: _, cap, body, is_pub } = s {
158            if *is_pub { Ok(format!("pub {} {:?} {:?} {:?}\n{}", name, ty, args, cap, body)) } else { Ok(format!("{} {:?} {:?} {:?}\n{}", name, ty, args, cap, body)) }
159        } else {
160            Err(anyhow!("未发现符号 {}", name))
161        }
162    }
163
164    pub fn get_field(&self, ty: &Type, name: &str) -> Result<(usize, Type)> {
165        //原生类型的函数 is_map is_list 或者 sqrt
166        let id = match ty {
167            Type::Any => {
168                if let Ok(id) = self.get_id("Any")
169                    && let Ok((_, Symbol::Struct(any_ty, _))) = self.get_symbol(id)
170                    && let Ok((idx, field_ty)) = any_ty.get_field(name)
171                {
172                    return Ok((idx, field_ty.clone()));
173                }
174                match name {
175                    "is_map" | "is_list" | "is_string" | "is_null" | "contains" | "starts_with" | "ends_with" => return Ok((usize::MAX, Type::Bool)),
176                    "len" => return Ok((usize::MAX, Type::I32)),
177                    _ => return Ok((usize::MAX, Type::Any)),
178                }
179            }
180            Type::Struct { params: _, fields: _ } => {
181                return ty.get_field(name).map(|(idx, ty)| (idx, ty.clone()));
182            }
183            Type::Str => {
184                let any_method = match name {
185                    "len" | "contains" | "split" | "starts_with" | "ends_with" | "is_string" | "is_null" => format!("Any::{}", name),
186                    _ => return Err(anyhow!("未发现 symbol {:?} {}", ty, name)),
187                };
188                return Ok((usize::MAX, Type::Symbol { id: self.get_id(&any_method)?, params: Vec::new() }));
189            }
190            Type::List(_) | Type::Array(_, _) => {
191                let any_method = match name {
192                    "len" | "push" | "pop" | "get_idx" | "set_idx" | "slice" | "contains" | "is_list" | "is_null" => format!("Any::{}", name),
193                    _ => return Err(anyhow!("未发现 symbol {:?} {}", ty, name)),
194                };
195                return Ok((usize::MAX, Type::Symbol { id: self.get_id(&any_method)?, params: Vec::new() }));
196            }
197            Type::Symbol { id, params: _ } => *id,
198            Type::Vec(_, _) => self.get_id("Vec")?,
199            Type::Fn { tys: _, ret } => {
200                return self.get_field(ret, name);
201            }
202            _ => {
203                //增加一个外部函数定义
204                if matches!(name, "is_map" | "is_list" | "is_string" | "is_null") {
205                    return Ok((usize::MAX, Type::Symbol { id: self.get_id(&format!("Any::{}", name))?, params: Vec::new() }));
206                }
207                return Err(anyhow!("未发现 symbol {:?} {}", ty, name));
208            }
209        };
210        let (_, s) = self.get_symbol(id)?;
211        if let Symbol::Struct(s, _) = s {
212            return s.get_field(name).and_then(|(idx, ty)| Ok((idx, ty.clone())));
213        };
214        Err(anyhow!("未发现 field {:?} {}", ty, name))
215    }
216
217    pub fn get_type(&self, ty: &Type) -> Result<Type> {
218        match ty {
219            Type::Ident { name, params } => {
220                let params = params.iter().map(|param| self.get_type(param)).collect::<Result<Vec<_>>>()?;
221                if name.as_str() == "Vec" && params.len() == 1 {
222                    return Ok(Type::Vec(Rc::new(params[0].clone()), 0));
223                }
224                if name.as_str() == "List" {
225                    return Ok(if params.is_empty() { Type::list_any() } else { Type::List(Rc::new(params[0].clone())) });
226                }
227                let id = self.get_id(&name)?;
228                if let (_, Symbol::Struct(ty, _)) = self.get_symbol(id)? {
229                    if let Type::Struct { params: generic_params, .. } = ty
230                        && !generic_params.is_empty()
231                        && generic_params.len() == params.len()
232                    {
233                        return self.get_type(&substitute_type(ty, generic_params, &params));
234                    }
235                    return self.get_type(ty);
236                }
237                return Ok(Type::Symbol { id, params });
238            }
239            Type::Symbol { id, params } => {
240                return match self.get_symbol(*id)? {
241                    (_, Symbol::Fn { ty, args: _, generic_params: _, cap: _, body: _, is_pub: _ }) => self.get_type(ty),
242                    (_, Symbol::Native(ty)) => self.get_type(ty),
243                    (_, Symbol::Struct(ty, _)) => {
244                        let params = params.iter().map(|param| self.get_type(param)).collect::<Result<Vec<_>>>()?;
245                        if let Type::Struct { params: generic_params, .. } = ty
246                            && !generic_params.is_empty()
247                            && generic_params.len() == params.len()
248                        {
249                            self.get_type(&substitute_type(ty, generic_params, &params))
250                        } else {
251                            self.get_type(ty)
252                        }
253                    }
254                    (_, s) => {
255                        log::debug!("s-> {:?}", s);
256                        Ok(Type::Symbol { id: *id, params: params.clone() })
257                    }
258                };
259            }
260            Type::Vec(elem, len) => {
261                return Ok(Type::Vec(Rc::new(self.get_type(elem)?), *len));
262            }
263            Type::List(elem) => {
264                return Ok(Type::List(Rc::new(self.get_type(elem)?)));
265            }
266            Type::Array(elem, len) => {
267                return Ok(Type::Array(Rc::new(self.get_type(elem)?), *len));
268            }
269            Type::ArrayParam(elem, len) => {
270                let elem = self.get_type(elem)?;
271                let len = self.get_type(len)?;
272                if let Some(len) = eval_const_int_type(&len) {
273                    let len = u32::try_from(len).map_err(|_| anyhow!("数组长度超出 u32 范围"))?;
274                    return Ok(Type::Array(Rc::new(elem), len));
275                }
276                return Ok(Type::ArrayParam(Rc::new(elem), Rc::new(len)));
277            }
278            Type::ConstBinary { op, left, right } => {
279                let left = self.get_type(left)?;
280                let right = self.get_type(right)?;
281                let ty = Type::ConstBinary { op: *op, left: Rc::new(left), right: Rc::new(right) };
282                return Ok(eval_const_int_type(&ty).map(Type::ConstInt).unwrap_or(ty));
283            }
284            Type::Fn { tys, ret } => {
285                return Ok(Type::Fn { tys: tys.iter().map(|ty| self.get_type(ty)).collect::<Result<Vec<_>>>()?, ret: Rc::new(self.get_type(ret)?) });
286            }
287            Type::Struct { params, fields } => {
288                return Ok(Type::Struct {
289                    params: params.iter().map(|param| self.get_type(param)).collect::<Result<Vec<_>>>()?,
290                    fields: fields.iter().map(|(name, ty)| if matches!(ty, Type::Symbol { .. }) { Ok((name.clone(), ty.clone())) } else { self.get_type(ty).map(|ty| (name.clone(), ty)) }).collect::<Result<Vec<_>>>()?,
291                });
292            }
293            _ => {}
294        }
295        Ok(ty.clone())
296    }
297
298    pub fn add_module(&mut self, name: SmolStr) {
299        self.invalidate_lookup_cache();
300        let len = self.roots.len();
301        if let Some(pos) = self.roots.iter().position(|r| r.as_str() == name.as_str()) {
302            if pos != len - 1 {
303                self.roots.swap(pos, len - 1);
304            }
305        } else {
306            self.roots.push(name.clone());
307        }
308        self.modules.insert(name, IndexMap::new());
309    }
310
311    pub fn push_module_scope(&mut self, name: SmolStr) {
312        self.invalidate_lookup_cache();
313        self.roots.push(name);
314    }
315
316    pub fn pop_module_scope(&mut self) {
317        self.invalidate_lookup_cache();
318        self.roots.pop();
319    }
320
321    pub fn pop_module(&mut self) {
322        self.invalidate_lookup_cache();
323        //如果不想模块成为全局的 add_module 之后调用 pop_module
324        if let Some(last) = self.roots.pop() {
325            if let Some(names) = self.modules.get(&last).map(|m| {
326                let kvs: Vec<(SmolStr, u32)> = m.iter().map(|kv| (kv.0.clone(), *kv.1)).collect();
327                kvs.iter().filter_map(|kv| if !self.get_symbol(kv.1).map(|s| s.1.is_pub()).unwrap_or(false) { Some(kv.0.clone()) } else { None }).collect::<Vec<_>>()
328            }) {
329                if let Some(m) = self.modules.get_mut(&last) {
330                    for name in names {
331                        m.shift_remove(&name); //删除非 pub 的符号;保持插入顺序
332                    }
333                }
334            }
335        }
336    }
337
338    pub fn get_id(&self, name: &str) -> Result<u32> {
339        if let Some(id) = self.cached_id(name) {
340            return Ok(id);
341        }
342        // 1) 全名命中 (含 :: 的 id 或 add_global 注册过的全局符号)。
343        if let Some(idx) = self.symbols.get_index_of(name) {
344            let id = idx as u32;
345            self.cache_id(name, id);
346            return Ok(id);
347        }
348        // 2) 含 :: 的名字拆分按模块查找,O(1)。
349        if let Some((mod_name, sym_name)) = name.split_once("::") {
350            if let Some(&id) = self.modules.get(mod_name).and_then(|m| m.get(sym_name)) {
351                self.cache_id(name, id);
352                return Ok(id);
353            }
354            // 即使 modules 里被 pop_module 移除,完整名仍可能在 symbols IndexMap 里。
355            let full = format!("{mod_name}::{sym_name}");
356            if let Some(idx) = self.symbols.get_index_of(full.as_str()) {
357                let id = idx as u32;
358                self.cache_id(name, id);
359                return Ok(id);
360            }
361        }
362        // 3) 不含 :: 的短名,按 roots 倒序查找第一个匹配的模块。O(M)。
363        //    pop_module 会清掉非 pub 项,所以 modules 可能没有;再退回按 root 拼全名查 symbols。
364        let short_name = name;
365        for root in self.roots.iter().rev() {
366            if let Some(&id) = self.modules.get(root).and_then(|m| m.get(short_name)) {
367                self.cache_id(name, id);
368                return Ok(id);
369            }
370            let full = format!("{root}::{short_name}");
371            if let Some(idx) = self.symbols.get_index_of(full.as_str()) {
372                let id = idx as u32;
373                self.cache_id(name, id);
374                return Ok(id);
375            }
376        }
377        // 4) 兜底:任意模块下的同名短符号。O(M*K),K 为模块内符号数。
378        for (_, m) in &self.modules {
379            if let Some(&id) = m.get(short_name) {
380                self.cache_id(name, id);
381                return Ok(id);
382            }
383        }
384        Err(anyhow!("{} 未发现", name))
385    }
386
387    pub fn add(&mut self, name: SmolStr, s: Symbol) -> u32 {
388        self.invalidate_lookup_cache();
389        let root = self.roots.last().cloned().unwrap();
390        let id = self.symbols.insert_full(format!("{}::{}", root, name).into(), s).0 as u32;
391        self.modules.get_mut(&root).map(|m| m.insert(name, id));
392        id
393    }
394
395    pub fn add_global(&mut self, name: SmolStr, s: Symbol) -> u32 {
396        if let Some(idx) = self.symbols.get_index_of(name.as_str()) {
397            return idx as u32;
398        }
399        if let Some((mod_name, symbol_name)) = name.as_str().split_once("::") {
400            if let Some(m) = self.modules.get_mut(mod_name) {
401                if let Some(&id) = m.get(symbol_name) {
402                    return id;
403                }
404            }
405        }
406        self.invalidate_lookup_cache();
407        let id = self.symbols.insert_full(name.clone(), s).0 as u32;
408        if let Some((mod_name, symbol_name)) = name.as_str().split_once("::") {
409            if let Some(m) = self.modules.get_mut(mod_name) {
410                m.insert(symbol_name.into(), id);
411            }
412        }
413        id
414    }
415
416    pub fn take(&mut self, id: u32) -> Option<Symbol> {
417        self.invalidate_lookup_cache();
418        self.symbols.get_index_mut(id as usize).map(|(_, s)| std::mem::take(s))
419    }
420}
compiler/symbol.rs

compiler/
symbol.rs
