open Iabsyn
open Types

let set_to_pred (s : string) = "is_" ^ s
let primed (s : string) = s ^ "'"


let string_of_value x =
  match x with
  | Int i -> Int32.to_string i
  | Bool true -> "true"
  | Bool false -> "false"
  | Float f -> string_of_float f
  | String s -> "\"" ^ (String.escaped s) ^ "\""
  | Byte c 
  | Char c -> "'" ^ (String.make 1 c) ^ "'"
  | Obj 0 -> "null"
  | Array _ -> "<array>"
  | Obj _ -> "<obj>"

let expand_new_statement (istr:string) (v:string) : string =
  istr ^ v ^" = fresh;\n" ^
  istr ^ "fresh = fresh.nextExternal;\n" ^
  istr ^ "fresh.prevExternal=null;\n" ^
  istr ^ v ^ ".nextExternal=null;\n" ^
  istr ^ v ^ ".prevExternal=null"

let string_of_expr tenv (e:Iabsyn.expr) =  
  let rec string_of_lvalue tenv (lv:Iabsyn.lvalue) =
    match lv with
      LocalLvalue v -> v
    | RefLvalue v -> v
    | FieldLvalue (e, f) -> (string_of_expr0 e) ^ "." ^ (Id.name_of_field f)
    | ArrayLvalue (e, i) -> (string_of_expr0 e) ^ "[" ^ (string_of_expr0 i) ^ "]"
  and string_of_expr0 (e:Iabsyn.expr) =  
  match e with
    LiteralExpr v -> string_of_value v
  | VarExpr v -> string_of_lvalue tenv v
  | FieldAccessExpr (e, f) -> (string_of_expr0 e) ^ "." ^ (Id.name_of_field f)
  | ArrayAccessExpr (e, i) -> 
      (string_of_expr0 e) ^ "[" ^ (string_of_expr0 i) ^ "]"
  | ArrayLengthExpr e -> 
      (string_of_expr0 e) ^ ".length"
  | NewExpr f -> "new "^f
  | NewArrayExpr (t, e, d) -> 
      let rec bprs i = if (i = 0) then "" else ("[]"^bprs (i-1)) in
      "new "^(string_of_value_type t)^
      String.concat "" (List.map (fun x -> "["^(string_of_expr0 x)^"]") e)^
      bprs d
  | InvokeExpr (p, []) -> (Id.module_of_proc p) ^ "." ^ 
      (Id.name_of_proc p) ^ "()"
  | InvokeExpr (p, args) -> (Id.module_of_proc p) ^ "." ^ 
      (Id.name_of_proc p) ^ "(" ^
      (List.fold_right (fun x y -> ((string_of_expr0 x) ^ ", " ^ y)) 
         (List.tl args) 
         (string_of_expr0 (List.hd args)) ^ ")")
  | AssignExpr (l, r) -> 
      let lt = match l with
        LocalLvalue v when List.mem_assoc v tenv -> List.assoc v tenv
      | RefLvalue v when List.mem_assoc v tenv -> List.assoc v tenv
      | _ -> TVoid in
      let rt = match r with
        VarExpr (LocalLvalue v) when List.mem_assoc v tenv -> List.assoc v tenv
      | VarExpr (RefLvalue v) when List.mem_assoc v tenv -> List.assoc v tenv
      | LiteralExpr (Obj _) -> TObj "null"
      | LiteralExpr (Bool _) -> TBool
      | _ -> TVoid in
      (match (lt, rt) with
      | (TObj _, _) | (_, TObj _) ->
          (let lstr = string_of_lvalue tenv l in
          match r with
          | NewExpr fmt ->
              expand_new_statement "    " lstr
          | _ -> lstr ^ " = " ^ (string_of_expr0 r))
      | (TBool, _) ->
          (string_of_lvalue tenv l) ^ " = " ^ (string_of_expr0 r)
      | (_, TBool) ->
          (string_of_lvalue tenv l) ^ " = " ^ (string_of_expr0 r)
      | _ -> "/* elided: "^(string_of_lvalue tenv l) ^ 
          " = " ^ (string_of_expr0 r)^" */")
  | PreIncExpr v -> "++" ^ (string_of_lvalue tenv v)
  | PostIncExpr v -> (string_of_lvalue tenv v) ^ "++"
  | PreDecExpr v -> "--" ^ (string_of_lvalue tenv v)
  | PostDecExpr v -> (string_of_lvalue tenv v) ^ "--"
  | NegExpr e -> "-" ^ (string_of_expr0 e)
  | PlusExpr (l, r) -> (string_of_expr0 l) ^ " + " ^ (string_of_expr0 r)
  | MinusExpr (l, r) -> (string_of_expr0 l) ^ " - " ^ (string_of_expr0 r)
  | MultExpr (l, r) -> (string_of_expr0 l) ^ " * " ^ (string_of_expr0 r)
  | DivExpr (l, r) -> (string_of_expr0 l) ^ " / " ^ (string_of_expr0 r)
  | ModExpr (l, r) -> (string_of_expr0 l) ^ " % " ^ (string_of_expr0 r)
  | AndExpr (l, r) -> (string_of_expr0 l) ^ " & " ^ (string_of_expr0 r)
  | OrExpr (l, r) -> (string_of_expr0 l) ^ " | " ^ (string_of_expr0 r)
  | NotExpr e -> " ! " ^ (string_of_expr0 e)
  | EqExpr (l, r) -> (string_of_expr0 l) ^ " = " ^ (string_of_expr0 r)
  | NeqExpr (l, r) -> (string_of_expr0 l) ^ " != " ^ (string_of_expr0 r)
  | LtExpr (l, r) -> "?"
  | GtExpr (l, r) -> "?"
  | LteqExpr (l, r) -> "?"
  | GteqExpr (l, r) -> "?"
  | BitAndExpr (l, r) -> (string_of_expr0 l) ^ " && " ^ (string_of_expr0 r)
  | BitOrExpr (l, r) -> (string_of_expr0 l) ^ " || " ^ (string_of_expr0 r)
  | BitXorExpr (l, r) -> (string_of_expr0 l) ^ " ^ " ^ (string_of_expr0 r)
  | ShiftLeftExpr (l, r) -> (string_of_expr0 l) ^ " << " ^ (string_of_expr0 r)
  | SignedShiftRightExpr (l, r) -> (string_of_expr0 l) ^ " >> " ^ (string_of_expr0 r)
  | UnsignedShiftRightExpr (l, r) -> 
      (string_of_expr0 l) ^ " >>> " ^ (string_of_expr0 r) in
  string_of_expr0 e

