use std::collections::HashMap;
use log::debug;
use petgraph::visit::DfsPostOrder;
use crate::{
bytecode::{Bytecode, Opcode, Operand, Register},
compiler::codegen::regalloc::Action,
ir::instruction::{Block, ControlFlowGraph, Instruction, Value},
};
use super::regalloc::RegAlloc;
type PatchFn = Box<dyn Fn(&mut Codegen)>;
pub struct Codegen {
reg_alloc: RegAlloc,
codes: Vec<Bytecode>,
block_map: HashMap<isize, isize>,
}
impl Codegen {
pub fn new(registers: &[Register]) -> Self {
Self {
reg_alloc: RegAlloc::new(registers),
codes: Vec::new(),
block_map: HashMap::new(),
}
}
pub fn generate_code(&mut self, cfg: ControlFlowGraph, is_func: bool) -> &[Bytecode] {
let mut cfg = cfg;
Self::resort_blocks(&mut cfg);
debug!("===IR===");
let mut pos = 0;
for block in &cfg.blocks {
debug!("block({}):", block.id.as_usize());
for inst in &block.instructions {
debug!("{pos}\t{inst}");
pos += 1;
}
}
debug!("===IR===");
self.reg_alloc.arrange(&cfg);
let mut index = 0;
let mut patchs: Vec<PatchFn> = Vec::new();
if is_func {
self.codes.push(Bytecode::single(
Opcode::PushC,
Operand::new_register(Register::RBP),
));
self.codes.push(Bytecode::double(
Opcode::MovC,
Operand::new_register(Register::RBP),
Operand::new_register(Register::RSP),
));
}
let pos = self.codes.len();
patchs.push(Box::new(move |this: &mut Self| {
this.codes[pos].operands[2] = Operand::new_immd(this.reg_alloc.stack_size() as isize)
}));
self.codes.push(Bytecode::triple(
Opcode::AddC,
Operand::Register(Register::RSP),
Operand::Register(Register::RSP),
Operand::new_immd(0),
));
for block in cfg.blocks.iter() {
self.block_map
.insert(block.id.as_usize() as isize, self.codes.len() as isize);
for inst in &block.instructions {
for action in self.reg_alloc.pre_allocate(index) {
match action {
Action::Restore { stack, register } => {
self.codes.push(Bytecode::double(
Opcode::Mov,
register.into(),
Operand::Stack(stack as isize),
));
}
_ => unreachable!(),
}
}
debug!("inst: {inst:?}");
debug!("register: {}", self.reg_alloc.reg_set);
match inst.clone() {
Instruction::Call { func, args, result } => {
self.gen_call(func, &args, result);
}
Instruction::CallEx {
callable,
args,
result,
} => {
self.gen_call_ex(callable, &args, result);
}
Instruction::CallNative { func, args, result } => {
self.gen_call_native(func, &args, result);
}
Instruction::LoadArg { dst, index } => {
let dst = self.gen_operand(dst);
let stack = self.reg_alloc.load_arg(index);
self.codes.push(Bytecode::double(
Opcode::Mov,
dst,
Operand::new_stack(stack),
));
}
Instruction::LoadConst { dst, const_id } => {
let dst = self.gen_operand(dst);
self.codes.push(Bytecode::double(
Opcode::LoadConst,
dst,
const_id.to_operand(),
));
}
Instruction::LoadEnv { dst, name } => {
let dst = self.gen_operand(dst);
self.codes
.push(Bytecode::double(Opcode::LoadEnv, dst, name.to_operand()));
}
Instruction::Move { dst, src } => {
let src = self.gen_operand(src);
let dst = self.gen_operand(dst);
self.codes.push(Bytecode::double(Opcode::Mov, dst, src));
}
Instruction::UnaryOp { op, dst, src } => {
let src = self.gen_operand(src);
let dst = self.gen_operand(dst);
self.codes.push(Bytecode::double(op, dst, src));
}
Instruction::BinaryOp { op, dst, lhs, rhs } => {
let src1 = self.gen_operand(lhs);
let src2 = self.gen_operand(rhs);
let dst = self.gen_operand(dst);
self.codes.push(Bytecode::triple(op, dst, src1, src2));
}
Instruction::MakeRange {
op,
begin,
end,
result,
} => match (begin, end) {
(Some(begin), Some(end)) => {
let src1 = self.gen_operand(begin);
let src2 = self.gen_operand(end);
let dst = self.gen_operand(result);
self.codes.push(Bytecode::triple(op, dst, src1, src2));
}
(Some(begin), None) => {
let src1 = self.gen_operand(begin);
let dst = self.gen_operand(result);
self.codes
.push(Bytecode::double(Opcode::RangeTo, dst, src1));
}
(None, Some(end)) => {
let src1 = self.gen_operand(end);
let dst = self.gen_operand(result);
match op {
Opcode::RangeInclusive => {
self.codes.push(Bytecode::double(
Opcode::RangeToInclusive,
dst,
src1,
));
}
Opcode::Range => {
self.codes
.push(Bytecode::double(Opcode::RangeTo, dst, src1));
}
_ => unreachable!("invalid op"),
}
}
(None, None) => {
let dst = self.gen_operand(result);
self.codes.push(Bytecode::single(Opcode::RangeFull, dst));
}
},
Instruction::MakeIterator {
src: iter,
dst: result,
} => {
let src = self.gen_operand(iter);
let dst = self.gen_operand(result);
self.codes
.push(Bytecode::double(Opcode::MakeIter, dst, src));
}
Instruction::IteratorHasNext { iter, dst: result } => {
let src = self.gen_operand(iter);
let dst = self.gen_operand(result);
self.codes
.push(Bytecode::double(Opcode::IterHasNext, dst, src));
}
Instruction::IterateNext { iter, dst: next } => {
let src = self.gen_operand(iter);
let dst = self.gen_operand(next);
self.codes
.push(Bytecode::double(Opcode::IterNext, dst, src));
}
Instruction::MakeArray { dst } => {
let dst = self.gen_operand(dst);
self.codes.push(Bytecode::single(Opcode::MakeArray, dst));
}
Instruction::ArrayPush { array, value } => {
let array = self.gen_operand(array);
let value = self.gen_operand(value);
self.codes
.push(Bytecode::double(Opcode::ArrayPush, array, value));
}
Instruction::MakeMap { dst } => {
let dst = self.gen_operand(dst);
self.codes.push(Bytecode::single(Opcode::MakeMap, dst));
}
Instruction::IndexSet {
object,
index: idx,
value,
} => {
let object = self.gen_operand(object);
let idx = self.gen_operand(idx);
let value = self.gen_operand(value);
self.codes
.push(Bytecode::triple(Opcode::IndexSet, object, idx, value));
}
Instruction::IndexGet {
dst,
object,
index: idx,
} => {
let dst = self.gen_operand(dst);
let object = self.gen_operand(object);
let idx = self.gen_operand(idx);
self.codes
.push(Bytecode::triple(Opcode::IndexGet, dst, object, idx));
}
Instruction::MakeSlice { dst, object, range } => {
let dst = self.gen_operand(dst);
let object = self.gen_operand(object);
let range = self.gen_operand(range);
self.codes
.push(Bytecode::triple(Opcode::MakeSlice, dst, object, range));
}
Instruction::PropertyCall {
object,
property,
args,
result,
} => {
self.gen_prop_call(object, property, &args, result);
}
Instruction::PropertyGet {
dst,
object,
property,
} => {
let dst = self.gen_operand(dst);
let object = self.gen_operand(object);
let property = self.gen_operand(property);
self.codes
.push(Bytecode::triple(Opcode::PropGet, dst, object, property));
}
Instruction::PropertySet {
object,
property,
value,
} => {
let object = self.gen_operand(object);
let property = self.gen_operand(property);
let value = self.gen_operand(value);
self.codes
.push(Bytecode::triple(Opcode::PropSet, object, property, value));
}
Instruction::Return { value } => {
if let Some(v) = value {
let ret = self.gen_operand(v);
self.codes.push(Bytecode::double(
Opcode::Mov,
Operand::new_register(Register::RV),
ret,
));
}
if is_func {
self.codes.push(Bytecode::double(
Opcode::MovC,
Operand::new_register(Register::RSP),
Operand::new_register(Register::RBP),
));
self.codes
.push(Bytecode::single(Opcode::PopC, Register::RBP.into()));
}
self.codes.push(Bytecode::empty(Opcode::Ret));
}
Instruction::Br { dst } => {
let dst = self.gen_operand(dst);
let pos = self.codes.len();
patchs.push(Box::new(move |this: &mut Self| {
let dst = this.codes[pos].operands[0].as_immd();
this.codes[pos].operands[0] =
Operand::new_immd(this.block_map[&dst] - pos as isize);
}));
self.codes.push(Bytecode::single(Opcode::Br, dst));
}
Instruction::BrIf {
condition,
true_blk,
false_blk,
} => {
let condition = self.gen_operand(condition);
let true_blk = self.gen_operand(true_blk);
let false_blk = self.gen_operand(false_blk);
let pos = self.codes.len();
patchs.push(Box::new(move |this: &mut Self| {
let true_blk = this.codes[pos].operands[1].as_immd();
this.codes[pos].operands[1] =
Operand::new_immd(this.block_map[&true_blk] - pos as isize);
let false_blk = this.codes[pos].operands[2].as_immd();
this.codes[pos].operands[2] =
Operand::new_immd(this.block_map[&false_blk] - pos as isize);
}));
self.codes.push(Bytecode::triple(
Opcode::BrIf,
condition,
true_blk,
false_blk,
));
}
Instruction::Halt => {
self.codes.push(Bytecode::empty(Opcode::Halt));
}
Instruction::Await { promise, dst } => {
let promise = self.gen_operand(promise);
let dst = self.gen_operand(dst);
self.codes
.push(Bytecode::double(Opcode::Await, dst, promise));
}
}
for action in self.reg_alloc.post_allocate(index) {
match action {
Action::Spill { stack, register } => {
self.codes.push(Bytecode::double(
Opcode::Mov,
Operand::Stack(stack as isize),
register.into(),
));
}
_ => unreachable!(),
}
}
index += 1;
}
}
for patch in patchs {
patch(self);
}
&self.codes
}
fn gen_call(&mut self, func: Value, args: &[Value], result: Value) {
let in_use_registers = self.reg_alloc.in_use_registers();
for reg in in_use_registers.iter().copied() {
self.codes.push(Bytecode::single(Opcode::Push, reg.into()));
}
for arg in args.iter().rev() {
let arg = self.gen_operand(*arg);
self.codes.push(Bytecode::single(Opcode::Push, arg));
}
self.codes
.push(Bytecode::single(Opcode::Call, func.to_operand()));
self.codes.push(Bytecode::triple(
Opcode::SubC,
Operand::Register(Register::RSP),
Operand::Register(Register::RSP),
Operand::new_immd(args.len() as isize),
));
for reg in in_use_registers.iter().rev().copied() {
self.codes.push(Bytecode::single(Opcode::Pop, reg.into()));
}
let result_reg = self.gen_operand(result);
self.codes.push(Bytecode::double(
Opcode::Mov,
result_reg,
Operand::new_register(Register::RV),
));
}
fn gen_call_ex(&mut self, func: Value, args: &[Value], result: Value) {
let in_use_registers = self.reg_alloc.in_use_registers();
for reg in in_use_registers.iter().copied() {
self.codes.push(Bytecode::single(Opcode::Push, reg.into()));
}
for arg in args.iter().rev() {
let arg = self.gen_operand(*arg);
self.codes.push(Bytecode::single(Opcode::Push, arg));
}
let callable = self.gen_operand(func);
self.codes.push(Bytecode::single(Opcode::CallEx, callable));
self.codes.push(Bytecode::triple(
Opcode::SubC,
Operand::Register(Register::RSP),
Operand::Register(Register::RSP),
Operand::new_immd(args.len() as isize),
));
for reg in in_use_registers.iter().rev().copied() {
self.codes.push(Bytecode::single(Opcode::Pop, reg.into()));
}
let result_reg = self.gen_operand(result);
self.codes.push(Bytecode::double(
Opcode::Mov,
result_reg,
Operand::new_register(Register::RV),
));
}
fn gen_call_native(&mut self, func: Value, args: &[Value], result: Value) {
for arg in args.iter().rev() {
let arg = self.gen_operand(*arg);
self.codes.push(Bytecode::single(Opcode::Push, arg));
}
self.codes.push(Bytecode::single(
Opcode::PushC,
Operand::new_register(Register::RBP),
));
self.codes.push(Bytecode::double(
Opcode::MovC,
Operand::new_register(Register::RBP),
Operand::new_register(Register::RSP),
));
let callable = self.gen_operand(func);
self.codes.push(Bytecode::double(
Opcode::CallNative,
callable,
Operand::new_immd(args.len() as isize),
));
self.codes.push(Bytecode::single(
Opcode::PopC,
Operand::new_register(Register::RBP),
));
self.codes.push(Bytecode::triple(
Opcode::SubC,
Operand::Register(Register::RSP),
Operand::Register(Register::RSP),
Operand::new_immd(args.len() as isize),
));
let result_reg = self.gen_operand(result);
self.codes.push(Bytecode::double(
Opcode::Mov,
result_reg,
Operand::new_register(Register::RV),
));
}
fn gen_prop_call(&mut self, object: Value, property: Value, args: &[Value], result: Value) {
for arg in args.iter().rev() {
let arg = self.gen_operand(*arg);
self.codes.push(Bytecode::single(Opcode::Push, arg));
}
self.codes.push(Bytecode::single(
Opcode::PushC,
Operand::new_register(Register::RBP),
));
self.codes.push(Bytecode::double(
Opcode::MovC,
Operand::new_register(Register::RBP),
Operand::new_register(Register::RSP),
));
let callable = self.gen_operand(object);
let prop = self.gen_operand(property);
self.codes.push(Bytecode::triple(
Opcode::PropCall,
callable,
prop,
Operand::new_immd(args.len() as isize),
));
self.codes.push(Bytecode::single(
Opcode::PopC,
Operand::new_register(Register::RBP),
));
self.codes.push(Bytecode::triple(
Opcode::SubC,
Operand::Register(Register::RSP),
Operand::Register(Register::RSP),
Operand::new_immd(args.len() as isize),
));
let result_reg = self.gen_operand(result);
self.codes.push(Bytecode::double(
Opcode::Mov,
result_reg,
Operand::new_register(Register::RV),
));
}
fn gen_operand(&mut self, value: Value) -> Operand {
match value {
Value::Primitive(v) => Operand::new_primitive(v),
Value::Constant(id) => Operand::new_immd(id.as_usize() as isize),
Value::Function(id) => Operand::new_symbol(id.as_usize() as u32),
Value::Block(id) => Operand::new_immd(id.as_usize() as isize),
Value::Variable(_) => {
let register = self.reg_alloc.alloc(value);
Operand::new_register(register)
}
}
}
fn resort_blocks(control_flow_graph: &mut ControlFlowGraph) {
let mut sorted = sort_graph_blocks(control_flow_graph);
let mut iter = sorted.iter_mut().peekable();
while let Some(block) = iter.next() {
let next_block_id = iter.peek().map(|b| b.id);
if let Some(Instruction::Br { dst }) = block.instructions.last() {
if next_block_id == dst.as_block() {
block.instructions.pop();
}
}
}
control_flow_graph.blocks = sorted;
}
}
fn sort_graph_blocks(control_flow_graph: &ControlFlowGraph) -> Vec<Block> {
let graph = &control_flow_graph.graph;
let entry = control_flow_graph.entry().expect("no entry block");
let start = control_flow_graph.block_node_map[&entry];
let mut dfs = DfsPostOrder::new(graph, start);
let mut sorted = Vec::new();
while let Some(node) = dfs.next(graph) {
let block_id = graph[node];
sorted.push(block_id);
}
sorted.reverse();
sorted
.into_iter()
.map(|block_id| {
let block = control_flow_graph
.get_block(block_id)
.expect("no such block");
block.clone()
})
.collect()
}
#[cfg(test)]
mod tests {
use crate::{
bytecode::{Opcode, Primitive},
ir::instruction::VariableId,
};
use super::*;
#[test]
fn test_register_allocator() {
let instructions = vec![
Instruction::Move {
dst: Value::Variable(VariableId::new(0)),
src: Value::Primitive(Primitive::from(3.141592653589793)),
},
Instruction::Move {
dst: Value::Variable(VariableId::new(1)),
src: Value::Primitive(Primitive::from(718281828459045)),
},
Instruction::BinaryOp {
op: Opcode::Mulx,
dst: Value::Variable(VariableId::new(2)),
lhs: Value::Variable(VariableId::new(0)),
rhs: Value::Variable(VariableId::new(1)),
},
Instruction::Return {
value: Some(Value::Variable(VariableId::new(2))),
},
];
let mut cfg = ControlFlowGraph::new();
let entry = cfg.create_block_with_instructions("test", instructions);
cfg.set_entry(entry);
let mut codegen = Codegen::new(&Register::small_general());
let bytecodes = codegen.generate_code(cfg, false);
for bytecode in bytecodes {
println!("{bytecode};");
}
}
#[test]
fn test_register_allocator2() {
let instructions = vec![
Instruction::BinaryOp {
op: Opcode::Addx,
dst: Value::Variable(VariableId::new(0)),
lhs: Value::Primitive(Primitive::from(1)),
rhs: Value::Primitive(Primitive::from(1)),
},
Instruction::BinaryOp {
op: Opcode::Mulx,
dst: Value::Variable(VariableId::new(1)),
lhs: Value::Primitive(Primitive::from(2)),
rhs: Value::Primitive(Primitive::from(2)),
},
Instruction::BinaryOp {
op: Opcode::Addx,
dst: Value::Variable(VariableId::new(2)),
lhs: Value::Variable(VariableId::new(0)),
rhs: Value::Variable(VariableId::new(1)),
},
Instruction::Return {
value: Some(Value::Variable(VariableId::new(2))),
},
];
let mut cfg = ControlFlowGraph::new();
let entry = cfg.create_block_with_instructions("test", instructions);
cfg.set_entry(entry);
let mut codegen = Codegen::new(&Register::small_general());
let bytecodes = codegen.generate_code(cfg, false);
for bytecode in bytecodes {
println!("{bytecode};");
}
}
#[test]
fn test_codegen() {
let mut cfg = ControlFlowGraph::new();
let block1 = cfg.create_block(None);
let block2 = cfg.create_block(None);
let block3 = cfg.create_block(None);
let block4 = cfg.create_block(None);
let block5 = cfg.create_block(None);
cfg.set_entry(block1);
cfg.switch_to_block(block1);
cfg.emit(Instruction::Move {
dst: Value::Variable(VariableId::new(0)),
src: Value::Primitive(Primitive::from(0)),
});
cfg.switch_to_block(block2);
cfg.emit(Instruction::MakeRange {
op: Opcode::RangeInclusive,
result: Value::Variable(VariableId::new(1)),
begin: Some(Value::Primitive(Primitive::from(0))),
end: Some(Value::Primitive(Primitive::from(2))),
});
cfg.emit(Instruction::MakeIterator {
src: Value::Variable(VariableId::new(1)),
dst: Value::Variable(VariableId::new(2)),
});
cfg.switch_to_block(block3);
cfg.emit(Instruction::IteratorHasNext {
iter: Value::Variable(VariableId::new(2)),
dst: Value::Variable(VariableId::new(3)),
});
cfg.emit(Instruction::BrIf {
condition: Value::Variable(VariableId::new(3)),
true_blk: Value::Block(block4),
false_blk: Value::Block(block5),
});
cfg.switch_to_block(block4);
cfg.emit(Instruction::IterateNext {
iter: Value::Variable(VariableId::new(2)),
dst: Value::Variable(VariableId::new(4)),
});
cfg.emit(Instruction::Move {
dst: Value::Variable(VariableId::new(5)),
src: Value::Variable(VariableId::new(4)),
});
cfg.emit(Instruction::BinaryOp {
op: Opcode::Addx,
dst: Value::Variable(VariableId::new(6)),
lhs: Value::Variable(VariableId::new(0)),
rhs: Value::Variable(VariableId::new(5)),
});
cfg.emit(Instruction::Move {
dst: Value::Variable(VariableId::new(0)),
src: Value::Variable(VariableId::new(6)),
});
cfg.emit(Instruction::Br {
dst: Value::Block(block3),
});
cfg.switch_to_block(block5);
cfg.emit(Instruction::Return {
value: Some(Value::Variable(VariableId::new(0))),
});
let mut codegen = Codegen::new(&Register::small_general());
let bytecodes = codegen.generate_code(cfg, false);
for bytecode in bytecodes {
println!("{bytecode};");
}
}
}