(** 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)