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| trans_expr(prog, env, func, &Ast::Var(var)))
 154                 .collect::<Result<_, String>>()?;
 155             func.block.insert(
 156                 ssa,
 157                 Code::MakeLam {
 158                     name: new_func.name,
 159                     upvars,
 160                     params: binds.len() as u16,
 161                 },
 162             );
 163             func.order.push(ssa);
 164             prog.functions.insert(new_func.name, new_func);
 165             Ok(ssa)
 166         }
 167         Ast::App(terms) => {
 168             let terms = terms
 169                 .iter()
 170                 .map(|term| trans_expr(prog, env, func, term))
 171                 .collect::<Result<_, _>>()?;
 172             let ssa = func.fresh_ssa();
 173             func.block.insert(ssa, Code::App(terms));
 174             func.order.push(ssa);
 175             Ok(ssa)
 176         }
 177         Ast::Var(name) => {
 178             match env.resolve(name) {
 179                 // Forwarding SSA variables skips unnecessary loads
 180                 Some(Var::Ssa(ssa)) => Ok(ssa),
 181                 Some(var) => {
 182                     if let Var::Global(g) = var {
 183                         prog.globals.insert(g);
 184                     }
 185                     let ssa = func.fresh_ssa();
 186                     func.block.insert(ssa, Code::Load(var));
 187                     // Inserting the result of the lookup into the environment deduplicates loads.
 188                     env.mapping.insert(*name, Var::Ssa(ssa));
 189                     func.order.push(ssa);
 190                     Ok(ssa)
 191                 }
 192                 None => Err(format!("Unable to resolve variable {:?}", name)),
 193             }
 194         }
 195         Ast::Num(n) => {
 196             let ssa = func.fresh_ssa();
 197             func.block.insert(ssa, Code::Num(*n));
 198             func.order.push(ssa);
 199             Ok(ssa)
 200         }
 201     }
 202 }
 203
 204 impl fmt::Debug for FnName {
 205     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 206         write!(fmt, "@fn{}", self.0)
 207     }
 208 }
 209
 210 impl fmt::Debug for SsaName {
 211     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 212         write!(fmt, "%{}", self.0)
 213     }
 214 }
 215
 216 impl fmt::Debug for Var {
 217     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 218         match *self {
 219             Var::Global(x) => write!(fmt, "@{}", x),
 220             Var::Param(x) => write!(fmt, "p{}", x),
 221             Var::Upvar(x) => write!(fmt, "c{}", x),
 222             Var::Ssa(x) => write!(fmt, "{:?}", x),
 223         }
 224     }
 225 }
 226
 227 impl fmt::Debug for Code {
 228     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 229         match self {
 230             Code::Load(var) => write!(fmt, "load {:?}", var),
 231             Code::StoreGlobal(g, ssa) => write!(fmt, "store [{:?}] = {:?}", Var::Global(*g), ssa),
 232             Code::MakeLam {
 233                 name,
 234                 upvars,
 235                 params,
 236             } => {
 237                 write!(fmt, "lam {:?} ({}) [", name, params)?;
 238                 for up in upvars {
 239                     write!(fmt, "{:?} ", up)?;
 240                 }
 241                 write!(fmt, "]")
 242             }
 243             Code::App(args) => {
 244                 write!(fmt, "app (")?;
 245                 for arg in args {
 246                     write!(fmt, "{:?} ", arg)?;
 247                 }
 248                 write!(fmt, ")")
 249             }
 250             Code::Num(n) => write!(fmt, "{}", n),
 251         }
 252     }
 253 }
 254
 255 impl fmt::Debug for Func {
 256     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
 257         write!(fmt, "fn {:?} ({}) ", self.name, self.params)?;
 258         fmt.debug_map()
 259             .entries(self.order.iter().map(|x| (x, &self.block[x])))
 260             .entry(&"ret", &self.result)
 261             .finish()
 262     }
 263 }