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 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 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, ¶ms));
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, ¶ms))
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 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); }
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}