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(
359 ResolvedLibrary {
360 id,
361 content: crate::libraries::ResolvedLibraryContent::Source(source),
362 },
363 gas,
364 )
365 .map_err(CompileError::from)
366 }
367
368 pub async fn compile_repl_program(
369 &mut self,
370 program: &Program,
371 state: &mut ReplState,
372 gas: &mut GasMeter,
373 ) -> Result<CompiledProgram, CompileError> {
374 self.compile_repl_program_internal(program, state, gas)
375 .await
376 .map_err(CompileError::from)
377 }
378
379 async fn compile_repl_program_internal(
380 &mut self,
381 program: &Program,
382 state: &mut ReplState,
383 gas: &mut GasMeter,
384 ) -> Result<CompiledProgram, EngineError> {
385 let importer = state
386 .importer_path
387 .as_ref()
388 .map(|p| LibraryId::Local { path: p.clone() });
389
390 let mut local_values = state.defined_values.clone();
391 local_values.extend(decl_value_names(&program.decls));
392 let existing_imported: HashSet<Symbol> = state.imported_values.keys().cloned().collect();
393 let import_bindings = self
394 .engine
395 .import_bindings_for_decls(
396 &program.decls,
397 importer.clone(),
398 &local_values,
399 Some(&existing_imported),
400 gas,
401 )
402 .await?;
403 state.alias_exports.extend(import_bindings.alias_exports);
404 state
405 .imported_values
406 .extend(import_bindings.imported_values);
407
408 let mut shadowed_values = state.defined_values.clone();
409 shadowed_values.extend(decl_value_names(&program.decls));
410
411 validate_import_uses(program, &state.alias_exports, Some(&shadowed_values))?;
412 let rewritten = rewrite_import_uses(
413 program,
414 &state.alias_exports,
415 &state.imported_values,
416 Some(&shadowed_values),
417 );
418
419 self.engine.inject_decls(&rewritten.decls)?;
420 state
421 .defined_values
422 .extend(decl_value_names(&program.decls));
423 self.compile_expr_internal(rewritten.expr.as_ref())
424 }
425
426 fn compile_library_source(
427 &mut self,
428 resolved: ResolvedLibrary,
429 gas: &mut GasMeter,
430 ) -> Result<CompiledProgram, EngineError> {
431 let mut loaded: HashMap<LibraryId, LibraryExports> = HashMap::new();
432 let mut loading: HashSet<LibraryId> = HashSet::new();
433
434 loading.insert(resolved.id.clone());
435
436 let prefix = prefix_for_library(&resolved.id);
437 let program = crate::libraries::program_from_resolved(&resolved, &mut *gas)?;
438 let rewritten = self.rewrite_and_inject_program(
439 &program,
440 Some(resolved.id.clone()),
441 &prefix,
442 gas,
443 &mut loaded,
444 &mut loading,
445 )?;
446
447 let exports = exports_from_program(&program, &prefix, &resolved.id);
448 loaded.insert(resolved.id.clone(), exports);
449 loading.remove(&resolved.id);
450
451 self.compile_expr_internal(rewritten.expr.as_ref())
452 }
453
454 fn compile_snippet_with_gas_and_importer(
455 &mut self,
456 source: &str,
457 gas: &mut GasMeter,
458 importer_path: Option<PathBuf>,
459 ) -> Result<CompiledProgram, EngineError> {
460 let program = parse_program_from_source(source, None, Some(&mut *gas))?;
461
462 let importer = importer_path.map(|p| LibraryId::Local { path: p });
463 let prefix = format!("@snippet{}", Uuid::new_v4());
464 let mut loaded: HashMap<LibraryId, LibraryExports> = HashMap::new();
465 let mut loading: HashSet<LibraryId> = HashSet::new();
466 let rewritten = self.rewrite_and_inject_program(
467 &program,
468 importer,
469 &prefix,
470 gas,
471 &mut loaded,
472 &mut loading,
473 )?;
474 self.compile_expr_internal(rewritten.expr.as_ref())
475 }
476}