arithmetic_eval/compiler/
expr.rs

1//! Compilation logic for `Expr`essions.
2
3use hashbrown::HashMap;
4
5use core::iter;
6
7use super::{captures::extract_vars_iter, CapturesExtractor, Compiler};
8use crate::{
9    alloc::{ToOwned, Vec},
10    error::RepeatedAssignmentContext,
11    executable::{Atom, Command, CompiledExpr, Executable, ExecutableFn, FieldName, SpannedAtom},
12    Error, ErrorKind,
13};
14use arithmetic_parser::{
15    grammars::Grammar, is_valid_variable_name, BinaryOp, Block, Expr, FnDefinition, MaybeSpanned,
16    ObjectExpr, Spanned, SpannedExpr, SpannedStatement, Statement,
17};
18
19impl Compiler {
20    fn compile_expr<'a, T: Grammar<'a>>(
21        &mut self,
22        executable: &mut Executable<'a, T::Lit>,
23        expr: &SpannedExpr<'a, T>,
24    ) -> Result<SpannedAtom<'a, T::Lit>, Error<'a>> {
25        let atom = match &expr.extra {
26            Expr::Literal(lit) => Atom::Constant(lit.clone()),
27
28            Expr::Variable => self.compile_var_access(expr)?,
29
30            Expr::TypeCast { value, .. } => {
31                // Just ignore the type annotation.
32                self.compile_expr(executable, value)?.extra
33            }
34
35            Expr::Tuple(tuple) => {
36                let registers = tuple
37                    .iter()
38                    .map(|elem| {
39                        self.compile_expr(executable, elem)
40                            .map(|spanned| spanned.extra)
41                    })
42                    .collect::<Result<Vec<_>, _>>()?;
43                let register =
44                    self.push_assignment(executable, CompiledExpr::Tuple(registers), expr);
45                Atom::Register(register)
46            }
47
48            Expr::Unary { op, inner } => {
49                let inner = self.compile_expr(executable, inner)?;
50                let register = self.push_assignment(
51                    executable,
52                    CompiledExpr::Unary {
53                        op: self.check_unary_op(op)?,
54                        inner,
55                    },
56                    expr,
57                );
58                Atom::Register(register)
59            }
60
61            Expr::Binary { op, lhs, rhs } => {
62                self.compile_binary_expr(executable, expr, op, lhs, rhs)?
63            }
64            Expr::Function { name, args } => self.compile_fn_call(executable, expr, name, args)?,
65
66            Expr::FieldAccess { name, receiver } => {
67                self.compile_field_access(executable, expr, name, receiver)?
68            }
69
70            Expr::Method {
71                name,
72                receiver,
73                args,
74            } => self.compile_method_call(executable, expr, name, receiver, args)?,
75
76            Expr::Block(block) => self.compile_block(executable, expr, block)?,
77            Expr::FnDefinition(def) => self.compile_fn_definition(executable, expr, def)?,
78            Expr::Object(object) => self.compile_object(executable, expr, object)?,
79
80            _ => {
81                let err = ErrorKind::unsupported(expr.extra.ty());
82                return Err(self.create_error(expr, err));
83            }
84        };
85
86        Ok(expr.copy_with_extra(atom).into())
87    }
88
89    fn compile_var_access<'a, T, A>(
90        &self,
91        var_span: &Spanned<'a, T>,
92    ) -> Result<Atom<A>, Error<'a>> {
93        let var_name = *var_span.fragment();
94        let register = self.vars_to_registers.get(var_name).ok_or_else(|| {
95            let err = ErrorKind::Undefined(var_name.to_owned());
96            self.create_error(var_span, err)
97        })?;
98        Ok(Atom::Register(*register))
99    }
100
101    fn compile_binary_expr<'a, T: Grammar<'a>>(
102        &mut self,
103        executable: &mut Executable<'a, T::Lit>,
104        binary_expr: &SpannedExpr<'a, T>,
105        op: &Spanned<'a, BinaryOp>,
106        lhs: &SpannedExpr<'a, T>,
107        rhs: &SpannedExpr<'a, T>,
108    ) -> Result<Atom<T::Lit>, Error<'a>> {
109        let lhs = self.compile_expr(executable, lhs)?;
110        let rhs = self.compile_expr(executable, rhs)?;
111
112        let compiled = CompiledExpr::Binary {
113            op: self.check_binary_op(op)?,
114            lhs,
115            rhs,
116        };
117
118        let register = self.push_assignment(executable, compiled, binary_expr);
119        Ok(Atom::Register(register))
120    }
121
122    fn compile_fn_call<'a, T: Grammar<'a>>(
123        &mut self,
124        executable: &mut Executable<'a, T::Lit>,
125        call_expr: &SpannedExpr<'a, T>,
126        name: &SpannedExpr<'a, T>,
127        args: &[SpannedExpr<'a, T>],
128    ) -> Result<Atom<T::Lit>, Error<'a>> {
129        let original_name = *name.fragment();
130        let original_name = if is_valid_variable_name(original_name) {
131            Some(original_name.to_owned())
132        } else {
133            None
134        };
135
136        let name = self.compile_expr(executable, name)?;
137
138        let args = args
139            .iter()
140            .map(|arg| self.compile_expr(executable, arg))
141            .collect::<Result<Vec<_>, _>>()?;
142        let function = CompiledExpr::Function {
143            name,
144            original_name,
145            args,
146        };
147        let register = self.push_assignment(executable, function, call_expr);
148        Ok(Atom::Register(register))
149    }
150
151    fn compile_field_access<'a, T: Grammar<'a>>(
152        &mut self,
153        executable: &mut Executable<'a, T::Lit>,
154        call_expr: &SpannedExpr<'a, T>,
155        name: &Spanned<'a>,
156        receiver: &SpannedExpr<'a, T>,
157    ) -> Result<Atom<T::Lit>, Error<'a>> {
158        let name_str = *name.fragment();
159        let field = name_str
160            .parse::<usize>()
161            .map(FieldName::Index)
162            .or_else(|_| {
163                if is_valid_variable_name(name_str) {
164                    Ok(FieldName::Name(name_str.to_owned()))
165                } else {
166                    let err = ErrorKind::InvalidFieldName(name_str.to_owned());
167                    Err(self.create_error(name, err))
168                }
169            })?;
170
171        let receiver = self.compile_expr(executable, receiver)?;
172        let field_access = CompiledExpr::FieldAccess { receiver, field };
173        let register = self.push_assignment(executable, field_access, call_expr);
174        Ok(Atom::Register(register))
175    }
176
177    fn compile_method_call<'a, T: Grammar<'a>>(
178        &mut self,
179        executable: &mut Executable<'a, T::Lit>,
180        call_expr: &SpannedExpr<'a, T>,
181        name: &Spanned<'a>,
182        receiver: &SpannedExpr<'a, T>,
183        args: &[SpannedExpr<'a, T>],
184    ) -> Result<Atom<T::Lit>, Error<'a>> {
185        let original_name = Some((*name.fragment()).to_owned());
186        let name: MaybeSpanned<'_, _> = name
187            .copy_with_extra(Atom::Register(self.vars_to_registers[*name.fragment()]))
188            .into();
189        let args = iter::once(receiver)
190            .chain(args)
191            .map(|arg| self.compile_expr(executable, arg))
192            .collect::<Result<Vec<_>, _>>()?;
193
194        let function = CompiledExpr::Function {
195            name,
196            original_name,
197            args,
198        };
199        let register = self.push_assignment(executable, function, call_expr);
200        Ok(Atom::Register(register))
201    }
202
203    fn compile_block<'r, 'a: 'r, T: Grammar<'a>>(
204        &mut self,
205        executable: &mut Executable<'a, T::Lit>,
206        block_expr: &SpannedExpr<'a, T>,
207        block: &Block<'a, T>,
208    ) -> Result<Atom<T::Lit>, Error<'a>> {
209        let backup_state = self.backup();
210        if self.scope_depth == 0 {
211            let command = Command::StartInnerScope;
212            executable.push_command(block_expr.copy_with_extra(command));
213        }
214        self.scope_depth += 1;
215
216        // Move the return value to the next register.
217        let return_value = self
218            .compile_block_inner(executable, block)?
219            .map_or(Atom::Void, |spanned| spanned.extra);
220
221        let new_register = if let Atom::Register(ret_register) = return_value {
222            let command = Command::Copy {
223                source: ret_register,
224                destination: backup_state.register_count,
225            };
226            executable.push_command(block_expr.copy_with_extra(command));
227            true
228        } else {
229            false
230        };
231
232        // Return to the previous state. This will erase register mapping
233        // for the inner scope and set the `scope_depth`.
234        *self = backup_state;
235        if new_register {
236            self.register_count += 1;
237        }
238        if self.scope_depth == 0 {
239            let command = Command::EndInnerScope;
240            executable.push_command(block_expr.copy_with_extra(command));
241        }
242        executable.push_command(
243            block_expr.copy_with_extra(Command::TruncateRegisters(self.register_count)),
244        );
245
246        Ok(if new_register {
247            Atom::Register(self.register_count - 1)
248        } else {
249            Atom::Void
250        })
251    }
252
253    pub(super) fn compile_block_inner<'a, T: Grammar<'a>>(
254        &mut self,
255        executable: &mut Executable<'a, T::Lit>,
256        block: &Block<'a, T>,
257    ) -> Result<Option<SpannedAtom<'a, T::Lit>>, Error<'a>> {
258        for statement in &block.statements {
259            self.compile_statement(executable, statement)?;
260        }
261
262        Ok(if let Some(return_value) = &block.return_value {
263            Some(self.compile_expr(executable, return_value)?)
264        } else {
265            None
266        })
267    }
268
269    #[allow(clippy::option_if_let_else)] // false positive
270    fn compile_object<'a, T: Grammar<'a>>(
271        &mut self,
272        executable: &mut Executable<'a, T::Lit>,
273        object_expr: &SpannedExpr<'a, T>,
274        object: &ObjectExpr<'a, T>,
275    ) -> Result<Atom<T::Lit>, Error<'a>> {
276        let fields = object.fields.iter().map(|(name, field_expr)| {
277            let name_str = *name.fragment();
278            if let Some(field_expr) = field_expr {
279                self.compile_expr(executable, field_expr)
280                    .map(|register| (name_str.to_owned(), register.extra))
281            } else {
282                self.compile_var_access(name)
283                    .map(|register| (name_str.to_owned(), register))
284            }
285        });
286        let obj_expr = CompiledExpr::Object(fields.collect::<Result<_, _>>()?);
287        let register = self.push_assignment(executable, obj_expr, object_expr);
288        Ok(Atom::Register(register))
289    }
290
291    fn compile_fn_definition<'a, T: Grammar<'a>>(
292        &mut self,
293        executable: &mut Executable<'a, T::Lit>,
294        def_expr: &SpannedExpr<'a, T>,
295        def: &FnDefinition<'a, T>,
296    ) -> Result<Atom<T::Lit>, Error<'a>> {
297        let module_id = self.module_id.clone_boxed();
298
299        let mut extractor = CapturesExtractor::new(module_id);
300        extractor.eval_function(def)?;
301        let captures = self.get_captures(extractor);
302
303        let fn_executable = self.compile_function(def, &captures)?;
304        let fn_executable = ExecutableFn {
305            inner: fn_executable,
306            def_span: def_expr.with_no_extra().into(),
307            arg_count: def.args.extra.len(),
308        };
309
310        let ptr = executable.push_child_fn(fn_executable);
311        let (capture_names, captures) = captures
312            .into_iter()
313            .map(|(name, value)| (name.to_owned(), value))
314            .unzip();
315        let register = self.push_assignment(
316            executable,
317            CompiledExpr::DefineFunction {
318                ptr,
319                captures,
320                capture_names,
321            },
322            def_expr,
323        );
324        Ok(Atom::Register(register))
325    }
326
327    fn get_captures<'a, T>(
328        &self,
329        extractor: CapturesExtractor<'a>,
330    ) -> HashMap<&'a str, SpannedAtom<'a, T>> {
331        extractor
332            .captures
333            .into_iter()
334            .map(|(var_name, var_span)| {
335                let register = self.get_var(var_name);
336                let capture = var_span.copy_with_extra(Atom::Register(register));
337                (var_name, capture.into())
338            })
339            .collect()
340    }
341
342    fn compile_function<'a, T: Grammar<'a>>(
343        &self,
344        def: &FnDefinition<'a, T>,
345        captures: &HashMap<&'a str, SpannedAtom<'a, T::Lit>>,
346    ) -> Result<Executable<'a, T::Lit>, Error<'a>> {
347        // Allocate registers for captures.
348        let mut this = Self::new(self.module_id.clone_boxed());
349        this.scope_depth = 1; // Disable generating variable annotations.
350
351        for (i, &name) in captures.keys().enumerate() {
352            this.vars_to_registers.insert(name.to_owned(), i);
353        }
354        this.register_count = captures.len() + 1; // one additional register for args
355
356        let mut executable = Executable::new(self.module_id.clone_boxed());
357        let args_span = def.args.with_no_extra();
358        this.destructure(&mut executable, &def.args.extra, args_span, captures.len())?;
359
360        for statement in &def.body.statements {
361            this.compile_statement(&mut executable, statement)?;
362        }
363        if let Some(return_value) = &def.body.return_value {
364            let return_atom = this.compile_expr(&mut executable, return_value)?;
365            let return_span = return_atom.with_no_extra();
366            let command = Command::Push(CompiledExpr::Atom(return_atom.extra));
367            executable.push_command(return_span.copy_with_extra(command));
368        }
369
370        executable.finalize_function(this.register_count);
371        Ok(executable)
372    }
373
374    fn compile_statement<'a, T: Grammar<'a>>(
375        &mut self,
376        executable: &mut Executable<'a, T::Lit>,
377        statement: &SpannedStatement<'a, T>,
378    ) -> Result<Option<SpannedAtom<'a, T::Lit>>, Error<'a>> {
379        Ok(match &statement.extra {
380            Statement::Expr(expr) => Some(self.compile_expr(executable, expr)?),
381
382            Statement::Assignment { lhs, rhs } => {
383                extract_vars_iter(
384                    self.module_id.as_ref(),
385                    &mut HashMap::new(),
386                    iter::once(lhs),
387                    RepeatedAssignmentContext::Assignment,
388                )?;
389
390                let rhs = self.compile_expr(executable, rhs)?;
391                // Allocate the register for the constant if necessary.
392                let rhs_register = match rhs.extra {
393                    Atom::Constant(_) | Atom::Void => {
394                        self.push_assignment(executable, CompiledExpr::Atom(rhs.extra), statement)
395                    }
396                    Atom::Register(register) => register,
397                };
398                self.assign(executable, lhs, rhs_register)?;
399                None
400            }
401
402            _ => {
403                let err = ErrorKind::unsupported(statement.extra.ty());
404                return Err(self.create_error(statement, err));
405            }
406        })
407    }
408}
arithmetic_eval/compiler/expr.rs

arithmetic_eval/compiler/
expr.rs
