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