Commits

Anonymous committed 756893e

Cleaned up version of Son of Blub that actually invokes the LLVM JIT compiler.

Comments (0)

Files changed (1)

part2/son_of_blub.ml

   | Ast_ref of variable * int * int  (* variable; frame #; offset # *)
   | Ast_cnd of ast * ast * ast  (* if part; then part; else part *)
   | Ast_app of ast * ast array
-  | Ast_abs of variable array * ast
+  | Ast_abs of lambda
 
 and lambda = {
   lam_ast : ast;
   lam_params : variable array;
+  lam_lltype : Llvm.lltype option;  (* this is a hack because we don't have a
+				       type system *)
 }
 
 and sval =
   close_lam : lambda;
   (* we could either have this as part of the lambda (compiled without a
      local environment) or as part of a closure (compiled in an environment) *)
-  mutable close_jitcode : (Llvm.lltype * Llvm.llvalue) option  (* return value; code *)
+  close_jitcode : (Llvm.lltype * Llvm.llvalue) Lazy.t  (* return value; code *)
 }
 
 and environment_frame = {
 exception Jit_failed
 open Llvm
 
-let make_bool value = 
-  Sllvm (Llvm.i1_type, GV.of_int Llvm.i1_type (if value then 1 else 0)) ;;
-let make_int value = Sllvm (Llvm.i64_type, GV.of_int Llvm.i64_type value) ;;
 
-let get_type ast (typeinfo:(ast * lltype) list) =
-  List.assq ast typeinfo
-;;
+let llvm_val_of_int t x = Sllvm (t, GV.of_int t x) ;;
+let make_bool value = llvm_val_of_int Llvm.i1_type (if value then 1 else 0) ;;
+let make_int value = llvm_val_of_int Llvm.i64_type value ;;
 
 let llvalue_of_gv t v =
   match t with
 let cur_module = Llvm.create_module "helloworld" ;;
 let jit = EE.create (ModuleProvider.create cur_module) ;;
 
-let compile_fn (env:environment) (ast:ast) typeinfo = 
-  let fn_type = try get_type ast typeinfo 
-  with Not_found -> Format.printf "OH OH\n%!"; raise Jit_failed in
+
+
+let compile_fn (env:environment) (lambda:lambda) = 
+  let fn_type = match lambda.lam_lltype with
+      Some t -> t | None -> raise Jit_failed
+  in
   let rettype = return_type fn_type in
   let argtypes = param_types fn_type in
 
-    let cur_fn = define_function "lambda" fn_type cur_module in
-    let builder = builder_at_end (entry_block cur_fn) in
-    
-    let params, body = match ast with
-	Ast_abs (params, body) -> params, body
-      | _ -> Format.printf "FAILED\n%!"; raise Jit_failed
-    in
-    let rec gen_llvm (builder:llbuilder) (ast:ast) : (llbuilder * llvalue) =
-      match ast with
-	  Ast_lit (Sllvm (t, v)) -> 
-	    let lit_val = llvalue_of_gv t v in
-	    (builder, lit_val)
-	| Ast_ref (var, 0, offset) ->
-	    assert (ExtArray.Array.mem var params);
-	    (builder, param cur_fn offset)
-	| Ast_ref (var, frame, offset) -> 
-	    Format.printf "(%d, %d) with %d frames" frame offset (List.length env);
-	    let v = match (List.nth env (frame-1)).env_frame_vals.(offset) with
-		Sllvm (t, v) -> llvalue_of_gv t v 
-	      | _ -> raise Jit_failed
-	    in
-	    (builder, v)
-	| Ast_cnd (pred, cons, alt) ->
+  let cur_fn = define_function "lambda" fn_type cur_module in
+  let builder = builder_at_end (entry_block cur_fn) in
+
+  let rec gen_llvm (builder:llbuilder) (ast:ast) : (llbuilder * llvalue) =
+    match ast with
+	Ast_lit (Sllvm (t, v)) -> 
+	  let lit_val = llvalue_of_gv t v in
+	  (builder, lit_val)
+      | Ast_ref (var, 0, offset) ->
+	  assert (ExtArray.Array.mem var lambda.lam_params);
+	  (builder, param cur_fn offset)
+      | Ast_ref (var, frame, offset) -> 
+	  (* frame 0 is found in the fn params, not in the environment *)
+	  let v = match (List.nth env (frame-1)).env_frame_vals.(offset) with
+	      Sllvm (t, v) -> llvalue_of_gv t v 
+	    | _ -> raise Jit_failed
+	  in
+	  (builder, v)
+      | Ast_cnd (pred, cons, alt) ->
 	  let builder, pred_val = gen_llvm builder pred in
 	  let test = build_icmp Icmp.Ne pred_val (const_int i1_type 0) 
 	    "test" builder in
-      
+	  
 	  let cons_block = append_block "true_branch" cur_fn in
 	  let alt_block = append_block "false_branch" cur_fn in
 	  let join_block = append_block "join_branches" cur_fn in
 	  (join_builder, return)
     in
 
-    let builder, retval = gen_llvm builder body in
+    let builder, retval = gen_llvm builder lambda.lam_ast in
     ignore (build_ret retval builder);
     (rettype, cur_fn)
 
 
 (* INTERPRETER SECTION *)
 
-let eval typeinfo env ast =
-  let rec eval env = function 
+let rec eval env = function 
     Ast_lit lit -> lit
   | Ast_ref (_, frame, offset) -> (List.nth env frame).env_frame_vals.(offset)
   | Ast_cnd (pred, cons, alt) -> begin
 	| _ -> eval env cons
     end
   | Ast_app (fn, args) -> apply (eval env fn) (Array.map (eval env) args)    
-  | Ast_abs (params, body) as ast ->
-      let jitcode = 
-	try
-	  Format.printf "Try compiling\n%!";
-	  let c = compile_fn env ast typeinfo in
-	  Format.printf "Got code\n%!";
-	  Some c
-	with Jit_failed -> None
-      in
+  | Ast_abs lambda ->
       Sclosure { close_env = env; 
-		 close_lam = { lam_ast = body;
-			       lam_params = params };
-		 close_jitcode = jitcode }
+		 close_lam = lambda;
+		 close_jitcode = lazy (compile_fn env lambda) }
 and apply fn args =
   match fn with
       Sclosure close -> begin
