open Std
module CFG = BStd.Graphs.Ir

type memory_offset = int64

type cfa_pos =
    RspOffset of memory_offset
  | RbpOffset of memory_offset
  | CfaLostTrack

type cfa_change = CfaChange of BStd.word * cfa_pos

type cfa_changes_fde = cfa_change list

module StrMap = Map.Make(String)
type cfa_changes = cfa_changes_fde StrMap.t

module TIdMap = Map.Make(BStd.Tid)

exception InvalidSub

let pp_cfa_pos ppx = function
  | RspOffset off -> Format.fprintf ppx "RSP + (%s)@." (Int64.to_string off)
  | RbpOffset off -> Format.fprintf ppx "RBP + (%s)@." (Int64.to_string off)
  | CfaLostTrack -> Format.fprintf ppx "??@."

let pp_cfa_change ppx = function CfaChange(addr, cfa_pos) ->
  Format.fprintf ppx "%a: %a" BStd.Word.pp_hex addr pp_cfa_pos cfa_pos

let pp_cfa_changes_fde ppx = List.iter (pp_cfa_change ppx)

let pp_cfa_changes ppx =
  StrMap.iter (fun fde_name entry ->
      Format.fprintf ppx "%s@\n====@\n@\n%a@." fde_name
        pp_cfa_changes_fde entry)

let pp_option_of sub_pp ppx = function
  | None -> Format.fprintf ppx "None"
  | Some x -> Format.fprintf ppx "Some %a" sub_pp x

let addr_of term =
  (** Get the address of a term *)
  match BStd.Term.get_attr term BStd.address with
  | None -> assert false
  | Some addr -> addr

let interpret_var_expr c_var offset expr = BStd.Bil.(
    let closed_form = BStd.Exp.substitute
      (var c_var)
      (int (BStd.Word.of_int64 offset))
      expr
    in
    let res = BStd.Exp.eval closed_form in
    match res with
    | Imm value ->
      Some (BStd.Word.to_int64_exn @@ BStd.Word.signed value)
    | _ -> None
  )

let process_def def (cur_offset: memory_offset)
  : ((cfa_change * memory_offset) option) =
  let lose_track addr =
    Some (CfaChange(addr, CfaLostTrack), Int64.zero)
  in
  (match Regs.X86_64.of_var (BStd.Def.lhs def) with
   | Some reg when reg = Regs.X86_64.rsp ->
     let exp = BStd.Def.rhs def in
     let free_vars = BStd.Exp.free_vars exp in
     let free_x86_regs = Regs.X86_64.map_varset free_vars in
     (match Regs.DwRegOptSet.cardinal free_x86_regs with
      | 1 ->
        let free_var = Regs.DwRegOptSet.choose free_x86_regs in
        (match free_var with
         | Some dw_var when dw_var = Regs.X86_64.rsp ->
           let bil_var = (match BStd.Var.Set.choose free_vars with
               | None -> assert false
               | Some x -> x) in
           let interpreted = interpret_var_expr bil_var cur_offset exp in
           (match interpreted with
            | None -> lose_track (addr_of def)
            | Some interp_val ->
              let gap = Int64.sub interp_val cur_offset in
              let new_offset = Int64.sub cur_offset gap in
              Some (CfaChange(addr_of def, RspOffset(new_offset)), new_offset)
           )
         | _ -> lose_track (addr_of def)
        )
      | _ -> lose_track @@ addr_of def
     )
   | _ -> None)

let process_jmp jmp (cur_offset: memory_offset)
  : ((cfa_change * memory_offset) option) =
  let gen_change off =
    let new_offset = Int64.add cur_offset (Int64.of_int off) in
    Some (CfaChange(addr_of jmp, RspOffset(new_offset)), new_offset)
  in

  match (BStd.Jmp.kind jmp) with
  | BStd.Call call -> gen_change (-8)
  | BStd.Ret ret -> gen_change (8)
  | _ -> None

