(**********************************************************************)
(* (c) Greg Morrisett, Steve Zdancewic                                *)
(*     June 1998, all rights reserved.                                *)
(**********************************************************************)

(* talutil.ml
 *
 * Utility functions for manipulating TAL types.  (Mainly stacks and registers)
 *)

open Tal

(* Some useful instruction macros *)
let mov_reg_slot r i = Mov (Reg r, (Prjr((Esp, []), i*4), []))
let mov_slot_reg i r = Mov (Prjr((Esp, []), i*4), (Reg r, []))
let stack_alloc slots = ArithBin (Sub, Reg Esp, Immed (slots*4))
let stack_free slots =  ArithBin (Add, Reg Esp, Immed (slots*4))

(* Some useful constructor constants -- left over from before the new TAL      *)
let junk_con = pcjunk 4
let var_con id = cvar id
let empty_stack = cempty
let code_con rs = ccode rs

(* Given a register state returns the current stack constructor *)
let get_stack_con rs =
  try let sptr = rs_get_reg rs Esp in
  match sptr.rcon with
    Csptr stack -> stack
  | _ -> failwith "talhelp.ml: get_stack_con Esp not a Csptr"
  with Dict.Absent -> failwith "talhelp.ml: get_stack_con Esp not in rs"

(* Given a register state and a constructor, sets the Esp register to point to *)
(* the constructor *)
let set_stack_con rs c =
  rs_set_reg rs Esp (csptr c)
    
(* The following functions are for stack-based programming at the constructor  *)
(* level.                                                                      *)

(* Cons the given constructor onto the top of the given stack constructor      *)
let cons_stack c stack = defcon (Ccons (c, stack))

exception Stack

(* Determines whether the given stack is empty.                                *)
let stack_is_empty stack = 
  match stack.rcon with
    Cempty -> true
  | Ccons _ -> false
  | Cappend _ -> false
  | _ -> failwith "talhelp.ml: empty_stack called on non-stack"

let stack_bottom stack = 
  match stack.rcon with
    Cempty -> true
  | Ccons _ -> false
  | Cappend _ -> false
  | Cvar _ -> true
  | _ -> failwith "talhelp.ml: empty_stack called on non-stack"

(* Returns the top of the stack.  If the stack has a stack variable on top,    *)
(* this function fails.                                                        *)
let rec hd_stack stack = 
  match stack.rcon with
    Cempty -> raise Stack
  | Ccons (c, _) -> c
  | Cappend (c, _) -> hd_stack c
  | _ -> failwith "talhelp.ml: hd_stack called on non-stack"

(* Returns the tail of the stack.  Fails if the top stack is a variable.       *)
let rec tl_stack stack =
  match stack.rcon with
    Cempty -> raise Stack
  | Ccons (_, c) -> c
  | Cappend (c1, c2) -> defcon (Cappend (tl_stack c1, c2))
  | _ -> failwith "talhelp.ml: tl_stack called on non-stack"

let rec pop_stack stack n =
  if n = 0 then stack
  else pop_stack (tl_stack stack) (n-1)

