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