(** Bombs-Must-Detonate: Final Compiler
    @author: Brian Go *)

open Bmdirc

type instruction = Bmdirc.instruction

(** Compilation environment location information for variables *)
type location = 
    StackLoc of int
  | GlobalLoc of int
  | FunctionLoc of string (** Function types are stored internally as strings.  The VM uses a reflection-esque mechanism to call a string.  This is safe b/c the typechecker is awesome. *)

(** Loop information for the environment to compile break & continue statements *)
type loop_info = 
    (** Break label, continue label *)
    LoopBreakContinue of string * string
  | LoopNoInfo

(** Environment keeps track of variables, stack restoration, and loop break/continue labels within different frames *)
type environment = 
    EnvTopLevel of int * (string * location) list
  | EnvFrame of int * loop_info * (string * location) list * environment


exception Compile_error of string

(** Raises a compiler error from the concatenation of the given string list *)
let raise_compile_error slist = raise (Compile_error (String.concat "" slist))

(** Makes a new compilation environment *)
let make_env () = (EnvTopLevel (0,[]))

(** Gets all variable definitions of the current stack frame and its parents*)
let rec get_cur_defs env = match env with
    EnvTopLevel (_,dlist) -> dlist
  | EnvFrame (_,_,dlist,parent) -> dlist @ (get_cur_defs parent)

(** Gets the location of the given variable *)
let get_loc v env = List.assoc v (get_cur_defs env)

(** Adds the location information of the given variable to the current stack frame *)
let add_loc (v,loc) env = match env with
    EnvTopLevel (n,dlist) -> EnvTopLevel (n,((v,loc)::dlist))
  | EnvFrame (n,info,dlist,parent) -> EnvFrame (n,info,((v,loc)::dlist),parent)

(** Checks if the given variable is defined in the current scope or one of its parent scopes *)
let isdef v env = List.mem_assoc v (get_cur_defs env)

