use std::io::Read;
use std::io::Write;
use macroassembler::assembler::abstract_macro_assembler::AbsoluteAddress;
use macroassembler::assembler::abstract_macro_assembler::Address;
use macroassembler::assembler::abstract_macro_assembler::Jump;
use macroassembler::assembler::abstract_macro_assembler::Label;
use macroassembler::assembler::link_buffer::LinkBuffer;
use macroassembler::assembler::RelationalCondition;
use macroassembler::assembler::*;
use macroassembler::jit::gpr_info::*;
use macroassembler::jit::helpers::AssemblyHelpers;
use macroassembler::wtf::executable_memory_handle::CodeRef;
pub struct BfJIT {
ctx: CGContext,
}
pub struct CGContext {
pub opt_level: u8,
}
#[derive(Copy, Clone, Debug)]
enum Token {
Forward(u32),
Backward(u32),
Add(u8),
Sub(u8),
Output,
Input,
LoopBegin,
LoopEnd,
LoopToZero,
LoopToAdd,
}
impl BfJIT {
pub fn new(ctx: CGContext) -> Self {
Self { ctx }
}
pub fn translate(&self, disasm: bool, input: &str) -> CodeRef {
let mut tokens: Vec<Token> = vec![];
let mut chars = input.chars().peekable();
while let Some(c) = chars.next() {
match c {
'>' => {
let mut n: u32 = 1;
if self.ctx.opt_level > 0 {
while chars.peek() == Some(&'>') {
n += 1;
chars.next().unwrap();
}
}
tokens.push(Token::Forward(n));
}
'<' => {
let mut n: u32 = 1;
if self.ctx.opt_level > 0 {
while chars.peek() == Some(&'<') {
n += 1;
chars.next().unwrap();
}
}
tokens.push(Token::Backward(n));
}
'+' => {
let mut n: u8 = 1;
if self.ctx.opt_level > 0 {
while chars.peek() == Some(&'+') {
n += 1;
chars.next().unwrap();
}
}
tokens.push(Token::Add(n));
}
'-' => {
let mut n: u8 = 1;
if self.ctx.opt_level > 0 {
while chars.peek() == Some(&'-') {
n += 1;
chars.next().unwrap();
}
}
tokens.push(Token::Sub(n));
}
'.' => tokens.push(Token::Output),
',' => tokens.push(Token::Input),
'[' => tokens.push(Token::LoopBegin),
']' => tokens.push(Token::LoopEnd),
_ => {}
};
}
if self.ctx.opt_level > 0 {
tokens = self.opt_inst_combine(&tokens);
}
self.do_translate(disasm, &tokens)
}
fn opt_inst_combine(&self, tokens: &[Token]) -> Vec<Token> {
let mut ret: Vec<Token> = vec![];
let mut i: usize = 0;
loop {
if i >= tokens.len() {
break;
}
match tokens[i..] {
[Token::LoopBegin, Token::Sub(1), Token::LoopEnd, ..] => {
ret.push(Token::LoopToZero);
i += 3;
}
[Token::LoopBegin, Token::Sub(1), Token::Forward(1), Token::Add(1), Token::Backward(1), Token::LoopEnd, ..] =>
{
ret.push(Token::LoopToAdd);
i += 6;
}
_ => {
ret.push(tokens[i]);
i += 1;
}
}
}
ret
}
fn do_translate(&self, disasm: bool, input: &[Token]) -> CodeRef {
let mut jmps_to_end: Vec<(Label, Jump)> = vec![];
let mut masm = TargetMacroAssembler::new();
masm.emit_function_prologue();
masm.mov(ARGUMENT_GPR0, NON_PRESERVED_NON_RETURN_GPR);
for t in input {
match *t {
Token::Forward(n) => {
masm.comment(format!("forward {}", n));
masm.add64(n as i32, NON_PRESERVED_NON_RETURN_GPR);
}
Token::Backward(n) => {
masm.comment(format!("backward {}", n));
masm.sub64(n as i32, NON_PRESERVED_NON_RETURN_GPR);
}
Token::Add(n) => {
masm.comment(format!("add {}", n));
masm.load8_signed_extend_to_32(
Address::new(NON_PRESERVED_NON_RETURN_GPR, 0),
T0,
);
masm.add32(n as i32, T0);
masm.store8(T0, Address::new(NON_PRESERVED_NON_RETURN_GPR, 0));
}
Token::Sub(n) => {
masm.comment(format!("sub {}", n));
masm.load8_signed_extend_to_32(
Address::new(NON_PRESERVED_NON_RETURN_GPR, 0),
T0,
);
masm.add32(-(n as i32), T0);
masm.store8(T0, Address::new(NON_PRESERVED_NON_RETURN_GPR, 0));
}
Token::Output => {
masm.comment("output");
masm.load8_signed_extend_to_32(
Address::new(NON_PRESERVED_NON_RETURN_GPR, 0),
ARGUMENT_GPR0,
);
masm.push_to_save(NON_PRESERVED_NON_RETURN_GPR);
masm.call_op(Some(AbsoluteAddress::new(putchar as _)));
masm.pop_to_restore(NON_PRESERVED_NON_RETURN_GPR);
}
Token::Input => {
masm.comment("input");
masm.call_op(Some(AbsoluteAddress::new(getchr as _)));
masm.store8(
RETURN_VALUE_GPR,
Address::new(NON_PRESERVED_NON_RETURN_GPR, 0),
);
}
Token::LoopBegin => {
masm.comment("loop begin");
masm.load8(Address::new(NON_PRESERVED_NON_RETURN_GPR, 0), T0);
let jend = masm.branch32(RelationalCondition::Equal, T0, 0i32);
let start = masm.label();
jmps_to_end.push((start, jend));
}
Token::LoopEnd => {
masm.comment("loop end");
let (start, jend) = jmps_to_end.pop().unwrap();
masm.load8_signed_extend_to_32(
Address::new(NON_PRESERVED_NON_RETURN_GPR, 0),
T0,
);
let j = masm.branch32(RelationalCondition::NotEqual, T0, 0i32);
j.link_to(&mut masm, start);
jend.link(&mut masm);
}
Token::LoopToZero => {
masm.comment("loop to zero");
masm.store8(0i32, Address::new(NON_PRESERVED_NON_RETURN_GPR, 0));
}
Token::LoopToAdd => {
masm.comment("loop to add");
masm.load8(Address::new(NON_PRESERVED_NON_RETURN_GPR, 0), T0);
masm.load8(Address::new(NON_PRESERVED_NON_RETURN_GPR, 1), T1);
masm.add32(T0, T1);
masm.store8(0i32, Address::new(NON_PRESERVED_NON_RETURN_GPR, 0));
masm.store8(T1, Address::new(NON_PRESERVED_NON_RETURN_GPR, 1));
}
}
}
masm.emit_function_epilogue();
masm.ret();
assert!(jmps_to_end.is_empty());
let mut buffer = LinkBuffer::from_macro_assembler(&mut masm).unwrap();
let mut fmt = String::new();
let code = buffer
.finalize_with_disassembly(disasm, "brainfuck", &mut fmt)
.unwrap();
println!("{}", fmt);
code
}
}
extern "C" fn putchar(x: u8) {
let mut out = ::std::io::stdout();
out.write_all(&[x]).unwrap();
out.flush().unwrap();
}
extern "C" fn getchr() -> u8 {
let mut buf = [0u8; 1];
std::io::stdin().read_exact(&mut buf).unwrap();
buf[0]
}
use pico_args::*;
use std::ffi::OsStr;
use std::path::PathBuf;
fn parse_path(s: &OsStr) -> Result<PathBuf, &'static str> {
Ok(s.into())
}
fn main() {
let mut args = Arguments::from_env();
let disasm = args.contains(["-d", "--disasm"]);
let opt = args.contains(["-O", "--optimize"]);
let input = args
.opt_value_from_os_str(["-i", "--input"], parse_path)
.unwrap();
if input.is_none() {
println!(
"No input provided. To compile brainfuck program run `brainfuck -i <program_name>"
);
return;
}
let input = input.unwrap();
let jit = BfJIT::new(CGContext {
opt_level: opt as u8,
});
let compile_start = std::time::Instant::now();
let code = jit.translate(disasm, &std::fs::read_to_string(input).unwrap());
println!(
"Compiled in {:.2}ms",
compile_start.elapsed().as_micros() as f64 / 1000.0
);
let fun = unsafe { std::mem::transmute::<_, extern "C" fn(*mut u8)>(code.start()) };
let mut mem = vec![0u8; 100 * 1024];
fun(mem.as_mut_ptr());
drop(code);
}