let expand_pragma (istr : string) (pragma : string) : string =
  let inserted_pragma v = 
    istr ^ "nxt77 = "^v^".nextExternal;\n" ^
    istr ^ "prv77 = "^v^".prevExternal;\n" ^
    istr ^ "if (prv77!=null) {\n" ^
    istr ^ "  prv77.nextExternal = nxt77;\n" ^
    istr ^ "}\n" ^
    istr ^ "if (externals="^v^") {\n" ^
    istr ^ "  externals = nxt77;\n" ^
    istr ^ "}\n" ^
    istr ^ "if (nxt77!=null) {\n" ^
    istr ^ "  nxt77.prevExternal = prv77;\n" ^
    istr ^ "}\n" ^
    istr ^ v ^ ".prevExternal = null;\n" ^
    istr ^ v ^ ".nextExternal = null;\n" ^
    istr ^ v ^ ".isExternal = false;\n" in
  let removed_pragma v =
    istr ^ v ^ ".nextExternal = externals;\n" ^
    istr ^ "if (externals!=null) {\n" ^
    istr ^ "externals.prevExternal = "^v^";\n" ^
    istr ^ "}\n" ^
    istr ^ "externals = " ^ v ^ ";\n" ^
    istr ^ v ^ ".isExternal = true;\n" in
  if (String.sub pragma 0 9)="inserted " then
    inserted_pragma (String.sub pragma 9 (String.length pragma - 9))
  else if (String.sub pragma 0 8)="removed " then
    removed_pragma (String.sub pragma 8 (String.length pragma - 8))
  else "ERROR"

