nandi.uk
Fork of daniellemaywood.uk/gleam — Wasm codegen work
9.9 kB
259 lines
1use std::collections::HashMap;
2
3use cranelift::{codegen::ir::BlockArg, prelude::*};
4use cranelift_jit::{JITBuilder, JITModule};
5use cranelift_module::{FuncOrDataId, Linkage, Module, default_libcall_names};
6use ecow::EcoString;
7use num_bigint::BigInt;
8
9pub mod mir;
10
11pub fn compile(m: mir::ast::Module) {
12 let mut flag_builder = settings::builder();
13
14 flag_builder
15 .set("is_pic", "false")
16 .expect("unable to set `is_pic` to false");
17
18 let isa = cranelift_native::builder()
19 .unwrap_or_else(|reason| panic!("host machine not support: {reason}"))
20 .finish(settings::Flags::new(flag_builder))
21 .expect("unable to build isa target");
22
23 let mut module = JITModule::new(JITBuilder::with_isa(isa, default_libcall_names()));
24 let mut context = module.make_context();
25 let mut function_context = FunctionBuilderContext::new();
26
27 for function in m.functions {
28 let sig = &mut context.func.signature;
29
30 for _ in &function.parameters {
31 sig.params.push(AbiParam::new(types::I64));
32 }
33
34 sig.returns.push(AbiParam::new(types::I64));
35
36 let mut builder = FunctionBuilder::new(&mut context.func, &mut function_context);
37
38 let entry = builder.create_block();
39 builder.append_block_params_for_function_params(entry);
40 builder.switch_to_block(entry);
41 builder.seal_block(entry);
42
43 let mut translator = FunctionTranslator {
44 builder,
45 module: &mut module,
46 variables: HashMap::new(),
47 };
48
49 for (index, parameter) in function.parameters.into_iter().enumerate() {
50 if let Some(name) = parameter.name {
51 let var = translator.make_var(name);
52 let val = translator.builder.block_params(entry)[index];
53
54 translator.builder.def_var(var, val);
55 }
56 }
57
58 let output = translator.translate_expression(function.body);
59
60 _ = translator.builder.ins().return_(&[output]);
61 translator.builder.finalize();
62
63 let id = module
64 .declare_function(&function.name, Linkage::Export, &context.func.signature)
65 .expect("unable to declare function");
66
67 module
68 .define_function(id, &mut context)
69 .expect("unable to define function");
70
71 module.clear_context(&mut context);
72 }
73
74 module
75 .finalize_definitions()
76 .expect("unable to finalize definitions");
77
78 let FuncOrDataId::Func(main) = module.get_name("main").expect("expected a main function")
79 else {
80 panic!("expected main to be a function")
81 };
82
83 let code = module.get_finalized_function(main);
84 #[allow(unsafe_code)]
85 let main = unsafe { std::mem::transmute::<_, extern "C" fn() -> i64>(code) };
86
87 println!("{}", main());
88}
89
90struct FunctionTranslator<'a, M: Module> {
91 builder: FunctionBuilder<'a>,
92 module: &'a mut M,
93 variables: HashMap<EcoString, Variable>,
94}
95
96impl<'a, M: Module> FunctionTranslator<'a, M> {
97 fn make_var(&mut self, name: EcoString) -> Variable {
98 let variable = self.builder.declare_var(types::I64);
99 _ = self.variables.insert(name, variable);
100 variable
101 }
102
103 fn get_var(&mut self, name: &EcoString) -> Variable {
104 self.variables
105 .get(name)
106 .cloned()
107 .unwrap_or_else(|| panic!("variable '{}' not found", name))
108 }
109
110 fn make_int(&mut self, value: BigInt) -> Value {
111 let value: i64 = value.try_into().expect("bigint out of representable range");
112
113 self.builder.ins().iconst(types::I64, value)
114 }
115
116 fn translate_expression(&mut self, expression: mir::ast::Expression) -> Value {
117 _ = self.module;
118
119 match expression {
120 mir::ast::Expression::Block(expressions) => expressions
121 .into_iter()
122 .map(|expression| self.translate_expression(expression))
123 .last()
124 .expect("block shouldn't be empty"),
125 mir::ast::Expression::FunctionRef {
126 module: _,
127 name,
128 arity,
129 } => {
130 let mut signature = self.module.make_signature();
131
132 for _ in 0..arity {
133 signature.params.push(AbiParam::new(types::I64));
134 }
135 signature.returns.push(AbiParam::new(types::I64));
136
137 let id = self
138 .module
139 .declare_function(&name, Linkage::Import, &signature)
140 .expect("unable to declare function for import");
141
142 let func = self.module.declare_func_in_func(id, self.builder.func);
143
144 self.builder.ins().func_addr(types::I64, func)
145 }
146 mir::ast::Expression::Var(var) => {
147 let var = self.get_var(&var.name);
148
149 self.builder.use_var(var)
150 }
151 mir::ast::Expression::Int { value } => self.make_int(value),
152 mir::ast::Expression::Float { value: _ } => todo!(),
153 mir::ast::Expression::Bool { value: _ } => todo!(),
154 mir::ast::Expression::String { value: _ } => todo!(),
155 mir::ast::Expression::Equals { lhs, rhs } => {
156 let lhs = self.translate_expression(*lhs);
157 let rhs = self.translate_expression(*rhs);
158
159 self.builder.ins().icmp(IntCC::Equal, lhs, rhs)
160 }
161 mir::ast::Expression::NotEquals { lhs: _, rhs: _ } => todo!(),
162 mir::ast::Expression::IntGt { lhs: _, rhs: _ } => todo!(),
163 mir::ast::Expression::IntGtEq { lhs: _, rhs: _ } => todo!(),
164 mir::ast::Expression::IntLt { lhs: _, rhs: _ } => todo!(),
165 mir::ast::Expression::IntLtEq { lhs: _, rhs: _ } => todo!(),
166 mir::ast::Expression::IntAdd { lhs, rhs } => {
167 let lhs = self.translate_expression(*lhs);
168 let rhs = self.translate_expression(*rhs);
169
170 self.builder.ins().iadd(lhs, rhs)
171 }
172 mir::ast::Expression::IntSub { lhs, rhs } => {
173 let lhs = self.translate_expression(*lhs);
174 let rhs = self.translate_expression(*rhs);
175
176 self.builder.ins().isub(lhs, rhs)
177 }
178 mir::ast::Expression::IntMul { lhs: _, rhs: _ } => todo!(),
179 mir::ast::Expression::IntDiv { lhs: _, rhs: _ } => todo!(),
180 mir::ast::Expression::IntRem { lhs: _, rhs: _ } => todo!(),
181 mir::ast::Expression::FloatGt { lhs: _, rhs: _ } => todo!(),
182 mir::ast::Expression::FloatGtEq { lhs: _, rhs: _ } => todo!(),
183 mir::ast::Expression::FloatLt { lhs: _, rhs: _ } => todo!(),
184 mir::ast::Expression::FloatLtEq { lhs: _, rhs: _ } => todo!(),
185 mir::ast::Expression::FloatAdd { lhs: _, rhs: _ } => todo!(),
186 mir::ast::Expression::FloatSub { lhs: _, rhs: _ } => todo!(),
187 mir::ast::Expression::FloatMul { lhs: _, rhs: _ } => todo!(),
188 mir::ast::Expression::FloatDiv { lhs: _, rhs: _ } => todo!(),
189 mir::ast::Expression::StringConcat { lhs: _, rhs: _ } => todo!(),
190 mir::ast::Expression::List { items: _, tail: _ } => todo!(),
191 mir::ast::Expression::Struct { tag: _, items: _ } => todo!(),
192 mir::ast::Expression::StructTag { value: _ } => todo!(),
193 mir::ast::Expression::StructAccess { value: _, index: _ } => todo!(),
194 mir::ast::Expression::Set { name, value } => {
195 let var = self.make_var(name.name);
196 let val = self.translate_expression(*value);
197
198 self.builder.def_var(var, val);
199 self.builder.use_var(var)
200 }
201 mir::ast::Expression::If { cond, then, else_ } => {
202 let cond = self.translate_expression(*cond);
203
204 let then_block = self.builder.create_block();
205 let else_block = self.builder.create_block();
206 let merge_block = self.builder.create_block();
207
208 _ = self.builder.append_block_param(merge_block, types::I64);
209
210 _ = self
211 .builder
212 .ins()
213 .brif(cond, then_block, &[], else_block, &[]);
214
215 self.builder.switch_to_block(then_block);
216 self.builder.seal_block(then_block);
217
218 let then_return = self.translate_expression(*then);
219 _ = self
220 .builder
221 .ins()
222 .jump(merge_block, &[BlockArg::Value(then_return)]);
223
224 self.builder.switch_to_block(else_block);
225 self.builder.seal_block(else_block);
226
227 let else_return = self.translate_expression(*else_);
228 _ = self
229 .builder
230 .ins()
231 .jump(merge_block, &[BlockArg::Value(else_return)]);
232
233 self.builder.switch_to_block(merge_block);
234 self.builder.seal_block(merge_block);
235
236 self.builder.block_params(merge_block)[0]
237 }
238 mir::ast::Expression::Call { target, args } => {
239 let mut signature = self.module.make_signature();
240
241 for _ in 0..args.len() {
242 signature.params.push(AbiParam::new(types::I64));
243 }
244 signature.returns.push(AbiParam::new(types::I64));
245
246 let target = self.translate_expression(*target);
247 let args: Vec<_> = args
248 .into_iter()
249 .map(|arg| self.translate_expression(arg))
250 .collect();
251
252 let sig = self.builder.import_signature(signature);
253 let call = self.builder.ins().call_indirect(sig, target, &args);
254
255 self.builder.inst_results(call)[0]
256 }
257 }
258 }
259}