Add OCaml tutorial to the examples.
[oota-llvm.git] / examples / OCaml-Kaleidoscope / Chapter7 / codegen.ml
1 (*===----------------------------------------------------------------------===
2  * Code Generation
3  *===----------------------------------------------------------------------===*)
4
5 open Llvm
6
7 exception Error of string
8
9 let context = global_context ()
10 let the_module = create_module context "my cool jit"
11 let builder = builder context
12 let named_values:(string, llvalue) Hashtbl.t = Hashtbl.create 10
13 let double_type = double_type context
14
15 (* Create an alloca instruction in the entry block of the function. This
16  * is used for mutable variables etc. *)
17 let create_entry_block_alloca the_function var_name =
18   let builder = builder_at context (instr_begin (entry_block the_function)) in
19   build_alloca double_type var_name builder
20
21 let rec codegen_expr = function
22   | Ast.Number n -> const_float double_type n
23   | Ast.Variable name ->
24       let v = try Hashtbl.find named_values name with
25         | Not_found -> raise (Error "unknown variable name")
26       in
27       (* Load the value. *)
28       build_load v name builder
29   | Ast.Unary (op, operand) ->
30       let operand = codegen_expr operand in
31       let callee = "unary" ^ (String.make 1 op) in
32       let callee =
33         match lookup_function callee the_module with
34         | Some callee -> callee
35         | None -> raise (Error "unknown unary operator")
36       in
37       build_call callee [|operand|] "unop" builder
38   | Ast.Binary (op, lhs, rhs) ->
39       begin match op with
40       | '=' ->
41           (* Special case '=' because we don't want to emit the LHS as an
42            * expression. *)
43           let name =
44             match lhs with
45             | Ast.Variable name -> name
46             | _ -> raise (Error "destination of '=' must be a variable")
47           in
48
49           (* Codegen the rhs. *)
50           let val_ = codegen_expr rhs in
51
52           (* Lookup the name. *)
53           let variable = try Hashtbl.find named_values name with
54           | Not_found -> raise (Error "unknown variable name")
55           in
56           ignore(build_store val_ variable builder);
57           val_
58       | _ ->
59           let lhs_val = codegen_expr lhs in
60           let rhs_val = codegen_expr rhs in
61           begin
62             match op with
63             | '+' -> build_add lhs_val rhs_val "addtmp" builder
64             | '-' -> build_sub lhs_val rhs_val "subtmp" builder
65             | '*' -> build_mul lhs_val rhs_val "multmp" builder
66             | '<' ->
67                 (* Convert bool 0/1 to double 0.0 or 1.0 *)
68                 let i = build_fcmp Fcmp.Ult lhs_val rhs_val "cmptmp" builder in
69                 build_uitofp i double_type "booltmp" builder
70             | _ ->
71                 (* If it wasn't a builtin binary operator, it must be a user defined
72                  * one. Emit a call to it. *)
73                 let callee = "binary" ^ (String.make 1 op) in
74                 let callee =
75                   match lookup_function callee the_module with
76                   | Some callee -> callee
77                   | None -> raise (Error "binary operator not found!")
78                 in
79                 build_call callee [|lhs_val; rhs_val|] "binop" builder
80           end
81       end
82   | Ast.Call (callee, args) ->
83       (* Look up the name in the module table. *)
84       let callee =
85         match lookup_function callee the_module with
86         | Some callee -> callee
87         | None -> raise (Error "unknown function referenced")
88       in
89       let params = params callee in
90
91       (* If argument mismatch error. *)
92       if Array.length params == Array.length args then () else
93         raise (Error "incorrect # arguments passed");
94       let args = Array.map codegen_expr args in
95       build_call callee args "calltmp" builder
96   | Ast.If (cond, then_, else_) ->
97       let cond = codegen_expr cond in
98
99       (* Convert condition to a bool by comparing equal to 0.0 *)
100       let zero = const_float double_type 0.0 in
101       let cond_val = build_fcmp Fcmp.One cond zero "ifcond" builder in
102
103       (* Grab the first block so that we might later add the conditional branch
104        * to it at the end of the function. *)
105       let start_bb = insertion_block builder in
106       let the_function = block_parent start_bb in
107
108       let then_bb = append_block context "then" the_function in
109
110       (* Emit 'then' value. *)
111       position_at_end then_bb builder;
112       let then_val = codegen_expr then_ in
113
114       (* Codegen of 'then' can change the current block, update then_bb for the
115        * phi. We create a new name because one is used for the phi node, and the
116        * other is used for the conditional branch. *)
117       let new_then_bb = insertion_block builder in
118
119       (* Emit 'else' value. *)
120       let else_bb = append_block context "else" the_function in
121       position_at_end else_bb builder;
122       let else_val = codegen_expr else_ in
123
124       (* Codegen of 'else' can change the current block, update else_bb for the
125        * phi. *)
126       let new_else_bb = insertion_block builder in
127
128       (* Emit merge block. *)
129       let merge_bb = append_block context "ifcont" the_function in
130       position_at_end merge_bb builder;
131       let incoming = [(then_val, new_then_bb); (else_val, new_else_bb)] in
132       let phi = build_phi incoming "iftmp" builder in
133
134       (* Return to the start block to add the conditional branch. *)
135       position_at_end start_bb builder;
136       ignore (build_cond_br cond_val then_bb else_bb builder);
137
138       (* Set a unconditional branch at the end of the 'then' block and the
139        * 'else' block to the 'merge' block. *)
140       position_at_end new_then_bb builder; ignore (build_br merge_bb builder);
141       position_at_end new_else_bb builder; ignore (build_br merge_bb builder);
142
143       (* Finally, set the builder to the end of the merge block. *)
144       position_at_end merge_bb builder;
145
146       phi
147   | Ast.For (var_name, start, end_, step, body) ->
148       (* Output this as:
149        *   var = alloca double
150        *   ...
151        *   start = startexpr
152        *   store start -> var
153        *   goto loop
154        * loop:
155        *   ...
156        *   bodyexpr
157        *   ...
158        * loopend:
159        *   step = stepexpr
160        *   endcond = endexpr
161        *
162        *   curvar = load var
163        *   nextvar = curvar + step
164        *   store nextvar -> var
165        *   br endcond, loop, endloop
166        * outloop: *)
167
168       let the_function = block_parent (insertion_block builder) in
169
170       (* Create an alloca for the variable in the entry block. *)
171       let alloca = create_entry_block_alloca the_function var_name in
172
173       (* Emit the start code first, without 'variable' in scope. *)
174       let start_val = codegen_expr start in
175
176       (* Store the value into the alloca. *)
177       ignore(build_store start_val alloca builder);
178
179       (* Make the new basic block for the loop header, inserting after current
180        * block. *)
181       let loop_bb = append_block context "loop" the_function in
182
183       (* Insert an explicit fall through from the current block to the
184        * loop_bb. *)
185       ignore (build_br loop_bb builder);
186
187       (* Start insertion in loop_bb. *)
188       position_at_end loop_bb builder;
189
190       (* Within the loop, the variable is defined equal to the PHI node. If it
191        * shadows an existing variable, we have to restore it, so save it
192        * now. *)
193       let old_val =
194         try Some (Hashtbl.find named_values var_name) with Not_found -> None
195       in
196       Hashtbl.add named_values var_name alloca;
197
198       (* Emit the body of the loop.  This, like any other expr, can change the
199        * current BB.  Note that we ignore the value computed by the body, but
200        * don't allow an error *)
201       ignore (codegen_expr body);
202
203       (* Emit the step value. *)
204       let step_val =
205         match step with
206         | Some step -> codegen_expr step
207         (* If not specified, use 1.0. *)
208         | None -> const_float double_type 1.0
209       in
210
211       (* Compute the end condition. *)
212       let end_cond = codegen_expr end_ in
213
214       (* Reload, increment, and restore the alloca. This handles the case where
215        * the body of the loop mutates the variable. *)
216       let cur_var = build_load alloca var_name builder in
217       let next_var = build_add cur_var step_val "nextvar" builder in
218       ignore(build_store next_var alloca builder);
219
220       (* Convert condition to a bool by comparing equal to 0.0. *)
221       let zero = const_float double_type 0.0 in
222       let end_cond = build_fcmp Fcmp.One end_cond zero "loopcond" builder in
223
224       (* Create the "after loop" block and insert it. *)
225       let after_bb = append_block context "afterloop" the_function in
226
227       (* Insert the conditional branch into the end of loop_end_bb. *)
228       ignore (build_cond_br end_cond loop_bb after_bb builder);
229
230       (* Any new code will be inserted in after_bb. *)
231       position_at_end after_bb builder;
232
233       (* Restore the unshadowed variable. *)
234       begin match old_val with
235       | Some old_val -> Hashtbl.add named_values var_name old_val
236       | None -> ()
237       end;
238
239       (* for expr always returns 0.0. *)
240       const_null double_type
241   | Ast.Var (var_names, body) ->
242       let old_bindings = ref [] in
243
244       let the_function = block_parent (insertion_block builder) in
245
246       (* Register all variables and emit their initializer. *)
247       Array.iter (fun (var_name, init) ->
248         (* Emit the initializer before adding the variable to scope, this
249          * prevents the initializer from referencing the variable itself, and
250          * permits stuff like this:
251          *   var a = 1 in
252          *     var a = a in ...   # refers to outer 'a'. *)
253         let init_val =
254           match init with
255           | Some init -> codegen_expr init
256           (* If not specified, use 0.0. *)
257           | None -> const_float double_type 0.0
258         in
259
260         let alloca = create_entry_block_alloca the_function var_name in
261         ignore(build_store init_val alloca builder);
262
263         (* Remember the old variable binding so that we can restore the binding
264          * when we unrecurse. *)
265         begin
266           try
267             let old_value = Hashtbl.find named_values var_name in
268             old_bindings := (var_name, old_value) :: !old_bindings;
269           with Not_found -> ()
270         end;
271
272         (* Remember this binding. *)
273         Hashtbl.add named_values var_name alloca;
274       ) var_names;
275
276       (* Codegen the body, now that all vars are in scope. *)
277       let body_val = codegen_expr body in
278
279       (* Pop all our variables from scope. *)
280       List.iter (fun (var_name, old_value) ->
281         Hashtbl.add named_values var_name old_value
282       ) !old_bindings;
283
284       (* Return the body computation. *)
285       body_val
286
287 let codegen_proto = function
288   | Ast.Prototype (name, args) | Ast.BinOpPrototype (name, args, _) ->
289       (* Make the function type: double(double,double) etc. *)
290       let doubles = Array.make (Array.length args) double_type in
291       let ft = function_type double_type doubles in
292       let f =
293         match lookup_function name the_module with
294         | None -> declare_function name ft the_module
295
296         (* If 'f' conflicted, there was already something named 'name'. If it
297          * has a body, don't allow redefinition or reextern. *)
298         | Some f ->
299             (* If 'f' already has a body, reject this. *)
300             if block_begin f <> At_end f then
301               raise (Error "redefinition of function");
302
303             (* If 'f' took a different number of arguments, reject. *)
304             if element_type (type_of f) <> ft then
305               raise (Error "redefinition of function with different # args");
306             f
307       in
308
309       (* Set names for all arguments. *)
310       Array.iteri (fun i a ->
311         let n = args.(i) in
312         set_value_name n a;
313         Hashtbl.add named_values n a;
314       ) (params f);
315       f
316
317 (* Create an alloca for each argument and register the argument in the symbol
318  * table so that references to it will succeed. *)
319 let create_argument_allocas the_function proto =
320   let args = match proto with
321     | Ast.Prototype (_, args) | Ast.BinOpPrototype (_, args, _) -> args
322   in
323   Array.iteri (fun i ai ->
324     let var_name = args.(i) in
325     (* Create an alloca for this variable. *)
326     let alloca = create_entry_block_alloca the_function var_name in
327
328     (* Store the initial value into the alloca. *)
329     ignore(build_store ai alloca builder);
330
331     (* Add arguments to variable symbol table. *)
332     Hashtbl.add named_values var_name alloca;
333   ) (params the_function)
334
335 let codegen_func the_fpm = function
336   | Ast.Function (proto, body) ->
337       Hashtbl.clear named_values;
338       let the_function = codegen_proto proto in
339
340       (* If this is an operator, install it. *)
341       begin match proto with
342       | Ast.BinOpPrototype (name, args, prec) ->
343           let op = name.[String.length name - 1] in
344           Hashtbl.add Parser.binop_precedence op prec;
345       | _ -> ()
346       end;
347
348       (* Create a new basic block to start insertion into. *)
349       let bb = append_block context "entry" the_function in
350       position_at_end bb builder;
351
352       try
353         (* Add all arguments to the symbol table and create their allocas. *)
354         create_argument_allocas the_function proto;
355
356         let ret_val = codegen_expr body in
357
358         (* Finish off the function. *)
359         let _ = build_ret ret_val builder in
360
361         (* Validate the generated code, checking for consistency. *)
362         Llvm_analysis.assert_valid_function the_function;
363
364         (* Optimize the function. *)
365         let _ = PassManager.run_function the_function the_fpm in
366
367         the_function
368       with e ->
369         delete_function the_function;
370         raise e