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





(** Final instruction *)


type instruction =


(** 0-arity *)


    Push
  | Pop
  | Read
  | Print
  | PrintLn
  | Swap
  | Stop
  | Return
  | Cons
  | Car
  | Cdr
  | IsNull
  | Nil
  | GetElem
  | SetElem
  | And
  | Or
  | Add
  | Sub
  | Mul
  | Div
  | DivI
  | RemI
  | Concat
  | Neg
  | Frac
  | Int
  | Lt
  | Gt
  | Lte
  | Gte
  | Eq
  | Neq
  | ConstUninit
  | Apply


(** 1-arity *)


  | ConstInt of int
  | ConstFloat of float
  | ConstString of string
  | PushSf of string
  | Assign of int
  | Acc of int
  | Rev of int
  | MakeBlock of int
  | MakeBlockFilled of int
  | AllocFields of int
  | GetField of int
  | SetField of int
  | Call of string
  | Jmp of string
  | Jz of string
  | Jnz of string  | RPC of string
  | Label of string
  | Comment of string
  | Annotation of string
  | GetElemStatic of int
  | SetElemStatic of int


(** 2-arity *)


  | MakeBlockStatic of int * int



(** IR Instruction *)


type ir_instruction =
    Instruction of instruction


(** Ir-Specific *)


  | GetVar of string
  | SetVar of string
  | NoteVar of string * int
  | NoteFunction of string
  | If of ir_buffer * ir_buffer * ir_buffer 

(** test, then, else *)


  | For of ir_buffer * ir_buffer * ir_buffer * ir_buffer 
  | While of ir_buffer * ir_buffer
  | DoWhile of ir_buffer * ir_buffer
  | Break
  | Continue
  | BeginScope
  | EndScope
  | EndScopeSf
  | EndScopeSfRPC
  | DeclareGlobal of string

and ir_buffer = ir_instruction list