(* Inserts the constructor into the stack beginning at the i'th position.      *)
(* For example: insert_stack c1::c2::c3::c4::... 2 s returns the stack :       *)
(*    c1::c2::s::c3::c4::...                                                   *)
(* Fails if the stack has fewer than i non-stack-variable elements on top.     *)
let rec insert_stack stack i c =
  if i=0 then cons_stack c stack else
  cons_stack (hd_stack stack) (insert_stack (tl_stack stack) (i-1) c)

(* Inserts an ordered sequence of constructors into the stack beginning at the *)
(* i'th position.  For example: insert_list_stack c1::c2::c3::... 2 [s1; s2]   *)
(* returns the stack: c1::c2::s1::s2::c3::...                                  *)
(* Fails if the stack has fewer than i non-stack-variable elements on top      *)
let rec insert_list_stack stack i clist =
  if i=0 then List.fold_right cons_stack clist stack else
  cons_stack (hd_stack stack) (insert_list_stack (tl_stack stack) (i-1) clist)

(* stack_size returns the number of non-stack variable slots on top of the     *)
(* given stack.                                                                *)
exception StackVar of int

let stack_size stack =
  let rec helper stack acc =
    match stack.rcon with
      Cempty -> acc
    | Ccons (_, tail) -> helper tail (acc+1)
    | Cappend (topstack, bottomstack) -> 
	(try 
	  let top = helper topstack acc in
	  let bot = try helper bottomstack 0 with StackVar sz -> sz in
	  top + bot
	with
	  StackVar sz -> sz)
    | Cvar _ -> raise (StackVar acc)
    | _ -> failwith "talhelp.ml: stack_size found non-stack constructor"
  in
  try helper stack 0 with StackVar sz -> sz

(* The following functions are for manipulating code types and register states *)

let rec get_code c =
  match c.rcon with
    Cforall (_, _, c') -> get_code c'
  | Cexist (_, _, c') -> get_code c'
  | Ccode _ -> c
  | _ -> failwith "talhelp.ml: get_code c doesn't have code type"

let set_code c code =
  let rec helper c =
    match c.rcon with
      Cforall (id, k, c') -> defcon (Cforall (id, k, helper c'))
    | Cexist (id, k, c') -> defcon (Cexist (id, k, helper c'))
    | Ccode rs -> code
    | _ -> failwith "talhelp.ml: set_code found unexpected constructor"
  in
  helper c


let get_rs c =
  match (get_code c).rcon with
    Ccode rs -> rs
  | _ -> failwith "talhelp.ml: get_rs shouldn't happen"

let get_rs_copt copt =
  match copt with
    None -> failwith "talhelp.ml: get_rs_copt function as no type"
  | Some c -> get_rs c


let rec get_type_arglist_h c l =
  match c.rcon with
    Cforall (id, _, c') -> get_type_arglist_h c' ((defcon (Cvar id))::l)
  | Cexist (id, _, c') -> get_type_arglist_h c' ((defcon (Cvar id))::l)
  | _ -> l
    
let get_type_arglist copt =
  match copt with
    None -> failwith "talhelp.ml: get_type_arglist function has no type" 
  | Some c -> get_type_arglist_h c []

let rec cons_n_junks stack n =
  if n=0 then stack else cons_stack junk_con (cons_n_junks stack (n-1))

let get_stack c = get_stack_con (get_rs c)

(* NOTE: The register allocator assumes that the return address is on the top *)
(* of the stack whenever a function call is made. *)
let get_return_rs c =
  let stack = get_stack c in
  let code = hd_stack stack in
  match code.rcon with
    Ccode rs -> rs
  | _ -> failwith "rewrite.ml: get_return_rs top of stack not code!"

let get_registers rs =
  rs_fold_reg (fun r _ set -> Set.insert set r) rs (Set.empty compare_regs)  

let get_arg_registers c = get_registers (get_rs c)
let get_ret_registers c = get_registers (get_return_rs c)

let delete_regset_from_rs = 
  Set.fold (fun r rs -> rs_del_reg rs r)

let make_var_subst con cl =
  let rconlist = List.map (fun c -> match c with Tapp con -> con.rcon 
  | _ -> failwith "talhelp.ml: make_var_subst found non-Tapp") cl
  in
  let rec helper c cl subst =
    match (c.rcon, cl) with
      (Cforall(id1, k, c'), (Cvar id2)::rest) ->
	helper c' rest (Dict.insert subst id2 (var_con id1))
    | (Cexist(id1, k, c'), (Cvar id2)::rest) ->
	helper c' rest (Dict.insert subst id2 (var_con id1))
    | (Ccode _, []) -> subst
    | _ -> failwith "talhelp.ml: make_var_subst"
  in
  helper con rconlist (Dict.empty Identifier.id_compare)

let con_eq ctxt c1 c2 = try begin
  Talcon.alphaeqcon ctxt c1 c2;
  true
end with Talctxt.Talverify (_, Talctxt.Neqcon _) -> false

let compare_stacks ctxt stack1 stack2 =
  let rec helper s1 s2 i =
    match (stack_bottom s1, stack_bottom s2) with
      (true, true) -> ([], s1)
    | (false, false) ->
	let (hd1, tl1) = (hd_stack s1, tl_stack s1) in
	let (hd2, tl2) = (hd_stack s2, tl_stack s2) in
	let (junk_slots, tail) = helper tl1 tl2 (i+1) in
	if con_eq ctxt hd1 hd2 then
	  (junk_slots, cons_stack hd1 tail)
	else
	  (i::junk_slots, cons_stack junk_con tail)
    | _ -> failwith "talhelp.ml: compare_stacks different lengths"
  in
  helper stack1 stack2 0
	

(* Adds the type-variables to ctxt and sets the register state to the one in c *)
let block_type_to_ctxt ctxt copt =
  let rec helper c =
    match c.rcon with
      Cforall (id, k, c') -> Talctxt.add_var (helper c') id k
    | Cexist (id, k, c') -> Talctxt.add_var (helper c') id k
    | Ccode rs -> Talctxt.set_register_state ctxt rs
    | _ -> failwith "talhelp.ml: block_type_to_ctxt found unexpected con"
  in
  match copt with
    None -> ctxt
  | Some con -> helper (snd (Talcon.check ctxt con))

let rec elim_stack_slots slots stack =
  match slots with
    [] -> stack
  | (true::rest) -> elim_stack_slots rest (tl_stack stack)
  | (false::rest) -> cons_stack (hd_stack stack)
	            (elim_stack_slots rest (tl_stack stack))

(* Applies function f to each constructor of kind Kstack in the coercion list *)
let fix_cl ctxt cl f =
  let fix_one_con c =
    match c.rcon with
      (Cempty | Ccons (_, _) | Cappend (_, _)) ->
	f c
    | Cvar id -> (match Talctxt.get_variable_kind ctxt id with
	Kstack -> f c
      |	_ -> c)
    | _ -> c
  in
  let rec helper list =
    match list with
      [] -> []
    | (Tapp c)::rest -> (Tapp (fix_one_con c))::(helper rest)
    | _ -> failwith "talhelp.ml: fix_cl found non-type-application"
  in
  helper cl

(* EOF talutil.ml *)