-	match close.close_jitcode with
-	    Some (rettype, fn) -> 
-	      let args = Array.map (fun (Sllvm (_, v)) -> v) args in
-	      let result = EE.run_function fn args jit in
-	      Format.printf "Ran code\n%!";
-	      Sllvm (rettype, result)
-
-	  | None ->
-	      let frame = { env_frame_vars = close.close_lam.lam_params;
-			    env_frame_vals = args } in
-	      let env' = frame :: close.close_env in
-	      eval env' close.close_lam.lam_ast
+	try
+	  let rettype, fn = Lazy.force close.close_jitcode in
+	  let args = Array.map (fun (Sllvm (_, v)) -> v) args in
+	  let result = EE.run_function fn args jit in
+	  Sllvm (rettype, result)
+	with Jit_failed ->
+	  let frame = { env_frame_vars = close.close_lam.lam_params;
+			env_frame_vals = args } in
+	  let env' = frame :: close.close_env in
+	  eval env' close.close_lam.lam_ast
       end
-  in
-  eval env ast
 ;;
 
-let make_var name = Variable name ;;
 
-let make_lit gv = Ast_lit gv ;;
-let make_ref var frame offset = Ast_ref (var, frame, offset) ;;
-let make_cnd pred cons alt = Ast_cnd (pred, cons, alt) ;;
-let make_abs params body = Ast_abs (Array.of_list params, body) ;;
-let make_app fn args = Ast_app (fn, Array.of_list args) ;;
+(* TEST SECTION *)
 
