]> Witch of Git - ivy/blob - src/ast.rs
Add skeleton of x64 compilation code
[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
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 }