use crate::{ast::{ASTType, BinOp, Lit, Name, UnOp, AST}, codegen_binary, codegen_unary, extension::Extension, opcodes::OpCode, vm::{Value, VM}};
pub struct Codegen {
pub vm: VM,
pub scopes: Vec<usize>, // number of variables declared in the scope
}
impl Codegen {
pub fn new(text: String) -> Self {
Codegen {
vm: VM::new(Vec::new(), Vec::new(),
Vec::new(), text),
scopes: vec![0],
}
}
pub fn visit(&mut self, ast: AST) {
use crate::ast::ASTType::*;
match ast.ast_type {
Program(asts) => {
for ast in asts {
self.visit(ast);
}
}
If(condition, block, else_expr) => {
self.visit_if(*condition, *block, *else_expr);
}
Block(ref stmts) => {
self.visit_block(&ast, stmts.clone());
}
Decl(ref name, ref expr) => {
self.visit_decl(&ast, name.clone(), *expr.clone());
}
Expr(ref expr) => {
self.visit_expr_stmt(&ast, *expr.clone());
}
// REMOVE LATER
Print(ref expr) => {
self.visit_print(&ast, *expr.clone());
}
Assign(ref name, ref expr) => {
self.visit_assign(&ast, name.clone(), *expr.clone());
}
Binary(left, oper, right) => {
self.visit_binary(*left, oper, *right);
}
Unary(oper, right) => {
self.visit_unary(oper, *right);
}
Lit(ref lit) => {
self.visit_lit(&ast, lit.clone());
}
_ => {}
}
}
pub fn visit_if(&mut self, condition: AST, block: AST, else_expr: Option<AST>) {
self.visit(condition.clone());
self.vm.chunk.push(OpCode::IF_NE.into());
self.vm.chunk.push(0xFF); // placeholder
self.vm.chunk.push(0xFF); // placeholder
let start_jmp_loc = self.vm.chunk.len()-2;
if self.vm.lines.last().unwrap().0 != condition.end_line {
self.vm.lines.push((condition.end_line, 3));
} else {
self.vm.lines.last_mut().unwrap().1 += 3;
}
self.visit(block);
if let Some(ast) = else_expr {
match ast.ast_type {
ASTType::If(_, _, _) => {
self.vm.chunk.push(OpCode::JMP.into());
self.vm.chunk.push(0xFF); // placeholder
self.vm.chunk.push(0xFF); // placeholder
let end_jmp_loc = self.vm.chunk.len()-2;
if self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 3));
} else {
self.vm.lines.last_mut().unwrap().1 += 3;
}
self.patch_jump(start_jmp_loc);
self.visit(ast);
self.patch_jump(end_jmp_loc);
}
ASTType::Block(_) => {
self.vm.chunk.push(OpCode::JMP.into());
self.vm.chunk.push(0xFF); // placeholder
self.vm.chunk.push(0xFF); // placeholder
if self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 3));
} else {
self.vm.lines.last_mut().unwrap().1 += 3;
}
let end_jmp_loc = self.vm.chunk.len()-2;
self.patch_jump(start_jmp_loc); // jmp to else
self.visit(ast);
self.patch_jump(end_jmp_loc); // jmp to after else
}
_ => { println!("unimplemented"); }
}
} else {
self.patch_jump(start_jmp_loc);
}
}
pub fn visit_block(&mut self, ast: &AST, stmts: Vec<AST>) {
self.scopes.push(0);
for stmt in stmts {
self.visit(stmt);
}
let variables = self.scopes.pop().expect("popped scope in block");
for _ in 0..variables {
self.vm.chunk.push(OpCode::POP.into());
if self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 1));
} else {
self.vm.lines.last_mut().unwrap().1 += 1;
}
}
}
pub fn visit_decl(&mut self, ast: &AST, _name: Name, expr: AST) {
self.visit(expr);
#[allow(irrefutable_let_patterns)]
if let Extension::Resolution(_scope, _ptr) = ast.extensions[0] {
*self.scopes.last_mut().expect("codegen scopes vec was empty") += 1;
}
}
pub fn visit_expr_stmt(&mut self, ast: &AST, expr: AST) {
self.visit(expr);
self.vm.chunk.push(OpCode::POP.into());
if self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 1));
} else {
self.vm.lines.last_mut().unwrap().1 += 1;
}
}
// REMOVE LATER
pub fn visit_print(&mut self, ast: &AST, expr: AST) {
self.visit(expr);
self.vm.chunk.push(OpCode::PRINT.into());
if self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 1));
} else {
self.vm.lines.last_mut().unwrap().1 += 1;
}
}
pub fn visit_assign(&mut self, ast: &AST, _name: Name, expr: AST) {
self.visit(expr);
// pop stack to get value. then find variable in stack. set variable to value.
self.vm.chunk.push(OpCode::ASSIGN.into());
#[allow(irrefutable_let_patterns)]
if let Extension::Resolution(_scope, ptr) = ast.extensions[0] {
self.vm.chunk.push(ptr as u8);
if self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 1));
} else {
self.vm.lines.last_mut().unwrap().1 += 1;
}
}
if self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 1));
} else {
self.vm.lines.last_mut().unwrap().1 += 1;
}
}
pub fn visit_binary(&mut self, left: AST, oper: BinOp, right: AST) {
use crate::ast::BinOp::*;
codegen_binary!(self, left, right, oper, Add, ADD, Sub, SUB, Mul, MUL,
Div, DIV, Mod, MOD, BitAnd, BITAND, BitXor, BITXOR, BitOr, BITOR,
Eq, EQ, Ne, NE, Gr, GR, Ge, GE, Ls, LS, Le, LE, Or, OR, And, AND);
}
pub fn visit_unary(&mut self, oper: UnOp, right: AST) {
use crate::ast::UnOp::*;
codegen_unary!(self, right, oper, Not, NOT, Neg, NEG);
}
pub fn visit_lit(&mut self, ast: &AST, lit: Lit) {
match lit {
Lit::Num(value) => {
self.vm.pool.push(Value::Num(value));
self.vm.chunk.push(OpCode::LOAD.into());
self.vm.chunk.push((self.vm.pool.len()-1) as u8);
if self.vm.lines.len() == 0 || self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 2));
} else {
self.vm.lines.last_mut().unwrap().1 += 2;
}
}
Lit::Str(value) => {
self.vm.pool.push(Value::Str(value.into()));
self.vm.chunk.push(OpCode::LOAD.into());
self.vm.chunk.push((self.vm.pool.len()-1) as u8);
if self.vm.lines.len() == 0 || self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 2));
} else {
self.vm.lines.last_mut().unwrap().1 += 2;
}
}
Lit::Bool(value) => {
self.vm.pool.push(Value::Bool(value));
self.vm.chunk.push(OpCode::LOAD.into());
self.vm.chunk.push((self.vm.pool.len()-1) as u8);
if self.vm.lines.len() == 0 || self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 2));
} else {
self.vm.lines.last_mut().unwrap().1 += 2;
}
}
Lit::Ident(_name) => {
self.vm.chunk.push(OpCode::VAR.into());
#[allow(irrefutable_let_patterns)]
if let Extension::Resolution(_scope, ptr) = ast.extensions[0] {
/*println!("line: {} ptr: {} locals: {:?}", ast.end_line, ptr, self.locals);
let (_scope, slot) = self.locals.get(ptr).unwrap();
println!("slot: {}", *slot);
self.vm.chunk.push(*slot);*/
self.vm.chunk.push(ptr as u8);
if self.vm.lines.len() == 0 || self.vm.lines.last().unwrap().0 != ast.end_line {
self.vm.lines.push((ast.end_line, 2));
} else {
self.vm.lines.last_mut().unwrap().1 += 2;
}
}
}
_ => {}
}
}
pub fn patch_jump(&mut self, loc: usize) {
let jump = self.vm.chunk.len() - loc - 2;
self.vm.chunk[loc] = ((jump >> 8) & 0xFF) as u8;
self.vm.chunk[loc + 1] = (jump & 0xFF) as u8;
}
}