]> Witch of Git - ivy/blob - src/ast.rs
[rt] Split the runtime into multiple modules
[ivy] / 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 }