-module Input = struct
-  type node =
-      Lit of myval
-    | Ref of string
-    | Cnd of node * node * node
-    | Abs of string list * node
-    | App of node * node list
-  and myval =
-      Bool of bool
-    | Int of int
-
-  let find var env =
-    let rec find frame_idx = function
-	[] -> assert false
-      | { env_frame_vars = vars } :: env -> 
-	  try 
-	    frame_idx, ExtArray.Array.findi ((=) var) vars
-	  with Not_found -> find (frame_idx+1) env
-    in
-    find 0 env
-  ;;
-
-  let rec to_ast env = function
-      Lit (Bool x) -> make_lit (make_bool x)
-    | Lit (Int x) -> make_lit (make_int x)
-    | Ref v -> 
-	let variable = Variable v in
-	let frame, offset = find variable env in
-	Ast_ref (variable, frame, offset)
-    | Cnd (pred, cons, alt) -> Ast_cnd (to_ast env pred,
-					to_ast env cons,
-					to_ast env alt)
-    | Abs (params, body) -> 
-	let vars = Array.of_list (List.map (fun v -> Variable v) params) in
-	let vals = Array.map (fun _ -> Sunbound) vars in
-	let frame = { env_frame_vars = vars;
-		      env_frame_vals = vals } in
-	let env = frame :: env in
-	Ast_abs (vars, to_ast env body)
-    | App (fn, args) ->
-	Ast_app (to_ast env fn, Array.of_list (List.map (to_ast env) args))
-end
+let mkvar name = Variable name ;;
+let lit gv = Ast_lit gv ;;
+let lit_bool x = lit (make_bool x) ;;
+let lit_int x = lit (make_int x) ;;
+let mkref var frame offset = Ast_ref (mkvar var, frame, offset) ;;
+let mkcnd pred cons alt = Ast_cnd (pred, cons, alt) ;;
+let mkabs params body lltype = 
+  Ast_abs { lam_params = Array.of_list (List.map mkvar params);
+	    lam_ast = body;
+	    lam_lltype = lltype } ;;
+let mkapp fn args = Ast_app (fn, Array.of_list args) ;;
 
 
 let () =
-  let lit_bool x = Input.Lit (Input.Bool x) in
-  let lit_int x = Input.Lit (Input.Int x) in
-  let app fn args = Input.App (fn, args) in
-  let abs params body = Input.Abs (params, body) in
-  let cnd pred cons alt = Input.Cnd (pred, cons, alt) in
-  let var v = Input.Ref v in
+  let lambda = mkabs ["a"; "x"; "y"] 
+    (mkcnd (mkref "a" 0 0) (mkref "x" 0 1) (mkref "y" 0 2))
+    (Some (function_type i64_type [| i1_type; i64_type; i64_type |]))
+  in
+  let expr = mkapp lambda [lit_bool true; lit_int 34; lit_int 47] in
+  let result = eval [] expr in
 
-  let lambda = abs ["a"; "x"; "y"] (cnd (var "a") (var "x") (var "y")) in
-  let expr = app lambda [lit_bool true; lit_int 34; lit_int 47] in
-  let expr = Input.to_ast [] expr in
-  let Ast_app (lambda, _) = expr in 
-
-  let result = eval [lambda, function_type i64_type 
-		       [| i1_type; i64_type; i64_type |] ]
-    [] expr
-  in
-
-  let () = match result with
+  match result with
       Sllvm (t, v) when t = Llvm.i64_type -> assert (GV.as_int v = 34)
     | _ -> assert false
+;;
+
+let () =
+  let lambda1 = mkabs ["a"; "x"] 
+    (mkcnd (mkref "a" 0 0) (mkref "x" 0 1) (mkref "y" 1 0))
+    (Some (function_type i64_type [| i1_type; i64_type |]))
   in
-  
+  let lambda2 = mkabs ["y"] lambda1 None in (* can't type this _yet_ *)
+  let expr = mkapp lambda2 [lit_int 47] in
+  let expr = mkapp expr [lit_bool true; lit_int 34] in
+  let result = eval [] expr in
 
-  let lambda1 = abs ["a"; "x"] (cnd (var "a") (var "x") (var "y")) in
-  let lambda2 = abs ["y"] lambda1 in
-  let expr = app lambda2 [lit_int 37] in
-  let expr = app expr [lit_bool true; lit_int 34] in
-  let expr = Input.to_ast [] expr in
-
-  let Ast_app (Ast_app (Ast_abs (_, lambda), _), _) = expr in
-
-  let result = eval [lambda, function_type i64_type
-		       [| i1_type; i64_type |] ] [] expr in
-  let () = match result with
+  match result with
       Sllvm (t, v) when t = Llvm.i64_type -> assert (GV.as_int v = 34)
     | _ -> assert false
-  in
-  dump_module cur_module
 ;;
 
+let () = Llvm.dump_module cur_module ;;
 
-(* add a Let statement *)