Source

blub / blub_vm.ml

open Printf
open Blub_ast
open Blub_common
open Blub_types
open Blub_environment

exception Var_not_found of variable * environment


module Vm = struct
  open Format

  let array_to_pair (lst:sval array) =
    Array.fold_right (fun v pr ->
		      Spair { car=v; cdr=pr }) lst Snull

  let pp_bc f = function
      B_lit (l, _) -> fprintf f "lit %a" pp_sval l
    | B_ref ((var, Gbinding idx), _) -> fprintf f "gref %s:%d" var.var_uname idx
    | B_ref ((var, Lbinding (frame, idx)), _) ->
	fprintf f "lref %s:%d.%d" var.var_uname frame idx
    | B_set ((var, Gbinding idx), _) -> fprintf f "gset %s:%d" var.var_uname idx
    | B_set ((var, Lbinding (frame, idx)), _) ->
	fprintf f "lset %s:%d.%d" var.var_uname frame idx
    | B_test _ -> fprintf f "test"
    | B_abs (abs, _) ->
	fprintf f "lambda/%d" (Array.length abs.lam_params)
    | B_ret -> fprintf f "ret"
    | B_apply -> fprintf f "apply"
    | B_arg (idx, _) -> fprintf f "arg %d" idx
    | B_tailframe (sz, _) -> fprintf f "tailframe %d" sz
    | B_frame (sz, _, _) -> fprintf f "frame %d" sz
    | B_conti _ -> fprintf f "conti"
    | B_nuate _ -> fprintf f "nuate"
    | B_halt -> fprintf f "halt"
  ;;

  (* TODO need to be able to invoke this from the eval loop; the main
     problem is the typechecking which has to be worked out.  Sometimes
     you do better with the full program; sometimes you do better with
     having local variables bound to values.  Is it possible to do an
     incrementally improving type inference? *)

  let llvm_compile tmap globals env lambda = 
    let ast = A_abs (lambda.lam_params, lambda.lam_varargs, lambda.lam_ast) in

    (* Gather up all the free variables. *)

    (* Parameters will be bound when the function is called, of course. *)
    let bound_vars = Array.fold_right 
      (fun var acc ->
	 Blub_closure.VarSet.add var acc)
      lambda.lam_params Blub_closure.VarSet.empty 
    in

    (* Should only add other vars if bound to values *)
    let add_bound_vars bound vars vals =
      let rec iterate idx =
	if idx < 0 then 
          bound
	else
          if vals.(idx) != Sunbound then
            Blub_closure.VarSet.add vars.(idx) (iterate (idx-1))
          else
            iterate (idx-1)
      in
      iterate (Array.length vars - 1)
    in
    
    let bound_vars = add_bound_vars bound_vars globals.vars globals.vals in

    let bound_vars = List.fold_right 
      (fun frame bnd ->
         add_bound_vars bnd frame.vars frame.vals)
      env bound_vars
    in  

    let () = try
      let free_vars = Blub_closure.find_free bound_vars lambda.lam_ast in
      let num_free = Blub_closure.VarSet.cardinal free_vars in
      printf "Function has %d free variables: %a\n%!" 
	num_free (pp_list pp_var) (Blub_closure.VarSet.elements free_vars);
      
      (* 
	 Conditions for invoking LLVM right now:
	 [1] no free variables
	 [2] all type variables are bound to concrete types 
      *)

      let unsolved_typevars = Blub_typecheck.all_unsolved tmap ast in
      let num_unsolved = Blub_typecheck.TypeVarSet.cardinal unsolved_typevars in
      printf "Function has %d unsolved type variables\n%!" num_unsolved;

      if (num_free = 0) && (num_unsolved = 0) then begin
	let llvm_code = Blub_llvm.jit_compile globals env ast tmap in
	let Fun (argtypes, _) = PMap.find ast tmap in
	
	let argtypes = Array.map Blub_typecheck.simplify argtypes in
	printf "Adding JITcode for argtypes %a\n%!"
	  (pp_array pp_type) argtypes;

	lambda.lam_jitcode <- 
	  (argtypes, llvm_code) :: lambda.lam_jitcode
      end;
      ()
    with 
	Assert_failure (filename, line, _) ->
	  printf "ASSERT FAILED in LLVM compile: %s:%d\n%!" 
	    filename line
      |	Not_found ->
	  printf "Llvm compilation not successful\n%!"
    in
    ()
  ;;

  let compile (globals:env_frame) (env:environment) (ast:ast) (next:bytecode) : bytecode =
    let tmap = try
      let tmap = Blub_typecheck.tcheck globals env ast in
      PMap.iter (fun ast typ ->
		    printf "type: %a => %a (%d unsolved typevars)\n%!" pp_ast ast pp_type typ (Blub_typecheck.TypeVarSet.cardinal (Blub_typecheck.unsolved typ)))
	tmap;
      tmap
    with Blub_typecheck.Tcheck_failed _ -> 
      printf "Failed to typecheck\n%!";
      PMap.empty
    in
    let rec bc_compile env ast next =
      printf "bc-compile: %a  in env: %a\n%!" pp_ast ast pp_env env;
      match ast with
	  A_lit l -> B_lit (l, next)
	| A_ref r -> 
	    B_ref ((r, Env.find globals env r), next)
	| A_set (r, expr) -> 
	    let e = B_set ((r, Env.find globals env r), next) in
	    bc_compile env expr e
	| A_cnd (e1, e2, e3) ->
	    let e2 = bc_compile env e2 next in
	    let e3 = bc_compile env e3 next in
	    bc_compile env e1 (B_test (e2, e3))
	| A_seq exprs ->
	    Array.fold_right (bc_compile env) exprs next
	| A_abs (params, varargs, body) ->
	    (* We will definitely compile the bytecode version of the function
	       now.  If conditions are right, we might compile the native
	       code version here as well. *)

	    let frame = 
	      { vars = params;
		vals = Array.map (fun _ -> Sunbound) params } in
	    let env = frame :: env in
	    let code = bc_compile env body B_ret in
	    let lambda = { lam_params = params;
			   lam_varargs = varargs;
			   lam_ast = body;
			   lam_bytecode = code;
			   lam_jitcode = [] } in

	    llvm_compile tmap globals env lambda;
	    B_abs (lambda , next)

	| A_app (fn, args) -> begin
	    let c = bc_compile env fn B_apply in
	    let _, n = Array.fold_left
	      (fun (idx, next) arg ->
		 let arg = bc_compile env arg (B_arg (idx, next)) in
		 (idx+1, arg))
	      (0, c) args
	    in
	    match next with
		B_ret -> B_tailframe (Array.length args, n)
	      | _ -> B_frame (Array.length args, next, n)
	  end
	    
	| A_let (bindings, body, LT_let) ->
	    let params = Array.map fst bindings in
	    let args = Array.map snd bindings in
	    bc_compile env (A_app (A_abs (params, false, body), args)) next
	      
	| A_let (bindings, body, _) ->
	    let params = Array.map fst bindings in
	    let args = Array.map (fun _ -> A_lit Snull) bindings in
	    
	    let init = A_seq 
	      (Array.map (fun (var, expr) -> A_set (var, expr)) bindings)
	    in
	    let body = A_seq [| init; body |] in
	    bc_compile env (A_app (A_abs (params, false, body), args)) next
	      
	| A_callcc expr -> begin
	    let c = B_conti (B_arg (0, bc_compile env expr B_apply)) in
	    match next with
		B_ret -> B_tailframe (1, c)
	      | _ -> B_frame (1, next, c)
	  end
    in
    bc_compile env ast next
  ;;

  let frame2 = Array.create 2 Sunbound;;

  let rec eval acc globals (env:environment) (rib:sval array) stack inst =
    printf "eval: %a %a  acc: %a  rib: %a\n%!" pp_bc inst pp_env env pp_sval acc (pp_array pp_sval) rib;
    match inst with
	B_lit (lit, next) -> eval lit globals env rib stack next
      | B_ref ((var, Gbinding idx), next) ->
	  let r = globals.vals.(idx) in
	  if (r = Sunbound) then begin
	    raise (Var_not_found (var, env))
	  end;
	  eval r globals env rib stack next
      | B_ref ((var, Lbinding (f, idx)), next) ->
	  let r = (List.nth env f).vals.(idx) in
	  if (r = Sunbound) then begin
	    raise (Var_not_found (var, env))
	  end;
	  eval r globals env rib stack next
      | B_set ((_, Gbinding idx), next) ->
	  globals.vals.(idx) <- acc;
	  eval Snull globals env rib stack next
      | B_set ((_, Lbinding (f, idx)), next) ->
	  (List.nth env f).vals.(idx) <- acc;
	  eval Snull globals env rib stack next    
      | B_test (trueb, falseb) ->
	  let next = if acc == Sfalse then falseb else trueb in
	  eval acc globals env rib stack next
      | B_abs (abs, next) -> begin
	  let closure = Sclosure (env, abs) in
	  eval closure globals env rib stack next
	end
      | B_ret ->
	  let frame :: stack = stack in
	  eval acc globals frame.env frame.rib stack frame.ret
	    
      | B_tailframe (sz, next) ->
	  let rib = Array.create sz Sunbound in
	  eval acc globals env rib stack next

      | B_frame (sz, ret, next) ->
	  let stack = { env=env; rib=rib; ret=ret } :: stack in
	  let rib = Array.create sz Sunbound in
	  eval acc globals env rib stack next
	    
      | B_apply -> 
	  apply acc globals env rib stack
	    
      | B_arg (idx, next) ->
	  rib.(idx) <- acc;
	  eval acc globals env rib stack next
	    
      | B_conti x ->
	  let var = mklocal "conti" in
	  let body = B_nuate (stack, (var, Lbinding (0, 0))) in
	  let abs = { lam_params = [| var |];
		      lam_varargs = false;
		      lam_ast = A_lit (Sunbound);
		      lam_bytecode = body;
		      lam_jitcode = [] } in
	  let acc = Sclosure ([], abs) in
	  eval acc globals env rib stack x
	    
      | B_nuate (stack, r) ->
	  let acc = (List.hd env).vals.(0) in
	  eval acc globals env rib stack B_ret
	    
      | B_halt -> acc
	  
  and apply acc globals env rib stack =
    match acc with
	Sclosure (env, abs) -> begin
	  let { lam_params = params;
		lam_varargs = varargs;
		lam_bytecode = body } = abs in

	  let frame = if varargs then
	    (* one param (the last) is the varargs, the rest are the
	       regular args *)
	    let nparams = Array.length params in
	    let nargs = nparams-1 in
	    let nvargs = (Array.length rib) - nargs in
	    
	    let vargs = Array.sub rib nargs nvargs in
	    let last_arg = array_to_pair vargs in
	    
	    let args = Array.create nparams Sunbound in
	    Array.blit rib 0 args 0 nargs;
	    args.(nargs) <- last_arg;
	    { vars = params; vals = args }
	  else
	    { vars = params; vals = rib }
	  in
	  let env = frame :: env in

	  (* Gather up the types of the arguments *)
	  let arg_types = Array.map Blub_typecheck.type_of_sval frame.vals in

	  (* Look for some LLVM code to save us *)
	  try
	    let code = List.assoc arg_types abs.lam_jitcode in
	    let result = Blub_llvm.execute_function code frame.vals in
	    eval result globals env rib stack B_ret
	  with Not_found ->
	    printf "Failed to find JITcode for argtypes %a (%d options)\n%!" 
	      (pp_array pp_type) arg_types (List.length abs.lam_jitcode);
	    eval acc globals env rib stack body
	end

      | Spfn { prim_bc = fn } ->
	  eval (fn rib) globals env rib stack B_ret
      | Spfe fn -> fn globals env rib stack
      | x ->
	  printf "Trying to apply non-function sval: %a\n%!" pp_sval x;
	  assert false
	    
  ;;
  let topeval globals code = 
    try
      eval Sunbound globals [] [||] [] code
    with
	Var_not_found (var, env) ->
	  printf "Variable@ %a@ not found@ in environment@ %a@\n%!"
	    pp_var var pp_env env;
	  Sunbound
  ;;

end