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