cas_compiler/expr/
assign.rs

1use cas_compute::numerical::{func::{Function, User}, value::Value};
2use cas_error::Error;
3use cas_parser::parser::{
4    ast::{assign::{Assign, AssignTarget}, LitSym, Param},
5    token::op::AssignOpKind,
6};
7use crate::{
8    error::OverrideBuiltinConstant,
9    item::Symbol,
10    Compile,
11    Compiler,
12    InstructionKind,
13    NewChunk,
14};
15
16/// Extracts the user symbol ID from a [`Symbol`], returning an error if the symbol is not a
17/// user-defined symbol.
18fn extract_user_symbol(lit: &LitSym, symbol: Symbol) -> Result<usize, Error> {
19    symbol.index()
20        .map_err(|name| Error::new(vec![lit.span.clone()], OverrideBuiltinConstant {
21            name: name.to_string(),
22        }))
23}
24
25impl Compile for Assign {
26    fn compile(&self, compiler: &mut Compiler) -> Result<(), Error> {
27        match &self.target {
28            AssignTarget::Symbol(symbol) => {
29                // variable assignment
30                match self.op.kind {
31                    AssignOpKind::Assign => {
32                        compiler.with_state(|state| {
33                            state.top_level_assign = false;
34                        }, |compiler| {
35                            self.value.compile(compiler)
36                        })?;
37
38                        let symbol_id = compiler.resolve_user_symbol_or_insert(symbol)?;
39                        compiler.add_instr(InstructionKind::StoreVar(symbol_id));
40                    },
41                    compound => {
42                        let resolved = compiler.resolve_symbol(symbol)?;
43                        compiler.add_instr_with_spans(
44                            InstructionKind::LoadVar(resolved),
45                            vec![symbol.span.clone()],
46                        );
47
48                        compiler.with_state(|state| {
49                            state.top_level_assign = false;
50                        }, |compiler| {
51                            self.value.compile(compiler)
52                        })?;
53
54                        compiler.add_instr(InstructionKind::Binary(compound.into()));
55                        let symbol_id = extract_user_symbol(symbol, compiler.resolve_symbol(symbol)?)?;
56                        compiler.add_instr(InstructionKind::StoreVar(symbol_id));
57                    }
58                }
59            },
60            AssignTarget::Index(index) => {
61                // assignment to index of list
62                match self.op.kind {
63                    AssignOpKind::Assign => {
64                        // load new value
65                        self.value.compile(compiler)?;
66                    },
67                    compound => {
68                        // load the (hopefully) list value
69                        index.target.compile(compiler)?;
70
71                        // load value to index by
72                        index.index.compile(compiler)?;
73
74                        // compute the new value to assign
75                        compiler.add_instr_with_spans(
76                            InstructionKind::LoadIndexed,
77                            vec![index.target.span(), index.index.span()],
78                        );
79                        self.value.compile(compiler)?;
80                        compiler.add_instr(InstructionKind::Binary(compound.into()));
81                    }
82                }
83
84                // load the (hopefully) list value
85                index.target.compile(compiler)?;
86
87                // load value to index by
88                index.index.compile(compiler)?;
89
90                compiler.add_instr_with_spans(
91                    InstructionKind::StoreIndexed,
92                    vec![index.target.span(), index.index.span()],
93                );
94            },
95            AssignTarget::Func(header) => {
96                // for function assignment, create a new chunk for the function body
97                let NewChunk { id, chunk, captures } = compiler.new_chunk(header, |compiler| {
98                    compiler.add_instr(InstructionKind::InitFunc(
99                        header.name.to_string(),
100                        header.to_string(),
101                        header.params.len(),
102                        header.params.iter().filter(|p| p.has_default()).count(),
103                    ));
104
105                    // arguments to function are placed on the stack, so we need to go through the
106                    // parameters in reverse order to store them in the correct order
107                    let mut iter = header.params.iter().rev().peekable();
108                    while let Some(Param::Default(sym, default_expr)) =
109                        iter.next_if(|p| p.has_default()) {
110                        let assign_param = compiler.new_unassociated_label();
111
112                        // are there currently enough arguments on the stack?
113                        // if so, no need to push a default value
114                        compiler.add_instr(InstructionKind::CheckExecReady);
115                        compiler.add_instr(InstructionKind::JumpIfTrue(assign_param));
116
117                        // if not, push a default value onto the stack, count that as one new
118                        // argument
119                        default_expr.compile(compiler)?;
120                        compiler.add_instr(InstructionKind::NextArg);
121
122                        compiler.set_end_label(assign_param);
123
124                        // assign the default (or already provided) value to the parameter
125                        let symbol_id = compiler.add_symbol(sym)?;
126                        compiler.add_instr(InstructionKind::AssignVar(symbol_id));
127                    }
128
129                    // if we get here and there are still fewer or more arguments than parameters,
130                    // there are missing or extra arguments; raise an error
131                    compiler.add_instr(InstructionKind::ErrorIfMissingArgs);
132
133                    // assign required arguments; we know there are enough arguments on the stack,
134                    // as proven by `CheckExecReady`
135                    for param in iter {
136                        let symbol_id = compiler.add_symbol(param.symbol())?;
137                        compiler.add_instr(InstructionKind::AssignVar(symbol_id));
138                    }
139
140                    self.value.compile(compiler)?;
141                    compiler.add_instr(InstructionKind::Return);
142
143                    Ok(())
144                })?;
145
146                let user_func = User::new(chunk, captures);
147                compiler.add_instr(InstructionKind::LoadConst(Value::Function(Function::User(user_func))));
148                compiler.add_instr(InstructionKind::StoreVar(id));
149            },
150        };
151
152        Ok(())
153    }
154}