sphinx/
codegen.rs

1#![allow(unused_variables)]
2
3use core::iter;
4
5use crate::language::{IntType, FloatType, InternSymbol, Access};
6use crate::parser::stmt::{StmtMeta, Stmt, Label, StmtList, ControlFlow};
7use crate::parser::expr::{Expr, ExprMeta, ExprBlock, ConditionalBranch};
8use crate::parser::primary::{Atom, Primary, AccessItem};
9use crate::parser::lvalue::{LValue, AssignType};
10use crate::parser::fundefs::{FunctionDef, SignatureDef};
11use crate::parser::operator::{UnaryOp, BinaryOp};
12use crate::runtime::strings::{StringInterner};
13use crate::runtime::errors::ErrorKind;
14use crate::debug::symbol::{DebugSymbol, ChunkSymbols, DebugSymbolTable};
15
16mod scope;
17
18pub mod chunk;
19pub mod consts;
20pub mod funproto;
21pub mod opcodes;
22pub mod errors;
23
24pub use opcodes::{OpCode, LocalIndex};
25pub use chunk::{UnloadedProgram, Program, ProgramData, Chunk};
26pub use consts::{ConstID, Constant};
27pub use funproto::{FunctionID, FunctionProto, UpvalueTarget};
28pub use errors::{CompileResult, CompileError};
29
30use scope::{ScopeTracker, ScopeTag, Scope, LocalName, InsertLocal, ControlFlowTarget};
31use chunk::{ChunkBuilder, ChunkInfo, ChunkBuf};
32use funproto::{UnloadedFunction, UnloadedSignature, UnloadedParam};
33
34
35// Helpers
36
37#[derive(Clone, Copy, PartialEq, Eq, Hash)]
38enum JumpOffset {
39    Short(i16),
40    Long(i32),
41}
42
43#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
44enum Jump {
45    Uncond,
46    IfFalse,
47    IfTrue,
48    PopIfFalse,
49    PopIfTrue,
50}
51
52impl Jump {
53    pub const fn dummy_width(&self) -> usize {
54        get_jump_opcode(Jump::Uncond, JumpOffset::Short(0)).instr_len()
55    }
56}
57
58const fn get_jump_opcode(jump: Jump, offset: JumpOffset) -> OpCode {
59    match (jump, offset) {
60        (Jump::Uncond,  JumpOffset::Short(..)) => OpCode::Jump,
61        (Jump::IfFalse, JumpOffset::Short(..)) => OpCode::JumpIfFalse,
62        (Jump::IfTrue,  JumpOffset::Short(..)) => OpCode::JumpIfTrue,
63        
64        (Jump::Uncond,  JumpOffset::Long(..))  => OpCode::LongJump,
65        (Jump::IfFalse, JumpOffset::Long(..))  => OpCode::LongJumpIfFalse,
66        (Jump::IfTrue,  JumpOffset::Long(..))  => OpCode::LongJumpIfTrue,
67        
68        (Jump::PopIfFalse, JumpOffset::Short(..))  => OpCode::PopJumpIfFalse,
69        (Jump::PopIfTrue,  JumpOffset::Short(..))  => OpCode::PopJumpIfTrue,
70        
71        (Jump::PopIfFalse, JumpOffset::Long(..))   => OpCode::PopLongJumpIfFalse,
72        (Jump::PopIfTrue,  JumpOffset::Long(..))   => OpCode::PopLongJumpIfTrue,
73    }
74}
75
76// represents the site of a dummy jump instruction that will be patched with a target later
77#[derive(Debug)]
78struct JumpSite {
79    jump: Jump,
80    offset: usize,
81    width: usize,
82}
83
84
85/// Output container
86#[derive(Debug)]
87pub struct CompiledProgram {
88    pub program: UnloadedProgram,
89    pub symbols: ChunkSymbols,
90}
91
92
93// Code Generator
94pub struct Compiler {
95    builder: ChunkBuilder,
96    scopes: ScopeTracker,
97    errors: Vec<CompileError>,
98    symbols: ChunkSymbols,
99}
100
101impl Compiler {
102    pub fn new(strings: StringInterner) -> Self {
103        // insert symbol container for main chunk
104        let mut symbols = ChunkSymbols::new();
105        symbols.insert(Chunk::Main, DebugSymbolTable::new());
106        
107        Self {
108            builder: ChunkBuilder::with_strings(strings),
109            scopes: ScopeTracker::new(),
110            errors: Vec::new(),
111            symbols,
112        }
113    }
114    
115    fn new_chunk(&mut self, info: ChunkInfo) -> CompileResult<Chunk> {
116        let chunk_id = self.builder.new_chunk(info)?;
117        self.symbols.entry(chunk_id)
118            .or_insert_with(DebugSymbolTable::new);
119        
120        Ok(chunk_id)
121    }
122    
123    fn get_chunk(&mut self, chunk_id: Chunk) -> CodeGenerator {
124        CodeGenerator {
125            compiler: self,
126            chunk_id,
127        }
128    }
129    
130    pub fn compile_program<'a>(mut self, program: impl Iterator<Item=&'a StmtMeta>) -> Result<CompiledProgram, Vec<CompileError>> {
131        for stmt in program {
132            self.push_stmt(stmt);
133        }
134        self.finish()
135    }
136    
137    pub fn push_stmt(&mut self, stmt: &StmtMeta) {
138        if let Err(error) = self.get_chunk(Chunk::Main).push_stmt(stmt) {
139            self.errors.push(error);
140        }
141    }
142    
143    pub fn finish(mut self) -> Result<CompiledProgram, Vec<CompileError>> {
144        if self.errors.is_empty() {
145            self.get_chunk(Chunk::Main)
146                .finish();
147            
148            let output = CompiledProgram {
149                program: self.builder.build(),
150                symbols: self.symbols,
151            };
152            
153            Ok(output)
154        } else {
155            Err(self.errors)
156        }
157    }
158}
159
160struct CodeGenerator<'c> {
161    compiler: &'c mut Compiler,
162    chunk_id: Chunk,
163}
164
165impl CodeGenerator<'_> {
166    
167    pub fn push_stmt(&mut self, stmt: &StmtMeta) -> CompileResult<()> {
168        let symbol = stmt.debug_symbol();
169        
170        let result = self.compile_stmt(Some(symbol), stmt.variant());
171        if let Err(error) = result {
172            Err(error.with_symbol(*symbol))
173        } else {
174            Ok(())
175        }
176    }
177    
178    pub fn finish(mut self) {
179        match self.chunk_id {
180            Chunk::Main => self.emit_instr(None, OpCode::Exit),
181            
182            Chunk::Function(..) => self.emit_instr(None, OpCode::Return),
183        }
184    }
185    
186    fn chunk_id(&self) -> Chunk { self.chunk_id }
187    
188    fn builder(&self) -> &ChunkBuilder { &self.compiler.builder }
189    fn builder_mut(&mut self) -> &mut ChunkBuilder { &mut self.compiler.builder }
190    
191    fn scopes(&self) -> &ScopeTracker { &self.compiler.scopes }
192    fn scopes_mut(&mut self) -> &mut ScopeTracker { &mut self.compiler.scopes }
193    
194    fn symbols(&self) -> &ChunkSymbols { &self.compiler.symbols }
195    fn symbols_mut(&mut self) -> &mut ChunkSymbols { &mut self.compiler.symbols }
196    
197    fn chunk(&self) -> &ChunkBuf {
198        self.builder().chunk(self.chunk_id)
199    }
200    
201    fn chunk_mut(&mut self) -> &mut ChunkBuf {
202        let chunk_id = self.chunk_id;
203        self.builder_mut().chunk_mut(chunk_id)
204    }
205    
206    fn current_offset(&self) -> usize {
207        self.chunk().len()
208    }
209    
210    fn push_symbol(&mut self, symbol: DebugSymbol) {
211        let chunk_id = self.chunk_id;
212        let offset = self.current_offset();
213        self.symbols_mut()
214            .get_mut(&chunk_id).unwrap()
215            .insert(offset, symbol)
216    }
217    
218    fn create_chunk(&mut self, metadata: ChunkInfo) -> CompileResult<CodeGenerator> {
219        let chunk_id = self.compiler.new_chunk(metadata)?;
220        Ok(self.compiler.get_chunk(chunk_id))
221    }
222}
223
224
225///////// Emitting Bytecode /////////
226impl CodeGenerator<'_> {
227    fn emit_instr(&mut self, symbol: Option<&DebugSymbol>, opcode: OpCode) {
228        debug_assert!(opcode.instr_len() == 1);
229        
230        if let Some(symbol) = symbol {
231            self.push_symbol(*symbol);
232        }
233        
234        self.chunk_mut().push_byte(opcode);
235    }
236    
237    fn emit_instr_byte(&mut self, symbol: Option<&DebugSymbol>, opcode: OpCode, byte: u8) {
238        debug_assert!(opcode.instr_len() == 2);
239        
240        if let Some(symbol) = symbol {
241            self.push_symbol(*symbol);
242        }
243        
244        self.chunk_mut().push_byte(opcode);
245        self.chunk_mut().push_byte(byte);
246    }
247    
248    fn emit_instr_data(&mut self, symbol: Option<&DebugSymbol>, opcode: OpCode, bytes: &[u8]) {
249        debug_assert!(opcode.instr_len() == 1 + bytes.len());
250        
251        if let Some(symbol) = symbol {
252            self.push_symbol(*symbol);
253        }
254        
255        self.chunk_mut().push_byte(opcode);
256        self.chunk_mut().extend_bytes(bytes);
257    }
258}
259
260///////// Patching Bytecode /////////
261impl CodeGenerator<'_> {
262    fn patch_instr_data<const N: usize>(&mut self, offset: usize, opcode: OpCode, bytes: &[u8; N]) {
263        debug_assert!(opcode.instr_len() == 1 + N);
264        
265        self.chunk_mut().as_mut_slice()[offset] = u8::from(opcode);
266        self.chunk_mut().patch_bytes(offset + 1, bytes);
267    }
268    
269    fn emit_dummy_instr(&mut self, symbol: Option<&DebugSymbol>, width: usize) {
270        if let Some(symbol) = symbol {
271            self.push_symbol(*symbol);
272        }
273        
274        for i in 0..width {
275            self.chunk_mut().push_byte(OpCode::Nop);
276        }
277    }
278}
279
280///////// Constants /////////
281impl CodeGenerator<'_> {
282    fn get_or_make_const(&mut self, value: Constant) -> CompileResult<ConstID> {
283        self.builder_mut().get_or_insert_const(value)
284    }
285    
286    fn emit_load_const(&mut self, symbol: Option<&DebugSymbol>, value: Constant) -> CompileResult<()> {
287        let cid = self.get_or_make_const(value)?;
288        
289        if cid <= u8::MAX.into() {
290            self.emit_instr_byte(symbol, OpCode::LoadConst, u8::try_from(cid).unwrap());
291        } else {
292            self.emit_instr_data(symbol, OpCode::LoadConst16, &cid.to_le_bytes());
293        }
294        Ok(())
295    }
296    
297    fn emit_load_error(&mut self, symbol: Option<&DebugSymbol>, error: ErrorKind, message: &str) -> CompileResult<()> {
298        let cid = self.builder_mut().get_or_insert_error(error, message)?;
299        
300        if cid <= u8::MAX.into() {
301            self.emit_instr_byte(symbol, OpCode::LoadConst, u8::try_from(cid).unwrap());
302        } else {
303            self.emit_instr_data(symbol, OpCode::LoadConst16, &cid.to_le_bytes());
304        }
305        Ok(())
306    }
307    
308    fn make_function(&mut self, function: UnloadedFunction) {
309        self.builder_mut().insert_function(function)
310    }
311    
312    fn emit_load_function(&mut self, symbol: Option<&DebugSymbol>, fun_id: FunctionID) {
313        if fun_id <= u8::MAX.into() {
314            self.emit_instr_byte(symbol, OpCode::LoadFunction, u8::try_from(fun_id).unwrap());
315        } else {
316            self.emit_instr_data(symbol, OpCode::LoadFunction16, &fun_id.to_le_bytes());
317        }
318    }
319}
320
321///////// Jumps /////////
322impl CodeGenerator<'_> {
323    fn emit_jump_instr(&mut self, symbol: Option<&DebugSymbol>, jump: Jump, target: usize) -> CompileResult<()> {
324        let jump_site = self.current_offset();
325        let guess_width = jump.dummy_width();  // guess the width of the jump instruction
326        
327        let mut jump_offset = Self::calc_jump_offset(jump_site + guess_width, target)?;
328        let mut jump_opcode = get_jump_opcode(jump, jump_offset);
329        
330        if guess_width != jump_opcode.instr_len() {
331            // guessed wrong, need to recalc offset with new width
332            let new_width = jump_opcode.instr_len();
333            let new_offset = Self::calc_jump_offset(jump_site + new_width, target)?;
334            let new_opcode = get_jump_opcode(jump, new_offset);
335            
336            // if we *still* don't have the right width, just abort
337            if new_width != new_opcode.instr_len() {
338                return Err("could not calculate jump offset".into());
339            }
340            
341            jump_offset = new_offset;
342            jump_opcode = new_opcode;
343        }
344        
345        match jump_offset {
346            JumpOffset::Short(offset) => self.emit_instr_data(symbol, jump_opcode, &offset.to_le_bytes()),
347            JumpOffset::Long(offset)  => self.emit_instr_data(symbol, jump_opcode, &offset.to_le_bytes()),
348        }
349        Ok(())
350    }
351    
352    fn emit_dummy_jump(&mut self, symbol: Option<&DebugSymbol>, jump: Jump) -> JumpSite {
353        let offset = self.current_offset();
354        let jump_site = JumpSite {
355            jump, offset,
356            width: jump.dummy_width(),
357        };
358        
359        self.emit_dummy_instr(symbol, jump_site.width);
360        
361        jump_site
362    }
363    
364    fn patch_jump_instr(&mut self, jump: &JumpSite, target: usize) -> CompileResult<()> {
365        let jump_type = jump.jump;
366        let jump_site = jump.offset;
367        let dummy_width = jump.width;
368        
369        let mut jump_offset = Self::calc_jump_offset(jump_site + dummy_width, target)?;
370        let mut jump_opcode = get_jump_opcode(jump_type, jump_offset);
371        
372        if dummy_width != jump_opcode.instr_len() {
373            // need to recalculate offset with the new width
374            let new_width = jump_opcode.instr_len();
375            let new_offset = Self::calc_jump_offset(jump_site + new_width, target)?;
376            let new_opcode = get_jump_opcode(jump_type, new_offset);
377            
378            // if we *still* don't have the right width, just abort
379            if new_width != new_opcode.instr_len() {
380                return Err("could not calculate jump offset".into());
381            }
382            
383            jump_offset = new_offset;
384            jump_opcode = new_opcode;
385            self.chunk_mut().resize_patch(jump_site, dummy_width, new_width);
386        }
387        
388        match jump_offset {
389            JumpOffset::Short(offset) => self.patch_instr_data(jump_site, jump_opcode, &offset.to_le_bytes()),
390            JumpOffset::Long(offset)  => self.patch_instr_data(jump_site, jump_opcode, &offset.to_le_bytes()),
391        }
392        Ok(())
393    }
394    
395    // Expects the *end* offset of the jump instruction
396    fn calc_jump_offset(jump_end_offset: usize, target: usize) -> CompileResult<JumpOffset> {
397        // inefficent, but this is compile time so that's okay
398        let target = i128::try_from(target).unwrap();
399        let jump_site = i128::try_from(jump_end_offset).unwrap();
400        
401        if let Ok(offset) = i16::try_from(target - jump_site) {
402            return Ok(JumpOffset::Short(offset));
403        }
404        
405        if let Ok(offset) = i32::try_from(target - jump_site) {
406            return Ok(JumpOffset::Long(offset));
407        }
408        
409        Err("could not calculate jump offset".into())
410    }
411}
412
413///////// Scopes /////////
414
415// container for data needed to drop a scope
416struct ScopeDrop {
417    tag: ScopeTag,
418    locals: usize,
419    close_upvals: Vec<LocalIndex>,
420}
421
422impl From<&Scope> for ScopeDrop {
423    fn from(scope: &Scope) -> Self {
424        ScopeDrop {
425            tag: scope.tag(),
426            locals: scope.locals().len(),
427            close_upvals: scope.locals().iter()
428                .filter_map(|local| if local.captured() { Some(local.index()) } else { None })
429                .collect(),
430        }
431    }
432}
433
434impl CodeGenerator<'_> {
435    fn emit_begin_scope(&mut self, symbol: Option<&DebugSymbol>, label: Option<&Label>, tag: ScopeTag) -> &mut Scope {
436        let chunk_id = self.chunk_id;
437        self.scopes_mut().push_scope(symbol, label.copied(), tag)
438    }
439    
440    fn emit_end_scope(&mut self) -> Scope {
441        let scope = self.scopes_mut().pop_scope();
442        self.emit_scope_drop(scope.debug_symbol(), &(&scope).into());
443        scope
444    }
445    
446    fn patch_break_sites(&mut self, scope: &Scope, break_target: usize) -> CompileResult<()> {
447        for break_site in scope.break_sites().iter() {
448            self.patch_jump_instr(break_site, break_target)?;
449        }
450        Ok(())
451    }
452    
453    fn patch_continue_sites(&mut self, scope: &Scope, continue_target: usize) -> CompileResult<()> {
454        for continue_site in scope.continue_sites().iter() {
455            self.patch_jump_instr(continue_site, continue_target)?;
456        }
457        Ok(())
458    }
459    
460    fn emit_scope_drop(&mut self, symbol: Option<&DebugSymbol>, scope: &ScopeDrop) {
461        // close all upvalues
462        for local_index in scope.close_upvals.iter() {
463            self.emit_close_upvalue(symbol, *local_index);
464        }
465        
466        // discard all the locals from the stack
467        let mut discard = scope.locals;
468        while discard > u8::MAX.into() {
469            self.emit_instr_byte(symbol, OpCode::DropLocals, u8::MAX);
470            discard -= usize::from(u8::MAX);
471        }
472        
473        if discard > 0 {
474            self.emit_instr_byte(symbol, OpCode::DropLocals, u8::try_from(discard).unwrap());
475        }
476    }
477    
478    // If the local name cannot be found, no instructions are emitted and None is returned
479    fn try_emit_load_local(&mut self, symbol: Option<&DebugSymbol>, name: &LocalName) -> Option<u16> {
480        if let Some(index) = self.scopes().resolve_local(name).map(|local| local.index()) {
481            self.emit_load_local_index(symbol, index);
482            Some(index)
483        } else{
484            None
485        }
486    }
487    
488    fn emit_load_local_index(&mut self, symbol: Option<&DebugSymbol>, index: LocalIndex) {
489        if let Ok(index) = u8::try_from(index) {
490            self.emit_instr_byte(symbol, OpCode::LoadLocal, index);
491        } else {
492            self.emit_instr_data(symbol, OpCode::LoadLocal16, &index.to_le_bytes());
493        }
494    }
495    
496    fn try_emit_load_upval(&mut self, symbol: Option<&DebugSymbol>, name: &LocalName) -> CompileResult<Option<u16>> {
497        if let Some(index) = self.scopes_mut().resolve_or_create_upval(name)?.map(|upval| upval.index()) {
498            if let Ok(index) = u8::try_from(index) {
499                self.emit_instr_byte(symbol, OpCode::LoadUpvalue, index);
500            } else {
501                self.emit_instr_data(symbol, OpCode::LoadUpvalue16, &index.to_le_bytes());
502            }
503            
504            Ok(Some(index))
505        } else {
506            Ok(None)
507        }
508    }
509    
510    fn emit_close_upvalue(&mut self, symbol: Option<&DebugSymbol>, index: LocalIndex) {
511        if let Ok(index) = u8::try_from(index) {
512            self.emit_instr_byte(symbol, OpCode::CloseUpvalue, index);
513        } else {
514            self.emit_instr_data(symbol, OpCode::CloseUpvalue16, &index.to_le_bytes());
515        }
516    }
517    
518    fn emit_create_temporary(&mut self, symbol: Option<&DebugSymbol>, access: Access) -> CompileResult<LocalIndex> {
519        debug_assert!(self.scopes().is_temporary_scope());
520        
521        match self.scopes_mut().insert_local(access, LocalName::Anonymous)? {
522            InsertLocal::CreateNew(local_index) => {
523                self.emit_instr(symbol, OpCode::InsertLocal);
524                Ok(local_index)
525            }
526            
527            InsertLocal::HideExisting(local_index) => {
528                self.emit_assign_local(symbol, local_index);
529                Ok(local_index)
530            }
531        }
532    }
533}
534
535
536///////// Statements /////////
537impl CodeGenerator<'_> {
538    fn compile_stmt_with_symbol(&mut self, stmt: &StmtMeta) -> CompileResult<()> {
539        let symbol = stmt.debug_symbol();
540        self.compile_stmt(Some(symbol), stmt.variant())
541            .map_err(|err| err.with_symbol(*symbol))
542    }
543    
544    fn compile_stmt(&mut self, symbol: Option<&DebugSymbol>, stmt: &Stmt) -> CompileResult<()> {
545        match stmt {
546            Stmt::Loop { label, body } => self.compile_loop(symbol, label.as_ref(), body)?,
547            
548            Stmt::WhileLoop { label, condition, body } => self.compile_while_loop(symbol, label.as_ref(), condition, body)?,
549            
550            Stmt::ForLoop { label, lvalue, iter, body } => self.compile_for_loop(symbol, label.as_ref(), lvalue, iter, body)?,
551            
552            Stmt::Assert(expr) => {
553                self.compile_expr(symbol, expr)?;
554                self.emit_instr(symbol, OpCode::Assert);
555                self.emit_instr(symbol, OpCode::Pop);
556            }
557            
558            Stmt::Expression(expr) => {
559                self.compile_expr(symbol, expr)?;
560                self.emit_instr(symbol, OpCode::Pop);
561            },
562        }
563        Ok(())
564    }
565    
566    fn compile_stmt_list(&mut self, stmt_list: &StmtList) -> CompileResult<()> {
567        // compile stmt suite
568        for stmt in stmt_list.iter() {
569            self.compile_stmt_with_symbol(stmt)?;
570        }
571        
572        Ok(())
573    }
574    
575    // compile a statment list that will not evaluate to a value
576    fn compile_stmt_block(&mut self, stmt_list: &StmtList) -> CompileResult<()> {
577        self.compile_stmt_list(stmt_list)?;
578        self.compile_end_control(stmt_list)?;
579        Ok(())
580    }
581    
582    fn compile_expr_block(&mut self, symbol: Option<&DebugSymbol>, suite: &ExprBlock) -> CompileResult<()> {
583        let stmt_list = suite.stmt_list();
584        self.compile_stmt_list(stmt_list)?;
585        
586        if stmt_list.end_control().is_none() {
587            // result expression
588            if let Some(expr) = suite.result() {
589                self.compile_expr_with_symbol(expr)?;
590            } else {
591                self.emit_instr(symbol, OpCode::Nil); // implicit nil
592            }
593        }
594        
595        self.compile_end_control(stmt_list)?;
596        
597        Ok(())
598    }
599    
600    fn compile_end_control(&mut self, stmt_list: &StmtList) -> CompileResult<()> {
601        // handle control flow
602        if let Some(control) = stmt_list.end_control() {
603            let result = self.compile_control_flow(control);
604            if let Some(symbol) = control.debug_symbol() {
605                result.map_err(|error| error.with_symbol(*symbol))?;
606            } else {
607                result?;
608            }
609        }
610        Ok(())
611    }
612    
613    fn compile_control_flow(&mut self, control_flow: &ControlFlow) -> CompileResult<()> {
614        match control_flow {
615            ControlFlow::Continue { label, symbol } => self.compile_continue_control(
616                symbol.as_ref(), label.as_ref()
617            )?,
618            
619            ControlFlow::Break { label, expr, symbol } => self.compile_break_control(
620                symbol.as_ref(), label.as_ref(), expr.as_deref()
621            )?,
622            
623            ControlFlow::Return { expr, symbol } => {
624                let symbol = symbol.as_ref();
625                if let Some(expr) = expr {
626                    self.compile_expr(symbol, expr)?;
627                } else {
628                    self.emit_instr(symbol, OpCode::Nil);
629                }
630                self.emit_instr(symbol, OpCode::Return);
631            }
632        }
633        Ok(())
634    }
635    
636    fn compile_break_control(&mut self, symbol: Option<&DebugSymbol>, label: Option<&Label>, expr: Option<&Expr>) -> CompileResult<()> {
637        // find the target scope
638        let target_depth = match self.scopes().resolve_control_flow(ControlFlowTarget::Break(label.copied())) {
639            Some(scope) => scope.depth(),
640            None => {
641                let message =
642                    if label.is_some() { "can't find loop or block with matching label for \"break\"" }
643                    else { "\"break\" outside of loop or block" };
644                return Err(message.into());
645            },
646        };
647        
648        // drop all scopes up to and including the target
649        let scope_drop: Vec<ScopeDrop> = self.scopes().iter_scopes()
650            .take_while(|scope| scope.depth() >= target_depth)
651            .map(ScopeDrop::from)
652            .collect();
653        
654        let (target, through_scopes) = scope_drop.split_last().unwrap();
655        for scope in through_scopes.iter() {
656            self.emit_scope_drop(symbol, scope);
657            
658            // expression blocks leave their value on the stack
659            // (this is helped by the fact that break/contine must come last in a list of statements)
660            // so if we break through an expression block we need to pop its value
661            if scope.tag.is_expr_block() {
662                self.emit_instr(symbol, OpCode::Pop);
663            }
664        }
665        self.emit_scope_drop(symbol, target); // drop target scope
666        
667        // if breaking from an expression block, emit the expression value before jumping
668        if target.tag.is_expr_block() {
669            if let Some(expr) = expr {
670                self.compile_expr(symbol, expr)?;
671            } else {
672                self.emit_instr(symbol, OpCode::Nil);
673            }
674        } else if expr.is_some() {
675            return Err("\"break\" with value outside of block expression".into())
676        }
677        
678        // emit jump site, register with scope
679        let break_site = self.emit_dummy_jump(symbol, Jump::Uncond);
680        
681        let target_scope = self.scopes_mut().iter_scopes_mut()
682            .find(|scope| scope.depth() == target_depth)
683            .unwrap();
684        target_scope.register_break(break_site);
685        
686        Ok(())
687        
688    }
689    
690    fn compile_continue_control(&mut self, symbol: Option<&DebugSymbol>, label: Option<&Label>) -> CompileResult<()> {
691        // find the target scope
692        let target_depth = match self.scopes().resolve_control_flow(ControlFlowTarget::Continue(label.copied())) {
693            Some(scope) => scope.depth(),
694            None => {
695                let message =
696                    if label.is_some() { "can't find loop with matching label for \"continue\"" }
697                    else { "\"continue\" outside of loop" };
698                return Err(message.into());
699            }
700        };
701        
702        // drop all scopes up to and including the target
703        let scope_drop: Vec<ScopeDrop> = self.scopes().iter_scopes()
704            .take_while(|scope| scope.depth() >= target_depth)
705            .map(ScopeDrop::from)
706            .collect();
707        
708        for scope in scope_drop.iter() {
709            self.emit_scope_drop(symbol, scope);
710            
711            // expression blocks leave their value on the stack
712            // (this is helped by the fact that break/contine must come last in a list of statements)
713            // so if we jump out of an expression block we need to pop its value
714            if scope.tag.is_expr_block() {
715                self.emit_instr(symbol, OpCode::Pop);
716            }
717        }
718        
719        // emit jump site, register with scope
720        let continue_site = self.emit_dummy_jump(symbol, Jump::Uncond);
721        
722        let target_scope = self.scopes_mut().iter_scopes_mut()
723            .find(|scope| scope.depth() == target_depth)
724            .unwrap();
725        target_scope.register_continue(continue_site);
726        
727        Ok(())
728    }
729    
730    fn compile_loop(&mut self, symbol: Option<&DebugSymbol>, label: Option<&Label>, body: &StmtList) -> CompileResult<()> {
731        
732        let loop_target = self.current_offset();
733        
734        self.emit_begin_scope(symbol, label, ScopeTag::Loop);
735        
736        self.compile_stmt_block(body)?;
737        let loop_scope = self.emit_end_scope();
738        
739        self.emit_jump_instr(symbol, Jump::Uncond, loop_target)?;
740        
741        // finalize scope
742        let break_target = self.current_offset();
743        self.patch_break_sites(&loop_scope, break_target)?;
744        self.patch_continue_sites(&loop_scope, loop_target)?;
745        
746        Ok(())
747    }
748    
749    fn compile_while_loop(&mut self, symbol: Option<&DebugSymbol>, label: Option<&Label>, condition: &Expr, body: &StmtList) -> CompileResult<()> {
750        
751        // first iteration conditional jump
752        let continue_target = self.current_offset();
753        self.compile_expr(symbol, condition)?;
754        
755        let end_jump_site = self.emit_dummy_jump(symbol, Jump::PopIfFalse);
756        
757        let loop_target = self.current_offset();
758        
759        self.emit_begin_scope(symbol, label, ScopeTag::Loop);
760        self.compile_stmt_block(body)?;
761        let loop_scope = self.emit_end_scope();
762        
763        // rest iteration conditional jump
764        self.compile_expr(symbol, condition)?;
765        self.emit_jump_instr(symbol, Jump::PopIfTrue, loop_target)?;
766        
767        self.patch_jump_instr(&end_jump_site, self.current_offset())?;
768        
769        // finalize scope
770        let break_target = self.current_offset();
771        self.patch_break_sites(&loop_scope, break_target)?;
772        self.patch_continue_sites(&loop_scope, continue_target)?;
773        
774        Ok(())
775    }
776    
777    fn compile_for_loop(&mut self, symbol: Option<&DebugSymbol>, label: Option<&Label>, lvalue: &LValue, iter: &Expr, body: &StmtList) -> CompileResult<()> {
778        
779        self.emit_begin_scope(symbol, label, ScopeTag::Loop);
780        
781        // initialize iterator
782        self.compile_expr(symbol, iter)?;
783        self.emit_instr(symbol, OpCode::IterInit);
784        
785        // first iteration conditional jump
786        let continue_target = self.current_offset();
787        let end_jump_site = self.emit_dummy_jump(symbol, Jump::IfFalse);
788        
789        let loop_target = self.current_offset();
790        
791        // advance iterator and assign value
792        // default to "let" for loop variables (unlike normal assignment, which defaults to "local")
793        self.emit_instr(symbol, OpCode::IterNext);
794        self.compile_assignment(symbol, AssignType::DeclImmutable, lvalue)?; 
795        self.emit_instr(symbol, OpCode::Pop);
796        
797        // compile body
798        self.compile_stmt_block(body)?;
799        let loop_scope = self.emit_end_scope();
800        
801        // rest iteration conditional jump
802        // should have just [ ... iter state[N] ] on the stack here
803        self.emit_jump_instr(symbol, Jump::IfTrue, loop_target)?;
804        
805        let break_target = self.current_offset();
806        self.emit_instr_byte(symbol, OpCode::Drop, 2); // drop [ iter state ]
807        
808        // finalize scope
809        self.patch_jump_instr(&end_jump_site, break_target)?;
810        self.patch_break_sites(&loop_scope, break_target)?;
811        self.patch_continue_sites(&loop_scope, continue_target)?;
812        
813        Ok(())
814    }
815}
816
817///////// Expressions /////////
818
819// helper to indicate if the length of a sequence is statically known or if it is stored in a local
820enum Unpack {
821    Empty,
822    Static(IntType),
823    Dynamic,
824}
825
826impl CodeGenerator<'_> {
827    fn compile_expr_with_symbol(&mut self, expr: &ExprMeta) -> CompileResult<()> {
828        let symbol = expr.debug_symbol();
829        self.compile_expr(Some(symbol), expr.variant())
830            .map_err(|err| err.with_symbol(*symbol))
831    }
832    
833    fn compile_expr(&mut self, symbol: Option<&DebugSymbol>, expr: &Expr) -> CompileResult<()> {
834        match expr {
835            Expr::Atom(atom) => self.compile_atom(symbol, atom)?,
836            
837            Expr::Primary(primary) => self.compile_primary(symbol, primary)?,
838            
839            Expr::UnaryOp(op, expr) => self.compile_unary_op(symbol, *op, expr)?,
840            
841            Expr::BinaryOp(op, exprs) => {
842                let (lhs, rhs) = &**exprs;
843                self.compile_binary_op(symbol, *op, lhs, rhs)?;
844            },
845            
846            Expr::Assignment(assign) => {
847                if let Some(op) = assign.op {
848                    self.compile_update_assignment(symbol, op, assign.modifier, &assign.lhs, &assign.rhs)?;
849                } else {
850                    self.compile_expr(symbol, &assign.rhs)?;
851                    self.compile_assignment(symbol, assign.modifier, &assign.lhs)?;
852                }
853            },
854            
855            Expr::Tuple(items) => self.compile_tuple(symbol, items)?,
856            
857            // unpacking is only allowed in invocation, tuple literals, and by itself in parentheses
858            // note: assignment uses *packing*, not unpacking, which is the LValue dual of packing.
859            Expr::Unpack(Some(..)) => return Err("unpack expression must be enclosed in parentheses".into()),
860            Expr::Unpack(None) => return Err("\"...\" is not allowed here".into()),
861            
862            Expr::Block { label, suite } => self.compile_block_expression(symbol, label.as_ref(), suite)?,
863            Expr::IfExpr { branches, else_clause } => self.compile_if_expression(symbol, branches, else_clause.as_ref().map(|expr| &**expr))?,
864            
865            Expr::FunctionDef(fundef) => self.compile_function_def(symbol, fundef)?,
866        }
867        Ok(())
868    }
869    
870    fn compile_tuple(&mut self, symbol: Option<&DebugSymbol>, expr_list: &[ExprMeta]) -> CompileResult<()> {
871        match self.compile_unpack_sequence(symbol, expr_list)? {
872            Unpack::Empty => self.emit_instr(symbol, OpCode::Empty),
873            
874            Unpack::Static(len) => {
875                if let Ok(len) = u8::try_from(len) {
876                    self.emit_instr_byte(symbol, OpCode::Tuple, len);
877                } else {
878                    self.compile_integer(symbol, len)?;
879                    self.emit_instr(symbol, OpCode::TupleN);
880                }
881            }
882            
883            Unpack::Dynamic => {
884                self.emit_instr(symbol, OpCode::TupleN);
885            }
886        }
887        Ok(())
888    }
889    
890    // compiles to a sequence of values
891    fn compile_unpack_sequence(&mut self, symbol: Option<&DebugSymbol>, seq: &[ExprMeta]) -> CompileResult<Unpack> {
892        if seq.is_empty() {
893            return Ok(Unpack::Empty)
894        }
895        
896        let (last, rest) = seq.split_last().unwrap();
897        
898        let mut static_len: IntType = 0;
899        let mut unpack_len = None; // accumulate length in an anonymous temporary if needed
900        
901        for expr in rest.iter() {
902            match expr.variant() {
903                Expr::Unpack(None) => return Err("need a value to unpack".into()),
904                
905                Expr::Unpack(Some(unpack)) => {
906                    let symbol = Some(expr.debug_symbol());
907                    
908                    self.compile_expr(symbol, unpack)?;
909                    self.emit_instr(symbol, OpCode::IterInit);
910                    self.emit_instr(symbol, OpCode::IterUnpack);
911                    
912                    // store unpack len in accumulator
913                    if let Some(local_index) = unpack_len {
914                        self.emit_load_local_index(symbol, local_index);
915                        self.emit_instr(symbol, OpCode::Add);
916                        self.emit_assign_local(symbol, local_index);
917                    } else {
918                        self.emit_begin_scope(symbol, None, ScopeTag::Temporary);
919                        let local_index = self.emit_create_temporary(symbol, Access::ReadWrite)?;
920                        unpack_len = Some(local_index);
921                    }
922                    
923                    self.emit_instr(symbol, OpCode::Pop);
924                }
925                
926                _ => {
927                    self.compile_expr_with_symbol(expr)?;
928                    static_len = static_len.checked_add(1)
929                        .ok_or("unpack length limit exceeded")?;
930                }
931            }
932        }
933        
934        // if the last item is an unpack expression, it does not need to use the local accumulator
935        // TODO should there be a dedicated accumulator register?
936        match last.variant() {
937            Expr::Unpack(None) => return Err("need a value to unpack".into()),
938            
939            Expr::Unpack(Some(unpack)) => {
940                let symbol = Some(last.debug_symbol());
941                
942                self.compile_expr(symbol, unpack)?;
943                self.emit_instr(symbol, OpCode::IterInit);
944                self.emit_instr(symbol, OpCode::IterUnpack);
945                
946                if let Some(local_index) = unpack_len {
947                    self.emit_load_local_index(symbol, local_index);
948                    self.emit_instr(symbol, OpCode::Add);
949                    
950                    debug_assert!(self.scopes().is_temporary_scope());
951                    self.emit_end_scope();
952                }
953                
954                if static_len > 0 {
955                    self.compile_integer(symbol, static_len)?;
956                    self.emit_instr(symbol, OpCode::Add);
957                }
958                
959                Ok(Unpack::Dynamic)
960            }
961            
962            _ => {
963                self.compile_expr_with_symbol(last)?;
964                static_len = static_len.checked_add(1)
965                    .ok_or("unpack length limit exceeded")?;
966                    
967                if let Some(local_index) = unpack_len {
968                    self.emit_load_local_index(symbol, local_index);
969                    
970                    if static_len > 0 {
971                        self.compile_integer(symbol, static_len)?;
972                        self.emit_instr(symbol, OpCode::Add);
973                    }
974                    
975                    debug_assert!(self.scopes().is_temporary_scope());
976                    self.emit_end_scope();
977                    
978                    return Ok(Unpack::Dynamic)
979                }
980                
981                Ok(Unpack::Static(static_len))
982            }
983        }
984    }
985
986    fn compile_atom(&mut self, symbol: Option<&DebugSymbol>, atom: &Atom) -> CompileResult<()> {
987        match atom {
988            Atom::Nil => self.emit_instr(symbol, OpCode::Nil),
989            Atom::EmptyTuple => self.emit_instr(symbol, OpCode::Empty),
990            Atom::BooleanLiteral(true) => self.emit_instr(symbol, OpCode::True),
991            Atom::BooleanLiteral(false) => self.emit_instr(symbol, OpCode::False),
992            
993            Atom::IntegerLiteral(value) => self.compile_integer(symbol, *value)?,
994            Atom::FloatLiteral(value) => self.compile_float(symbol, *value)?,
995            
996            Atom::StringLiteral(value) => self.emit_load_const(symbol, Constant::from(*value))?,
997            Atom::Identifier(name) => self.compile_name_lookup(symbol, name)?,
998            
999            // Atom::Self_ => unimplemented!(),
1000            // Atom::Super => unimplemented!(),
1001            
1002            Atom::Group { modifier, inner } => {
1003                // modifiers are not allowed outside of assignment
1004                if let Some(modifier) = modifier {
1005                    return Err("assignment modifiers are not allowed outside of an assignment expression".into())
1006                }
1007                
1008                match &**inner {
1009                    // tuple constructor
1010                    Expr::Unpack(None) => return Err("need a value to unpack".into()),
1011                    Expr::Unpack(Some(iter)) => {
1012                        self.compile_expr(symbol, iter)?;
1013                        self.emit_instr(symbol, OpCode::IterInit);
1014                        self.emit_instr(symbol, OpCode::IterUnpack);
1015                        self.emit_instr(symbol, OpCode::TupleN);
1016                    }
1017                    
1018                    // parenthesized group
1019                    expr => self.compile_expr(symbol, inner)?,
1020                }
1021                
1022            },
1023        }
1024        Ok(())
1025    }
1026    
1027    fn compile_integer(&mut self, symbol: Option<&DebugSymbol>, value: IntType) -> CompileResult<()> {
1028        if let Ok(value) = u8::try_from(value) {
1029            self.emit_instr_byte(symbol, OpCode::UInt8, value);
1030        } else if let Ok(value) = i8::try_from(value) {
1031            self.emit_instr_byte(symbol, OpCode::Int8, value.to_le_bytes()[0]);
1032        } else if let Ok(value) = i16::try_from(value) {
1033            self.emit_instr_data(symbol, OpCode::Int16, &value.to_le_bytes());
1034        } else {
1035            self.emit_load_const(symbol, Constant::from(value))?;
1036        }
1037        Ok(())
1038    }
1039    
1040    fn compile_float(&mut self, symbol: Option<&DebugSymbol>, value: FloatType) -> CompileResult<()> {
1041        self.emit_load_const(symbol, Constant::from(value))
1042    }
1043    
1044    fn compile_name_lookup(&mut self, symbol: Option<&DebugSymbol>, name: &InternSymbol) -> CompileResult<()> {
1045        let local_name = LocalName::Symbol(*name);
1046        
1047        // Try loading a Local variable
1048        if self.try_emit_load_local(symbol, &local_name).is_some() {
1049            return Ok(());
1050        }
1051        
1052        // Next, try loading an upvalue
1053        if self.try_emit_load_upval(symbol, &local_name)?.is_some() {
1054            return Ok(());
1055        }
1056        
1057        // Otherwise, it must be a Global variable
1058        self.emit_load_const(symbol, Constant::from(*name))?;
1059        self.emit_instr(symbol, OpCode::LoadGlobal);
1060        Ok(())
1061    }
1062    
1063    fn compile_primary(&mut self, symbol: Option<&DebugSymbol>, primary: &Primary) -> CompileResult<()> {
1064        self.compile_atom(symbol, primary.atom())?;
1065        
1066        for item in primary.path().iter() {
1067            match item {
1068                AccessItem::Attribute(name) => unimplemented!(),
1069                AccessItem::Index(index) => unimplemented!(),
1070                AccessItem::Invoke(args) => self.compile_invocation(symbol, args)?,
1071            }
1072        }
1073        
1074        Ok(())
1075    }
1076    
1077    fn compile_invocation(&mut self, symbol: Option<&DebugSymbol>, args: &[ExprMeta]) -> CompileResult<()> {
1078        // prepare argument list:
1079        // [ callobj arg[0] ... arg[n] nargs ] => [ ret_value ] 
1080
1081        // process argument unpacking
1082        match self.compile_unpack_sequence(symbol, args)? {
1083            Unpack::Empty => self.emit_instr_byte(symbol, OpCode::UInt8, 0),
1084            Unpack::Static(len) => self.compile_integer(symbol, len)?,
1085            Unpack::Dynamic => { } // nothing to do
1086        }
1087
1088        self.emit_instr(symbol, OpCode::Call);
1089
1090        Ok(())
1091    }
1092    
1093    fn compile_unary_op(&mut self, symbol: Option<&DebugSymbol>, op: UnaryOp, expr: &Expr) -> CompileResult<()> {
1094        self.compile_expr(symbol, expr)?;
1095        match op {
1096            UnaryOp::Neg => self.emit_instr(symbol, OpCode::Neg),
1097            UnaryOp::Pos => self.emit_instr(symbol, OpCode::Pos),
1098            UnaryOp::Inv => self.emit_instr(symbol, OpCode::Inv),
1099            UnaryOp::Not => self.emit_instr(symbol, OpCode::Not),
1100        };
1101        Ok(())
1102    }
1103    
1104    fn compile_binary_op(&mut self, symbol: Option<&DebugSymbol>, op: BinaryOp, lhs: &Expr, rhs: &Expr) -> CompileResult<()> {
1105        
1106        if matches!(op, BinaryOp::And) {
1107            return self.compile_shortcircuit_and(symbol, lhs, rhs);
1108        }
1109        if matches!(op, BinaryOp::Or) {
1110            return self.compile_shortcircuit_or(symbol, lhs, rhs);
1111        }
1112        
1113        self.compile_expr(symbol, lhs)?;
1114        self.compile_expr(symbol, rhs)?;
1115        self.emit_binary_op(symbol, op);
1116        
1117        Ok(())
1118    }
1119    
1120    fn emit_binary_op(&mut self, symbol: Option<&DebugSymbol>, op: BinaryOp) {
1121        match op {
1122            BinaryOp::And | BinaryOp::Or => unreachable!(),
1123            
1124            BinaryOp::Mul => self.emit_instr(symbol, OpCode::Mul),
1125            BinaryOp::Div => self.emit_instr(symbol, OpCode::Div),
1126            BinaryOp::Mod => self.emit_instr(symbol, OpCode::Mod),
1127            BinaryOp::Add => self.emit_instr(symbol, OpCode::Add),
1128            BinaryOp::Sub => self.emit_instr(symbol, OpCode::Sub),
1129            
1130            BinaryOp::BitAnd => self.emit_instr(symbol, OpCode::And),
1131            BinaryOp::BitXor => self.emit_instr(symbol, OpCode::Xor),
1132            BinaryOp::BitOr  => self.emit_instr(symbol, OpCode::Or),
1133            
1134            BinaryOp::LShift  => self.emit_instr(symbol, OpCode::Shl),
1135            BinaryOp::RShift => self.emit_instr(symbol, OpCode::Shr),
1136            
1137            BinaryOp::LT => self.emit_instr(symbol, OpCode::LT),
1138            BinaryOp::GT => self.emit_instr(symbol, OpCode::GT),
1139            BinaryOp::LE => self.emit_instr(symbol, OpCode::LE),
1140            BinaryOp::GE => self.emit_instr(symbol, OpCode::GE),
1141            BinaryOp::EQ => self.emit_instr(symbol, OpCode::EQ),
1142            BinaryOp::NE => self.emit_instr(symbol, OpCode::NE),
1143        };
1144    }
1145}
1146
1147///////// Declarations and Assignments /////////
1148impl CodeGenerator<'_> {
1149
1150    
1151    fn compile_update_assignment(&mut self, symbol: Option<&DebugSymbol>, op: BinaryOp, assign: AssignType, lhs: &LValue, rhs: &Expr) -> CompileResult<()> {
1152        
1153        let local_only = match assign {
1154            AssignType::AssignLocal => true,
1155            AssignType::AssignNonLocal => false,
1156            
1157            AssignType::DeclImmutable | AssignType::DeclMutable
1158                => return Err("update-assignment is invalid when declaring a variable".into()),
1159        };
1160        
1161        // TODO suport Attribute and Index LValues as well
1162        match lhs {
1163            LValue::Identifier(name) => {
1164                self.compile_name_lookup(symbol, name)?;
1165                self.compile_expr(symbol, rhs)?;
1166                self.emit_binary_op(symbol, op);
1167                
1168                self.compile_assign_identifier(symbol, name, local_only)
1169            },
1170            
1171            LValue::Attribute(target) => unimplemented!(),
1172            
1173            LValue::Index(target) => unimplemented!(),
1174            
1175            LValue::Tuple {..} | LValue::Pack(..)
1176                => Err("can't update-assign to this".into()),
1177            
1178            LValue::Modifier {..} => unreachable!(),
1179        }
1180    }
1181    
1182    fn compile_assignment(&mut self, symbol: Option<&DebugSymbol>, mut assign: AssignType, mut lhs: &LValue) -> CompileResult<()> {
1183        
1184        while let LValue::Modifier { modifier, lvalue } = lhs {
1185            assign = *modifier;
1186            lhs = lvalue;
1187        }
1188        
1189        match lhs {
1190            LValue::Tuple(items) => self.compile_assign_tuple(symbol, assign, items),
1191            
1192            LValue::Pack(..) => {
1193                let item = std::slice::from_ref(lhs);
1194                self.compile_assign_tuple(symbol, assign, item)
1195            }
1196            
1197            lhs => {
1198                match assign {
1199                    AssignType::AssignLocal => self.compile_assign_variable(symbol, lhs, false),
1200                    AssignType::AssignNonLocal => self.compile_assign_variable(symbol, lhs, true),
1201                    AssignType::DeclImmutable => self.compile_decl_variable(symbol, Access::ReadOnly, lhs),
1202                    AssignType::DeclMutable => self.compile_decl_variable(symbol, Access::ReadWrite, lhs),
1203                }
1204            }
1205        }
1206    }
1207
1208    fn compile_decl_variable(&mut self, symbol: Option<&DebugSymbol>, access: Access, lhs: &LValue) -> CompileResult<()> {
1209        match lhs {
1210            LValue::Identifier(name) => if self.scopes().is_global_scope() {
1211                self.compile_decl_global_name(symbol, access, *name)
1212            } else {
1213                self.compile_decl_local_name(symbol, access, *name)
1214            },
1215            
1216            LValue::Tuple {..} => unreachable!(),
1217            LValue::Modifier {..} => unreachable!(),
1218            
1219            _ => Err("not a variable name".into()),
1220        }
1221    }
1222    
1223    fn compile_decl_global_name(&mut self, symbol: Option<&DebugSymbol>, access: Access, name: InternSymbol) -> CompileResult<()> {
1224        
1225        self.emit_load_const(symbol, Constant::from(name))?;
1226        match access {
1227            Access::ReadOnly => self.emit_instr(symbol, OpCode::InsertGlobal),
1228            Access::ReadWrite => self.emit_instr(symbol, OpCode::InsertGlobalMut),
1229        }
1230        Ok(())
1231    }
1232    
1233    fn compile_decl_local_name(&mut self, symbol: Option<&DebugSymbol>, access: Access, name: InternSymbol) -> CompileResult<()> {
1234        
1235        match self.scopes_mut().insert_local(access, LocalName::Symbol(name))? {
1236            InsertLocal::CreateNew(..) => 
1237                self.emit_instr(symbol, OpCode::InsertLocal),
1238            
1239            InsertLocal::HideExisting(local_index) =>
1240                self.emit_assign_local(symbol, local_index),
1241        }
1242        
1243        Ok(())
1244    }
1245    
1246    fn compile_assign_variable(&mut self, symbol: Option<&DebugSymbol>, lhs: &LValue, allow_nonlocal: bool) -> CompileResult<()> {
1247        
1248        match lhs {
1249            LValue::Identifier(name) => self.compile_assign_identifier(symbol, name, allow_nonlocal),
1250            
1251            LValue::Attribute(target) => unimplemented!(),
1252            
1253            LValue::Index(target) => unimplemented!(),
1254            
1255            _ => panic!("invalid assignment target"),
1256        }
1257    }
1258    
1259    fn compile_assign_identifier(&mut self, symbol: Option<&DebugSymbol>, name: &InternSymbol, allow_nonlocal: bool) -> CompileResult<()> {
1260        
1261        // Generate assignment
1262        
1263        if !self.scopes().is_global_scope() {
1264            
1265            let local_name = LocalName::Symbol(*name);
1266            
1267            // check if the name is found in the local scope...
1268            let result = self.scopes().resolve_local(&local_name);
1269            
1270            if let Some(local) = result.cloned() {
1271                if !local.mode().can_write() {
1272                    return Err("can't assign to immutable local variable".into());
1273                }
1274                
1275                self.emit_assign_local(symbol, local.index());
1276                
1277                return Ok(());
1278            }
1279            
1280            // nonlocal keyword is not required in the global frame
1281            if !allow_nonlocal && self.scopes().is_call_frame() {
1282                return Err("can't assign to a non-local variable without the \"nonlocal\" keyword".into());
1283            }
1284            
1285            // check if an upvalue is found or can be created...
1286            if self.scopes().is_call_frame() {
1287                if let Some(upval) = self.scopes_mut().resolve_or_create_upval(&local_name)? {
1288                    if !upval.mode().can_write() {
1289                        return Err("can't assign to immutable local variable".into());
1290                    }
1291                    
1292                    let index = upval.index();
1293                    if let Ok(index) = u8::try_from(index) {
1294                        self.emit_instr_byte(symbol, OpCode::StoreUpvalue, index);
1295                    } else {
1296                        self.emit_instr_data(symbol, OpCode::StoreUpvalue16, &index.to_le_bytes());
1297                    }
1298                    
1299                    return Ok(());
1300                }
1301            }
1302        }
1303
1304        // ...finally, try to assign to a global, which are late bound
1305        self.emit_load_const(symbol, Constant::from(*name))?;
1306        self.emit_instr(symbol, OpCode::StoreGlobal);
1307        Ok(())
1308    }
1309    
1310    fn emit_assign_local(&mut self, symbol: Option<&DebugSymbol>, offset: LocalIndex) {
1311        if let Ok(offset) = u8::try_from(offset) {
1312            self.emit_instr_byte(symbol, OpCode::StoreLocal, offset);
1313        } else {
1314            self.emit_instr_data(symbol, OpCode::StoreLocal16, &offset.to_le_bytes());
1315        }
1316    }
1317    
1318    fn compile_assign_tuple(&mut self, symbol: Option<&DebugSymbol>, assign: AssignType, item_targets: &[LValue]) -> CompileResult<()> {
1319        // process tuple packing patterns
1320        
1321        let mut pack_targets = item_targets.iter().enumerate()
1322            .filter_map(|(idx, target)| match target {
1323                LValue::Pack(pack_target) => Some((idx, pack_target.as_deref())),
1324                _ => None,
1325            });
1326            
1327        let (idx, pack_target) = match pack_targets.next() {
1328            Some(pack_target) => pack_target,
1329            None => return self.compile_assign_tuple_nopack(symbol, assign, item_targets),
1330        };
1331        
1332        if !pack_targets.next().is_none() {
1333            return Err("tuple assignment may contain only one \"...\" pack pattern".into())
1334        }
1335        
1336        let (pre_pack, rest) = item_targets.split_at(idx);
1337        let (_, post_pack) = rest.split_at(1);
1338        
1339        self.compile_assign_tuple_pack(symbol, assign, pack_target, pre_pack, post_pack)
1340    }
1341    
1342    fn compile_assign_tuple_pack(&mut self, symbol: Option<&DebugSymbol>, assign: AssignType, pack: Option<&LValue>, pre_pack: &[LValue], post_pack: &[LValue]) -> CompileResult<()> {
1343        let mut error_jump_sites = Vec::new();
1344        
1345        // assignment needs to preserve original value for expression result
1346        self.emit_instr(symbol, OpCode::Clone);
1347        
1348        // iterate to yield values
1349        self.emit_instr(symbol, OpCode::IterInit);
1350        
1351        // compile to unrolled iteration for pre-pack items
1352        for target in pre_pack.iter() {
1353            debug_assert!(!matches!(target, LValue::Pack(..)));
1354            
1355            // check if there is an item left for this target
1356            let error_jump = self.emit_dummy_jump(symbol, Jump::IfFalse);
1357            error_jump_sites.push(error_jump);
1358            
1359            // advance the iterator and put the item on the stack
1360            self.emit_instr(symbol, OpCode::IterNext);
1361            
1362            self.compile_assignment(symbol, assign, target)?;
1363            self.emit_instr(symbol, OpCode::Pop);
1364        }
1365        
1366        let temp_scope;
1367        match (pack, post_pack) {
1368            (None, []) => {
1369                // if there are no post-pack items and no pack target just discard the iterator
1370                self.emit_instr_byte(symbol, OpCode::Drop, 2);
1371                temp_scope = false;
1372            }
1373            
1374            (Some(pack_target), []) => {
1375                // exhaust the rest of the iterator and assign to pack_target
1376                self.emit_instr(symbol, OpCode::IterUnpack);
1377                self.emit_instr(symbol, OpCode::TupleN);
1378                
1379                self.compile_assignment(symbol, assign, pack_target)?;
1380                self.emit_instr(symbol, OpCode::Pop);
1381                temp_scope = false;
1382            }
1383            
1384            (pack_target, post_pack) => {
1385                // exhaust the rest of the iterator and assign the last items to post_pack
1386                // then assign whatever is left to the pack_target
1387                self.emit_instr(symbol, OpCode::IterUnpack);
1388                
1389                // calculate the pack length and store it
1390                let post_len = IntType::try_from(post_pack.len())
1391                    .map_err(|_| "unpack length limit exceeded")?;
1392                
1393                self.compile_integer(symbol, post_len)?;
1394                self.emit_instr(symbol, OpCode::Sub);
1395                
1396                temp_scope = true;
1397                self.emit_begin_scope(symbol, None, ScopeTag::Temporary);
1398                let pack_len = self.emit_create_temporary(symbol, Access::ReadOnly)?;
1399                
1400                // check if there are enough items
1401                self.emit_instr_byte(symbol, OpCode::UInt8, 0);
1402                self.emit_instr(symbol, OpCode::LT);
1403                let error_jump = self.emit_dummy_jump(symbol, Jump::PopIfTrue);
1404                error_jump_sites.push(error_jump);
1405                
1406                // assign post-pack items
1407                for target in post_pack.iter().rev() {
1408                    debug_assert!(!matches!(target, LValue::Pack(..)));
1409                    self.compile_assignment(symbol, assign, target)?;
1410                    self.emit_instr(symbol, OpCode::Pop);
1411                }
1412                
1413                // load the stored pack len and assign to pack target or discard
1414                self.emit_load_local_index(symbol, pack_len);
1415                if let Some(pack_target) = pack_target {
1416                    self.emit_instr(symbol, OpCode::TupleN);
1417                    self.compile_assignment(symbol, assign, pack_target)?;
1418                    self.emit_instr(symbol, OpCode::Pop);
1419                } else {
1420                    self.emit_instr(symbol, OpCode::DropN);
1421                }
1422            }
1423        }
1424        
1425        let done_jump_site = self.emit_dummy_jump(symbol, Jump::Uncond);
1426        
1427        // not enough items
1428        let error_target = self.current_offset();
1429        for jump_site in error_jump_sites.iter() {
1430            self.patch_jump_instr(jump_site, error_target)?;
1431        }
1432        
1433        let message = format!("not enough values to unpack (expected at least {})", pre_pack.len() + post_pack.len());
1434        self.emit_load_error(symbol, ErrorKind::UnpackError, message.as_str())?;
1435        self.emit_instr(symbol, OpCode::Error);
1436        
1437        // cleanup
1438        self.patch_jump_instr(&done_jump_site, self.current_offset())?;
1439        
1440        if temp_scope {
1441            debug_assert!(self.scopes().is_temporary_scope());
1442            self.emit_end_scope();
1443        }
1444        
1445        Ok(())
1446    }
1447    
1448    fn compile_assign_tuple_nopack(&mut self, symbol: Option<&DebugSymbol>, assign: AssignType, items: &[LValue]) -> CompileResult<()> {
1449        let mut error_jump_sites = Vec::new();
1450        
1451        // assignment needs to preserve original value for expression result
1452        self.emit_instr(symbol, OpCode::Clone);
1453        
1454        // iterate to yield values
1455        self.emit_instr(symbol, OpCode::IterInit);
1456        
1457        // compile to unrolled iteration
1458        for target in items.iter() {
1459            debug_assert!(!matches!(target, LValue::Pack(..)));
1460            
1461            // check if there is an item left for this target
1462            let error_jump = self.emit_dummy_jump(symbol, Jump::IfFalse);
1463            error_jump_sites.push(error_jump);
1464            
1465            // advance the iterator and put the item on the stack
1466            self.emit_instr(symbol, OpCode::IterNext);
1467            
1468            self.compile_assignment(symbol, assign, target)?;
1469            self.emit_instr(symbol, OpCode::Pop);
1470        }
1471        
1472        // if the iterator is finished we've succeeded
1473        let done_jump_site = self.emit_dummy_jump(symbol, Jump::PopIfFalse);
1474        
1475        // too many items
1476        let message = format!("too many values to unpack (expected {})", items.len());
1477        self.emit_load_error(symbol, ErrorKind::UnpackError, message.as_str())?;
1478        self.emit_instr(symbol, OpCode::Error);
1479        
1480        // not enough items
1481        let error_target = self.current_offset();
1482        for jump_site in error_jump_sites.iter() {
1483            self.patch_jump_instr(jump_site, error_target)?;
1484        }
1485        
1486        let message = format!("not enough values to unpack (expected {})", items.len());
1487        self.emit_load_error(symbol, ErrorKind::UnpackError, message.as_str())?;
1488        self.emit_instr(symbol, OpCode::Error);
1489        
1490        // cleanup
1491        self.patch_jump_instr(&done_jump_site, self.current_offset())?;
1492        
1493        self.emit_instr(symbol, OpCode::Pop);  // pop the iterator
1494        
1495        Ok(())
1496    }
1497}
1498
1499///////// Blocks and If-Expressions /////////
1500impl CodeGenerator<'_> {
1501
1502    fn compile_block_expression(&mut self, symbol: Option<&DebugSymbol>, label: Option<&Label>, suite: &ExprBlock) -> CompileResult<()> {
1503        
1504        self.emit_begin_scope(symbol, label, ScopeTag::Block);
1505        self.compile_expr_block(symbol, suite)?;
1506        let block_scope = self.emit_end_scope();
1507        
1508        // finalize scope
1509        let break_target = self.current_offset();
1510        self.patch_break_sites(&block_scope, break_target)?;
1511        
1512        Ok(())
1513    }
1514    
1515    fn compile_if_expression(&mut self, symbol: Option<&DebugSymbol>, branches: &[ConditionalBranch], else_clause: Option<&ExprBlock>) -> CompileResult<()> {
1516        debug_assert!(!branches.is_empty());
1517        
1518        // track the sites where we jump to the end, so we can patch them later
1519        let mut end_jump_sites = Vec::new();
1520        
1521        // if there is no else branch, the last non-else branch won't have a jump to end, and should not pop the condition
1522        // this is because if-expressions without an else clause evaluate to their condition when not entered
1523        let (last_branch, rest) = branches.split_last().unwrap();
1524        let iter_branches = rest.iter()
1525            .map(|branch| (false, branch))
1526            .chain(iter::once((true, last_branch)));
1527        
1528        for (is_last, branch) in iter_branches {
1529            let is_final_branch = is_last && else_clause.is_none();
1530            
1531            self.compile_expr(symbol, branch.condition())?;
1532            
1533            // need to keep condition value on the stack in case there is a break/continue
1534            // inside the statement list
1535            let branch_jump_site = self.emit_dummy_jump(symbol, Jump::IfFalse);
1536            
1537            self.emit_begin_scope(symbol, None, ScopeTag::Branch);
1538            self.compile_expr_block(symbol, branch.suite())?;
1539            self.emit_end_scope();
1540            
1541            // site for the jump to the end of if-expression
1542            if !is_final_branch {
1543                self.emit_instr(symbol, OpCode::Pop);
1544                let jump_site = self.emit_dummy_jump(symbol, Jump::Uncond);
1545                end_jump_sites.push(jump_site);
1546            }
1547            
1548            // target for the jump from the conditional of the now compiled branch
1549            self.patch_jump_instr(&branch_jump_site, self.current_offset())?;
1550        }
1551        
1552        // else clause
1553        if let Some(suite) = else_clause {
1554            
1555            self.emit_begin_scope(symbol, None, ScopeTag::Branch);
1556            self.compile_expr_block(symbol, suite)?;
1557            self.emit_end_scope();
1558            
1559        }
1560        
1561        // patch all of the end jump sites
1562        let end_target = self.current_offset();
1563        for jump_site in end_jump_sites.iter() {
1564            self.patch_jump_instr(jump_site, end_target)?;
1565        }
1566        
1567        Ok(())
1568    }
1569    
1570    fn compile_shortcircuit_and(&mut self, symbol: Option<&DebugSymbol>, lhs: &Expr, rhs: &Expr) -> CompileResult<()> {
1571        self.compile_expr(symbol, lhs)?;
1572        
1573        let shortcircuit = self.emit_dummy_jump(symbol, Jump::IfFalse);
1574        
1575        self.emit_instr(symbol, OpCode::Pop);
1576        self.compile_expr(symbol, rhs)?;
1577        
1578        self.patch_jump_instr(&shortcircuit, self.current_offset())?;
1579        
1580        Ok(())
1581    }
1582    
1583    fn compile_shortcircuit_or(&mut self, symbol: Option<&DebugSymbol>, lhs: &Expr, rhs: &Expr) -> CompileResult<()> {
1584        self.compile_expr(symbol, lhs)?;
1585        
1586        let shortcircuit = self.emit_dummy_jump(symbol, Jump::IfTrue);
1587        
1588        self.emit_instr(symbol, OpCode::Pop);
1589        self.compile_expr(symbol, rhs)?;
1590        
1591        self.patch_jump_instr(&shortcircuit, self.current_offset())?;
1592        
1593        Ok(())
1594    }
1595}
1596
1597///////// Function Definitions /////////
1598impl CodeGenerator<'_> {
1599    fn compile_function_def(&mut self, symbol: Option<&DebugSymbol>, fundef: &FunctionDef) -> CompileResult<()> {
1600        // create a new chunk for the function
1601        let info = ChunkInfo::Function {
1602            symbol: symbol.copied(),
1603        };
1604
1605        let mut chunk_gen = self.create_chunk(info)?;
1606        
1607        let chunk_id = chunk_gen.chunk_id();
1608        let fun_id = match chunk_id {
1609            Chunk::Function(fun_id) => fun_id,
1610            _ => panic!("chunk {:?} is not valid for function", chunk_id),
1611        };
1612        
1613        // and a new local scope
1614        // don't need to emit new scope instructions, should handled by function call
1615        chunk_gen.scopes_mut().push_frame(symbol);
1616        
1617        // don't need to generate IN_LOCAL instructions for these, the VM should include them automatically
1618        chunk_gen.scopes_mut().insert_local(Access::ReadOnly, LocalName::Receiver)?;
1619        chunk_gen.scopes_mut().insert_local(Access::ReadOnly, LocalName::NArgs)?;
1620        
1621        // prepare argument list
1622        chunk_gen.compile_function_preamble(symbol, fundef)?;
1623        
1624        // function body
1625        chunk_gen.compile_stmt_block(fundef.body.stmt_list())?;
1626        
1627        // function result
1628        if let Some(expr) = fundef.body.result() {
1629            chunk_gen.compile_expr_with_symbol(expr)?;
1630        } else {
1631            chunk_gen.emit_instr(symbol, OpCode::Nil);
1632        }
1633        
1634        // end the function scope
1635        // don't need to drop locals explicitly, that will be done when the VMCallFrame returns
1636        let frame = chunk_gen.scopes_mut().pop_frame();
1637        
1638        // however we do still need to close upvalues before we return
1639        for local in frame.iter_locals().filter(|local| local.captured()) {
1640            chunk_gen.emit_close_upvalue(None, local.index());
1641        }
1642        
1643        chunk_gen.finish();
1644        
1645        // compile the function signature
1646        let signature = self.compile_function_signature(symbol, &fundef.signature)?;
1647        
1648        // compile upvalues
1649        let upvalues = frame.upvalues().iter()
1650            .map(|upval| upval.target())
1651            .collect::<Vec<UpvalueTarget>>()
1652            .into_boxed_slice();
1653        
1654        let function = UnloadedFunction {
1655            signature, upvalues, fun_id,
1656        };
1657        
1658        // load the function object as the expression result
1659        self.make_function(function);
1660        self.emit_load_function(symbol, fun_id);
1661        
1662        Ok(())
1663    }
1664    
1665    fn compile_function_preamble(&mut self, symbol: Option<&DebugSymbol>, fundef: &FunctionDef) -> CompileResult<()> {
1666        // process default and variadic arguments
1667        // this ensures that exactly `signature.param_count()` values are on the stack
1668        
1669        let signature = &fundef.signature;
1670        if signature.param_count() == 0 {
1671            return Ok(())
1672        }
1673        
1674        for param in signature.required.iter() {
1675            self.scopes_mut().insert_local(param.mode, LocalName::Symbol(param.name))?;
1676        }
1677        
1678        if !signature.default.is_empty() {
1679            self.compile_default_args(signature)?;
1680            for param in signature.default.iter() {
1681                self.scopes_mut().insert_local(param.mode, LocalName::Symbol(param.name))?;
1682            }
1683        }
1684        
1685        if let Some(param) = &signature.variadic {
1686            self.compile_variadic_arg(signature)?;
1687            self.scopes_mut().insert_local(param.mode, LocalName::Symbol(param.name))?;
1688        }
1689        
1690        self.emit_instr(symbol, OpCode::InsertArgs);
1691
1692        Ok(())
1693    }
1694    
1695    fn compile_default_args(&mut self, signature: &SignatureDef) -> CompileResult<()> {
1696        debug_assert!(!signature.default.is_empty());
1697        
1698        let mut jump_sites = Vec::new();
1699        let mut jump_targets = Vec::new();
1700        
1701        // depending on the number of arguments, jump into the default argument sequence
1702        let required_count = u8::try_from(signature.required.len())
1703            .map_err(|_| "parameter count limit exceeded")?;
1704        let default_count = u8::try_from(signature.default.len())
1705            .map_err(|_| "parameter count limit exceeded")?;
1706        
1707        // "defaults passed" = NArgs - required_count
1708        self.try_emit_load_local(None, &LocalName::NArgs).unwrap();
1709        if required_count != 0 {
1710            self.emit_instr_byte(None, OpCode::UInt8, required_count);
1711            self.emit_instr(None, OpCode::Sub);
1712        }
1713        
1714        /*
1715            NArgs - required count:
1716            
1717            <  0   => impossible, missing required argument
1718               0   => jump to default[0]
1719               1   => jump to default[1]
1720               ...
1721               N-1 => jump to default[N-1]
1722            >= N   => jump to end
1723        */
1724        
1725        for count in 0..default_count {
1726            self.emit_instr(None, OpCode::Clone);
1727            self.emit_instr_byte(None, OpCode::UInt8, count);
1728            self.emit_instr(None, OpCode::EQ);
1729            
1730            let jump_site = self.emit_dummy_jump(None, Jump::PopIfTrue);
1731            jump_sites.insert(count.into(), jump_site);
1732        }
1733        
1734        // if we get here, jump unconditionally to the end
1735        let jump_site = self.emit_dummy_jump(None, Jump::Uncond);
1736        jump_sites.insert(default_count.into(), jump_site);
1737        
1738        // generate default arguments
1739        for (idx, param) in signature.default.iter().enumerate() {
1740            jump_targets.insert(idx, self.current_offset());
1741            
1742            let symbol = Some(param.default.debug_symbol());
1743            let expr = param.default.variant();
1744            self.compile_expr(symbol, expr)?;
1745            // self.emit_instr(None, OpCode::InsertLocal);
1746        }
1747        
1748        jump_targets.insert(default_count.into(), self.current_offset());
1749        self.emit_instr(None, OpCode::Pop);  // drop "defaults passed"
1750        
1751        // patch all jumps
1752        debug_assert!(jump_sites.len() == jump_targets.len());
1753        for (jump_site, jump_target) in jump_sites.iter().zip(jump_targets.into_iter()) {
1754            self.patch_jump_instr(jump_site, jump_target)?;
1755        }
1756        
1757        Ok(())
1758    }
1759    
1760    fn compile_variadic_arg(&mut self, signature: &SignatureDef) -> CompileResult<()> {
1761        debug_assert!(signature.variadic.is_some());
1762        
1763        let positional_count = u8::try_from(signature.required.len() + signature.default.len())
1764            .map_err(|_| "parameter count limit exceeded")?;
1765        
1766        // "variadic count" = NArgs - required_count - default_count
1767        self.try_emit_load_local(None, &LocalName::NArgs).unwrap();
1768        if positional_count != 0 {
1769            self.emit_instr_byte(None, OpCode::UInt8, positional_count);
1770            self.emit_instr(None, OpCode::Sub);
1771        }
1772        
1773        // check if "variadic count" > 0
1774        self.emit_instr(None, OpCode::Clone);
1775        self.emit_instr_byte(None, OpCode::UInt8, 0);
1776        self.emit_instr(None, OpCode::GT);
1777        
1778        let tuple_jump_site = self.emit_dummy_jump(None, Jump::PopIfTrue);
1779        
1780        // if we get here then the variadic arg is empty
1781        self.emit_instr(None, OpCode::Pop);
1782        self.emit_instr(None, OpCode::Empty);
1783        
1784        let end_jump_site = self.emit_dummy_jump(None, Jump::Uncond);
1785        
1786        self.patch_jump_instr(&tuple_jump_site, self.current_offset())?;
1787        self.emit_instr(None, OpCode::TupleN);
1788        // self.emit_instr(None, OpCode::InsertLocal);
1789        
1790        self.patch_jump_instr(&end_jump_site, self.current_offset())?;
1791        
1792        Ok(())
1793    }
1794    
1795    fn compile_function_signature(&mut self, symbol: Option<&DebugSymbol>, signature: &SignatureDef) -> CompileResult<UnloadedSignature> {
1796        let name = 
1797            if let Some(name) = signature.name {
1798                Some(self.get_or_make_const(Constant::from(name))?)
1799            } else { None };
1800        
1801        let mut required = Vec::new();
1802        for param in signature.required.iter() {
1803            let name = self.get_or_make_const(Constant::from(param.name))?;
1804            required.push(UnloadedParam { name, mode: param.mode });
1805        }
1806        
1807        let mut default = Vec::new();
1808        for param in signature.default.iter() {
1809            let name = self.get_or_make_const(Constant::from(param.name))?;
1810            default.push(UnloadedParam { name, mode: param.mode });
1811        }
1812        
1813        let mut variadic = None;
1814        if let Some(param) = &signature.variadic {
1815            let name = self.get_or_make_const(Constant::from(param.name))?;
1816            variadic.replace(UnloadedParam { name, mode: param.mode });
1817        }
1818        
1819        let signature = UnloadedSignature {
1820            name,
1821            required: required.into_boxed_slice(),
1822            default: default.into_boxed_slice(),
1823            variadic,
1824        };
1825        
1826        Ok(signature)
1827    }
1828    
1829}