examples/OCaml-Kaleidoscope/Chapter5/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 let rec codegen_expr = function
  16   | Ast.Number n -> const_float double_type n
  17   | Ast.Variable name ->
  18       (try Hashtbl.find named_values name with
  19         | Not_found -> raise (Error "unknown variable name"))
  20   | Ast.Binary (op, lhs, rhs) ->
  21       let lhs_val = codegen_expr lhs in
  22       let rhs_val = codegen_expr rhs in
  23       begin
  24         match op with
  25         | '+' -> build_fadd lhs_val rhs_val "addtmp" builder
  26         | '-' -> build_fsub lhs_val rhs_val "subtmp" builder
  27         | '*' -> build_fmul lhs_val rhs_val "multmp" builder
  28         | '<' ->
  29             (* Convert bool 0/1 to double 0.0 or 1.0 *)
  30             let i = build_fcmp Fcmp.Ult lhs_val rhs_val "cmptmp" builder in
  31             build_uitofp i double_type "booltmp" builder
  32         | _ -> raise (Error "invalid binary operator")
  33       end
  34   | Ast.Call (callee, args) ->
  35       (* Look up the name in the module table. *)
  36       let callee =
  37         match lookup_function callee the_module with
  38         | Some callee -> callee
  39         | None -> raise (Error "unknown function referenced")
  40       in
  41       let params = params callee in
  42
  43       (* If argument mismatch error. *)
  44       if Array.length params == Array.length args then () else
  45         raise (Error "incorrect # arguments passed");
  46       let args = Array.map codegen_expr args in
  47       build_call callee args "calltmp" builder
  48   | Ast.If (cond, then_, else_) ->
  49       let cond = codegen_expr cond in
  50
  51       (* Convert condition to a bool by comparing equal to 0.0 *)
  52       let zero = const_float double_type 0.0 in
  53       let cond_val = build_fcmp Fcmp.One cond zero "ifcond" builder in
  54
  55       (* Grab the first block so that we might later add the conditional branch
  56        * to it at the end of the function. *)
  57       let start_bb = insertion_block builder in
  58       let the_function = block_parent start_bb in
  59
  60       let then_bb = append_block context "then" the_function in
  61
  62       (* Emit 'then' value. *)
  63       position_at_end then_bb builder;
  64       let then_val = codegen_expr then_ in
  65
  66       (* Codegen of 'then' can change the current block, update then_bb for the
  67        * phi. We create a new name because one is used for the phi node, and the
  68        * other is used for the conditional branch. *)
  69       let new_then_bb = insertion_block builder in
  70
  71       (* Emit 'else' value. *)
  72       let else_bb = append_block context "else" the_function in
  73       position_at_end else_bb builder;
  74       let else_val = codegen_expr else_ in
  75
  76       (* Codegen of 'else' can change the current block, update else_bb for the
  77        * phi. *)
  78       let new_else_bb = insertion_block builder in
  79
  80       (* Emit merge block. *)
  81       let merge_bb = append_block context "ifcont" the_function in
  82       position_at_end merge_bb builder;
  83       let incoming = [(then_val, new_then_bb); (else_val, new_else_bb)] in
  84       let phi = build_phi incoming "iftmp" builder in
  85
  86       (* Return to the start block to add the conditional branch. *)
  87       position_at_end start_bb builder;
  88       ignore (build_cond_br cond_val then_bb else_bb builder);
  89
  90       (* Set a unconditional branch at the end of the 'then' block and the
  91        * 'else' block to the 'merge' block. *)
  92       position_at_end new_then_bb builder; ignore (build_br merge_bb builder);
  93       position_at_end new_else_bb builder; ignore (build_br merge_bb builder);
  94
  95       (* Finally, set the builder to the end of the merge block. *)
  96       position_at_end merge_bb builder;
  97
  98       phi
  99   | Ast.For (var_name, start, end_, step, body) ->
 100       (* Emit the start code first, without 'variable' in scope. *)
 101       let start_val = codegen_expr start in
 102
 103       (* Make the new basic block for the loop header, inserting after current
 104        * block. *)
 105       let preheader_bb = insertion_block builder in
 106       let the_function = block_parent preheader_bb in
 107       let loop_bb = append_block context "loop" the_function in
 108
 109       (* Insert an explicit fall through from the current block to the
 110        * loop_bb. *)
 111       ignore (build_br loop_bb builder);
 112
 113       (* Start insertion in loop_bb. *)
 114       position_at_end loop_bb builder;
 115
 116       (* Start the PHI node with an entry for start. *)
 117       let variable = build_phi [(start_val, preheader_bb)] var_name builder in
 118
 119       (* Within the loop, the variable is defined equal to the PHI node. If it
 120        * shadows an existing variable, we have to restore it, so save it
 121        * now. *)
 122       let old_val =
 123         try Some (Hashtbl.find named_values var_name) with Not_found -> None
 124       in
 125       Hashtbl.add named_values var_name variable;
 126
 127       (* Emit the body of the loop.  This, like any other expr, can change the
 128        * current BB.  Note that we ignore the value computed by the body, but
 129        * don't allow an error *)
 130       ignore (codegen_expr body);
 131
 132       (* Emit the step value. *)
 133       let step_val =
 134         match step with
 135         | Some step -> codegen_expr step
 136         (* If not specified, use 1.0. *)
 137         | None -> const_float double_type 1.0
 138       in
 139
 140       let next_var = build_add variable step_val "nextvar" builder in
 141
 142       (* Compute the end condition. *)
 143       let end_cond = codegen_expr end_ in
 144
 145       (* Convert condition to a bool by comparing equal to 0.0. *)
 146       let zero = const_float double_type 0.0 in
 147       let end_cond = build_fcmp Fcmp.One end_cond zero "loopcond" builder in
 148
 149       (* Create the "after loop" block and insert it. *)
 150       let loop_end_bb = insertion_block builder in
 151       let after_bb = append_block context "afterloop" the_function in
 152
 153       (* Insert the conditional branch into the end of loop_end_bb. *)
 154       ignore (build_cond_br end_cond loop_bb after_bb builder);
 155
 156       (* Any new code will be inserted in after_bb. *)
 157       position_at_end after_bb builder;
 158
 159       (* Add a new entry to the PHI node for the backedge. *)
 160       add_incoming (next_var, loop_end_bb) variable;
 161
 162       (* Restore the unshadowed variable. *)
 163       begin match old_val with
 164       | Some old_val -> Hashtbl.add named_values var_name old_val
 165       | None -> ()
 166       end;
 167
 168       (* for expr always returns 0.0. *)
 169       const_null double_type
 170
 171 let codegen_proto = function
 172   | Ast.Prototype (name, args) ->
 173       (* Make the function type: double(double,double) etc. *)
 174       let doubles = Array.make (Array.length args) double_type in
 175       let ft = function_type double_type doubles in
 176       let f =
 177         match lookup_function name the_module with
 178         | None -> declare_function name ft the_module
 179
 180         (* If 'f' conflicted, there was already something named 'name'. If it
 181          * has a body, don't allow redefinition or reextern. *)
 182         | Some f ->
 183             (* If 'f' already has a body, reject this. *)
 184             if block_begin f <> At_end f then
 185               raise (Error "redefinition of function");
 186
 187             (* If 'f' took a different number of arguments, reject. *)
 188             if element_type (type_of f) <> ft then
 189               raise (Error "redefinition of function with different # args");
 190             f
 191       in
 192
 193       (* Set names for all arguments. *)
 194       Array.iteri (fun i a ->
 195         let n = args.(i) in
 196         set_value_name n a;
 197         Hashtbl.add named_values n a;
 198       ) (params f);
 199       f
 200
 201 let codegen_func the_fpm = function
 202   | Ast.Function (proto, body) ->
 203       Hashtbl.clear named_values;
 204       let the_function = codegen_proto proto in
 205
 206       (* Create a new basic block to start insertion into. *)
 207       let bb = append_block context "entry" the_function in
 208       position_at_end bb builder;
 209
 210       try
 211         let ret_val = codegen_expr body in
 212
 213         (* Finish off the function. *)
 214         let _ = build_ret ret_val builder in
 215
 216         (* Validate the generated code, checking for consistency. *)
 217         Llvm_analysis.assert_valid_function the_function;
 218
 219         (* Optimize the function. *)
 220         let _ = PassManager.run_function the_function the_fpm in
 221
 222         the_function
 223       with e ->
 224         delete_function the_function;
 225         raise e