let sym_of_blk blk : cfa_changes_fde =
  (** Extracts the symbolic CFA changes of a block. These changes assume that
      at the beginning of the block, CFA = RspOffset(0) and will be offset
      after *)

  let apply_offset (accu, cur_offset) = function
    | None -> (accu, cur_offset)
    | Some (change, n_offset) -> (change :: accu, n_offset)
  in

  let fold_elt (accu, cur_offset) elt = match elt with
    | `Def(def) ->
      apply_offset (accu, cur_offset) @@ process_def def cur_offset
    | `Jmp(jmp) ->
      apply_offset (accu, cur_offset) @@ process_jmp jmp cur_offset
    | _ -> (accu, cur_offset)
  in

  let elts_seq = BStd.Blk.elts blk in
  let out, end_offset = BStd.Seq.fold elts_seq
    ~init:([], Int64.zero)
    ~f:fold_elt in
  out

let end_offset (changelist: cfa_changes_fde): memory_offset option =
  match changelist with
  | CfaChange(_, RspOffset(x)) :: _ -> Some x
  | _ -> None

exception Inconsistent of BStd.tid

let rec dfs_propagate changemap propagated parent_val node graph =
  let c_tid = BStd.Term.tid @@ CFG.Node.label node in
  match TIdMap.find_opt c_tid propagated with
  | Some x ->
    if x = parent_val then
      (* Already propagated and consistent, all fine *)
      propagated
    else
      (* Already propagated with a different value: inconsistency *)
      raise (Inconsistent c_tid)
  | None ->
    let n_propagated = TIdMap.add c_tid parent_val propagated in
    let outwards = CFG.Node.outputs node graph in
    let self_entry = TIdMap.find c_tid changemap in
    let offset = (match end_offset self_entry with
          Some x -> x
        | None -> Int64.zero) in
    let cur_val = Int64.add parent_val offset in
    BStd.Seq.fold outwards
      ~init:n_propagated
      ~f:(fun accu edge ->
          dfs_propagate changemap accu cur_val (CFG.Edge.dst edge) graph)

let get_entry_blk graph =
  let entry =
    BStd.Seq.find (CFG.nodes graph)
      (fun node -> BStd.Seq.is_empty @@ CFG.Node.inputs node graph)
  in match entry with
  | None -> assert false
  | Some x -> x

let term_addr term =
  BStd.Term.get_attr term BStd.address

let same_keys map1 map2 =
  let exists_in_other other key _ =
    TIdMap.mem key other in

  TIdMap.for_all (exists_in_other map2) map1
  && TIdMap.for_all (exists_in_other map1) map2

