1use std::collections::{HashMap, HashSet};
2use std::path::{Path, PathBuf};
3
4use rexlang_ast::expr::{Expr, Program, Symbol};
5use rexlang_typesystem::{TypedExpr, TypedExprKind};
6use rexlang_util::GasMeter;
7use uuid::Uuid;
8
9use crate::engine::{
10 ClassMethodRequirement, CompiledExterns, CompiledProgram, Engine, NativeRequirement,
11 RUNTIME_LINK_ABI_VERSION, RuntimeLinkContract, collect_pattern_bindings, type_check_engine,
12};
13use crate::libraries::{
14 LibraryExports, LibraryId, ReplState, ResolvedLibrary, decl_value_names, exports_from_program,
15 parse_program_from_source, prefix_for_library, rewrite_import_uses, validate_import_uses,
16};
17use crate::{CompileError, EngineError, Env};
18
19#[derive(Clone)]
20pub struct Compiler<State = ()>
21where
22 State: Clone + Send + Sync + 'static,
23{
24 pub(crate) engine: Engine<State>,
25}
26
27impl<State> Compiler<State>
28where
29 State: Clone + Send + Sync + 'static,
30{
31 pub fn new(engine: Engine<State>) -> Self {
32 Self { engine }
33 }
34
35 pub fn compile_expr(&mut self, expr: &Expr) -> Result<CompiledProgram, CompileError> {
36 self.compile_expr_internal(expr).map_err(CompileError::from)
37 }
38
39 pub(crate) fn compile_expr_internal(
40 &mut self,
41 expr: &Expr,
42 ) -> Result<CompiledProgram, EngineError> {
43 let typed = self.type_check(expr)?;
44 let env = self.engine.env_snapshot();
45 let externs = self.collect_externs(&typed, &env);
46 let link_contract = self.link_contract(&typed, &env);
47 Ok(CompiledProgram::new(externs, link_contract, env, typed))
48 }
49
50 pub(crate) fn type_check(&mut self, expr: &Expr) -> Result<TypedExpr, EngineError> {
51 type_check_engine(&mut self.engine, expr)
52 }
53
54 fn collect_externs(&self, expr: &TypedExpr, env: &Env) -> CompiledExterns {
55 enum Frame<'b> {
56 Expr(&'b TypedExpr),
57 Push(Symbol),
58 PushMany(Vec<Symbol>),
59 Pop(usize),
60 }
61
62 let mut natives = HashSet::new();
63 let mut class_methods = HashSet::new();
64 let mut bound: Vec<Symbol> = Vec::new();
65 let mut stack = vec![Frame::Expr(expr)];
66 while let Some(frame) = stack.pop() {
67 match frame {
68 Frame::Expr(expr) => match &expr.kind {
69 TypedExprKind::Var { name, .. } => {
70 if bound.iter().any(|sym| sym == name) || env.get(name).is_some() {
71 continue;
72 }
73 if self.engine.type_system.class_methods.contains_key(name) {
74 class_methods.insert(name.clone());
75 } else if self.engine.has_native_name(name) {
76 natives.insert(name.clone());
77 }
78 }
79 TypedExprKind::Tuple(elems) | TypedExprKind::List(elems) => {
80 for elem in elems.iter().rev() {
81 stack.push(Frame::Expr(elem));
82 }
83 }
84 TypedExprKind::Dict(kvs) => {
85 for v in kvs.values().rev() {
86 stack.push(Frame::Expr(v));
87 }
88 }
89 TypedExprKind::RecordUpdate { base, updates } => {
90 for v in updates.values().rev() {
91 stack.push(Frame::Expr(v));
92 }
93 stack.push(Frame::Expr(base));
94 }
95 TypedExprKind::App(f, x) => {
96 stack.push(Frame::Expr(x));
97 stack.push(Frame::Expr(f));
98 }
99 TypedExprKind::Project { expr, .. } => stack.push(Frame::Expr(expr)),
100 TypedExprKind::Lam { param, body } => {
101 stack.push(Frame::Pop(1));
102 stack.push(Frame::Expr(body));
103 stack.push(Frame::Push(param.clone()));
104 }
105 TypedExprKind::Let { name, def, body } => {
106 stack.push(Frame::Pop(1));
107 stack.push(Frame::Expr(body));
108 stack.push(Frame::Push(name.clone()));
109 stack.push(Frame::Expr(def));
110 }
111 TypedExprKind::LetRec { bindings, body } => {
112 if !bindings.is_empty() {
113 stack.push(Frame::Pop(bindings.len()));
114 stack.push(Frame::Expr(body));
115 for (_, def) in bindings.iter().rev() {
116 stack.push(Frame::Expr(def));
117 }
118 stack.push(Frame::PushMany(
119 bindings.iter().map(|(name, _)| name.clone()).collect(),
120 ));
121 } else {
122 stack.push(Frame::Expr(body));
123 }
124 }
125 TypedExprKind::Ite {
126 cond,
127 then_expr,
128 else_expr,
129 } => {
130 stack.push(Frame::Expr(else_expr));
131 stack.push(Frame::Expr(then_expr));
132 stack.push(Frame::Expr(cond));
133 }
134 TypedExprKind::Match { scrutinee, arms } => {
135 for (pat, arm_expr) in arms.iter().rev() {
136 let mut bindings = Vec::new();
137 collect_pattern_bindings(pat, &mut bindings);
138 let count = bindings.len();
139 if count != 0 {
140 stack.push(Frame::Pop(count));
141 stack.push(Frame::Expr(arm_expr));
142 stack.push(Frame::PushMany(bindings));
143 } else {
144 stack.push(Frame::Expr(arm_expr));
145 }
146 }
147 stack.push(Frame::Expr(scrutinee));
148 }
149 TypedExprKind::Bool(..)
150 | TypedExprKind::Uint(..)
151 | TypedExprKind::Int(..)
152 | TypedExprKind::Float(..)
153 | TypedExprKind::String(..)
154 | TypedExprKind::Uuid(..)
155 | TypedExprKind::DateTime(..)
156 | TypedExprKind::Hole => {}
157 },
158 Frame::Push(sym) => bound.push(sym),
159 Frame::PushMany(syms) => bound.extend(syms),
160 Frame::Pop(count) => bound.truncate(bound.len().saturating_sub(count)),
161 }
162 }
163
164 let mut natives = natives.into_iter().collect::<Vec<_>>();
165 let mut class_methods = class_methods.into_iter().collect::<Vec<_>>();
166 natives.sort();
167 class_methods.sort();
168 CompiledExterns {
169 natives,
170 class_methods,
171 }
172 }
173
174 fn link_contract(&self, expr: &TypedExpr, env: &Env) -> RuntimeLinkContract {
175 enum Frame<'b> {
176 Expr(&'b TypedExpr),
177 Push(Symbol),
178 PushMany(Vec<Symbol>),
179 Pop(usize),
180 }
181
182 let mut native_requirements = HashSet::new();
183 let mut class_method_requirements = HashSet::new();
184 let mut bound: Vec<Symbol> = Vec::new();
185 let mut stack = vec![Frame::Expr(expr)];
186 while let Some(frame) = stack.pop() {
187 match frame {
188 Frame::Expr(expr) => match &expr.kind {
189 TypedExprKind::Var { name, .. } => {
190 if bound.iter().any(|sym| sym == name) || env.get(name).is_some() {
191 continue;
192 }
193 if self.engine.type_system.class_methods.contains_key(name) {
194 class_method_requirements.insert(ClassMethodRequirement {
195 name: name.clone(),
196 typ: expr.typ.clone(),
197 });
198 } else if self.engine.has_native_name(name) {
199 native_requirements.insert(NativeRequirement {
200 name: name.clone(),
201 typ: expr.typ.clone(),
202 });
203 }
204 }
205 TypedExprKind::Tuple(elems) | TypedExprKind::List(elems) => {
206 for elem in elems.iter().rev() {
207 stack.push(Frame::Expr(elem));
208 }
209 }
210 TypedExprKind::Dict(kvs) => {
211 for v in kvs.values().rev() {
212 stack.push(Frame::Expr(v));
213 }
214 }
215 TypedExprKind::RecordUpdate { base, updates } => {
216 for v in updates.values().rev() {
217 stack.push(Frame::Expr(v));
218 }
219 stack.push(Frame::Expr(base));
220 }
221 TypedExprKind::App(f, x) => {
222 stack.push(Frame::Expr(x));
223 stack.push(Frame::Expr(f));
224 }
225 TypedExprKind::Project { expr, .. } => stack.push(Frame::Expr(expr)),
226 TypedExprKind::Lam { param, body } => {
227 stack.push(Frame::Pop(1));
228 stack.push(Frame::Expr(body));
229 stack.push(Frame::Push(param.clone()));
230 }
231 TypedExprKind::Let { name, def, body } => {
232 stack.push(Frame::Pop(1));
233 stack.push(Frame::Expr(body));
234 stack.push(Frame::Push(name.clone()));
235 stack.push(Frame::Expr(def));
236 }
237 TypedExprKind::LetRec { bindings, body } => {
238 if !bindings.is_empty() {
239 stack.push(Frame::Pop(bindings.len()));
240 stack.push(Frame::Expr(body));
241 for (_, def) in bindings.iter().rev() {
242 stack.push(Frame::Expr(def));
243 }
244 stack.push(Frame::PushMany(
245 bindings.iter().map(|(name, _)| name.clone()).collect(),
246 ));
247 } else {
248 stack.push(Frame::Expr(body));
249 }
250 }
251 TypedExprKind::Ite {
252 cond,
253 then_expr,
254 else_expr,
255 } => {
256 stack.push(Frame::Expr(else_expr));
257 stack.push(Frame::Expr(then_expr));
258 stack.push(Frame::Expr(cond));
259 }
260 TypedExprKind::Match { scrutinee, arms } => {
261 for (pat, arm_expr) in arms.iter().rev() {
262 let mut bindings = Vec::new();
263 collect_pattern_bindings(pat, &mut bindings);
264 let count = bindings.len();
265 if count != 0 {
266 stack.push(Frame::Pop(count));
267 stack.push(Frame::Expr(arm_expr));
268 stack.push(Frame::PushMany(bindings));
269 } else {
270 stack.push(Frame::Expr(arm_expr));
271 }
272 }
273 stack.push(Frame::Expr(scrutinee));
274 }
275 TypedExprKind::Bool(..)
276 | TypedExprKind::Uint(..)
277 | TypedExprKind::Int(..)
278 | TypedExprKind::Float(..)
279 | TypedExprKind::String(..)
280 | TypedExprKind::Uuid(..)
281 | TypedExprKind::DateTime(..)
282 | TypedExprKind::Hole => {}
283 },
284 Frame::Push(sym) => bound.push(sym),
285 Frame::PushMany(syms) => bound.extend(syms),
286 Frame::Pop(count) => bound.truncate(bound.len().saturating_sub(count)),
287 }
288 }
289
290 let mut natives = native_requirements.into_iter().collect::<Vec<_>>();
291 let mut class_methods = class_method_requirements.into_iter().collect::<Vec<_>>();
292 natives.sort_by(|a, b| {
293 a.name
294 .cmp(&b.name)
295 .then_with(|| a.typ.to_string().cmp(&b.typ.to_string()))
296 });
297 class_methods.sort_by(|a, b| {
298 a.name
299 .cmp(&b.name)
300 .then_with(|| a.typ.to_string().cmp(&b.typ.to_string()))
301 });
302 RuntimeLinkContract {
303 abi_version: RUNTIME_LINK_ABI_VERSION,
304 natives,
305 class_methods,
306 }
307 }
308
309 fn rewrite_and_inject_program(
310 &mut self,
311 program: &Program,
312 importer: Option<LibraryId>,
313 prefix: &str,
314 gas: &mut GasMeter,
315 loaded: &mut HashMap<LibraryId, LibraryExports>,
316 loading: &mut HashSet<LibraryId>,
317 ) -> Result<Program, EngineError> {
318 let rewritten = self
319 .engine
320 .rewrite_program_with_imports(program, importer, prefix, gas, loaded, loading)?;
321 self.engine.inject_decls(&rewritten.decls)?;
322 Ok(rewritten)
323 }
324
325 pub fn compile_snippet(
326 &mut self,
327 source: &str,
328 gas: &mut GasMeter,
329 ) -> Result<CompiledProgram, CompileError> {
330 self.compile_snippet_with_gas_and_importer(source, gas, None)
331 .map_err(CompileError::from)
332 }
333
334 pub fn compile_snippet_at(
335 &mut self,
336 source: &str,
337 importer_path: impl AsRef<Path>,
338 gas: &mut GasMeter,
339 ) -> Result<CompiledProgram, CompileError> {
340 let path = importer_path.as_ref().to_path_buf();
341 self.compile_snippet_with_gas_and_importer(source, gas, Some(path))
342 .map_err(CompileError::from)
343 }
344
345 pub fn compile_library_file(
346 &mut self,
347 path: impl AsRef<Path>,
348 gas: &mut GasMeter,
349 ) -> Result<CompiledProgram, CompileError> {
350 let (id, bytes) = self
351 .engine
352 .read_local_library_bytes(path.as_ref())
353 .map_err(CompileError::from)?;
354 let source = self
355 .engine
356 .decode_local_library_source(&id, bytes)
357 .map_err(CompileError::from)?;
358 self.compile_library_source(ResolvedLibrary { id, source }, gas)
359 .map_err(CompileError::from)
360 }
361
362 pub async fn compile_repl_program(
363 &mut self,
364 program: &Program,
365 state: &mut ReplState,
366 gas: &mut GasMeter,
367 ) -> Result<CompiledProgram, CompileError> {
368 self.compile_repl_program_internal(program, state, gas)
369 .await
370 .map_err(CompileError::from)
371 }
372
373 async fn compile_repl_program_internal(
374 &mut self,
375 program: &Program,
376 state: &mut ReplState,
377 gas: &mut GasMeter,
378 ) -> Result<CompiledProgram, EngineError> {
379 let importer = state
380 .importer_path
381 .as_ref()
382 .map(|p| LibraryId::Local { path: p.clone() });
383
384 let mut local_values = state.defined_values.clone();
385 local_values.extend(decl_value_names(&program.decls));
386 let existing_imported: HashSet<Symbol> = state.imported_values.keys().cloned().collect();
387 let import_bindings = self
388 .engine
389 .import_bindings_for_decls(
390 &program.decls,
391 importer.clone(),
392 &local_values,
393 Some(&existing_imported),
394 gas,
395 )
396 .await?;
397 state.alias_exports.extend(import_bindings.alias_exports);
398 state
399 .imported_values
400 .extend(import_bindings.imported_values);
401
402 let mut shadowed_values = state.defined_values.clone();
403 shadowed_values.extend(decl_value_names(&program.decls));
404
405 validate_import_uses(program, &state.alias_exports, Some(&shadowed_values))?;
406 let rewritten = rewrite_import_uses(
407 program,
408 &state.alias_exports,
409 &state.imported_values,
410 Some(&shadowed_values),
411 );
412
413 self.engine.inject_decls(&rewritten.decls)?;
414 state
415 .defined_values
416 .extend(decl_value_names(&program.decls));
417 self.compile_expr_internal(rewritten.expr.as_ref())
418 }
419
420 fn compile_library_source(
421 &mut self,
422 resolved: ResolvedLibrary,
423 gas: &mut GasMeter,
424 ) -> Result<CompiledProgram, EngineError> {
425 let mut loaded: HashMap<LibraryId, LibraryExports> = HashMap::new();
426 let mut loading: HashSet<LibraryId> = HashSet::new();
427
428 loading.insert(resolved.id.clone());
429
430 let prefix = prefix_for_library(&resolved.id);
431 let program =
432 parse_program_from_source(&resolved.source, Some(&resolved.id), Some(&mut *gas))?;
433 let rewritten = self.rewrite_and_inject_program(
434 &program,
435 Some(resolved.id.clone()),
436 &prefix,
437 gas,
438 &mut loaded,
439 &mut loading,
440 )?;
441
442 let exports = exports_from_program(&program, &prefix, &resolved.id);
443 loaded.insert(resolved.id.clone(), exports);
444 loading.remove(&resolved.id);
445
446 self.compile_expr_internal(rewritten.expr.as_ref())
447 }
448
449 fn compile_snippet_with_gas_and_importer(
450 &mut self,
451 source: &str,
452 gas: &mut GasMeter,
453 importer_path: Option<PathBuf>,
454 ) -> Result<CompiledProgram, EngineError> {
455 let program = parse_program_from_source(source, None, Some(&mut *gas))?;
456
457 let importer = importer_path.map(|p| LibraryId::Local { path: p });
458 let prefix = format!("@snippet{}", Uuid::new_v4());
459 let mut loaded: HashMap<LibraryId, LibraryExports> = HashMap::new();
460 let mut loading: HashSet<LibraryId> = HashSet::new();
461 let rewritten = self.rewrite_and_inject_program(
462 &program,
463 importer,
464 &prefix,
465 gas,
466 &mut loaded,
467 &mut loading,
468 )?;
469 self.compile_expr_internal(rewritten.expr.as_ref())
470 }
471}