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

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