let of_sub sub : cfa_changes_fde =
  (** Extracts the `cfa_changes_fde` of a subroutine *)

  Format.eprintf "Sub %s...@." @@ BStd.Sub.name sub ;

  let initial_cfa_rsp_offset = Int64.of_int 8 in
  let store_sym accu blk =
    let blk = CFG.Node.label blk in
    let res = sym_of_blk blk in
    TIdMap.add (BStd.Term.tid blk) res accu
  in

  let node_addr nd = term_addr @@ CFG.Node.label nd in

  let merge_corrected blk_tid changes offset = match (changes, offset) with
    | Some changes, Some offset ->
      Some (
        List.map (fun (CfaChange(addr, pos)) -> match pos with
              RspOffset(off) -> CfaChange(addr,
                                          RspOffset(Int64.add off offset))
            | RbpOffset(off) -> CfaChange(addr, RbpOffset(off))
            | CfaLostTrack -> CfaChange(addr, CfaLostTrack)
          )
          changes
      )
    | _ -> None
  in

  let cfg = BStd.Sub.to_cfg sub in

  let tid_match = BStd.Seq.fold (CFG.nodes cfg)
      ~init:TIdMap.empty
      ~f:(fun accu node ->
          let tid = BStd.Term.tid @@ CFG.Node.label node in
          TIdMap.add tid node accu)
  in

  let blk_sym = BStd.Seq.fold
      ~init:TIdMap.empty
      ~f:store_sym
      @@ CFG.nodes cfg
  in

  let entry_blk = get_entry_blk cfg in
  let offset_map = dfs_propagate
      blk_sym (TIdMap.empty) initial_cfa_rsp_offset entry_blk cfg in

  let corrected = TIdMap.merge merge_corrected blk_sym offset_map in

  let is_connex = same_keys tid_match corrected in
  if not is_connex then
    raise InvalidSub ;

  let tid_list = TIdMap.bindings tid_match in
  let sorted_blk = List.sort (fun (tid1, bl1) (tid2, bl2) ->
      let res = match (node_addr bl1, node_addr bl2) with
        | Some addr1, Some addr2 -> compare addr1 addr2
        | Some _, None -> 1
        | None, Some _ -> -1
        | None, None -> compare tid1 tid2
      in
      -res)
      tid_list
  in

  let out = List.fold_left
      (fun accu blk ->
         let changes = TIdMap.find blk corrected in
         List.fold_left (fun accu chg -> chg::accu)
           accu
           changes
      )
      []
      (List.map (fun (x, y) -> x) sorted_blk) in

  let sub_addr = (match term_addr sub with
      | Some x -> x
      | None -> assert false) in
  let init = [
    CfaChange(sub_addr, RspOffset(initial_cfa_rsp_offset))
  ] in
  init @ out

let cleanup_fde (fde_changes: cfa_changes_fde) : cfa_changes_fde =
  (** Cleanup the result of `of_sub`.

      Merges entries at the same address, propagates track lost *)

  let fold_one (accu, (last_commit:cfa_pos option), in_flight, lost_track) = function
    | CfaChange(addr, cfa_change) as cur_change -> (
        match lost_track, in_flight, cfa_change with
        | true, _, _ ->
          (* Already lost track: give up *)
          (accu, last_commit, None, lost_track)
        | false, _, CfaLostTrack ->
          (* Just lost track: give up the operation on flight as well *)
          (cur_change :: accu, None, None, true)
        | _, Some CfaChange(flight_addr, flight_chg), _
          when flight_addr = addr ->
          (* On flight address matches current address: continue flying *)
            accu, last_commit, Some cur_change, false
        | _, Some CfaChange(_, in_flight_inner_pos), _
          when last_commit = Some in_flight_inner_pos ->
          (* Doesn't match anymore, but there was some operation in flight,
             which has the same result as what was last committed. Discard. *)
          (accu, last_commit, Some cur_change, false)
        | _, Some (CfaChange(_, in_flight_inner_pos) as in_flight_inner), _ ->
          (* Doesn't match anymore, but there was some operation in flight:
             commit it, put the new one in flight *)
          (in_flight_inner :: accu, Some in_flight_inner_pos,
           Some cur_change, false)
        | _, None, _ ->
          (* No operation in flight: put the new one in flight *)
          (accu, last_commit, Some cur_change, false)
    )
  in

  let extract_end_value (accu, _, in_flight, lost_track) =
    List.rev @@ match lost_track, in_flight with
    | true, _ -> accu
    | false, None -> accu
    | false, Some x -> x :: accu
  in

  extract_end_value
    @@ List.fold_left fold_one ([], None, None, false) fde_changes


let of_prog prog : cfa_changes =
  (** Extracts the `cfa_changes` of a program *)
  let fold_step accu sub =
    (try
       let res = (*cleanup_fde @@ *)of_sub sub in
       StrMap.add (BStd.Sub.name sub) res accu
     with InvalidSub -> accu)
  in
  let subroutines = BStd.Term.enum BStd.sub_t prog in
  BStd.Seq.fold subroutines
    ~init:StrMap.empty
    ~f:fold_step

let of_proj proj : cfa_changes =
  (** Extracts the `cfa_changes` of a project *)
  let prog = BStd.Project.program proj in
  of_prog prog