Witch of Git - ivy/blob - src/trans/x64.rs

   1 use crate::ast;
   2 use crate::trans::{Code, FnName, Func, Program, SsaName, Var};
   3 use std::{
   4     borrow::Cow,
   5     collections::HashMap,
   6     fmt,
   7     io::{self, Write},
   8 };
   9
  10 #[cfg(target_os = "macos")]
  11 macro_rules! extern_label {
  12     ($l:literal) => {
  13         Label::Extern(concat!("_", $l))
  14     };
  15 }
  16 #[cfg(not(target_os = "macos"))]
  17 macro_rules! extern_label {
  18     ($l:literal) => {
  19         Label::Extern($l)
  20     };
  21 }
  22
  23 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
  24 pub enum Label {
  25     Global(u32),
  26     Function(FnName),
  27     Extern(&'static str),
  28     Builtin(&'static str),
  29 }
  30
  31 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
  32 pub enum Addr {
  33     Rip(Label),
  34     Off(Reg, i32),
  35 }
  36
  37 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
  38 #[allow(unused)]
  39 pub enum Reg {
  40     Rax,
  41     Rbx,
  42     Rcx,
  43     Rdx,
  44     Rsi,
  45     Rdi,
  46     Rbp,
  47     Rsp,
  48     R8,
  49     R9,
  50     R10,
  51     R11,
  52     R12,
  53     R13,
  54     R14,
  55     R15,
  56 }
  57
  58 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
  59 pub enum Cfi {
  60     Def(Reg, i32),
  61 }
  62
  63 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
  64 pub enum Inst {
  65     Push(Reg),
  66     Pop(Reg),
  67     Cfi(Cfi),
  68     Load(Reg, Addr),
  69     Store(Addr, Reg),
  70     Mov(Reg, Reg),
  71     Imm(Reg, i64),
  72     Lea(Reg, Addr),
  73     Call(Label),
  74     Jmp(Label),
  75     AddI(Reg, u32),
  76     SubI(Reg, u32),
  77     Ret,
  78 }
  79
  80 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
  81 pub struct Definition {
  82     name: Label,
  83     body: Cow<'static, [Inst]>,
  84 }
  85
  86 pub struct Global {
  87     name: Label,
  88     value: i64,
  89 }
  90
  91 pub fn write_compile(
  92     builtins: &HashMap<String, ast::Name>,
  93     out: &mut impl Write,
  94     prog: &Program,
  95 ) -> io::Result<()> {
  96     let (entry, globals, defns) = compile(&prog);
  97     struct Entry {
  98         number: u32,
  99         defn: Definition,
 100         params: u16,
 101     }
 102     let builtins = builtins
 103         .iter()
 104         .map(|(name, id)| {
 105             let number = id.global_number().unwrap();
 106             let (params, name, code) = builtin(&name);
 107             Entry {
 108                 number,
 109                 defn: Definition {
 110                     name: Label::Builtin(name),
 111                     body: Cow::Borrowed(code),
 112                 },
 113                 params,
 114             }
 115         })
 116         .collect::<Vec<_>>();
 117     writeln!(out, ".data")?;
 118     for global in globals {
 119         writeln!(out, "{}", global)?;
 120     }
 121     writeln!(out)?;
 122     writeln!(
 123         out,
 124         r#".text
 125 .global _start
 126 _start:
 127     sub rsp, 15
 128     and spl, 0xF0"#
 129     )?;
 130
 131     for builtin in &builtins {
 132         writeln!(
 133             out,
 134             r#"
 135     lea rdi, [rip + {}]
 136     mov rsi, {}
 137     mov rdx, 0
 138     call {}
 139     mov [rip + {}], rax"#,
 140             builtin.defn.name,
 141             builtin.params,
 142             extern_label!("ivy_make_lam"),
 143             Label::Global(builtin.number),
 144         )?;
 145     }
 146
 147     writeln!(
 148         out,
 149         r#"
 150     call {}
 151     mov rdi, 0
 152     call {}
 153 "#,
 154         entry,
 155         extern_label!("ivy_exit"),
 156     )?;
 157     for builtin in &builtins {
 158         writeln!(out, "{}", builtin.defn)?;
 159         writeln!(out)?;
 160     }
 161     for defn in defns {
 162         writeln!(out, "{}", defn)?;
 163         writeln!(out)?;
 164     }
 165     Ok(())
 166 }
 167
 168 pub fn compile(prog: &Program) -> (Label, Vec<Global>, Vec<Definition>) {
 169     let globals = prog
 170         .globals
 171         .iter()
 172         .map(|&global| Global {
 173             name: Label::Global(global),
 174             value: 0,
 175         })
 176         .collect();
 177     let defns = prog
 178         .functions
 179         .values()
 180         .map(|func| compile_func(func))
 181         .collect();
 182     (Label::Function(prog.top), globals, defns)
 183 }
 184
 185 fn even_up(x: usize) -> usize {
 186     x + (x & 1)
 187 }
 188
 189 fn compile_func(func: &Func) -> Definition {
 190     fn stack(x: &SsaName) -> Addr {
 191         Addr::Off(Reg::Rsp, x.0 as i32 * 8)
 192     }
 193
 194     const HEADER_SIZE: i32 = 0x18;
 195
 196     fn param(x: u16) -> Addr {
 197         Addr::Off(Reg::Rbx, HEADER_SIZE + x as i32 * 8)
 198     }
 199
 200     let upvar = move |x: u16| {
 201         let params = func.params as i32 * 8;
 202         Addr::Off(Reg::Rbx, HEADER_SIZE + params + x as i32 * 8)
 203     };
 204
 205     let mut body = Vec::new();
 206     let stack_slots = even_up(func.order.len());
 207     body.push(Inst::Push(Reg::Rbx));
 208     body.push(Inst::Cfi(Cfi::Def(Reg::Rsp, 16)));
 209     body.push(Inst::Mov(Reg::Rbx, Reg::Rdi));
 210     body.push(Inst::SubI(Reg::Rsp, stack_slots as u32 * 8));
 211     body.push(Inst::Cfi(Cfi::Def(Reg::Rsp, stack_slots as i32 * 8 + 16)));
 212     let address = |v| match v {
 213         Var::Global(x) => Addr::Rip(Label::Global(x)),
 214         Var::Param(x) => param(x),
 215         Var::Upvar(x) => upvar(x),
 216         Var::Ssa(s) => stack(&s),
 217     };
 218     let get = |r, s| match func.block[s] {
 219         Code::Load(v) => Inst::Load(r, address(v)),
 220         Code::Num(v) => Inst::Imm(r, (v << 1) | 1),
 221         _ => Inst::Load(r, stack(s)),
 222     };
 223     for ssa in &func.order {
 224         match &func.block[ssa] {
 225             // We don't need to generate code for loads or nums themselves, they're consulted by other accesses
 226             Code::Load(_) => (),
 227             Code::Num(_) => (),
 228             Code::StoreGlobal(x, s) => {
 229                 body.push(get(Reg::Rax, s));
 230                 body.push(Inst::Store(Addr::Rip(Label::Global(*x)), Reg::Rax));
 231             }
 232             Code::MakeLam {
 233                 name,
 234                 upvars,
 235                 params,
 236             } => {
 237                 body.push(Inst::Lea(Reg::Rdi, Addr::Rip(Label::Function(*name))));
 238                 body.push(Inst::Imm(Reg::Rsi, *params as i64));
 239                 body.push(Inst::Imm(Reg::Rdx, upvars.len() as i64));
 240                 body.push(Inst::Call(extern_label!("ivy_make_lam")));
 241                 body.push(Inst::Store(stack(ssa), Reg::Rax));
 242                 for (i, s) in upvars.iter().enumerate() {
 243                     body.push(get(Reg::Rdi, s));
 244                     body.push(Inst::Store(
 245                         Addr::Off(Reg::Rax, HEADER_SIZE + (*params as i32) * 8 + i as i32 * 8),
 246                         Reg::Rdi,
 247                     ));
 248                     body.push(Inst::Call(extern_label!("ivy_incref")));
 249                     body.push(Inst::Load(Reg::Rax, stack(ssa)));
 250                 }
 251             }
 252             Code::App(terms) => {
 253                 let mut terms = terms.iter();
 254                 body.push(get(Reg::Rdi, terms.next().expect("a function to apply")));
 255                 match terms.next() {
 256                     Some(s) => body.push(get(Reg::Rsi, s)),
 257                     None => body.push(Inst::Imm(Reg::Rsi, 0)),
 258                 }
 259                 body.push(Inst::Call(extern_label!("ivy_app")));
 260                 while let Some(term) = terms.next() {
 261                     body.push(Inst::Mov(Reg::Rdi, Reg::Rax));
 262                     body.push(get(Reg::Rsi, term));
 263                     body.push(Inst::Call(extern_label!("ivy_app")));
 264                 }
 265                 body.push(Inst::Store(stack(ssa), Reg::Rax));
 266             }
 267         }
 268     }
 269     match func.block[&func.result] {
 270         Code::Load(v) => body.push(Inst::Load(Reg::Rax, address(v))),
 271         Code::Num(n) => body.push(Inst::Imm(Reg::Rax, (n << 1) | 1i64)),
 272         _ => (),
 273     }
 274     body.push(Inst::AddI(Reg::Rsp, stack_slots as u32 * 8));
 275     body.push(Inst::Cfi(Cfi::Def(Reg::Rsp, 16)));
 276     body.push(Inst::Pop(Reg::Rbx));
 277     body.push(Inst::Cfi(Cfi::Def(Reg::Rsp, 8)));
 278     body.push(Inst::Ret);
 279     Definition {
 280         name: Label::Function(func.name),
 281         body: Cow::Owned(body),
 282     }
 283 }
 284
 285 impl fmt::Display for Definition {
 286     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 287         writeln!(fmt, "    .globl {}", self.name)?;
 288         match self.body.len() {
 289             0 => write!(fmt, "{}:", self.name),
 290             1 => write!(fmt, "{}: {}", self.name, self.body[0]),
 291             _ => {
 292                 writeln!(fmt, "    .cfi_startproc")?;
 293                 writeln!(fmt, "{}:", self.name)?;
 294                 for inst in self.body.as_ref() {
 295                     writeln!(fmt, "    {}", inst)?;
 296                 }
 297                 write!(fmt, "    .cfi_endproc")
 298             }
 299         }
 300     }
 301 }
 302
 303 impl fmt::Display for Global {
 304     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 305         write!(fmt, "{}: .quad {}", self.name, self.value)
 306     }
 307 }
 308
 309 impl fmt::Display for Label {
 310     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 311         match self {
 312             Label::Global(g) => write!(fmt, "IVY_GLOBAL${}", g),
 313             Label::Function(f) => write!(fmt, "ivy_fn${}", f.0),
 314             Label::Extern(l) => write!(fmt, "{}", l),
 315             Label::Builtin(l) => write!(fmt, "ivy_builtin${}", l),
 316         }
 317     }
 318 }
 319
 320 impl fmt::Display for Reg {
 321     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 322         match self {
 323             Reg::Rax => write!(fmt, "rax"),
 324             Reg::Rbx => write!(fmt, "rbx"),
 325             Reg::Rcx => write!(fmt, "rcx"),
 326             Reg::Rdx => write!(fmt, "rdx"),
 327             Reg::Rsi => write!(fmt, "rsi"),
 328             Reg::Rdi => write!(fmt, "rdi"),
 329             Reg::Rbp => write!(fmt, "rbp"),
 330             Reg::Rsp => write!(fmt, "rsp"),
 331             Reg::R8 => write!(fmt, "r8"),
 332             Reg::R9 => write!(fmt, "r9"),
 333             Reg::R10 => write!(fmt, "r10"),
 334             Reg::R11 => write!(fmt, "r11"),
 335             Reg::R12 => write!(fmt, "r12"),
 336             Reg::R13 => write!(fmt, "r13"),
 337             Reg::R14 => write!(fmt, "r14"),
 338             Reg::R15 => write!(fmt, "r15"),
 339         }
 340     }
 341 }
 342
 343 impl fmt::Display for Addr {
 344     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 345         match self {
 346             Addr::Rip(l) => write!(fmt, "[rip + {}]", l),
 347             Addr::Off(r, o) => write!(fmt, "[{} + 0x{:02x}]", r, o),
 348         }
 349     }
 350 }
 351
 352 impl fmt::Display for Cfi {
 353     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 354         match self {
 355             Cfi::Def(r, i) => write!(fmt, ".cfi_def_cfa {}, {}", r, i),
 356         }
 357     }
 358 }
 359
 360 impl fmt::Display for Inst {
 361     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 362         match self {
 363             Inst::Push(r) => write!(fmt, "push {}", r),
 364             Inst::Pop(r) => write!(fmt, "pop {}", r),
 365             Inst::Cfi(c) => write!(fmt, "{}", c),
 366             Inst::Mov(d, s) => write!(fmt, "mov {}, {}", d, s),
 367             Inst::Load(r, a) => write!(fmt, "mov {}, {}", r, a),
 368             Inst::Store(a, r) => write!(fmt, "mov {}, {}", a, r),
 369             Inst::Imm(r, i) => write!(fmt, "mov {}, {}", r, i),
 370             Inst::Lea(r, a) => write!(fmt, "lea {}, {}", r, a),
 371             Inst::Call(l) => write!(fmt, "call {}", l),
 372             Inst::Jmp(l) => write!(fmt, "jmp {}", l),
 373             Inst::AddI(r, x) => write!(fmt, "add {}, {}", r, x),
 374             Inst::SubI(r, x) => write!(fmt, "sub {}, {}", r, x),
 375             Inst::Ret => write!(fmt, "ret"),
 376         }
 377     }
 378 }
 379
 380 fn builtin(name: &str) -> (u16, &'static str, &'static [Inst]) {
 381     match name {
 382         "debug" => (
 383             1,
 384             "debug",
 385             &[
 386                 Inst::Load(Reg::Rdi, Addr::Off(Reg::Rdi, 0x18)),
 387                 Inst::Jmp(extern_label!("ivy_debug")),
 388             ],
 389         ),
 390         name => panic!("Unsupported builtin '{}' for x64", name),
 391     }
 392 }