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 226