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

   1 use ess::{self, Sexp};
   2 use pretty::RcDoc;
   3 use std::{
   4     cell::RefCell,
   5     collections::HashMap,
   6     sync::atomic::{AtomicU32, Ordering},
   7 };
   8
   9 #[derive(Hash, PartialEq, Eq, Clone, Copy, Debug)]
  10 pub enum Name {
  11     Local(u32),
  12     Global(u32),
  13 }
  14
  15 impl Name {
  16     pub fn local(&self) -> bool {
  17         match self {
  18             Name::Local(_) => true,
  19             Name::Global(_) => false,
  20         }
  21     }
  22
  23     pub fn global_number(&self) -> Option<u32> {
  24         match self {
  25             Name::Local(_) => None,
  26             Name::Global(g) => Some(*g),
  27         }
  28     }
  29
  30     fn to_doc(&self) -> RcDoc {
  31         match self {
  32             Name::Local(n) => RcDoc::text(format!("x{}", n)),
  33             Name::Global(n) => RcDoc::text(format!("g{}", n)),
  34         }
  35     }
  36 }
  37
  38 #[derive(PartialEq, Clone, Debug)]
  39 pub enum Ast {
  40     Let(Vec<(Name, Ast)>, Box<Ast>),
  41     Lam {
  42         binds: Vec<Name>,
  43         upvars: Vec<Name>,
  44         body: Box<Ast>,
  45     },
  46     App(Vec<Ast>),
  47     Var(Name),
  48     Num(i64),
  49 }
  50
  51 impl Ast {
  52     pub fn to_doc(&self) -> RcDoc {
  53         match self {
  54             Ast::Let(binds, body) => RcDoc::text("let")
  55                 .append(
  56                     RcDoc::line()
  57                         .append(RcDoc::intersperse(
  58                             binds.iter().map(|(n, v)| {
  59                                 n.to_doc()
  60                                     .append(RcDoc::space())
  61                                     .append(RcDoc::text("="))
  62                                     .append(RcDoc::line().append(v.to_doc()).nest(2).group())
  63                             }),
  64                             RcDoc::text(";").append(RcDoc::line()),
  65                         ))
  66                         .nest(2),
  67                 )
  68                 .append(RcDoc::line())
  69                 .append(RcDoc::text("in"))
  70                 .append(RcDoc::line())
  71                 .append(body.to_doc()),
  72             Ast::Lam { binds, body, .. } => RcDoc::text("lam").append(
  73                 RcDoc::line()
  74                     .append(RcDoc::text("("))
  75                     .append(
  76                         RcDoc::intersperse(binds.iter().map(|x| x.to_doc()), RcDoc::line())
  77                             .nest(2)
  78                             .group(),
  79                     )
  80                     .append(RcDoc::text(") =>"))
  81                     .append(RcDoc::line())
  82                     .append(body.to_doc())
  83                     .nest(2)
  84                     .group(),
  85             ),
  86             Ast::App(terms) => RcDoc::intersperse(
  87                 terms.iter().map(|x| match x {
  88                     Ast::App(..) | Ast::Lam { .. } => {
  89                         RcDoc::text("(").append(x.to_doc()).append(RcDoc::text(")"))
  90                     }
  91                     _ => x.to_doc(),
  92                 }),
  93                 RcDoc::line(),
  94             )
  95             .nest(2)
  96             .group(),
  97             Ast::Var(name) => name.to_doc(),
  98             Ast::Num(n) => RcDoc::as_string(n),
  99         }
 100     }
 101 }
 102
 103 struct Env<'e> {
 104     fresh: AtomicU32,
 105     vars: HashMap<String, Name>,
 106     upvars: RefCell<HashMap<String, Name>>,
 107     up: Option<&'e Env<'e>>,
 108 }
 109
 110 impl Env<'_> {
 111     fn new() -> Env<'static> {
 112         Env {
 113             fresh: AtomicU32::new(0),
 114             vars: HashMap::new(),
 115             upvars: RefCell::new(HashMap::new()),
 116             up: None,
 117         }
 118     }
 119
 120     fn wrapping<'e, 'r: 'e>(env: &'r Env<'e>) -> Env<'r> {
 121         Env {
 122             fresh: AtomicU32::new(env.fresh.load(Ordering::Relaxed)),
 123             vars: HashMap::new(),
 124             upvars: RefCell::new(HashMap::new()),
 125             up: Some(env),
 126         }
 127     }
 128
 129     fn fresh(&mut self, text: &str) -> Name {
 130         let name = if self.up.is_some() {
 131             Name::Local(*self.fresh.get_mut())
 132         } else {
 133             Name::Global(*self.fresh.get_mut())
 134         };
 135         self.vars.insert(text.into(), name);
 136         *self.fresh.get_mut() += 1;
 137         name
 138     }
 139
 140     fn _get(&self, name: &str) -> Option<Name> {
 141         if let Some(&resolved) = self.vars.get(name) {
 142             return Some(resolved);
 143         }
 144         self.up.and_then(|up| up._get(name)).map(|resolved| {
 145             if resolved.local() {
 146                 self.upvars.borrow_mut().insert(name.into(), resolved);
 147             }
 148             resolved
 149         })
 150     }
 151
 152     fn lookup(&mut self, name: &str) -> Option<Name> {
 153         if let Some(&resolved) = self.vars.get(name) {
 154             return Some(resolved);
 155         }
 156         self.up.and_then(|up| up._get(name)).map(|resolved| {
 157             self.vars.insert(name.into(), resolved);
 158             if resolved.local() {
 159                 self.upvars.get_mut().insert(name.into(), resolved);
 160             }
 161             resolved
 162         })
 163     }
 164
 165     fn upvars(&self) -> Vec<Name> {
 166         self.upvars.borrow().values().map(|&x| x).collect()
 167     }
 168 }
 169
 170 impl Drop for Env<'_> {
 171     fn drop(&mut self) {
 172         if let Some(up) = &self.up {
 173             up.fresh.store(*self.fresh.get_mut(), Ordering::Relaxed);
 174         }
 175     }
 176 }
 177
 178 impl Ast {
 179     pub fn parse(
 180         builtins: &[impl AsRef<str>],
 181         sexp: &Sexp,
 182     ) -> Result<(Ast, HashMap<String, Name>), &'static str> {
 183         let mut env = Env::new();
 184         let mut builtin_map = HashMap::new();
 185         for builtin in builtins {
 186             let builtin = builtin.as_ref();
 187             builtin_map.insert(builtin.into(), env.fresh(builtin));
 188         }
 189         Ok((Ast::from_sexp(&mut env, sexp)?, builtin_map))
 190     }
 191
 192     fn from_sexp<'r, 'e: 'r>(env: &'r mut Env<'e>, sexp: &Sexp) -> Result<Ast, &'static str> {
 193         match sexp {
 194             Sexp::List(xs, _) => match xs.first() {
 195                 Some(Sexp::Sym(s, _)) if s == "let" => {
 196                     if xs.len() > 3 {
 197                         return Err("too many body expressions in 'let'");
 198                     }
 199                     let binds = match xs.get(1) {
 200                         Some(Sexp::List(binds, _)) => binds
 201                             .iter()
 202                             .map(|bind| match bind {
 203                                 Sexp::List(bind, _) => {
 204                                     if bind.len() != 2 {
 205                                         return Err("binding lists must have two elements");
 206                                     }
 207                                     // We do the body before the name to avoid being "letrec"
 208                                     let body = Ast::from_sexp(env, &bind[1])?;
 209                                     let name = match &bind[0] {
 210                                         Sexp::Sym(name, _) => env.fresh(&name),
 211                                         _ => {
 212                                             return Err("binding lists must start with a variable")
 213                                         }
 214                                     };
 215                                     Ok((name, body))
 216                                 }
 217                                 _ => Err("bindings must be lists"),
 218                             })
 219                             .collect::<Result<_, _>>()?,
 220                         Some(_) => return Err("the binding in 'let' is not a list"),
 221                         None => return Err("no binding form in 'let'"),
 222                     };
 223                     let body = match xs.get(2) {
 224                         Some(x) => Box::new(Ast::from_sexp(env, x)?),
 225                         None => return Err("no body expression in 'let'"),
 226                     };
 227                     Ok(Ast::Let(binds, body))
 228                 }
 229                 Some(Sexp::Sym(s, _)) if s == "lam" => {
 230                     if xs.len() > 3 {
 231                         return Err("too many body expressions in 'lam'");
 232                     }
 233                     let mut env = Env::wrapping(env);
 234                     let binds = match xs.get(1) {
 235                         Some(Sexp::List(names, _)) => names
 236                             .iter()
 237                             .map(|name| match name {
 238                                 Sexp::Sym(s, _) => Ok(env.fresh(&s)),
 239                                 _ => Err("non-symbol in 'lam' arguments"),
 240                             })
 241                             .collect::<Result<_, _>>()?,
 242                         Some(_) => return Err("'lam' arguments weren't a list"),
 243                         None => return Err("no arguments in 'lam'"),
 244                     };
 245                     let body = match xs.get(2) {
 246                         Some(x) => Box::new(Ast::from_sexp(&mut env, x)?),
 247                         None => return Err("no body expression in 'lam'"),
 248                     };
 249                     let upvars = env.upvars();
 250                     Ok(Ast::Lam {
 251                         binds,
 252                         upvars,
 253                         body,
 254                     })
 255                 }
 256                 Some(_) => Ok(Ast::App(
 257                     xs.iter()
 258                         .map(|x| Ast::from_sexp(env, x))
 259                         .collect::<Result<_, _>>()?,
 260                 )),
 261                 None => Err("expected a non-empty list for an application"),
 262             },
 263             Sexp::Sym(s, _) if s == "let" => Err("unexpected id 'let'"),
 264             Sexp::Sym(s, _) if s == "lam" => Err("unexpected id 'lam'"),
 265             Sexp::Sym(s, _) => Ok(Ast::Var(
 266                 env.lookup(s).ok_or("variable used but not defined")?,
 267             )),
 268             Sexp::Int(i, _) => Ok(Ast::Num(*i)),
 269             _ => Err("unexpected s-expression type"),
 270         }
 271     }
 272 }