(** Offsets every stack location in the environment by i (positive = push, negative = pop *)
let rec offset_locs i env = 
  let helper (x,y) = match y with
      StackLoc loc -> (x,StackLoc (loc+i))
    | z -> x,z
  in
    match env with 
	EnvTopLevel (n,vlist) -> 
	  EnvTopLevel (n+i, List.map helper vlist)
      | EnvFrame (n,info, vlist,parent) -> 
	  EnvFrame (n+i,info,List.map helper vlist, offset_locs i parent) 

(** Adds a frame to the environment *)
let add_frame env = EnvFrame (0,LoopNoInfo,[],env)

(** Removes a frame from the envrionment.  Offsets the parent assuming all remaining elements 
    in the frame were popped (see get_restore_code *)
let remove_frame env = match env with
    EnvTopLevel _ -> raise_compile_error ["Unexpected compilation error: attempted to pop toplevel compilation scope."]
  | EnvFrame (n,_,_,parent) -> offset_locs (-n) parent

(** Repeats the given list n times and concatenates the repetitions *)
let rec repeat_list lst n = 
  if n = 0 then []
  else if n = 1 then lst
  else if n > 1 then lst @ (repeat_list lst (n-1))
  else raise (Invalid_argument "repeat_list must have a nonnegative argument")

(** Gets the number of pops needed to restore the current stack frame to its parent *)
let get_restore_code env = match env with
    EnvTopLevel _ -> []
  | EnvFrame (n,_,_,_) -> repeat_list [Pop] n

(** Sets the break/continue information of the current environment frame *)
let set_loop_info info env = match env with
    EnvTopLevel _ -> raise_compile_error ["Unexpected compilation error: attempted to set loop information in top level."]
  | EnvFrame (n,_,defs,parent) -> EnvFrame (n,info,defs,parent)

(** Gets the relevant break label *)
let rec get_break_label env = match env with
    EnvTopLevel _ -> raise_compile_error ["Unexpected compilation error: attempted to break from a loop outside a loop."]
  | EnvFrame (_,info,_,parent) ->
      match info with
	  LoopNoInfo -> get_break_label parent
	| LoopBreakContinue (b,_) -> b

(** Gets the restore code needed to exit all frames that have been entered since the loop frame
    before breaking/continuing *)
let rec get_breakcontinue_restore env = match env with
    EnvTopLevel _ -> raise_compile_error ["Unexpected compilation error: attempted to break from a loop outside a loop."]
  | EnvFrame (n,info,_,parent) ->
      match info with
	  LoopNoInfo -> (repeat_list [Pop] n) @ get_breakcontinue_restore parent
	| LoopBreakContinue (_,_) -> repeat_list [Pop] n

(** Gets the relevant continue label *)
let rec get_continue_label env = match env with
    EnvTopLevel _ -> raise_compile_error ["Unexpected compilation error: attempted to continue a loop outside a loop."]
  | EnvFrame (_,info,_,parent) ->
      match info with
	  LoopNoInfo -> get_continue_label parent
	| LoopBreakContinue (_,c) -> c
 

    
(** See Bmdirc.make_label *)
let make_label str = Bmdirc.make_label str

(** Gets the change in stack pointer resulting from the given instruction.
    Note that this from a instruction-by-instruction point of view, not a code flow point of view. *)
let get_stack_offset inst = match inst with 
    Push -> 1
  | Pop -> -1
  | Read -> 0
  | Print -> 0
  | PrintLn -> 0
  | Swap -> 0
  | Stop -> 0
  | Return -> 0
  | Cons -> -1
  | Car -> 0
  | Cdr -> 0
  | IsNull -> 0
  | Nil -> 0
  | GetElem -> -1
  | SetElem -> -2
  | And -> -1
  | Or -> -1
  | Add -> -1
  | Sub -> -1
  | Mul -> -1
  | Div -> -1
  | DivI -> -1
  | RemI -> -1
  | Concat -> -1
  | Neg -> 0
  | Frac -> 0
  | Int -> 0
  | Lt -> -1
  | Gt -> -1
  | Lte -> -1
  | Gte -> -1
  | Eq -> -1
  | Neq -> -1
  | ConstUninit -> 0
  | ConstInt _ -> 0
  | ConstFloat _ -> 0
  | ConstString _ -> 0
  | PushSf _ -> 2
  | Assign _ -> 0
  | Acc _ -> 0
  | Rev _ -> 0
  | MakeBlock _ -> 0
  | AllocFields _ -> 0
  | GetField _ -> 0
  | SetField _ -> 0
  | Call _ -> 0
  | Apply -> 0
  | Jmp _ -> 0
  | Jz _ -> 0
  | Jnz _ -> 0
  | RPC _ -> 0
  | Label _ -> 0
  | Comment _ -> 0
  | Annotation _ -> 0
  | GetElemStatic _ -> 0
  | SetElemStatic _ -> -1
  | MakeBlockStatic (_,n) -> -n
  | MakeBlockFilled _ -> -1

(** Gets the number of global fields that need to be allocated *)
let rec n_globals env = 
  let rec helper lst = match lst with
      (_,GlobalLoc _)::rest -> 1 + (helper rest)
    | _::rest -> helper rest
    | [] -> 0
  in
    match env with 
	EnvTopLevel (_,vlist) -> helper vlist
      | EnvFrame (_,_,vlist,parent) -> (helper vlist) + (n_globals parent)

(** Compiles the given IR instruction *)
let rec compile_ir_instruction inst env = match inst with
    Instruction i -> [i], offset_locs (get_stack_offset i) env
  | GetVar v -> 
      if isdef v env then
	(match get_loc v env with
	     StackLoc loc ->
	       [Acc loc], env
	   | GlobalLoc loc ->
	       [GetField loc], env
	   | FunctionLoc funname ->
	       [ConstString funname], env)
      else
	raise_compile_error ["Error in instruction GetVar, ";v;" was not defined."]
  | SetVar v ->
      if isdef v env then
	(match get_loc v env with
	     StackLoc loc ->
	       [Assign loc], env
	   | GlobalLoc loc ->
	       [SetField loc], env
	   | FunctionLoc name ->
	       raise_compile_error ["Unexpected compile error: attempted to re-define function ";name;"."])
      else 
	[Push], add_loc (v,StackLoc 0) (offset_locs 1 env)
  | NoteVar (v,i) -> [], add_loc (v,StackLoc i) env
  | NoteFunction name -> [], add_loc (name,FunctionLoc name) env
  | If (test,thenbody,elsebody) ->
      let else_label = make_label "else" in
      let endif_label = make_label "endif" in
      let test_result,_ = compile_ir_instruction_list test env in
      let then_result,then_env = compile_ir_instruction_list thenbody (add_frame env) in
      let else_result,else_env = compile_ir_instruction_list elsebody (add_frame env) in
      let restore_code_then = get_restore_code then_env in
      let restore_code_else = get_restore_code else_env in
	[Comment "if"]
	@ test_result
	@ [Jz else_label]
	@ then_result
	@ restore_code_then
	@ [Jmp endif_label;
	   Label else_label]
	@ else_result
	@ restore_code_else
	@ [Label endif_label]
	@ [Comment "endif"], env
  | For (init, test, step, body) ->
      let test_label = make_label "for_test" in
      let endfor_label = make_label "for_end" in
      let init_result,init_env = compile_ir_instruction_list init (add_frame env) in
      let restore_init = get_restore_code init_env in
      let test_result,_ = compile_ir_instruction_list test init_env in
      let env_for_body = set_loop_info (LoopBreakContinue (endfor_label,test_label)) (add_frame init_env) in
      let body_result,body_env = compile_ir_instruction_list body env_for_body in
      let restore_body = get_restore_code body_env in
      let step_result,step_env = compile_ir_instruction_list step (add_frame init_env) in
      let restore_step = get_restore_code step_env in
	[Comment "for"]
	@ init_result
	@ [Label test_label]
	@ test_result
	@ [Jz endfor_label]
	@ body_result
	@ restore_body
	@ step_result
	@ restore_step
	@ [Jmp test_label]
	@ [Label endfor_label]
	@ restore_init
	@ [Comment "end for"], env
  | While (test, body) ->
      let test_label = make_label "while" in
      let wend_label = make_label "wend" in
      let test_result,_ = compile_ir_instruction_list test env in 
      let env_for_body = set_loop_info (LoopBreakContinue (wend_label,test_label)) (add_frame env) in
      let body_result,body_env = compile_ir_instruction_list body env_for_body in
      let restore_body = get_restore_code body_env in
	[Comment "while"]
	@ [Label test_label]
	@ test_result
	@ [Jz wend_label]
	@ body_result
	@ restore_body
	@ [Jmp test_label]
	@ [Label wend_label]
	@ [Comment "wend"], env
  | DoWhile (body,test) ->
      let do_label = make_label "do" in
      let end_label = make_label "end_dowhile" in
      let test_result,_ = compile_ir_instruction_list test env in 
      let env_for_body = set_loop_info (LoopBreakContinue (end_label,do_label)) (add_frame env) in
      let body_result,body_env = compile_ir_instruction_list body env_for_body in
      let restore_body = get_restore_code body_env in
	[Comment "dowhile"]
	@ [Label do_label]
	@ body_result
	@ restore_body
	@ [Comment "dowhile_test"]
	@ test_result
	@ [Jnz do_label]
	@ [Label end_label]
	@ [Comment "end dowhile"], env
  | Break -> 
      let restore_code = get_breakcontinue_restore env in
	restore_code @ [Jmp (get_break_label env)], env
  | Continue -> 
      let restore_code = get_breakcontinue_restore env in
	restore_code @ [Jmp (get_continue_label env)], env
  | BeginScope -> [], add_frame env
  | EndScope -> [], remove_frame env
  | EndScopeSf -> [], offset_locs (-2) (remove_frame env)
  | EndScopeSfRPC -> [], offset_locs (-1) (remove_frame env)
  | DeclareGlobal s ->
      let n = (n_globals env) in
	[], add_loc (s,GlobalLoc n) env

(** Compiles a list of IR instructions *)
and compile_ir_instruction_list ir_buf env = match ir_buf with
    cur::rest -> 
      let cur_result,new_env = compile_ir_instruction cur env in
      let rest_result,final_env = compile_ir_instruction_list rest new_env in
	cur_result @ rest_result, final_env
  | [] -> [], env

(** Compiles a complete buffer of IR instructions, including the compiler signal to allocate global fields *)
let compile_buffer ir_buf = 
  try
    let start_env = make_env () in
    let result,result_env = compile_ir_instruction_list ir_buf start_env in
      (AllocFields (n_globals result_env))::result
  with 
      Compile_error s ->
	let _ = output_string stderr s in
	let _ = output_string stderr "\n" in
	let _ = flush stderr in
	  raise (Compile_error s)

(** Produces a string representation of the given instruction buffer *)
let rec string_of_buffer buf = match buf with
    cur::rest -> String.concat "" [(Bmdirc.string_of_instruction cur);"\n";(string_of_buffer rest)]
  | [] -> ""

(** Prints the given instruction buffer to the given output channel *)
let print_buffer chan buf = 
  output_string chan (string_of_buffer buf)