let makeLoopInv p ms (contentSet,fmt,retparam,singletons,obj_locals) givenInv =
  let global_sets = Spec.collect_sets_types ms in
  let setnames = List.map Util.replace_dot_with_uscore
      (List.map fst global_sets) in
  let wr_quant s = "existset " ^ primed s ^ " of " ^ fmt ^ ":" in
  let wr_setdef s = set_to_pred s  ^ "(" ^ primed s ^ ") &\n" in
  let invs = try Util.split_by "<&>" givenInv with Not_found -> [] in
  let (pale_part,ba_part) = match invs with
  | [x;y] -> (x,y)
  | _ -> 
      Util.error "syntax error trying to separate pale and boolean algebra parts of loop invariant";
      ("true","true") in     
  let get_param_name (v,t) = v in   
  let str_external_general_case v =
    "(!"^v^ " in " ^ 
    contentSet ^
    " => "^v^".nextExternal!=null & "^v^".prevExternal!=null) & \n" ^
    "!"^v^".isFresh" in
  let parsed_ba_part = Spec.parse_formula
      ms.Sabsyn.module_name { Spec.s = []; Spec.b = [] } ba_part in
  (String.concat "" (List.map wr_quant setnames @
                     List.map wr_setdef setnames)) ^ 
  "(" ^ pale_part  ^ ") & (" ^ 
  Paleba.pale_string fmt (retparam @ singletons @ obj_locals) parsed_ba_part ^ ") & externalInv() & " ^
  "fresh != null & fresh.nextExternal!=null" ^
  (if (p.Sabsyn.formals != []) then
    let v = get_param_name (List.hd p.Sabsyn.formals) in
    (" &"^(str_external_general_case v))
  else "")


(* Note that we can now somewhat simplify this, because we can assume
 * that all locals have been pulled up. *)
let rec string_of_stmt (ms:Sabsyn.spec_module) loopInvData tenv (ind:int) (p: Sabsyn.proc_spec) (s0:string Iabsyn.stmt) =
  let istr = (String.make ind ' ') in
  let istr' = (String.make (ind-2) ' ') in
  let string_of_local_decl_stmts (ind:int) (s:string Iabsyn.stmt) =
    match s with
    | Iabsyn.LocalDeclStmt (v, vt, e) ->
        istr ^ 
        (match vt with
        | TBool -> "bool " ^ v ^ ";\n"
        | TObj t -> "pointer " ^ v ^ ": " ^ t ^ ";\n"
        | _ -> "")    
    | _ -> ""
  in
  let rec string_of_nondecl_stmts (ind:int) (s:string Iabsyn.stmt) =
    let pale_string_of_expr e =
