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 Cond {
  60     Le,
  61 }
  62
  63 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
  64 pub enum Cfi {
  65     Def(Reg, i32),
  66 }
  67
  68 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
  69 pub enum Inst {
  70     Push(Reg),
  71     Pop(Reg),
  72     Cfi(Cfi),
  73     Load(Reg, Addr),
  74     Store(Addr, Reg),
  75     Mov(Reg, Reg),
  76     Imm(Reg, i64),
  77     Lea(Reg, Addr),
  78     Call(Label),
  79     Jmp(Label),
  80     Add(Reg, Reg),
  81     Sub(Reg, Reg),
  82     AddI(Reg, u32),
  83     SubI(Reg, u32),
  84     Cmp(Reg, Reg),
  85     CLoad(Cond, Reg, Addr),
  86     CMov(Cond, Reg, Reg),
  87     Ret,
  88 }
  89
  90 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
  91 pub struct Definition {
  92     name: Label,
  93     body: Cow<'static, [Inst]>,
  94 }
  95
  96 pub struct Global {
  97     name: Label,
  98     value: i64,
  99 }
 100
 101 pub fn write_compile(
 102     builtins: &HashMap<String, ast::Name>,
 103     out: &mut impl Write,
 104     prog: &Program,
 105 ) -> io::Result<()> {
 106     let (entry, globals, defns) = compile(&prog);
 107     struct Entry {
 108         number: u32,
 109         defn: Definition,
 110         params: u16,
 111     }
 112     let builtins = builtins
 113         .iter()
 114         .map(|(name, id)| {
 115             let number = id.global_number().unwrap();
 116             let (params, name, code) = builtin(&name);
 117             Entry {
 118                 number,
 119                 defn: Definition {
 120                     name: Label::Builtin(name),
 121                     body: Cow::Borrowed(code),
 122                 },
 123                 params,
 124             }
 125         })
 126         .collect::<Vec<_>>();
 127     writeln!(out, ".data")?;
 128     for global in globals {
 129         writeln!(out, "{}", global)?;
 130     }
 131     writeln!(out)?;
 132     writeln!(
 133         out,
 134         r#".text
 135 .global _start
 136 _start:
 137     sub rsp, 15
 138     and spl, 0xF0"#
 139     )?;
 140
 141     for builtin in &builtins {
 142         writeln!(
 143             out,
 144             r#"
 145     lea rdi, [rip + {}]
 146     mov rsi, {}
 147     mov rdx, 0
 148     call {}
 149     mov [rip + {}], rax"#,
 150             builtin.defn.name,
 151             builtin.params,
 152             extern_label!("ivy_make_lam"),
 153             Label::Global(builtin.number),
 154         )?;
 155     }
 156
 157     writeln!(
 158         out,
 159         r#"
 160     call {}
 161     mov rdi, 0
 162     call {}
 163 "#,
 164         entry,
 165         extern_label!("ivy_exit"),
 166     )?;
 167     for builtin in &builtins {
 168         writeln!(out, "{}", builtin.defn)?;
 169         writeln!(out)?;
 170     }
 171     for defn in defns {
 172         writeln!(out, "{}", defn)?;
 173         writeln!(out)?;
 174     }
 175     Ok(())
 176 }
 177
 178 pub fn compile(prog: &Program) -> (Label, Vec<Global>, Vec<Definition>) {
 179     let globals = prog
 180         .globals
 181         .iter()
 182         .map(|&global| Global {
 183             name: Label::Global(global),
 184             value: 0,
 185         })
 186         .collect();
 187     let defns = prog
 188         .functions
 189         .values()
 190         .map(|func| compile_func(func))
 191         .collect();
 192     (Label::Function(prog.top), globals, defns)
 193 }
 194
 195 fn even_up(x: usize) -> usize {
 196     x + (x & 1)
 197 }
 198
 199 fn compile_func(func: &Func) -> Definition {
 200     fn stack(x: &SsaName) -> Addr {
 201         Addr::Off(Reg::Rsp, x.0 as i32 * 8)
 202     }
 203
 204     const HEADER_SIZE: i32 = 0x18;
 205
 206     fn param(x: u16) -> Addr {
 207         Addr::Off(Reg::Rbx, HEADER_SIZE + x as i32 * 8)
 208     }
 209
 210     let upvar = move |x: u16| {
 211         let params = func.params as i32 * 8;
 212         Addr::Off(Reg::Rbx, HEADER_SIZE + params + x as i32 * 8)
 213     };
 214
 215     let mut body = Vec::new();
 216     let stack_slots = even_up(func.order.len());
 217     body.push(Inst::Push(Reg::Rbx));
 218     body.push(Inst::Cfi(Cfi::Def(Reg::Rsp, 16)));
 219     body.push(Inst::Mov(Reg::Rbx, Reg::Rdi));
 220     body.push(Inst::SubI(Reg::Rsp, stack_slots as u32 * 8));
 221     body.push(Inst::Cfi(Cfi::Def(Reg::Rsp, stack_slots as i32 * 8 + 16)));
 222     let address = |v| match v {
 223         Var::Global(x) => Addr::Rip(Label::Global(x)),
 224         Var::Param(x) => param(x),
 225         Var::Upvar(x) => upvar(x),
 226         Var::Ssa(s) => stack(&s),
 227     };
 228     let get = |r, s| match func.block[s] {
 229         Code::Load(v) => Inst::Load(r, address(v)),
 230         Code::Num(v) => Inst::Imm(r, (v << 1) | 1),
 231         _ => Inst::Load(r, stack(s)),
 232     };
 233     for ssa in &func.order {
 234         match &func.block[ssa] {
 235             // We don't need to generate code for loads or nums themselves, they're consulted by other accesses
 236             Code::Load(_) => (),
 237             Code::Num(_) => (),
 238             Code::StoreGlobal(x, s) => {
 239                 body.push(get(Reg::Rax, s));
 240                 body.push(Inst::Store(Addr::Rip(Label::Global(*x)), Reg::Rax));
 241             }
 242             Code::MakeLam {
 243                 name,
 244                 upvars,
 245                 params,
 246             } => {
 247                 body.push(Inst::Lea(Reg::Rdi, Addr::Rip(Label::Function(*name))));
 248                 body.push(Inst::Imm(Reg::Rsi, *params as i64));
 249                 body.push(Inst::Imm(Reg::Rdx, upvars.len() as i64));
 250                 body.push(Inst::Call(extern_label!("ivy_make_lam")));
 251                 body.push(Inst::Store(stack(ssa), Reg::Rax));
 252                 for (i, s) in upvars.iter().enumerate() {
 253                     body.push(get(Reg::Rdi, s));
 254                     body.push(Inst::Store(
 255                         Addr::Off(Reg::Rax, HEADER_SIZE + (*params as i32) * 8 + i as i32 * 8),
 256                         Reg::Rdi,
 257                     ));
 258                     body.push(Inst::Call(extern_label!("ivy_incref")));
 259                     body.push(Inst::Load(Reg::Rax, stack(ssa)));
 260                 }
 261             }
 262             Code::App(terms) => {
 263                 let mut terms = terms.iter();
 264                 body.push(get(Reg::Rdi, terms.next().expect("a function to apply")));
 265                 match terms.next() {
 266                     Some(s) => body.push(get(Reg::Rsi, s)),
 267                     None => body.push(Inst::Imm(Reg::Rsi, 0)),
 268                 }
 269                 body.push(Inst::Call(extern_label!("ivy_app")));
 270                 while let Some(term) = terms.next() {
 271                     body.push(Inst::Mov(Reg::Rdi, Reg::Rax));
 272                     body.push(get(Reg::Rsi, term));
 273                     body.push(Inst::Call(extern_label!("ivy_app")));
 274                 }
 275                 body.push(Inst::Store(stack(ssa), Reg::Rax));
 276             }
 277         }
 278     }
 279     match func.block[&func.result] {
 280         Code::Load(v) => body.push(Inst::Load(Reg::Rax, address(v))),
 281         Code::Num(n) => body.push(Inst::Imm(Reg::Rax, (n << 1) | 1i64)),
 282         _ => (),
 283     }
 284     body.push(Inst::AddI(Reg::Rsp, stack_slots as u32 * 8));
 285     body.push(Inst::Cfi(Cfi::Def(Reg::Rsp, 16)));
 286     body.push(Inst::Pop(Reg::Rbx));
 287     body.push(Inst::Cfi(Cfi::Def(Reg::Rsp, 8)));
 288     body.push(Inst::Ret);
 289     Definition {
 290         name: Label::Function(func.name),
 291         body: Cow::Owned(body),
 292     }
 293 }
 294
 295 impl fmt::Display for Definition {
 296     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 297         writeln!(fmt, "    .globl {}", self.name)?;
 298         match self.body.len() {
 299             0 => write!(fmt, "{}:", self.name),
 300             1 => write!(fmt, "{}: {}", self.name, self.body[0]),
 301             _ => {
 302                 writeln!(fmt, "    .cfi_startproc")?;
 303                 writeln!(fmt, "{}:", self.name)?;
 304                 for inst in self.body.as_ref() {
 305                     writeln!(fmt, "    {}", inst)?;
 306                 }
 307                 write!(fmt, "    .cfi_endproc")
 308             }
 309         }
 310     }
 311 }
 312
 313 impl fmt::Display for Global {
 314     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 315         write!(fmt, "{}: .quad {}", self.name, self.value)
 316     }
 317 }
 318
 319 impl fmt::Display for Label {
 320     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 321         match self {
 322             Label::Global(g) => write!(fmt, "IVY_GLOBAL${}", g),
 323             Label::Function(f) => write!(fmt, "ivy_fn${}", f.0),
 324             Label::Extern(l) => write!(fmt, "{}", l),
 325             Label::Builtin(l) => write!(fmt, "ivy_builtin${}", l),
 326         }
 327     }
 328 }
 329
 330 impl fmt::Display for Reg {
 331     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 332         match self {
 333             Reg::Rax => write!(fmt, "rax"),
 334             Reg::Rbx => write!(fmt, "rbx"),
 335             Reg::Rcx => write!(fmt, "rcx"),
 336             Reg::Rdx => write!(fmt, "rdx"),
 337             Reg::Rsi => write!(fmt, "rsi"),
 338             Reg::Rdi => write!(fmt, "rdi"),
 339             Reg::Rbp => write!(fmt, "rbp"),
 340             Reg::Rsp => write!(fmt, "rsp"),
 341             Reg::R8 => write!(fmt, "r8"),
 342             Reg::R9 => write!(fmt, "r9"),
 343             Reg::R10 => write!(fmt, "r10"),
 344             Reg::R11 => write!(fmt, "r11"),
 345             Reg::R12 => write!(fmt, "r12"),
 346             Reg::R13 => write!(fmt, "r13"),
 347             Reg::R14 => write!(fmt, "r14"),
 348             Reg::R15 => write!(fmt, "r15"),
 349         }
 350     }
 351 }
 352
 353 impl fmt::Display for Addr {
 354     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 355         match self {
 356             Addr::Rip(l) => write!(fmt, "[rip + {}]", l),
 357             Addr::Off(r, o) => write!(fmt, "[{} + 0x{:02x}]", r, o),
 358         }
 359     }
 360 }
 361
 362 impl fmt::Display for Cond {
 363     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 364         match self {
 365             Cond::Le => write!(fmt, "le"),
 366         }
 367     }
 368 }
 369
 370 impl fmt::Display for Cfi {
 371     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 372         match self {
 373             Cfi::Def(r, i) => write!(fmt, ".cfi_def_cfa {}, {}", r, i),
 374         }
 375     }
 376 }
 377
 378 impl fmt::Display for Inst {
 379     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 380         match self {
 381             Inst::Push(r) => write!(fmt, "push {}", r),
 382             Inst::Pop(r) => write!(fmt, "pop {}", r),
 383             Inst::Cfi(c) => write!(fmt, "{}", c),
 384             Inst::Mov(d, s) => write!(fmt, "mov {}, {}", d, s),
 385             Inst::Load(r, a) => write!(fmt, "mov {}, {}", r, a),
 386             Inst::Store(a, r) => write!(fmt, "mov {}, {}", a, r),
 387             Inst::Imm(r, i) => write!(fmt, "mov {}, {}", r, i),
 388             Inst::Lea(r, a) => write!(fmt, "lea {}, {}", r, a),
 389             Inst::Call(l) => write!(fmt, "call {}", l),
 390             Inst::Jmp(l) => write!(fmt, "jmp {}", l),
 391             Inst::Add(d, s) => write!(fmt, "add {}, {}", d, s),
 392             Inst::Sub(d, s) => write!(fmt, "sub {}, {}", d, s),
 393             Inst::AddI(r, x) => write!(fmt, "add {}, {}", r, x),
 394             Inst::SubI(r, x) => write!(fmt, "sub {}, {}", r, x),
 395             Inst::Cmp(a, b) => write!(fmt, "cmp {}, {}", a, b),
 396             Inst::CLoad(c, r, a) => write!(fmt, "cmov{} {}, {}", c, r, a),
 397             Inst::CMov(c, d, s) => write!(fmt, "cmov{} {}, {}", c, d, s),
 398             Inst::Ret => write!(fmt, "ret"),
 399         }
 400     }
 401 }
 402
 403 fn builtin(name: &str) -> (u16, &'static str, &'static [Inst]) {
 404     match name {
 405         "debug" => (
 406             1,
 407             "debug",
 408             &[
 409                 Inst::Load(Reg::Rdi, Addr::Off(Reg::Rdi, 0x18)),
 410                 Inst::Jmp(extern_label!("ivy_debug")),
 411             ],
 412         ),
 413         "+" => (
 414             2,
 415             "add",
 416             &[
 417                 Inst::Load(Reg::Rax, Addr::Off(Reg::Rdi, 0x18)),
 418                 Inst::Load(Reg::Rdi, Addr::Off(Reg::Rdi, 0x18 + 8)),
 419                 Inst::Add(Reg::Rax, Reg::Rdi),
 420                 Inst::SubI(Reg::Rax, 1),
 421                 Inst::Ret,
 422             ],
 423         ),
 424         "-" => (
 425             2,
 426             "sub",
 427             &[
 428                 Inst::Load(Reg::Rax, Addr::Off(Reg::Rdi, 0x18)),
 429                 Inst::Load(Reg::Rdi, Addr::Off(Reg::Rdi, 0x18 + 8)),
 430                 Inst::Sub(Reg::Rax, Reg::Rdi),
 431                 Inst::AddI(Reg::Rax, 1),
 432                 Inst::Ret,
 433             ],
 434         ),
 435         "<=" => (
 436             2,
 437             "leq",
 438             &[
 439                 Inst::Load(Reg::Rax, Addr::Off(Reg::Rdi, 0x18)),
 440                 Inst::Load(Reg::Rdi, Addr::Off(Reg::Rdi, 0x18 + 8)),
 441                 Inst::Cmp(Reg::Rax, Reg::Rdi),
 442                 Inst::Load(Reg::Rax, Addr::Rip(Label::Global(2))), // false
 443                 Inst::CLoad(Cond::Le, Reg::Rax, Addr::Rip(Label::Global(1))), // true
 444                 Inst::Ret,
 445             ],
 446         ),
 447         "true" => (
 448             2,
 449             "true",
 450             &[Inst::Load(Reg::Rax, Addr::Off(Reg::Rdi, 0x18)), Inst::Ret],
 451         ),
 452         "false" => (
 453             2,
 454             "false",
 455             &[
 456                 Inst::Load(Reg::Rax, Addr::Off(Reg::Rdi, 0x18 + 8)),
 457                 Inst::Ret,
 458             ],
 459         ),
 460         name => panic!("Unsupported builtin '{}' for x64", name),
 461     }
 462 }