Skip to main content

compiler/
symbol.rs

1use dynamic::{ConstIntOp, Dynamic, Type};
2use parser::Stmt;
3use smol_str::SmolStr;
4use std::{collections::BTreeMap, 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;
60
61pub fn eval_const_int_type(ty: &Type) -> Option<i64> {
62    match ty {
63        Type::ConstInt(value) => Some(*value),
64        Type::ConstBinary { op, left, right } => {
65            let left = eval_const_int_type(left)?;
66            let right = eval_const_int_type(right)?;
67            match op {
68                ConstIntOp::Add => Some(left + right),
69                ConstIntOp::Sub => Some(left - right),
70                ConstIntOp::Mul => Some(left * right),
71                ConstIntOp::Div => (right != 0).then_some(left / right),
72                ConstIntOp::Mod => (right != 0).then_some(left % right),
73            }
74        }
75        _ => None,
76    }
77}
78
79pub fn substitute_type(ty: &Type, params: &[Type], args: &[Type]) -> Type {
80    match ty {
81        Type::Ident { name, params: nested } if nested.is_empty() => {
82            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())
83        }
84        Type::Ident { name, params: nested } => Type::Ident { name: name.clone(), params: nested.iter().map(|param| substitute_type(param, params, args)).collect() },
85        Type::Struct { params: struct_params, fields } => Type::Struct {
86            params: struct_params.iter().map(|param| substitute_type(param, params, args)).collect(),
87            fields: fields.iter().map(|(name, field_ty)| (name.clone(), substitute_type(field_ty, params, args))).collect(),
88        },
89        Type::Vec(elem, len) => Type::Vec(Rc::new(substitute_type(elem, params, args)), *len),
90        Type::Array(elem, len) => Type::Array(Rc::new(substitute_type(elem, params, args)), *len),
91        Type::ArrayParam(elem, len) => Type::ArrayParam(Rc::new(substitute_type(elem, params, args)), Rc::new(substitute_type(len, params, args))),
92        Type::ConstBinary { op, left, right } => {
93            let left = substitute_type(left, params, args);
94            let right = substitute_type(right, params, args);
95            let ty = Type::ConstBinary { op: *op, left: Rc::new(left), right: Rc::new(right) };
96            eval_const_int_type(&ty).map(Type::ConstInt).unwrap_or(ty)
97        }
98        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)) },
99        Type::Symbol { id, params: nested } => Type::Symbol { id: *id, params: nested.iter().map(|param| substitute_type(param, params, args)).collect() },
100        Type::Tuple(items) => Type::Tuple(items.iter().map(|item| substitute_type(item, params, args)).collect()),
101        _ => ty.clone(),
102    }
103}
104
105#[derive(Clone, Default)]
106pub struct SymbolTable {
107    pub symbols: IndexMap<SmolStr, Symbol>,
108    modules: BTreeMap<SmolStr, BTreeMap<SmolStr, u32>>,
109    pub roots: Vec<SmolStr>,
110}
111
112impl SymbolTable {
113    pub fn add_to_module(&mut self, module: &str, name: SmolStr, s: Symbol) -> Result<u32> {
114        let full_name: SmolStr = format!("{}::{}", module, name).into();
115        let id = self.symbols.insert_full(full_name, s).0 as u32;
116        let module_symbols = self.modules.get_mut(module).ok_or_else(|| anyhow!("模块 {} 不存在", module))?;
117        module_symbols.insert(name, id);
118        Ok(id)
119    }
120    pub fn get_symbol(&self, idx: u32) -> Result<(&SmolStr, &Symbol)> {
121        self.symbols.get_index(idx as usize).ok_or(anyhow!("未发现符号 {}", idx))
122    }
123
124    pub fn get_symbol_mut(&mut self, idx: u32) -> Option<(&SmolStr, &mut Symbol)> {
125        self.symbols.get_index_mut(idx as usize)
126    }
127
128    pub fn symbol(&self, name: &str) -> Vec<(SmolStr, u32)> {
129        self.modules.get(name).map(|m| m.iter().map(|(name, id)| (name.clone(), *id)).collect()).unwrap_or(Vec::new())
130    }
131
132    pub fn disassemble(&self, name: &str) -> Result<String> {
133        let id = self.get_id(name)?;
134        let (name, s) = self.get_symbol(id)?;
135        if let Symbol::Fn { ty, args, generic_params: _, cap, body, is_pub } = s {
136            if *is_pub { Ok(format!("pub {} {:?} {:?} {:?}\n{}", name, ty, args, cap, body)) } else { Ok(format!("{} {:?} {:?} {:?}\n{}", name, ty, args, cap, body)) }
137        } else {
138            Err(anyhow!("未发现符号 {}", name))
139        }
140    }
141
142    pub fn get_field(&self, ty: &Type, name: &str) -> Result<(usize, Type)> {
143        //原生类型的函数 is_map is_list 或者 sqrt
144        let id = match ty {
145            Type::Any => {
146                if let Ok(id) = self.get_id("Any")
147                    && let Ok((_, Symbol::Struct(any_ty, _))) = self.get_symbol(id)
148                    && let Ok((idx, field_ty)) = any_ty.get_field(name)
149                {
150                    return Ok((idx, field_ty.clone()));
151                }
152                return Ok((usize::MAX, Type::Any));
153            }
154            Type::Struct { params: _, fields: _ } => {
155                return ty.get_field(name).map(|(idx, ty)| (idx, ty.clone()));
156            }
157            Type::Symbol { id, params: _ } => *id,
158            Type::Vec(_, _) => self.get_id("Vec")?,
159            Type::Fn { tys: _, ret } => {
160                return self.get_field(ret, name);
161            }
162            _ => {
163                //增加一个外部函数定义
164                if name == "is_map" || name == "is_list" {
165                    return Ok((usize::MAX, Type::Bool));
166                }
167                return Err(anyhow!("未发现 symbol {:?} {}", ty, name));
168            }
169        };
170        let (_, s) = self.get_symbol(id)?;
171        if let Symbol::Struct(s, _) = s {
172            return s.get_field(name).and_then(|(idx, ty)| Ok((idx, ty.clone())));
173        };
174        Err(anyhow!("未发现 field {:?} {}", ty, name))
175    }
176
177    pub fn get_type(&self, ty: &Type) -> Result<Type> {
178        match ty {
179            Type::Ident { name, params } => {
180                let params = params.iter().map(|param| self.get_type(param)).collect::<Result<Vec<_>>>()?;
181                if name.as_str() == "Vec" && params.len() == 1 {
182                    return Ok(Type::Vec(Rc::new(params[0].clone()), 0));
183                }
184                let id = self.get_id(&name)?;
185                if let (_, Symbol::Struct(ty, _)) = self.get_symbol(id)? {
186                    if let Type::Struct { params: generic_params, .. } = ty
187                        && !generic_params.is_empty()
188                        && generic_params.len() == params.len()
189                    {
190                        return self.get_type(&substitute_type(ty, generic_params, &params));
191                    }
192                    return Ok(ty.clone());
193                }
194                return Ok(Type::Symbol { id, params });
195            }
196            Type::Symbol { id, params } => {
197                return match self.get_symbol(*id)? {
198                    (_, Symbol::Fn { ty, args: _, generic_params: _, cap: _, body: _, is_pub: _ }) => Ok(ty.clone()),
199                    (_, Symbol::Native(ty)) => Ok(ty.clone()),
200                    (_, Symbol::Struct(ty, _)) => {
201                        let params = params.iter().map(|param| self.get_type(param)).collect::<Result<Vec<_>>>()?;
202                        if let Type::Struct { params: generic_params, .. } = ty
203                            && !generic_params.is_empty()
204                            && generic_params.len() == params.len()
205                        {
206                            self.get_type(&substitute_type(ty, generic_params, &params))
207                        } else {
208                            Ok(ty.clone())
209                        }
210                    }
211                    (_, s) => {
212                        println!("s-> {:?}", s);
213                        Ok(Type::Symbol { id: *id, params: params.clone() })
214                    }
215                };
216            }
217            Type::Vec(elem, len) => {
218                return Ok(Type::Vec(Rc::new(self.get_type(elem)?), *len));
219            }
220            Type::Array(elem, len) => {
221                return Ok(Type::Array(Rc::new(self.get_type(elem)?), *len));
222            }
223            Type::ArrayParam(elem, len) => {
224                let elem = self.get_type(elem)?;
225                let len = self.get_type(len)?;
226                if let Some(len) = eval_const_int_type(&len) {
227                    let len = u32::try_from(len).map_err(|_| anyhow!("数组长度超出 u32 范围"))?;
228                    return Ok(Type::Array(Rc::new(elem), len));
229                }
230                return Ok(Type::ArrayParam(Rc::new(elem), Rc::new(len)));
231            }
232            Type::ConstBinary { op, left, right } => {
233                let left = self.get_type(left)?;
234                let right = self.get_type(right)?;
235                let ty = Type::ConstBinary { op: *op, left: Rc::new(left), right: Rc::new(right) };
236                return Ok(eval_const_int_type(&ty).map(Type::ConstInt).unwrap_or(ty));
237            }
238            Type::Fn { tys, ret } => {
239                return Ok(Type::Fn { tys: tys.iter().map(|ty| self.get_type(ty)).collect::<Result<Vec<_>>>()?, ret: Rc::new(self.get_type(ret)?) });
240            }
241            Type::Struct { params, fields } => {
242                return Ok(Type::Struct {
243                    params: params.iter().map(|param| self.get_type(param)).collect::<Result<Vec<_>>>()?,
244                    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<_>>>()?,
245                });
246            }
247            _ => {}
248        }
249        Ok(ty.clone())
250    }
251
252    pub fn add_module(&mut self, name: SmolStr) {
253        let len = self.roots.len();
254        if let Some(pos) = self.roots.iter().position(|r| r.as_str() == name.as_str()) {
255            if pos != len - 1 {
256                self.roots.swap(pos, len - 1);
257            }
258        } else {
259            self.roots.push(name.clone());
260        }
261        self.modules.insert(name, BTreeMap::new());
262    }
263
264    pub fn pop_module(&mut self) {
265        //如果不想模块成为全局的 add_module 之后调用 pop_module
266        if let Some(last) = self.roots.pop() {
267            if let Some(names) = self.modules.get(&last).map(|m| {
268                let kvs: Vec<(SmolStr, u32)> = m.iter().map(|kv| (kv.0.clone(), *kv.1)).collect();
269                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<_>>()
270            }) {
271                if let Some(m) = self.modules.get_mut(&last) {
272                    for name in names {
273                        m.remove(&name); //删除非 pub 的符号
274                    }
275                }
276            }
277        }
278    }
279
280    pub fn get_id(&self, name: &str) -> Result<u32> {
281        if let Some(id) = self.roots.iter().rev().find_map(|r| self.modules.get(r).and_then(|m| m.get(name))) {
282            return Ok(*id);
283        }
284        if let Some(id) = self.modules.values().find_map(|m| m.get(name)) {
285            return Ok(*id);
286        }
287        if let Some((mod_name, symbol_name)) = name.split_once("::") {
288            if let Some(m) = self.modules.get(mod_name) {
289                return m.get(symbol_name).copied().ok_or(anyhow!("{} 未发现", name));
290            }
291        }
292        self.roots.iter().find_map(|r| self.modules.get(r).and_then(|m| m.get(name))).copied().ok_or(anyhow!("{} 未发现", name))
293    }
294
295    pub fn add(&mut self, name: SmolStr, s: Symbol) -> u32 {
296        let root = self.roots.last().cloned().unwrap();
297        let id = self.symbols.insert_full(format!("{}::{}", root, name).into(), s).0 as u32;
298        self.modules.get_mut(&root).map(|m| m.insert(name, id));
299        id
300    }
301
302    pub fn add_global(&mut self, name: SmolStr, s: Symbol) -> u32 {
303        if let Some((mod_name, symbol_name)) = name.as_str().split_once("::") {
304            if let Some(m) = self.modules.get_mut(mod_name) {
305                if let Some(&id) = m.get(symbol_name) {
306                    return id;
307                }
308            }
309        }
310        let id = self.symbols.insert_full(name.clone(), s).0 as u32;
311        if let Some((mod_name, symbol_name)) = name.as_str().split_once("::") {
312            if let Some(m) = self.modules.get_mut(mod_name) {
313                m.insert(symbol_name.into(), id);
314            }
315        }
316        id
317    }
318
319    pub fn take(&mut self, id: u32) -> Option<Symbol> {
320        self.symbols.get_index_mut(id as usize).map(|(_, s)| std::mem::take(s))
321    }
322}