(* so we could have integer compares, but don't handle them yet *)
      (match e with
      | LtExpr _ | LteqExpr _ | GteqExpr _ | GtExpr _ -> "?"
      | _ -> string_of_expr tenv e) in
    match s with
      Iabsyn.EmptyStmt -> istr ^ ";\n"
    | Iabsyn.ChoiceStmt _ -> failwith "hey, don't paleprint choice stmts!"
    | Iabsyn.HavocStmt _ -> failwith "hey, don't paleprint havoc stmts!"
    | Iabsyn.CompoundStmt cs -> 
        (match cs with 
        | [Iabsyn.CompoundStmt cs'] -> string_of_stmt ms loopInvData tenv ind p (Iabsyn.CompoundStmt cs')
        | _ ->
            istr' ^ "{\n" ^
            (String.concat "" (List.map (string_of_local_decl_stmts ind) cs)) ^
            (String.concat "" (List.map (string_of_nondecl_stmts ind) cs)) ^
            istr' ^ "}\n")
    | Iabsyn.ExprStmt e -> 
        istr ^ (string_of_expr tenv e) ^ ";\n"
    | Iabsyn.ReturnStmt None ->
        istr ^ "return;\n"
    | Iabsyn.ReturnStmt Some r ->        
        (* istr ^ "return "^(string_of_expr tenv r)^";\n" *)
        (match p.Sabsyn.ret_val with
        | Some (v,_) -> istr ^ v ^ " = "^(string_of_expr tenv r)^";\n"

        | _ -> "")
    | Iabsyn.WhileStmt (Some a, e, ws) -> 
        let es = pale_string_of_expr e in
        istr ^ "while [" ^ 
          makeLoopInv p ms loopInvData (Util.trim_quotes a)
                    ^"]\n" ^
        istr ^ "("^es^")\n" ^
        string_of_stmt ms loopInvData tenv (ind+2) p ws
    | Iabsyn.WhileStmt (None, e, ws) -> 
        let (fn, p) = Ast.lookup_expr_pos e in
        Util.err (fn ^ " " ^ p) "PALE needs your loop invariants";
        Util.print_errors ();
        istr ^ ";\n"
    | Iabsyn.AssertStmt (_, a) -> 
        istr ^ "split [" ^ a ^ "];\n"
    | Iabsyn.AssumeStmt (_, a) -> 
        failwith "don't know what to print for assume stmts";
        istr ^ "split [" ^ a ^ "];\n"
    | Iabsyn.PragmaStmt ps ->
        expand_pragma istr ps
    | Iabsyn.IfStmt (e, 
                     (Iabsyn.CompoundStmt _ as t), 
                     (Iabsyn.CompoundStmt _ as f)) ->
        let es = pale_string_of_expr e in
        istr ^ "if ("^es^")\n" ^
        string_of_stmt ms loopInvData tenv (ind+2) p t ^
        istr ^ "else\n" ^
        string_of_stmt ms loopInvData tenv (ind+2) p f 
    | Iabsyn.IfStmt (e, (Iabsyn.CompoundStmt _ as t), Iabsyn.EmptyStmt) ->
        let es = pale_string_of_expr e in
        istr ^ "if ("^es^")\n" ^
        string_of_stmt ms loopInvData tenv (ind+2) p t
    | Iabsyn.IfStmt (e, t, Iabsyn.EmptyStmt) -> 
        string_of_stmt ms loopInvData tenv (ind+2) p
          (Iabsyn.IfStmt (e, 
                          (Itrans.flatten_compound_stmt 
                             (Iabsyn.CompoundStmt [t])), 
                          Iabsyn.EmptyStmt))
    | Iabsyn.IfStmt (e, t, f) -> 
        string_of_stmt ms loopInvData tenv (ind+2) p
          (Iabsyn.IfStmt (e, 
                          Itrans.flatten_compound_stmt 
                            (Iabsyn.CompoundStmt [t]), 
                          Itrans.flatten_compound_stmt
                            (Iabsyn.CompoundStmt [f])))
    | Iabsyn.LocalDeclStmt (v, vt, e) -> 
        string_of_nondecl_stmts ind 
          (Iabsyn.ExprStmt (Iabsyn.AssignExpr (Iabsyn.LocalLvalue v, e))) in
  ((string_of_local_decl_stmts ind s0) ^ (string_of_nondecl_stmts ind s0))

(* assume that locals have been pulled to top level *)
let fetch_lv_types p =
  let rec fetch_lv_types_stmt s = 
    match s with
    | CompoundStmt cs -> 
        List.concat (List.map fetch_lv_types_stmt cs)
    | LocalDeclStmt (v, t, e) -> 
        [(v, t)]
    | _ -> [] in
  List.concat [(match p.ret_val with None -> []
                                   | Some x -> [x]) @ p.formals @ 
               fetch_lv_types_stmt p.proc_body]
    
let print (file : out_channel)
    ((ma : string Aabsyn.abst_body),
     (ms : Sabsyn.spec_module),
     (mi : string Iabsyn.impl_module))
    loopInvData
    (ps : Sabsyn.proc_spec) (pi : string Iabsyn.proc_def) : unit = 
  let tenv = (fetch_lv_types pi @ mi.references) in
  let b = string_of_stmt ms loopInvData tenv 2 ps pi.Iabsyn.proc_body in
  let b' = Str.string_after b ((String.index b '{')+2) in
  let b'' = Str.string_before b' (String.rindex b' '}') in
  output_string file b''