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