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     pub 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 resolved = env
 156                         .resolve(var)
 157                         .ok_or_else(|| format!("Unable to resolve variable {:?}", var))?;
 158                     let load = Code::Load(resolved);
 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                     if let Var::Global(g) = var {
 192                         prog.globals.insert(g);
 193                     }
 194                     let ssa = func.fresh_ssa();
 195                     func.block.insert(ssa, Code::Load(var));
 196                     // Inserting the result of the lookup into the environment deduplicates loads.
 197                     env.mapping.insert(*name, Var::Ssa(ssa));
 198                     func.order.push(ssa);
 199                     Ok(ssa)
 200                 }
 201                 None => Err(format!("Unable to resolve variable {:?}", name)),
 202             }
 203         }
 204         Ast::Num(n) => {
 205             let ssa = func.fresh_ssa();
 206             func.block.insert(ssa, Code::Num(*n));
 207             func.order.push(ssa);
 208             Ok(ssa)
 209         }
 210     }
 211 }
 212
 213 impl fmt::Debug for FnName {
 214     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 215         write!(fmt, "@fn{}", self.0)
 216     }
 217 }
 218
 219 impl fmt::Debug for SsaName {
 220     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 221         write!(fmt, "%{}", self.0)
 222     }
 223 }
 224
 225 impl fmt::Debug for Var {
 226     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 227         match *self {
 228             Var::Global(x) => write!(fmt, "@{}", x),
 229             Var::Param(x) => write!(fmt, "p{}", x),
 230             Var::Upvar(x) => write!(fmt, "c{}", x),
 231             Var::Ssa(x) => write!(fmt, "{:?}", x),
 232         }
 233     }
 234 }
 235
 236 impl fmt::Debug for Code {
 237     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 238         match self {
 239             Code::Load(var) => write!(fmt, "load {:?}", var),
 240             Code::StoreGlobal(g, ssa) => write!(fmt, "store [{:?}] = {:?}", Var::Global(*g), ssa),
 241             Code::MakeLam {
 242                 name,
 243                 upvars,
 244                 params,
 245             } => {
 246                 write!(fmt, "lam {:?} ({}) [", name, params)?;
 247                 for up in upvars {
 248                     write!(fmt, "{:?} ", up)?;
 249                 }
 250                 write!(fmt, "]")
 251             }
 252             Code::App(args) => {
 253                 write!(fmt, "app (")?;
 254                 for arg in args {
 255                     write!(fmt, "{:?} ", arg)?;
 256                 }
 257                 write!(fmt, ")")
 258             }
 259             Code::Num(n) => write!(fmt, "{}", n),
 260         }
 261     }
 262 }
 263
 264 impl fmt::Debug for Func {
 265     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 266         write!(fmt, "fn {:?} ({}) ", self.name, self.params)?;
 267         fmt.debug_map()
 268             .entries(self.order.iter().map(|x| (x, &self.block[x])))
 269             .entry(&"ret", &self.result)
 270             .finish()
 271     }
 272 }