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