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
12pub struct CodeGenerator {
14 pub chunk: Chunk,
16 line: usize,
18 variables: Vec<String>,
20 functions: Vec<String>,
22 local_scopes: Vec<Vec<String>>,
24 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 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 fn set_current_location(&mut self, location: &Location) {
48 self.line = location.line;
49 }
50
51 fn add_error(&mut self, message: String) {
53 let error = CompilerError::new(
54 ErrorCode::GenericCompileError,
55 message,
56 self.line,
57 0, 0, None, );
61 self.errors.push(error);
62 }
63
64 fn compile(mut self, statements: &[Statement]) -> CompileResult<Chunk> {
74 for stmt in statements {
75 stmt.accept(&mut self).unwrap_or_else(|_| {
76 });
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 pub fn compile_statements(&mut self, statements: &[Statement]) -> CompileResult<()> {
99 for stmt in statements {
100 stmt.accept(self).unwrap_or_else(|_| {
101 });
103 }
104
105 if self.errors.is_empty() {
106 Ok(())
107 } else {
108 Err(std::mem::take(&mut self.errors))
109 }
110 }
111
112 pub fn get_chunk(&self) -> &Chunk {
114 &self.chunk
115 }
116
117 fn emit_byte(&mut self, byte: u8) {
123 self.chunk.write_byte(byte, self.line);
124 }
125
126 fn emit_op(&mut self, op: OpCode) {
132 self.chunk.write_op(op, self.line);
133 } 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 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 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 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 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 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 Ok(())
546 }
547}