let error_prefix = "ASMCosts"
let warning s = prerr_endline (error_prefix ^ s)


type instruction_nature =
  | Goto of BitVectors.word | Branch of BitVectors.word
  | Direct_fun_call of BitVectors.word | Return
  | Other

let inst_infos mem pc =
  let (inst, next_pc, inst_cost) = ASMInterpret.fetch mem pc in
  let (nature, next_pcs) = match inst with
    | `LCALL (`ADDR16 addr16) -> (Direct_fun_call addr16, [next_pc])
    | `ACALL (`ADDR11 addr11) ->
      (Direct_fun_call (Physical.addr16_of_addr11 pc addr11), [next_pc])
    | `LJMP (`ADDR16 addr16) -> (Goto addr16, [addr16])
    | `AJMP (`ADDR11 addr11) ->
      let addr = Physical.addr16_of_addr11 pc addr11 in
      (Goto addr, [addr])
    | `SJMP (`REL addr) ->
      let _, addr =
	BitVectors.half_add next_pc (BitVectors.sign_extension addr) in
      (Goto addr, [addr])
    | `JMP idptr ->
      (Other, [next_pc]) (* Indirect jump; precondition: every possible
                            destination should start with its own label *)
    | `JC addr
    | `JNC addr
    | `JB (_,addr)
    | `JNB (_,addr)
    | `JBC (_,addr)
    | `JZ addr
    | `JNZ addr
    | `CJNE (_,addr)
    | `DJNZ (_,addr) ->
      let `REL addr = addr in
      let _, addr =
	BitVectors.half_add next_pc (BitVectors.sign_extension addr) in
      (Branch addr, [next_pc ; addr])
    | `RET | `RETI -> (Return, [])
    | _ -> (Other, [next_pc]) in
  (nature, next_pc, next_pcs, inst_cost)


let rec compare = function
  | [] -> assert false (* do not use an this argument *)
  | [(_, cost)] -> cost
  | (pc1, cost1) :: (pc2, cost2) :: l when cost1 <> cost2 ->
    warning
      (Printf.sprintf
	 "Warning: branching to %s has cost %d, branching to %s has cost %d"
	 (string_of_int (BitVectors.int_of_vect pc1)) cost1
	 (string_of_int (BitVectors.int_of_vect pc2)) cost2) ;
    max cost1 (compare ((pc2, cost2) :: l))
  | _ :: l -> compare l

let rec block_costl mem costs = function
  | [] -> 0
  | [pc] when BitVectors.WordMap.mem pc costs -> 0
  | [pc] -> block_cost mem costs pc
  | next_pcs ->
    compare (List.map (fun pc -> (pc, block_costl mem costs [pc])) next_pcs)

and block_cost mem costs pc =
  let (_, _, next_pcs, cost) = inst_infos mem pc in
  cost + (block_costl mem costs next_pcs)


let traverse_code mem p =
  let rec aux pc code =
    let (_, newpc, _, _) = inst_infos mem pc in
    match code with
      | [] -> CostLabel.Map.empty
      | _::tl when BitVectors.WordMap.mem pc p.ASM.cost_labels ->
	let lbl = BitVectors.WordMap.find pc p.ASM.cost_labels in
	let cost = block_cost mem p.ASM.cost_labels pc in
	let costs_mapping = aux newpc tl in
	CostLabel.Map.add lbl cost costs_mapping
      | _::tl -> aux newpc tl
  in
  aux (BitVectors.zero `Sixteen) p.ASM.code


let first_cost_label mem costs =
  let rec aux oldpc =
    if BitVectors.WordMap.mem oldpc costs then
      (BitVectors.WordMap.find oldpc costs, 0)
    else
      let (nature, pc, _, inst_cost) = inst_infos mem oldpc in
      match nature with
	| Direct_fun_call pc ->
	  let (lbl, cost) = aux pc in
	  (lbl, inst_cost + cost)
	| Return
	| Goto _
	| Branch _ ->
	  assert false (* no such instructions before calling main *)
	| Other ->
	  let (lbl, cost) = aux pc in
	  (lbl, inst_cost + cost)
  in
  aux (BitVectors.zero `Sixteen)


let initialize_cost mem costs costs_mapping =
  let (lbl, cost) = first_cost_label mem costs in
  let old_cost =
    if CostLabel.Map.mem lbl costs_mapping then
      CostLabel.Map.find lbl costs_mapping
    else 0 in
  let new_cost = old_cost + cost in
  CostLabel.Map.add lbl new_cost costs_mapping


let compute p =
  let mem = ASMInterpret.load_code_memory p.ASM.code in
  let costs_mapping = traverse_code mem p in
  if p.ASM.has_main then initialize_cost mem p.ASM.cost_labels costs_mapping
  else costs_mapping