]> Witch of Git - ivy/blob - src/trans.rs
[rt] Split the runtime into multiple modules
[ivy] / src / trans.rs
1 use crate::ast::{self, Ast};
2 use std::{
3 collections::{HashMap, HashSet},
4 fmt,
5 };
6
7 pub mod x64;
8
9 #[derive(Hash, PartialEq, Eq, Clone, Copy)]
10 pub struct FnName(u32);
11
12 #[derive(Hash, PartialEq, Eq, Clone, Copy)]
13 pub struct SsaName(u32);
14
15 #[derive(Hash, PartialEq, Eq, Clone, Copy)]
16 pub enum Var {
17 Global(u32),
18 Param(u16),
19 Upvar(u16),
20 Ssa(SsaName),
21 }
22
23 #[derive(PartialEq, Eq, Clone)]
24 pub enum Code {
25 Load(Var),
26 StoreGlobal(u32, SsaName),
27 MakeLam {
28 name: FnName,
29 upvars: Vec<SsaName>,
30 params: u16,
31 },
32 App(Vec<SsaName>),
33 Num(i64),
34 }
35
36 #[derive(Debug, PartialEq, Eq, Clone)]
37 pub struct Program {
38 pub globals: HashSet<u32>,
39 functions: HashMap<FnName, Func>,
40 top: FnName,
41 fresh: u32,
42 }
43
44 impl Program {
45 fn fresh(&mut self) -> u32 {
46 let res = self.fresh;
47 self.fresh += 1;
48 res
49 }
50 }
51
52 #[derive(PartialEq, Eq, Clone)]
53 pub struct Func {
54 name: FnName,
55 result: SsaName,
56 params: u16,
57 block: HashMap<SsaName, Code>,
58 order: Vec<SsaName>,
59 fresh: u32,
60 }
61
62 impl Func {
63 fn fresh_ssa(&mut self) -> SsaName {
64 let fresh = SsaName(self.fresh);
65 self.fresh += 1;
66 fresh
67 }
68 }
69
70 struct Env {
71 mapping: HashMap<ast::Name, Var>,
72 }
73
74 impl Env {
75 fn resolve(&self, name: &ast::Name) -> Option<Var> {
76 self.mapping.get(&name).copied().or_else(|| match name {
77 ast::Name::Global(x) => Some(Var::Global(*x)),
78 _ => None,
79 })
80 }
81 }
82
83 pub fn translate(ast: &Ast) -> Result<Program, String> {
84 let mut prog = Program {
85 globals: HashSet::new(),
86 functions: HashMap::new(),
87 top: FnName(0),
88 fresh: 0,
89 };
90 let mut env = Env {
91 mapping: HashMap::new(),
92 };
93 let func = make_func(0, &mut prog, &mut env, ast)?;
94 prog.top = func.name;
95 prog.functions.insert(func.name, func);
96 Ok(prog)
97 }
98
99 fn make_func(params: u16, prog: &mut Program, env: &mut Env, ast: &Ast) -> Result<Func, String> {
100 let mut func = Func {
101 name: FnName(prog.fresh()),
102 result: SsaName(0),
103 params,
104 block: HashMap::new(),
105 order: Vec::new(),
106 fresh: 0,
107 };
108 func.result = trans_expr(prog, env, &mut func, ast)?;
109 Ok(func)
110 }
111
112 fn trans_expr(
113 prog: &mut Program,
114 env: &mut Env,
115 func: &mut Func,
116 ast: &Ast,
117 ) -> Result<SsaName, String> {
118 match ast {
119 Ast::Let(binds, body) => {
120 for (name, init) in binds {
121 let ssa = trans_expr(prog, env, func, init)?;
122 if let ast::Name::Global(g) = name {
123 let ssa2 = func.fresh_ssa();
124 prog.globals.insert(*g);
125 func.block.insert(ssa2, Code::StoreGlobal(*g, ssa));
126 func.order.push(ssa2);
127 }
128 // Inserting this unconditionally means that lookups within the
129 // local function will see the SSA name, and potentially avoid
130 // having to do a lookup on a global.
131 env.mapping.insert(*name, Var::Ssa(ssa));
132 }
133 trans_expr(prog, env, func, body)
134 }
135 Ast::Lam {
136 binds,
137 upvars,
138 body,
139 } => {
140 let mut new_env = Env {
141 mapping: HashMap::with_capacity(binds.len() + upvars.len()),
142 };
143 for (i, &name) in binds.iter().enumerate() {
144 new_env.mapping.insert(name, Var::Param(i as u16));
145 }
146 for (i, &name) in upvars.iter().enumerate() {
147 new_env.mapping.insert(name, Var::Upvar(i as u16));
148 }
149 let new_func = make_func(binds.len() as u16, prog, &mut new_env, body)?;
150 let ssa = func.fresh_ssa();
151 let upvars = upvars
152 .iter()
153 .map(|&var| trans_expr(prog, env, func, &Ast::Var(var)))
154 .collect::<Result<_, String>>()?;
155 func.block.insert(
156 ssa,
157 Code::MakeLam {
158 name: new_func.name,
159 upvars,
160 params: binds.len() as u16,
161 },
162 );
163 func.order.push(ssa);
164 prog.functions.insert(new_func.name, new_func);
165 Ok(ssa)
166 }
167 Ast::App(terms) => {
168 let terms = terms
169 .iter()
170 .map(|term| trans_expr(prog, env, func, term))
171 .collect::<Result<_, _>>()?;
172 let ssa = func.fresh_ssa();
173 func.block.insert(ssa, Code::App(terms));
174 func.order.push(ssa);
175 Ok(ssa)
176 }
177 Ast::Var(name) => {
178 match env.resolve(name) {
179 // Forwarding SSA variables skips unnecessary loads
180 Some(Var::Ssa(ssa)) => Ok(ssa),
181 Some(var) => {
182 if let Var::Global(g) = var {
183 prog.globals.insert(g);
184 }
185 let ssa = func.fresh_ssa();
186 func.block.insert(ssa, Code::Load(var));
187 // Inserting the result of the lookup into the environment deduplicates loads.
188 env.mapping.insert(*name, Var::Ssa(ssa));
189 func.order.push(ssa);
190 Ok(ssa)
191 }
192 None => Err(format!("Unable to resolve variable {:?}", name)),
193 }
194 }
195 Ast::Num(n) => {
196 let ssa = func.fresh_ssa();
197 func.block.insert(ssa, Code::Num(*n));
198 func.order.push(ssa);
199 Ok(ssa)
200 }
201 }
202 }
203
204 impl fmt::Debug for FnName {
205 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
206 write!(fmt, "@fn{}", self.0)
207 }
208 }
209
210 impl fmt::Debug for SsaName {
211 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
212 write!(fmt, "%{}", self.0)
213 }
214 }
215
216 impl fmt::Debug for Var {
217 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
218 match *self {
219 Var::Global(x) => write!(fmt, "@{}", x),
220 Var::Param(x) => write!(fmt, "p{}", x),
221 Var::Upvar(x) => write!(fmt, "c{}", x),
222 Var::Ssa(x) => write!(fmt, "{:?}", x),
223 }
224 }
225 }
226
227 impl fmt::Debug for Code {
228 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
229 match self {
230 Code::Load(var) => write!(fmt, "load {:?}", var),
231 Code::StoreGlobal(g, ssa) => write!(fmt, "store [{:?}] = {:?}", Var::Global(*g), ssa),
232 Code::MakeLam {
233 name,
234 upvars,
235 params,
236 } => {
237 write!(fmt, "lam {:?} ({}) [", name, params)?;
238 for up in upvars {
239 write!(fmt, "{:?} ", up)?;
240 }
241 write!(fmt, "]")
242 }
243 Code::App(args) => {
244 write!(fmt, "app (")?;
245 for arg in args {
246 write!(fmt, "{:?} ", arg)?;
247 }
248 write!(fmt, ")")
249 }
250 Code::Num(n) => write!(fmt, "{}", n),
251 }
252 }
253 }
254
255 impl fmt::Debug for Func {
256 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
257 write!(fmt, "fn {:?} ({}) ", self.name, self.params)?;
258 fmt.debug_map()
259 .entries(self.order.iter().map(|x| (x, &self.block[x])))
260 .entry(&"ret", &self.result)
261 .finish()
262 }
263 }