slang_backend/
codegen.rs

1use crate::bytecode::{Chunk, Function, OpCode};
2use crate::value::Value;
3use slang_error::{CompilerError, CompileResult, ErrorCode};
4use slang_ir::Visitor;
5use slang_ir::ast::{
6    BinaryExpr, BinaryOperator, BlockExpr, ConditionalExpr, Expression, FunctionCallExpr,
7    FunctionDeclarationStmt, FunctionTypeExpr, IfStatement, LetStatement, LiteralExpr, Statement, TypeDefinitionStmt,
8    UnaryExpr, UnaryOperator,
9};
10use slang_ir::location::Location;
11
12/// Compiles AST nodes into bytecode instructions
13pub struct CodeGenerator {
14    /// The bytecode chunk being constructed
15    pub chunk: Chunk,
16    /// Current line number for debugging information
17    line: usize,
18    /// Global variable names
19    variables: Vec<String>,
20    /// Function names for tracking declarations
21    functions: Vec<String>,
22    /// Stack of scopes for tracking local variables
23    local_scopes: Vec<Vec<String>>,
24    /// Accumulated errors during compilation
25    errors: Vec<CompilerError>,
26}
27
28pub fn generate_bytecode(statements: &[Statement]) -> CompileResult<Chunk> {
29    let compiler = CodeGenerator::new();
30    compiler.compile(statements)
31}
32
33impl CodeGenerator {
34    /// Creates a new compiler with an empty chunk
35    pub fn new() -> Self {
36        CodeGenerator {
37            chunk: Chunk::new(),
38            line: 1,
39            variables: Vec::new(),
40            functions: Vec::new(),
41            local_scopes: Vec::new(),
42            errors: Vec::new(),
43        }
44    }
45
46    /// Updates the current line from a source location
47    fn set_current_location(&mut self, location: &Location) {
48        self.line = location.line;
49    }
50
51    /// Creates a CompilerError with the current location and adds it to the error list
52    fn add_error(&mut self, message: String) {
53        let error = CompilerError::new(
54            ErrorCode::GenericCompileError,
55            message,
56            self.line,
57            0, // column - we don't track this in codegen currently
58            0, // position - we don't track this in codegen currently  
59            None, // token_length - not applicable for codegen errors
60        );
61        self.errors.push(error);
62    }
63
64    /// Compiles a list of statements into bytecode
65    ///
66    /// ### Arguments
67    ///
68    /// * `statements` - The statements to compile
69    ///
70    /// ### Returns
71    ///
72    /// A CompileResult containing the compiled bytecode chunk or errors
73    fn compile(mut self, statements: &[Statement]) -> CompileResult<Chunk> {
74        for stmt in statements {
75            stmt.accept(&mut self).unwrap_or_else(|_| {
76                // Error already added to self.errors
77            });
78        }
79
80        self.emit_op(OpCode::Return);
81
82        if self.errors.is_empty() {
83            Ok(self.chunk)
84        } else {
85            Err(self.errors)
86        }
87    }
88
89    /// Compiles statements and appends them to the existing chunk
90    ///
91    /// ### Arguments
92    ///
93    /// * `statements` - The statements to compile
94    ///
95    /// ### Returns
96    ///
97    /// CompileResult indicating success or containing errors
98    pub fn compile_statements(&mut self, statements: &[Statement]) -> CompileResult<()> {
99        for stmt in statements {
100            stmt.accept(self).unwrap_or_else(|_| {
101                // Error already added to self.errors
102            });
103        }
104        
105        if self.errors.is_empty() {
106            Ok(())
107        } else {
108            Err(std::mem::take(&mut self.errors))
109        }
110    }
111
112    /// Gets a reference to the current chunk
113    pub fn get_chunk(&self) -> &Chunk {
114        &self.chunk
115    }
116
117    /// Emits a single byte to the bytecode chunk
118    ///
119    /// ### Arguments
120    ///
121    /// * `byte` - The byte to emit
122    fn emit_byte(&mut self, byte: u8) {
123        self.chunk.write_byte(byte, self.line);
124    }
125
126    /// Emits an opcode to the bytecode chunk
127    ///
128    /// ### Arguments
129    ///
130    /// * `op` - The opcode to emit
131    fn emit_op(&mut self, op: OpCode) {
132        self.chunk.write_op(op, self.line);
133    }    /// Adds a constant value to the chunk and emits code to load it
134    ///
135    /// ### Arguments
136    ///
137    /// * `value` - The constant value to add
138    fn emit_constant(&mut self, value: Value) -> Result<(), ()> {
139        let constant_index = self.chunk.add_constant(value);
140        if constant_index > 255 {
141            self.add_error("Too many constants in one chunk".to_string());
142            return Err(());
143        }
144        self.emit_op(OpCode::Constant);
145        self.emit_byte(constant_index as u8);
146        Ok(())
147    }
148
149    /// Emits a jump instruction with placeholder offset
150    ///
151    /// ### Arguments
152    ///
153    /// * `op` - The jump opcode (Jump or JumpIfFalse)
154    ///
155    /// ### Returns
156    ///
157    /// The position where the jump offset needs to be patched later
158    fn emit_jump(&mut self, op: OpCode) -> usize {
159        self.emit_op(op);
160        self.emit_byte(0xFF);
161        self.emit_byte(0xFF);
162        self.chunk.code.len() - 2
163    }
164
165    /// Patches a previously emitted jump instruction with the actual offset
166    ///
167    /// ### Arguments
168    ///
169    /// * `offset` - The position of the jump offset to patch
170    fn patch_jump(&mut self, offset: usize) {
171        let jump = self.chunk.code.len() - offset - 2;
172        if jump > 0xFFFF {
173            panic!("Jump too far");
174        }
175
176        self.chunk.code[offset] = ((jump >> 8) & 0xFF) as u8;
177        self.chunk.code[offset + 1] = (jump & 0xFF) as u8;
178    }
179
180    fn begin_scope(&mut self) {
181        self.local_scopes.push(Vec::new());
182        self.emit_op(OpCode::BeginScope);
183    }
184
185    fn end_scope(&mut self) {
186        self.local_scopes.pop();
187        self.emit_op(OpCode::EndScope);
188    }
189}
190
191impl Visitor<Result<(), ()>> for CodeGenerator {
192    fn visit_statement(&mut self, stmt: &Statement) -> Result<(), ()> {
193        // Update current line from the statement's location
194        let location = match stmt {
195            Statement::Let(let_stmt) => let_stmt.location,
196            Statement::Assignment(assign_stmt) => assign_stmt.location,
197            Statement::TypeDefinition(type_stmt) => type_stmt.location,
198            Statement::Expression(expr) => expr.location(),
199            Statement::FunctionDeclaration(fn_decl) => fn_decl.location,
200            Statement::Return(return_stmt) => return_stmt.location,
201            Statement::If(if_stmt) => if_stmt.location,
202        };
203        self.set_current_location(&location);
204        
205        match stmt {
206            Statement::Let(let_stmt) => self.visit_let_statement(let_stmt),
207            Statement::Assignment(assign_stmt) => self.visit_assignment_statement(assign_stmt),
208            Statement::TypeDefinition(type_stmt) => self.visit_type_definition_statement(type_stmt),
209            Statement::Expression(expr) => self.visit_expression_statement(expr),
210            Statement::FunctionDeclaration(fn_decl) => {
211                self.visit_function_declaration_statement(fn_decl)
212            }
213            Statement::Return(expr) => self.visit_return_statement(expr),
214            Statement::If(if_stmt) => self.visit_if_statement(if_stmt),
215        }
216    }
217
218    fn visit_expression(&mut self, expr: &Expression) -> Result<(), ()> {
219        // Update current line from the expression's location
220        self.set_current_location(&expr.location());
221        
222        match expr {
223            Expression::Literal(lit_expr) => self.visit_literal_expression(lit_expr),
224            Expression::Binary(bin_expr) => self.visit_binary_expression(bin_expr),
225            Expression::Variable(var) => self.visit_variable_expression(var),
226            Expression::Unary(unary_expr) => self.visit_unary_expression(unary_expr),
227            Expression::Call(call_expr) => self.visit_call_expression(call_expr),
228            Expression::Conditional(cond_expr) => self.visit_conditional_expression(cond_expr),
229            Expression::Block(block_expr) => self.visit_block_expression(block_expr),
230            Expression::FunctionType(func_type_expr) => self.visit_function_type_expression(func_type_expr),
231        }
232    }
233
234    fn visit_function_declaration_statement(
235        &mut self,
236        fn_decl: &FunctionDeclarationStmt,
237    ) -> Result<(), ()> {
238        self.functions.push(fn_decl.name.clone());
239        let function_name_idx = self.chunk.add_identifier(fn_decl.name.clone());
240
241        let jump_over = self.emit_jump(OpCode::Jump);
242
243        let code_offset = self.chunk.code.len();
244        let mut locals = Vec::new();
245
246        self.begin_scope();
247        for param in &fn_decl.parameters {
248            locals.push(param.name.clone());
249            if let Some(current_scope) = self.local_scopes.last_mut() {
250                current_scope.push(param.name.clone());
251            }
252        }
253
254        for stmt in &fn_decl.body.statements {
255            stmt.accept(self)?;
256        }
257
258        if let Some(return_expr) = &fn_decl.body.return_expr {
259            return_expr.accept(self)?;
260        } else {
261            self.emit_constant(Value::Unit(()))?;
262        }
263
264        self.emit_op(OpCode::Return);
265
266        self.end_scope();
267        self.patch_jump(jump_over);
268
269        let function = Value::Function(Box::new(Function {
270            name: fn_decl.name.clone(),
271            arity: fn_decl.parameters.len() as u8,
272            code_offset,
273            locals,
274        }));
275        let fn_constant = self.chunk.add_constant(function);
276
277        self.emit_op(OpCode::DefineFunction);
278        self.emit_byte(function_name_idx as u8);
279        self.emit_byte(fn_constant as u8);
280
281        Ok(())
282    }
283
284    fn visit_return_statement(&mut self, return_stmt: &slang_ir::ast::ReturnStatement) -> Result<(), ()> {
285        if let Some(expr) = &return_stmt.value {
286            self.visit_expression(expr)?;
287        } else {
288            self.emit_constant(Value::Unit(()))?;
289        }
290        self.emit_op(OpCode::Return);
291        Ok(())
292    }
293
294    fn visit_expression_statement(&mut self, expr: &Expression) -> Result<(), ()> {
295        self.visit_expression(expr)?;
296        Ok(())
297    }
298
299    fn visit_let_statement(&mut self, let_stmt: &LetStatement) -> Result<(), ()> {
300        let is_local = !self.local_scopes.is_empty();
301
302        if is_local {
303            if let Some(current_scope) = self.local_scopes.last_mut() {
304                current_scope.push(let_stmt.name.clone());
305            }
306        } else {
307            self.variables.push(let_stmt.name.clone());
308        }
309
310        self.visit_expression(&let_stmt.value)?;
311
312        let var_index = self.chunk.add_identifier(let_stmt.name.clone());
313        if var_index > 255 {
314            self.add_error("Too many variables in one scope".to_string());
315            return Err(());
316        }
317
318        self.emit_op(OpCode::SetVariable);
319        self.emit_byte(var_index as u8);
320
321        self.emit_op(OpCode::Pop);
322
323        Ok(())
324    }
325
326    fn visit_assignment_statement(
327        &mut self,
328        assign_stmt: &slang_ir::ast::AssignmentStatement,
329    ) -> Result<(), ()> {
330        self.visit_expression(&assign_stmt.value)?;
331        let var_index = self.chunk.add_identifier(assign_stmt.name.clone());
332        if var_index > 255 {
333            self.add_error("Too many variables in one scope".to_string());
334            return Err(());
335        }
336        self.emit_op(OpCode::SetVariable);
337        self.emit_byte(var_index as u8);
338
339        Ok(())
340    }
341
342    fn visit_call_expression(&mut self, call_expr: &FunctionCallExpr) -> Result<(), ()> {
343        for arg in &call_expr.arguments {
344            self.visit_expression(arg)?;
345        }
346
347        let fn_name_idx = self.chunk.add_identifier(call_expr.name.clone());
348        self.emit_op(OpCode::GetVariable);
349        self.emit_byte(fn_name_idx as u8);
350
351        self.emit_op(OpCode::Call);
352        self.emit_byte(call_expr.arguments.len() as u8);
353
354        Ok(())
355    }
356
357    fn visit_literal_expression(&mut self, lit_expr: &LiteralExpr) -> Result<(), ()> {
358        match &lit_expr.value {
359            slang_ir::ast::LiteralValue::I32(i) => {
360                self.emit_constant(Value::I32(*i))?;
361            }
362            slang_ir::ast::LiteralValue::I64(i) => {
363                self.emit_constant(Value::I64(*i))?;
364            }
365            slang_ir::ast::LiteralValue::U32(i) => {
366                self.emit_constant(Value::U32(*i))?;
367            }
368            slang_ir::ast::LiteralValue::U64(i) => {
369                self.emit_constant(Value::U64(*i))?;
370            }
371            slang_ir::ast::LiteralValue::UnspecifiedInteger(i) => {
372                self.emit_constant(Value::I64(*i))?;
373            }
374            slang_ir::ast::LiteralValue::F32(f) => {
375                self.emit_constant(Value::F32(*f))?;
376            }
377            slang_ir::ast::LiteralValue::F64(f) => {
378                self.emit_constant(Value::F64(*f))?;
379            }
380            slang_ir::ast::LiteralValue::UnspecifiedFloat(f) => {
381                self.emit_constant(Value::F64(*f))?;
382            }
383            slang_ir::ast::LiteralValue::String(s) => {
384                self.emit_constant(Value::String(Box::new(s.clone())))?;
385            }
386            slang_ir::ast::LiteralValue::Boolean(b) => {
387                self.emit_constant(Value::Boolean(*b))?;
388            }
389            slang_ir::ast::LiteralValue::Unit => {
390                self.emit_constant(Value::Unit(()))?;
391            }
392        }
393
394        Ok(())
395    }
396
397    fn visit_binary_expression(&mut self, bin_expr: &BinaryExpr) -> Result<(), ()> {
398        match bin_expr.operator {
399            BinaryOperator::And => {
400                self.visit_expression(&bin_expr.left)?;
401                let jump_if_false = self.emit_jump(OpCode::JumpIfFalse);
402                self.emit_op(OpCode::Pop);
403                self.visit_expression(&bin_expr.right)?;
404                self.patch_jump(jump_if_false);
405                return Ok(());
406            }
407
408            BinaryOperator::Or => {
409                self.visit_expression(&bin_expr.left)?;
410                let jump_if_true = self.emit_jump(OpCode::JumpIfFalse);
411                let jump_to_end = self.emit_jump(OpCode::Jump);
412                self.patch_jump(jump_if_true);
413                self.emit_op(OpCode::Pop);
414                self.visit_expression(&bin_expr.right)?;
415                self.patch_jump(jump_to_end);
416                return Ok(());
417            }
418
419            _ => {
420                self.visit_expression(&bin_expr.left)?;
421                self.visit_expression(&bin_expr.right)?;
422
423                match bin_expr.operator {
424                    BinaryOperator::Add => self.emit_op(OpCode::Add),
425                    BinaryOperator::Subtract => self.emit_op(OpCode::Subtract),
426                    BinaryOperator::Multiply => self.emit_op(OpCode::Multiply),
427                    BinaryOperator::Divide => self.emit_op(OpCode::Divide),
428                    BinaryOperator::GreaterThan => self.emit_op(OpCode::Greater),
429                    BinaryOperator::LessThan => self.emit_op(OpCode::Less),
430                    BinaryOperator::GreaterThanOrEqual => self.emit_op(OpCode::GreaterEqual),
431                    BinaryOperator::LessThanOrEqual => self.emit_op(OpCode::LessEqual),
432                    BinaryOperator::Equal => self.emit_op(OpCode::Equal),
433                    BinaryOperator::NotEqual => self.emit_op(OpCode::NotEqual),
434                    _ => {
435                        self.add_error(format!(
436                            "Unsupported binary operator: {:?}",
437                            bin_expr.operator
438                        ));
439                        return Err(());
440                    }
441                }
442            }
443        }
444
445        Ok(())
446    }
447
448    fn visit_unary_expression(&mut self, unary_expr: &UnaryExpr) -> Result<(), ()> {
449        self.visit_expression(&unary_expr.right)?;
450
451        match unary_expr.operator {
452            UnaryOperator::Negate => self.emit_op(OpCode::Negate),
453            UnaryOperator::Not => self.emit_op(OpCode::BoolNot),
454        }
455
456        Ok(())
457    }
458
459    fn visit_variable_expression(
460        &mut self,
461        var_expr: &slang_ir::ast::VariableExpr,
462    ) -> Result<(), ()> {
463        let var_index = self.chunk.add_identifier(var_expr.name.clone());
464        if var_index > 255 {
465            self.add_error("Too many variables".to_string());
466            return Err(());
467        }
468        self.emit_op(OpCode::GetVariable);
469        self.emit_byte(var_index as u8);
470        Ok(())
471    }
472
473    fn visit_type_definition_statement(
474        &mut self,
475        _stmt: &TypeDefinitionStmt,
476    ) -> Result<(), ()> {
477        // Type definitions don't generate code at runtime
478        // They're just used by the semantic analyzer
479        Ok(())
480    }
481
482    fn visit_conditional_expression(&mut self, cond_expr: &ConditionalExpr) -> Result<(), ()> {
483        self.visit_expression(&cond_expr.condition)?;
484
485        let jump_to_else = self.emit_jump(OpCode::JumpIfFalse);
486        self.emit_op(OpCode::Pop);
487        self.visit_expression(&cond_expr.then_branch)?;
488
489        let jump_over_else = self.emit_jump(OpCode::Jump);
490        self.patch_jump(jump_to_else);
491        self.emit_op(OpCode::Pop);
492        self.visit_expression(&cond_expr.else_branch)?;
493
494        self.patch_jump(jump_over_else);
495
496        Ok(())
497    }
498
499    fn visit_if_statement(&mut self, if_stmt: &IfStatement) -> Result<(), ()> {
500        self.visit_expression(&if_stmt.condition)?;
501
502        let jump_to_else = self.emit_jump(OpCode::JumpIfFalse);
503        self.emit_op(OpCode::Pop);
504
505        self.visit_block_expression(&if_stmt.then_branch)?;
506
507        if let Some(else_branch) = &if_stmt.else_branch {
508            let jump_over_else = self.emit_jump(OpCode::Jump);
509
510            self.patch_jump(jump_to_else);
511            self.emit_op(OpCode::Pop);
512
513            self.visit_block_expression(else_branch)?;
514
515            self.patch_jump(jump_over_else);
516        } else {
517            self.patch_jump(jump_to_else);
518            self.emit_op(OpCode::Pop);
519        }
520
521        Ok(())
522    }
523
524    fn visit_block_expression(&mut self, block_expr: &BlockExpr) -> Result<(), ()> {
525        self.begin_scope();
526
527        for stmt in &block_expr.statements {
528            self.visit_statement(stmt)?;
529        }
530
531        if let Some(return_expr) = &block_expr.return_expr {
532            self.visit_expression(return_expr)?;
533        } else {
534            self.emit_constant(Value::Unit(()))?;
535        }
536
537        self.end_scope();
538
539        Ok(())
540    }
541
542    fn visit_function_type_expression(&mut self, _func_type_expr: &FunctionTypeExpr) -> Result<(), ()> {
543        // Function type expressions are compile-time constructs that don't generate runtime bytecode
544        // They are used for type checking and don't produce any values at runtime
545        Ok(())
546    }
547}