arithmetic_eval/compiler/
expr.rs

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

arithmetic_eval/compiler/
expr.rs
