examples/OCaml-Kaleidoscope/Chapter3/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_add lhs_val rhs_val "addtmp" builder
  26         | '-' -> build_sub lhs_val rhs_val "subtmp" builder
  27         | '*' -> build_mul 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
  49 let codegen_proto = function
  50   | Ast.Prototype (name, args) ->
  51       (* Make the function type: double(double,double) etc. *)
  52       let doubles = Array.make (Array.length args) double_type in
  53       let ft = function_type double_type doubles in
  54       let f =
  55         match lookup_function name the_module with
  56         | None -> declare_function name ft the_module
  57
  58         (* If 'f' conflicted, there was already something named 'name'. If it
  59          * has a body, don't allow redefinition or reextern. *)
  60         | Some f ->
  61             (* If 'f' already has a body, reject this. *)
  62             if block_begin f <> At_end f then
  63               raise (Error "redefinition of function");
  64
  65             (* If 'f' took a different number of arguments, reject. *)
  66             if element_type (type_of f) <> ft then
  67               raise (Error "redefinition of function with different # args");
  68             f
  69       in
  70
  71       (* Set names for all arguments. *)
  72       Array.iteri (fun i a ->
  73         let n = args.(i) in
  74         set_value_name n a;
  75         Hashtbl.add named_values n a;
  76       ) (params f);
  77       f
  78
  79 let codegen_func = function
  80   | Ast.Function (proto, body) ->
  81       Hashtbl.clear named_values;
  82       let the_function = codegen_proto proto in
  83
  84       (* Create a new basic block to start insertion into. *)
  85       let bb = append_block context "entry" the_function in
  86       position_at_end bb builder;
  87
  88       try
  89         let ret_val = codegen_expr body in
  90
  91         (* Finish off the function. *)
  92         let _ = build_ret ret_val builder in
  93
  94         (* Validate the generated code, checking for consistency. *)
  95         Llvm_analysis.assert_valid_function the_function;
  96
  97         the_function
  98       with e ->
  99         delete_function the_function;
 100         raise e