(** Shorthand for String.concat " *)


let concat s_list = String.concat "" s_list

exception Ircompile_error of string



(** Keeps track of the buffer of code used to initialize structs/arrays/lists. We need to initialize values of these types as null but of the correct structure. i.e. someStructType someStruct; stomeStruct.someField = someValue; should be valid. *)


let init_buffer = ref []



(** Raises an IR compile error by concatenating s_list *)


let raise_compile_error s_list = 
  raise (Ircompile_error (concat s_list))



(** Concatenates n repetitions of the given list *)


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")



(** Used by make_label *)


let label_id = ref (-1)



(** Generate labels with unique prefixes *)


let make_label str = 
  let _ = label_id := !label_id + 1 in
    concat ["_";string_of_int !label_id;"_";str]



(** Takes an argument list and notes the positions *)


let rec note_arg_list lst index = match lst with
    (_,varname)::rest -> (NoteVar (varname,index))::(note_arg_list rest (index+1))
  | [] -> []



(** Compiles the initialization code of data type dt *)


let rec compile_datatype_init dt = match dt with
    (Ast.SynIntType
    | Ast.SynFloatType
    | Ast.SynStringType
    | Ast.SynBoolType
    | Ast.SynVoidType
    | Ast.SynRefType _) -> [Instruction ConstUninit]
  | Ast.SynEnumOrStructType (name,field_types) -> 

(** field_types is empty if its an enum/template *)


      if List.length !field_types > 0 then
        let n_fields = List.length !field_types in
        let fields_result = compile_datatypes_pushing !field_types in
          fields_result
          @ [Instruction (Rev n_fields)]
          @ [Instruction (MakeBlockStatic (1,n_fields))]
      else
        [Instruction ConstUninit]
  | Ast.SynArrayType (t_arr,vprod) ->
      let dt_code = compile_datatype_init t_arr in
      let vprod_code = compile_value_producer vprod in
        dt_code 
        @ [Instruction Push]
        @ vprod_code
        @ [Instruction (MakeBlockFilled 0)]
  | Ast.SynListType _ -> 
      [Instruction Nil]
  | Ast.SynArrowType (_,_) -> [Instruction (ConstString "")]



(** Compiles each data type in dt_list, inserting push instructions between each compilation result *)


and compile_datatypes_pushing dt_list = match dt_list with
    dt::rest -> (compile_datatype_init dt)@[Instruction Push]@(compile_datatypes_pushing rest)
  | [] -> []



(** Compiles an expression *)


and compile_expression expr = match expr with
    Ast.SynVarDeclare vdecl -> compile_variable_declaration vdecl
  | Ast.SynEnumDeclare edecl -> []
  | Ast.SynVarAssign vasgn -> compile_variable_assignment vasgn
  | Ast.SynCond cond -> compile_conditional cond
  | Ast.SynLoop loop -> compile_loop loop
  | Ast.SynFunctionCall funcall -> compile_function_call funcall
  | Ast.SynReturnStatement ret -> compile_return ret
  | Ast.SynBreak -> [Break]
  | Ast.SynContinue -> [Continue]



(** Compiles a list of expressions *)


and compile_expression_list expr_list = match expr_list with
    cur::rest -> 
      let curbuf = compile_expression cur in
      let restbuf = compile_expression_list rest in
        curbuf@restbuf
  | [] -> []




(** Compiles a variable declarations *)


and compile_variable_declaration vdecl = match vdecl with
    Ast.SynVarDeclareNoInit (dt, varname) -> 
      let dt_init = compile_datatype_init dt in
        dt_init @ [SetVar varname]
  | Ast.SynVarDeclareWithInit(_, varname, vprod) -> 
      let valbuf = compile_value_producer vprod in
        valbuf@[SetVar varname]



(** Compiles a conditional. Would be nice to add short-circuiting in the future. *)


and compile_conditional cond = match cond with
    Ast.SynIf (vp,expr_list) -> 
      let vp_result = compile_value_producer vp in
      let expr_result = compile_expression_list expr_list in
        [If (vp_result, expr_result, [])]
  | Ast.SynIfCase (vp,expr_list,continued_cond) ->
      let vp_result = compile_value_producer vp in
      let expr_result = compile_expression_list expr_list in
      let continued_result = compile_continued_conditional continued_cond in
        [If (vp_result, expr_result, continued_result)]



(** Compiles an else case *)


and compile_continued_conditional cond = match cond with
    Ast.SynFinalElse expr_list -> compile_expression_list expr_list
  | Ast.SynElse cond -> compile_conditional cond



(** Compiles a loop *)


and compile_loop loop = match loop with
    Ast.SynWhile (vprod,lexpr_list) -> 
      let vp_result = compile_value_producer vprod in
      let lexpr_result = compile_expression_list lexpr_list in
        [While (vp_result,lexpr_result)]
  | Ast.SynFor (init,test,step,body) -> 
      let init_result = compile_expression init in
      let test_result = compile_value_producer test in 
      let step_result = compile_expression step in
      let body_result = compile_expression_list body in
        [For (init_result, test_result, step_result, body_result)]
  | Ast.SynDoWhile (lexpr_list, vprod) ->
      let vp_result = compile_value_producer vprod in
      let lexpr_result  = compile_expression_list lexpr_list in
        [While (lexpr_result,vp_result)]



(** Compiles a value *)


and compile_value value = match value with 
    Ast.SynIntValue i -> [Instruction (ConstInt i)]
  | Ast.SynFloatValue f -> [Instruction (ConstFloat f)]
  | Ast.SynBoolValue b -> [Instruction (ConstInt (if b then 1 else 0))]
  | Ast.SynStringValue s -> [Instruction (ConstString s)]



(** Compiles a value produder *)


and compile_value_producer vprod = match vprod with 
    Ast.SynValue value -> compile_value value
  | Ast.SynFunctionCallValue funcall -> compile_function_call funcall
  | Ast.SynVarIdentifier vident -> compile_get_variable_identifier vident
  | Ast.SynBinop (vp1,op,vp2) -> compile_binary_operation vp1 op vp2
  | Ast.SynPrefixUnop (op,vp) -> compile_prefix_unary_operation op vp
  | Ast.SynParenthesized vprod -> compile_value_producer vprod
  | Ast.SynArrayValueProducer vprod_list ->
      let n = List.length vprod_list in
      let values_result = compile_value_producer_list_pushing vprod_list in
        values_result 
        @ [Instruction (Rev n);
           Instruction (MakeBlockStatic (0,n))]
  | Ast.SynListValueProducer lvprod -> compile_list_value_producer lvprod



(** Compiles each value producer in order, pushing the results onto the stack *)


and compile_value_producer_list_pushing vp_list = match vp_list with
    cur::rest -> 
      let cur_result = compile_value_producer cur in
      let rest_result= compile_value_producer_list_pushing rest in
        cur_result
        @ [Instruction Push]
        @ rest_result
  | [] -> []



(** Compiles a list value producer *)


and compile_list_value_producer lvprod = match lvprod with
    Ast.SynListNil _ -> [Instruction Nil]
  | Ast.SynListList vprod_list -> 
      let result = compile_value_producer_list_pushing vprod_list in
        result
        @ [Instruction Nil]
        @ (repeat_list [Instruction Cons] (List.length vprod_list))
  | Ast.SynListCons (vp1,vp2) ->
      let result_car = compile_value_producer vp1 in
      let result_cdr = compile_value_producer vp2 in
        result_car
        @ [Instruction Push]
        @ result_cdr
        @ [Instruction Cons]



(** Compiles a variable assignment *)


and compile_variable_assignment vasgn = match vasgn with 
    Ast.SynVarAssignment (vident,vprod) ->
      let vp_result = compile_value_producer vprod in
      let setvar_result = compile_set_variable_identifier vident in
        vp_result @ setvar_result
  | Ast.SynVarModify (vident, op, vprod) ->
      let vp1, vp2 = Ast.SynVarIdentifier vident, vprod
      in
      let op_result = compile_value_producer 
        (Ast.SynBinop (vp1, op, vp2)) in
      let setvar_result = compile_set_variable_identifier vident in
        op_result @ setvar_result



(** Produces code that sets the variable identifier to the value in the accumulator. Leaves the value in the accumulator *)


and compile_set_variable_identifier vident =
  match vident with
      Ast.SynVarName s -> [SetVar s]
    | Ast.SynStructOrEnumValue (vid1,s2,ftype) ->
        (match !ftype with
             Ast.VIdStructField index ->
               let getstruct = compile_get_variable_identifier vid1 in
                 [Instruction Push]
                 @ getstruct
                 @ [Instruction (SetElemStatic index)] 

(** structs are reference types; we're done at this point *)


           | Ast.VIdEnumValue _ ->
               raise_compile_error ["Unexpected compile error: cannot set an enum value."]
           | Ast.VIdUnspecified ->
               raise_compile_error ["Unexpected compile error: field index ";s2;" was not determined by type checker."])
    | Ast.SynArrayCell (vid,vprod) ->
        let vid_result = compile_get_variable_identifier vid in
        let vprod_result = compile_value_producer vprod in
          [Instruction Push]
          @ vprod_result
          @ [Instruction Push]
          @ vid_result
          @ [Instruction SetElem]



(** Get a variable identifier into the accumulator *)


and compile_get_variable_identifier vident =
  match vident with
      Ast.SynVarName s -> [GetVar s]
    | Ast.SynStructOrEnumValue (vid1,s2,ftype) ->  
        (match !ftype with
             Ast.VIdStructField index ->
               let getstruct = compile_get_variable_identifier vid1 in
                 getstruct @ [Instruction (GetElemStatic index)]
           | Ast.VIdEnumValue index ->
               [Instruction (ConstInt index)]
           | Ast.VIdUnspecified -> 
               raise_compile_error ["Unexpected compile error: field index ";s2;" was not determined by type checker."])
    | Ast.SynArrayCell (vid,vprod) ->
        let vid_result = compile_get_variable_identifier vid in
        let vprod_result = compile_value_producer vprod in
          vprod_result 
          @ [Instruction Push]
          @ vid_result
          @ [Instruction GetElem]



(** Compiles a binary operation expression *)


and compile_binary_operation vp1 op vp2 = 
  let vp1_result = compile_value_producer vp1 in
  let vp2_result = compile_value_producer vp2 in
  let op_command = match op with
      Ast.SynBinopAnd -> [Instruction And]
    | Ast.SynBinopOr -> [Instruction Or]
    | Ast.SynBinopAdd -> [Instruction Add]
    | Ast.SynBinopSub -> [Instruction Sub]
    | Ast.SynBinopMul -> [Instruction Mul]
    | Ast.SynBinopDiv -> [Instruction Div]
    | Ast.SynBinopIDiv ->[Instruction DivI]
    | Ast.SynBinopMod -> [Instruction RemI]
    | Ast.SynBinopConcat -> [Instruction Concat]
    | Ast.SynBinopCons -> [Instruction Cons]
    | Ast.SynCompLt -> [Instruction Lt]
    | Ast.SynCompGt -> [Instruction Gt]
    | Ast.SynCompLte ->[Instruction Lte]
    | Ast.SynCompGte ->[Instruction Gte]
    | Ast.SynCompEq -> [Instruction Eq]
    | Ast.SynCompNeq ->[Instruction Neq]
  in
    vp2_result 
    @ [Instruction Push]
    @ vp1_result
    @ op_command
                         


(** Compiles unary operations *)

   
and compile_prefix_unary_operation op vp =
    let vp_result = compile_value_producer vp in
    let op_command = match op with 
        Ast.SynUnopNot -> [Instruction Push;
                           Instruction (ConstInt 1);
                           Instruction Add;
                           Instruction Push;
                           Instruction (ConstInt 2);
                           Instruction Swap;
                           Instruction RemI]
      | Ast.SynUnopCar -> [Instruction Car]
      | Ast.SynUnopCdr -> [Instruction Cdr]
      | Ast.SynUnopTrunc -> [Instruction Int]
      | Ast.SynUnopDeref -> [Instruction (GetElemStatic 0)]
      | Ast.SynUnopNeg -> [Instruction Neg]
      | Ast.SynUnopNull -> [Instruction IsNull]
    in
      vp_result @ op_command



(** Compiles a function call. Doesn't use tail-recursion -- this would be nice to add in the future. *)


and compile_function_call funcall = match funcall with
    Ast.SynLocalCall (fname,vprod_list,isvar) -> 
      let n_args = List.length vprod_list in
      let args_result = compile_value_producer_list_pushing vprod_list in
      let ret_label = make_label (concat ["returnfrom_local_";fname]) in
      let call_instruction = 
        if !isvar then 
          [GetVar fname;
           Instruction Apply]
        else
          [Instruction (Call fname)]
      in
        [Instruction (Comment (concat ["Local call to ";fname]));
         Instruction (PushSf ret_label);
         BeginScope]
        @ args_result
        @ [Instruction (Rev n_args)]
        @ call_instruction
        @ [EndScopeSf;
           Instruction (Label ret_label)]
  | Ast.SynRemoteCall (fname,vprod_list,ret_var) ->
      let n_args = List.length vprod_list in
      let args_result = compile_value_producer_list_pushing vprod_list in
      let ret_label = make_label (concat ["returnfrom_remote_withresult_";fname]) in
      let done_label = make_label (concat ["done_interpreting_remote_results_";fname]) in
      let vident_set_result = compile_set_variable_identifier ret_var in
        [Instruction (Comment (concat ["Remote call with result to ";fname]));
         Instruction (PushSf ret_label);
         BeginScope]
        @ args_result
        @ [Instruction (Rev n_args);
           Instruction (RPC fname);
           EndScopeSfRPC;
           Instruction (Label ret_label); 

(** Acc || Stack = success || result *)


           Instruction (Jz done_label);
           Instruction Swap] 

(** result || success *)


        @ vident_set_result 

(** ? || success *)


        @ [Instruction Swap; 

(** success || ? *)


           Instruction (Label done_label); 

(** success || ? *)


           Instruction Swap; 

(** ? || success *)


           Instruction (Acc 0); 

(** success || success *)


           Instruction Pop] 

(** success || *)


           
  | Ast.SynRemoteCallNoResult (fname, vprod_list) ->
      let n_args = List.length vprod_list in
      let args_result = compile_value_producer_list_pushing vprod_list in
      let ret_label = make_label (concat ["returnfrom_local_";fname]) in
        [Instruction (Comment (concat ["Remote call with no result to ";fname]));
         Instruction (PushSf ret_label);
         BeginScope]
        @ args_result
        @ [Instruction (Rev n_args);
           Instruction (RPC fname);
           EndScopeSfRPC;
           Instruction (Label ret_label); 

(** success || result *)


           Instruction Swap; 

(** result || success *)


           Instruction (Acc 0); 

(** success || success *)


           Instruction Pop] 

(** success || *)


              


(** Compiles a return statement *)


and compile_return ret = match ret with
    Ast.SynVoidReturn ->
      [Instruction ConstUninit; Instruction Return] 
  | Ast.SynValueReturn vprod ->
      let vp_result = compile_value_producer vprod in
        vp_result@[Instruction Return]



(** Compiles a function definition. Maybe this should have gone to IR first, but it works nicely as is. *)


let rec compile_function_definition name arg_list expr_list = 
  let label = make_label name in
  let args = note_arg_list arg_list 0 in
  let body_code = compile_expression_list expr_list in
    [NoteFunction name;
     Instruction (Annotation (concat ["function ";name;" ";string_of_int (List.length arg_list)])); 

(** annotate *)


     Instruction (Label label);
     BeginScope]
    @ args
    @ body_code
    @ [Instruction (Comment (concat ["end function ";name])); 
       Instruction ConstUninit; 
       Instruction Return;
       EndScope]



(** Compile global expression *)


let rec compile_global_expression expr = match expr with
    Ast.SynFunctionDeclare (dt, name, arglist) -> [NoteFunction name]
  | Ast.SynRemotableFunctionDeclare (dt, name, arglist) -> []
  | Ast.SynFunctionDefine ((dt, name, arglist), expr_list) -> 
        compile_function_definition name arglist expr_list
  | Ast.SynTemplatedDeclare (_,expr) -> compile_global_expression expr
  | Ast.SynTemplatedDefine (_,expr) -> compile_global_expression expr
  | Ast.SynStructDeclare sdecl -> []
  | Ast.SynGlobalEnumDeclare edecl -> []
  | Ast.SynGlobalVarDeclare vdecl -> 
      (match vdecl with
           Ast.SynVarDeclareNoInit (dt,v) -> 
             let dt_result = compile_datatype_init dt in
             let _ = init_buffer := !init_buffer @ dt_result @ [SetVar v] in
               [DeclareGlobal v]
         | _ -> raise_compile_error ["Unexpected compile error: cannot initialize a variable in the global scope."])                          
  | Ast.SynStateMachine _ ->
      raise_compile_error ["Unexpected compile error: source-to-source compilation of state machine failed."]
  | Ast.SynInclude data -> 
      match !data with
          Ast.IncludeAst ast -> compile_global_expression_list ast
        | Ast.IncludeFileName s -> 
            raise_compile_error ["Unexpected compile error: include ";s;" was not typechecked properly."]
            


(** Compile a global expression list *)


and compile_global_expression_list src_prog = 
  match src_prog with
      cur::rest -> 
        let curbuf = compile_global_expression cur in
        let restbuf = compile_global_expression_list rest in
          curbuf@restbuf
    | [] -> []



(** Compile a source program to IR representation, including the init buffer *)


and compile_ir src_prog =
  try
    let _ = init_buffer := [] in
    let init_label = make_label "INIT" in
    let program_result = compile_global_expression_list src_prog in
      [Instruction (Annotation (concat ["init ";init_label]))]
      @ program_result
      @ [Instruction (Annotation (concat ["function ";init_label;" 0"]));
         Instruction (Label (init_label))]
      @ !init_buffer
      @ [Instruction Return]
  with
      Ircompile_error s ->
        let _ = output_string stderr s in
        let _ = output_string stderr "\n" in
        let _ = flush stderr in
          raise (Ircompile_error s)



(** Produces a string representation of the given instruction *)


let string_of_instruction inst = match inst with
    Push -> "Push"
  | Pop -> "Pop"
  | Read -> "Read"
  | Print -> "Print"
  | PrintLn -> "PrintLn"
  | Swap -> "Swap"
  | Stop -> "Stop"
  | Return -> "Return"
  | Cons -> "Cons"
  | Car -> "Car"
  | Cdr -> "Cdr"
  | IsNull -> "IsNull"
  | Nil -> "Nil"
  | GetElem -> "GetElem"
  | SetElem -> "SetElem"
  | And -> "And"
  | Or -> "Or"
  | Add -> "Add"
  | Sub -> "Sub"
  | Mul -> "Mul"
  | Div -> "Div"
  | DivI -> "DivI"
  | RemI -> "RemI"
  | Concat -> "Concat"
  | Neg -> "Neg"
  | Frac -> "Frac"
  | Int -> "Int"
  | Lt -> "Lt"
  | Gt -> "Gt"
  | Lte -> "Lte"
  | Gte -> "Gte"
  | Eq -> "Eq"
  | Neq -> "Neq"
  | ConstUninit -> "Const"
  | ConstInt i -> concat ["Const ";string_of_int i]
  | ConstFloat f -> concat ["Const ";string_of_float f]
  | ConstString s -> concat ["Const \"";s;"\""]
  | PushSf s -> concat ["PushSf ";s]
  | Assign i -> concat ["Assign ";string_of_int i] 
  | Acc i -> concat ["Acc ";string_of_int i]
  | Rev i -> concat ["Rev ";string_of_int i]
  | MakeBlock i -> concat ["MakeBlock ";string_of_int i]
  | AllocFields i -> concat ["-allocfields, ";string_of_int i]
  | GetField i -> concat ["GetField ";string_of_int i]
  | SetField i -> concat ["SetField ";string_of_int i]
  | Call s -> concat ["Call ";s]
  | Jmp s -> concat ["Jmp ";s]
  | Jz s -> concat ["Jz ";s]
  | Jnz s -> concat ["Jnz ";s]
  | RPC s -> concat ["RPC ";s]
  | Label s -> concat [s;": "]
  | Comment s -> concat ["//";s]
  | Annotation s -> concat ["-";s]
  | GetElemStatic i -> concat ["GetElem ";string_of_int i]
  | SetElemStatic i -> concat ["SetElem ";string_of_int i]
  | MakeBlockStatic (i,j) -> concat ["MakeBlockStatic ";string_of_int i;" ";string_of_int j]
  | MakeBlockFilled n -> concat ["MakeFilledBlock ";string_of_int n]
  | Apply -> "Apply"



(** Produces a string representation of the given IR instruction *)


let rec string_of_ir_instruction inst = match inst with
    Instruction i -> string_of_instruction i
  | GetVar s -> concat ["GetVar ";s]
  | SetVar s -> concat ["SetVar ";s]
  | NoteVar (s,i) -> concat ["NoteVar ";s;" ";string_of_int i]
  | If (test,body,elsebody) -> concat 
      ["IF\n";
       string_of_ir_buffer test;
       "\nTHEN\n";
       string_of_ir_buffer body;
       "\nELSE\n";
       string_of_ir_buffer elsebody;
       "\nENDIF\n"]
  | For (init,test,step,body) -> concat
      ["FOR\nINIT\n";
       string_of_ir_buffer init;
       "\nTEST IF TRUE\n";
       string_of_ir_buffer test;
       "\nSTEP\n";
       string_of_ir_buffer step;
       "\nDO\n";
       string_of_ir_buffer body;
       "\nENDFOR\n"]
  | While (test,body) -> concat
      ["WHILE\nTEST IF TRUE\n";
       string_of_ir_buffer test;
       "\nDO\n";
       string_of_ir_buffer body;
       "\nENDWHILE\n"]
  | DoWhile (body,test) -> concat
      ["WHILE\nTEST IF TRUE\n";
       string_of_ir_buffer test;
       "\nDO\n";
       string_of_ir_buffer body;
       "\nENDWHILE\n"]
  | Break -> "Break"
  | Continue -> "Continue"
  | DeclareGlobal s -> concat ["DeclareGlobal ";s]
  | EndScope -> "EndScope"
  | EndScopeSf -> "EndScopeSf"
  | BeginScope -> "BeginScope"
  | EndScopeSfRPC -> "EndScopeSfRPC"
  | NoteFunction s -> concat ["NoteFunction ";s]
  



(** Produces a string representation of the given IR instruction buffer *)


and string_of_ir_buffer buf = match buf with
    cur::rest -> concat [(string_of_ir_instruction cur);"\n";(string_of_ir_buffer rest)]
  | [] -> ""



(** Prints the IR buffer to the given output channel *)


and print_ir_buffer chan buf =
  output_string chan (string_of_ir_buffer buf)