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

   1 use crate::ast::{self, Ast};
   2 use std::{
   3     collections::{HashMap, HashSet},
   4     fmt,
   5 };
   6
   7 pub mod x64;
   8
   9 #[derive(Hash, PartialEq, Eq, Clone, Copy)]
  10 pub struct FnName(u32);
  11
  12 #[derive(Hash, PartialEq, Eq, Clone, Copy)]
  13 pub struct SsaName(u32);
  14
  15 #[derive(Hash, PartialEq, Eq, Clone, Copy)]
  16 pub enum Var {
  17     Global(u32),
  18     Param(u16),
  19     Upvar(u16),
  20     Ssa(SsaName),
  21 }
  22
  23 #[derive(PartialEq, Eq, Clone)]
  24 pub enum Code {
  25     Load(Var),
  26     StoreGlobal(u32, SsaName),
  27     MakeLam {
  28         name: FnName,
  29         upvars: Vec<SsaName>,
  30         params: u16,
  31     },
  32     App(Vec<SsaName>),
  33     Num(i64),
  34 }
  35
  36 #[derive(Debug, PartialEq, Eq, Clone)]
  37 pub struct Program {
  38     globals: HashSet<u32>,
  39     functions: HashMap<FnName, Func>,
  40     top: FnName,
  41     fresh: u32,
  42 }
  43
  44 impl Program {
  45     fn fresh(&mut self) -> u32 {
  46         let res = self.fresh;
  47         self.fresh += 1;
  48         res
  49     }
  50 }
  51
  52 #[derive(PartialEq, Eq, Clone)]
  53 pub struct Func {
  54     name: FnName,
  55     result: SsaName,
  56     params: u16,
  57     block: HashMap<SsaName, Code>,
  58     order: Vec<SsaName>,
  59     fresh: u32,
  60 }
  61
  62 impl Func {
  63     fn fresh_ssa(&mut self) -> SsaName {
  64         let fresh = SsaName(self.fresh);
  65         self.fresh += 1;
  66         fresh
  67     }
  68 }
  69
  70 struct Env {
  71     mapping: HashMap<ast::Name, Var>,
  72 }
  73
  74 impl Env {
  75     fn resolve(&self, name: &ast::Name) -> Option<Var> {
  76         self.mapping.get(&name).copied().or_else(|| match name {
  77             ast::Name::Global(x) => Some(Var::Global(*x)),
  78             _ => None,
  79         })
  80     }
  81 }
  82
  83 pub fn translate(ast: &Ast) -> Result<Program, String> {
  84     let mut prog = Program {
  85         globals: HashSet::new(),
  86         functions: HashMap::new(),
  87         top: FnName(0),
  88         fresh: 0,
  89     };
  90     let mut env = Env {
  91         mapping: HashMap::new(),
  92     };
  93     let func = make_func(0, &mut prog, &mut env, ast)?;
  94     prog.top = func.name;
  95     prog.functions.insert(func.name, func);
  96     Ok(prog)
  97 }
  98
  99 fn make_func(params: u16, prog: &mut Program, env: &mut Env, ast: &Ast) -> Result<Func, String> {
 100     let mut func = Func {
 101         name: FnName(prog.fresh()),
 102         result: SsaName(0),
 103         params,
 104         block: HashMap::new(),
 105         order: Vec::new(),
 106         fresh: 0,
 107     };
 108     func.result = trans_expr(prog, env, &mut func, ast)?;
 109     Ok(func)
 110 }
 111
 112 fn trans_expr(
 113     prog: &mut Program,
 114     env: &mut Env,
 115     func: &mut Func,
 116     ast: &Ast,
 117 ) -> Result<SsaName, String> {
 118     match ast {
 119         Ast::Let(binds, body) => {
 120             for (name, init) in binds {
 121                 let ssa = trans_expr(prog, env, func, init)?;
 122                 if let ast::Name::Global(g) = name {
 123                     let ssa2 = func.fresh_ssa();
 124                     prog.globals.insert(*g);
 125                     func.block.insert(ssa2, Code::StoreGlobal(*g, ssa));
 126                     func.order.push(ssa2);
 127                 }
 128                 // Inserting this unconditionally means that lookups within the
 129                 // local function will see the SSA name, and potentially avoid
 130                 // having to do a lookup on a global.
 131                 env.mapping.insert(*name, Var::Ssa(ssa));
 132             }
 133             trans_expr(prog, env, func, body)
 134         }
 135         Ast::Lam {
 136             binds,
 137             upvars,
 138             body,
 139         } => {
 140             let mut new_env = Env {
 141                 mapping: HashMap::with_capacity(binds.len() + upvars.len()),
 142             };
 143             for (i, &name) in binds.iter().enumerate() {
 144                 new_env.mapping.insert(name, Var::Param(i as u16));
 145             }
 146             for (i, &name) in upvars.iter().enumerate() {
 147                 new_env.mapping.insert(name, Var::Upvar(i as u16));
 148             }
 149             let new_func = make_func(binds.len() as u16, prog, &mut new_env, body)?;
 150             let ssa = func.fresh_ssa();
 151             let upvars = upvars
 152                 .iter()
 153                 .map(|var| {
 154                     let ssa = func.fresh_ssa();
 155                     let load = Code::Load(
 156                         env.resolve(var)
 157                             .ok_or_else(|| format!("Unable to resolve variable {:?}", var))?,
 158                     );
 159                     func.block.insert(ssa, load);
 160                     func.order.push(ssa);
 161                     Ok(ssa)
 162                 })
 163                 .collect::<Result<_, String>>()?;
 164             func.block.insert(
 165                 ssa,
 166                 Code::MakeLam {
 167                     name: new_func.name,
 168                     upvars,
 169                     params: binds.len() as u16,
 170                 },
 171             );
 172             func.order.push(ssa);
 173             prog.functions.insert(new_func.name, new_func);
 174             Ok(ssa)
 175         }
 176         Ast::App(terms) => {
 177             let terms = terms
 178                 .iter()
 179                 .map(|term| trans_expr(prog, env, func, term))
 180                 .collect::<Result<_, _>>()?;
 181             let ssa = func.fresh_ssa();
 182             func.block.insert(ssa, Code::App(terms));
 183             func.order.push(ssa);
 184             Ok(ssa)
 185         }
 186         Ast::Var(name) => {
 187             match env.resolve(name) {
 188                 // Forwarding SSA variables skips unnecessary loads
 189                 Some(Var::Ssa(ssa)) => Ok(ssa),
 190                 Some(var) => {
 191                     let ssa = func.fresh_ssa();
 192                     func.block.insert(ssa, Code::Load(var));
 193                     // Inserting the result of the lookup into the environment deduplicates loads.
 194                     env.mapping.insert(*name, Var::Ssa(ssa));
 195                     func.order.push(ssa);
 196                     Ok(ssa)
 197                 }
 198                 None => Err(format!("Unable to resolve variable {:?}", name)),
 199             }
 200         }
 201         Ast::Num(n) => {
 202             let ssa = func.fresh_ssa();
 203             func.block.insert(ssa, Code::Num(*n));
 204             func.order.push(ssa);
 205             Ok(ssa)
 206         }
 207     }
 208 }
 209
 210 impl fmt::Debug for FnName {
 211     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 212         write!(fmt, "@fn{}", self.0)
 213     }
 214 }
 215
 216 impl fmt::Debug for SsaName {
 217     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 218         write!(fmt, "%{}", self.0)
 219     }
 220 }
 221
 222 impl fmt::Debug for Var {
 223     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 224         match *self {
 225             Var::Global(x) => write!(fmt, "@{}", x),
 226             Var::Param(x) => write!(fmt, "p{}", x),
 227             Var::Upvar(x) => write!(fmt, "c{}", x),
 228             Var::Ssa(x) => write!(fmt, "{:?}", x),
 229         }
 230     }
 231 }
 232
 233 impl fmt::Debug for Code {
 234     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 235         match self {
 236             Code::Load(var) => write!(fmt, "load {:?}", var),
 237             Code::StoreGlobal(g, ssa) => write!(fmt, "store [{:?}] = {:?}", Var::Global(*g), ssa),
 238             Code::MakeLam {
 239                 name,
 240                 upvars,
 241                 params,
 242             } => {
 243                 write!(fmt, "lam {:?} ({}) [", name, params)?;
 244                 for up in upvars {
 245                     write!(fmt, "{:?} ", up)?;
 246                 }
 247                 write!(fmt, "]")
 248             }
 249             Code::App(args) => {
 250                 write!(fmt, "app (")?;
 251                 for arg in args {
 252                     write!(fmt, "{:?} ", arg)?;
 253                 }
 254                 write!(fmt, ")")
 255             }
 256             Code::Num(n) => write!(fmt, "{}", n),
 257         }
 258     }
 259 }
 260
 261 impl fmt::Debug for Func {
 262     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 263         write!(fmt, "fn {:?} ({}) ", self.name, self.params)?;
 264         fmt.debug_map()
 265             .entries(self.order.iter().map(|x| (x, &self.block[x])))
 266             .entry(&"ret", &self.result)
 267             .finish()
 268     }
 269 }