1use std::collections::{HashMap, HashSet};
2
3use id_arena::Arena;
4use indexmap::IndexSet;
5use rue_types::{Type, TypeId};
6
7use crate::{Hir, HirId, Scope, ScopeId, Symbol, SymbolId};
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub enum Declaration {
11 Symbol(SymbolId),
12 Type(TypeId),
13}
14
15#[derive(Debug, Default, Clone)]
16pub struct Database {
17 hir: Arena<Hir>,
18 scopes: Arena<Scope>,
19 symbols: Arena<Symbol>,
20 types: Arena<Type>,
21 relevant_declarations: IndexSet<Declaration>,
22 declared_by: HashMap<Declaration, HashSet<Declaration>>,
23 referenced_by: HashMap<Declaration, HashSet<Declaration>>,
24 tests: Vec<SymbolId>,
25}
26
27impl Database {
28 pub fn new() -> Self {
29 Self::default()
30 }
31
32 pub fn alloc_hir(&mut self, hir: Hir) -> HirId {
33 self.hir.alloc(hir)
34 }
35
36 pub fn hir(&self, id: HirId) -> &Hir {
37 &self.hir[id]
38 }
39
40 pub fn hir_mut(&mut self, id: HirId) -> &mut Hir {
41 &mut self.hir[id]
42 }
43
44 pub fn alloc_scope(&mut self, scope: Scope) -> ScopeId {
45 self.scopes.alloc(scope)
46 }
47
48 pub fn scope(&self, id: ScopeId) -> &Scope {
49 &self.scopes[id]
50 }
51
52 pub fn scope_mut(&mut self, id: ScopeId) -> &mut Scope {
53 &mut self.scopes[id]
54 }
55
56 pub fn alloc_symbol(&mut self, symbol: Symbol) -> SymbolId {
57 self.symbols.alloc(symbol)
58 }
59
60 pub fn symbol(&self, id: SymbolId) -> &Symbol {
61 &self.symbols[id]
62 }
63
64 pub fn symbol_mut(&mut self, id: SymbolId) -> &mut Symbol {
65 &mut self.symbols[id]
66 }
67
68 pub fn alloc_type(&mut self, ty: Type) -> TypeId {
69 self.types.alloc(ty)
70 }
71
72 pub fn ty(&self, id: TypeId) -> &Type {
73 &self.types[id]
74 }
75
76 pub fn ty_mut(&mut self, id: TypeId) -> &mut Type {
77 &mut self.types[id]
78 }
79
80 pub fn types(&self) -> &Arena<Type> {
81 &self.types
82 }
83
84 pub fn types_mut(&mut self) -> &mut Arena<Type> {
85 &mut self.types
86 }
87
88 pub fn add_relevant_declaration(&mut self, declaration: Declaration) {
89 self.relevant_declarations.insert(declaration);
90 }
91
92 pub fn relevant_declarations(&self) -> impl Iterator<Item = Declaration> {
93 self.relevant_declarations.iter().copied()
94 }
95
96 pub fn add_reference(&mut self, outer: Declaration, inner: Declaration) {
97 self.referenced_by.entry(inner).or_default().insert(outer);
98 }
99
100 pub fn add_declaration(&mut self, outer: Declaration, inner: Declaration) {
101 self.declared_by.entry(inner).or_default().insert(outer);
102 }
103
104 pub fn reference_parents(&self, declaration: Declaration) -> Vec<Declaration> {
105 self.referenced_by
106 .get(&declaration)
107 .cloned()
108 .unwrap_or_default()
109 .into_iter()
110 .collect()
111 }
112
113 pub fn declaration_parents(&self, declaration: Declaration) -> Vec<Declaration> {
114 self.declared_by
115 .get(&declaration)
116 .cloned()
117 .unwrap_or_default()
118 .into_iter()
119 .collect()
120 }
121
122 pub fn add_test(&mut self, symbol: SymbolId) {
123 self.tests.push(symbol);
124 }
125
126 pub fn tests(&self) -> impl Iterator<Item = SymbolId> {
127 self.tests.iter().copied()
128 }
129
130 pub fn debug_symbol(&self, id: SymbolId) -> String {
131 let name = match self.symbol(id) {
132 Symbol::Unresolved => None,
133 Symbol::Module(module) => module.name.as_ref().map(|name| name.text().to_string()),
134 Symbol::Function(function) => {
135 function.name.as_ref().map(|name| name.text().to_string())
136 }
137 Symbol::Builtin(builtin) => Some(format!("{builtin:?}")),
138 Symbol::Parameter(parameter) => {
139 parameter.name.as_ref().map(|name| name.text().to_string())
140 }
141 Symbol::Constant(constant) => {
142 constant.name.as_ref().map(|name| name.text().to_string())
143 }
144 Symbol::Binding(binding) => binding.name.as_ref().map(|name| name.text().to_string()),
145 };
146
147 if let Some(name) = name {
148 format!("{}<{}>", name.replace('"', ""), id.index())
149 } else {
150 format!("<{}>", id.index())
151 }
152 }
153
154 pub fn debug_hir(&self, id: HirId) -> String {
155 match self.hir(id) {
156 Hir::Unresolved => "{unknown}".to_string(),
157 Hir::Nil => "nil".to_string(),
158 Hir::String(value) => format!("\"{value}\""),
159 Hir::Int(value) => format!("{value}"),
160 Hir::Bytes(value) => {
161 if value.is_empty() {
162 "nil".to_string()
163 } else {
164 format!("0x{}", hex::encode(value))
165 }
166 }
167 Hir::Bool(value) => format!("{value}"),
168 Hir::Pair(first, rest) => {
169 format!("({}, {})", self.debug_hir(*first), self.debug_hir(*rest))
170 }
171 Hir::Reference(symbol) => self.debug_symbol(*symbol),
172 Hir::Block(block) => block
173 .body
174 .map_or("{empty}".to_string(), |body| self.debug_hir(body)),
175 Hir::Lambda(lambda) => self.debug_symbol(*lambda),
176 Hir::If(condition, then, else_, inline) => format!(
177 "{}if {} {{ {} }} else {{ {} }}",
178 if *inline { "inline " } else { "" },
179 self.debug_hir(*condition),
180 self.debug_hir(*then),
181 self.debug_hir(*else_)
182 ),
183 Hir::FunctionCall(call) => format!(
184 "{}({})",
185 self.debug_hir(call.function),
186 call.args
187 .iter()
188 .enumerate()
189 .map(|(i, arg)| {
190 let arg = self.debug_hir(*arg);
191
192 if i == call.args.len() - 1 && !call.nil_terminated {
193 format!("...{arg}")
194 } else {
195 arg
196 }
197 })
198 .collect::<Vec<_>>()
199 .join(", ")
200 ),
201 Hir::Unary(op, hir) => format!("({op} {})", self.debug_hir(*hir)),
202 Hir::Binary(op, left, right) => {
203 format!(
204 "({} {} {})",
205 self.debug_hir(*left),
206 op,
207 self.debug_hir(*right)
208 )
209 }
210 Hir::CoinId(parent, puzzle, amount) => {
211 format!(
212 "coinid({}, {}, {})",
213 self.debug_hir(*parent),
214 self.debug_hir(*puzzle),
215 self.debug_hir(*amount)
216 )
217 }
218 Hir::Substr(hir, start, end) => {
219 if let Some(end) = end {
220 format!(
221 "substr({}, {}, {})",
222 self.debug_hir(*hir),
223 self.debug_hir(*start),
224 self.debug_hir(*end)
225 )
226 } else {
227 format!(
228 "substr({}, {})",
229 self.debug_hir(*hir),
230 self.debug_hir(*start)
231 )
232 }
233 }
234 Hir::G1Map(data, dst) => {
235 if let Some(dst) = dst {
236 format!(
237 "g1_map({}, {})",
238 self.debug_hir(*data),
239 self.debug_hir(*dst)
240 )
241 } else {
242 format!("g1_map({})", self.debug_hir(*data))
243 }
244 }
245 Hir::G2Map(data, dst) => {
246 if let Some(dst) = dst {
247 format!(
248 "g2_map({}, {})",
249 self.debug_hir(*data),
250 self.debug_hir(*dst)
251 )
252 } else {
253 format!("g2_map({})", self.debug_hir(*data))
254 }
255 }
256 Hir::Modpow(base, exponent, modulus) => {
257 format!(
258 "modpow({}, {}, {})",
259 self.debug_hir(*base),
260 self.debug_hir(*exponent),
261 self.debug_hir(*modulus)
262 )
263 }
264 Hir::BlsPairingIdentity(args) => {
265 format!(
266 "bls_pairing_identity({})",
267 args.iter()
268 .map(|arg| self.debug_hir(*arg))
269 .collect::<Vec<_>>()
270 .join(", ")
271 )
272 }
273 Hir::BlsVerify(sig, args) => {
274 format!(
275 "bls_verify({}, {})",
276 self.debug_hir(*sig),
277 args.iter()
278 .map(|arg| self.debug_hir(*arg))
279 .collect::<Vec<_>>()
280 .join(", ")
281 )
282 }
283 Hir::Secp256K1Verify(sig, pk, msg) => {
284 format!(
285 "secp256k1_verify({}, {}, {})",
286 self.debug_hir(*sig),
287 self.debug_hir(*pk),
288 self.debug_hir(*msg)
289 )
290 }
291 Hir::Secp256R1Verify(sig, pk, msg) => {
292 format!(
293 "secp256r1_verify({}, {}, {})",
294 self.debug_hir(*sig),
295 self.debug_hir(*pk),
296 self.debug_hir(*msg)
297 )
298 }
299 Hir::InfinityG1 => "INFINITY_G1".to_string(),
300 Hir::InfinityG2 => "INFINITY_G2".to_string(),
301 Hir::ClvmOp(op, args) => {
302 format!("{:?}({})", op, self.debug_hir(*args))
303 }
304 }
305 }
306}