include "ASM/ASM.ma".
include "ASM/Arithmetic.ma".
include "ASM/Fetch.ma".
include "ASM/Interpret.ma".
include "common/StructuredTraces.ma".

definition current_instruction0 ≝
  λmem,pc. \fst (\fst (fetch … mem pc)).

definition current_instruction ≝
 λs:Status. current_instruction0 (code_memory … s) (program_counter … s).

definition ASM_classify0: instruction → status_class ≝
 λi.
  match i with
   [ RealInstruction pre ⇒
     match pre with
      [ RET ⇒ cl_return
      | JZ _ ⇒ cl_jump
      | JNZ _ ⇒ cl_jump
      | JC _ ⇒ cl_jump
      | JNC _ ⇒ cl_jump
      | JB _ _ ⇒ cl_jump
      | JNB _ _ ⇒ cl_jump
      | JBC _ _ ⇒ cl_jump
      | CJNE _ _ ⇒ cl_jump
      | DJNZ _ _ ⇒ cl_jump
      | _ ⇒ cl_other
      ]
   | ACALL _ ⇒ cl_call
   | LCALL _ ⇒ cl_call
   | JMP _ ⇒ cl_call
   | AJMP _ ⇒ cl_jump
   | LJMP _ ⇒ cl_jump
   | SJMP _ ⇒ cl_jump
   | _ ⇒ cl_other
   ].

definition ASM_classify: Status → status_class ≝
 λs.ASM_classify0 (current_instruction s).

definition current_instruction_is_labelled ≝
  λcost_labels: BitVectorTrie costlabel 16.
  λs: Status.
  let pc ≝ program_counter … s in
    match lookup_opt … pc cost_labels with
    [ None ⇒ False
    | _    ⇒ True
    ].

definition label_of_current_instruction:
 BitVectorTrie costlabel 16 → Status → list costlabel
 ≝
  λcost_labels,s.
  let pc ≝ program_counter … s in
    match lookup_opt … pc cost_labels with
    [ None ⇒ []
    | Some l ⇒ [l]
    ].

definition next_instruction_properly_relates_program_counters ≝
  λbefore: Status.
  λafter : Status.
  let size ≝ current_instruction_cost before in
  let pc_before ≝ program_counter … before in
  let pc_after ≝ program_counter … after in
  let sum ≝ \snd (half_add … pc_before (bitvector_of_nat … size)) in
    sum = pc_after.

definition ASM_abstract_status: BitVectorTrie costlabel 16 → abstract_status ≝
 λcost_labels.
  mk_abstract_status
   Status
   (λs,s'. (execute_1 s) = s')
   (λs,class. ASM_classify s = class)
   (current_instruction_is_labelled cost_labels)
   next_instruction_properly_relates_program_counters.

(* To be moved in ASM/arithmetic.ma *)
definition addr16_of_addr11: Word → Word11 → Word ≝
  λpc: Word.
  λa: Word11.
  let 〈pc_upper, ignore〉 ≝ split … 8 8 pc in
  let 〈n1, n2〉 ≝ split … 4 4 pc_upper in
  let 〈b123, b〉 ≝ split … 3 8 a in
  let b1 ≝ get_index_v … b123 0 ? in
  let b2 ≝ get_index_v … b123 1 ? in
  let b3 ≝ get_index_v … b123 2 ? in
  let p5 ≝ get_index_v … n2 0 ? in
    (n1 @@ [[ p5; b1; b2; b3 ]]) @@ b.
  //
qed.
  
axiom minus_m_n_le_minus_m_So_to_Sm_minus_n_le_minus_m_o:
  ∀m, n, o: nat.
    m - n ≤ m - S o → S m - n ≤ m - o.

lemma strengthened_invariant:
  ∀code_memory: BitVectorTrie Byte 16.
  ∀total_program_size: nat.
  ∀new_program_counter: Word.
  ∀program_counter: Word.
    total_program_size ≤ code_memory_size →
      let end_instr ≝ \fst (\fst (fetch … code_memory (bitvector_of_nat … total_program_size))) in
        Or (ASM_classify0 end_instr = cl_jump) (ASM_classify0 end_instr = cl_return) →
          new_program_counter = \snd (\fst (fetch … code_memory program_counter)) →
            nat_of_bitvector … program_counter ≤ total_program_size →
              nat_of_bitvector … new_program_counter ≤ total_program_size.
  #code_memory #total_program_size #new_program_counter #program_counter
  #relation_total_program_code_memory_size #end_instruction_is_ok
  #new_program_counter_from_fetch #program_counter_within_program
  >new_program_counter_from_fetch
  lapply (sig2 ? ? (fetch code_memory program_counter)) #assm
    
        
(* XXX: use memoisation here in the future *)
let rec block_cost
  (code_memory: BitVectorTrie Byte 16) (cost_labels: BitVectorTrie costlabel 16)
    (program_counter: Word) (program_size: nat) (total_program_size: nat)
      (size_invariant: total_program_size ≤ code_memory_size)
        (pc_invariant: nat_of_bitvector … program_counter < total_program_size)
          (final_instr_invariant:
            let instr ≝ \fst (\fst (fetch … code_memory (bitvector_of_nat … total_program_size))) in
              Or (ASM_classify0 instr = cl_jump) (ASM_classify0 instr = cl_return))
                on program_size: total_program_size - nat_of_bitvector … program_counter < program_size → nat ≝
  match program_size return λprogram_size: nat. total_program_size - nat_of_bitvector … program_counter < program_size → nat with
  [ O ⇒ λabsurd. ⊥
  | S program_size ⇒ λrecursive_case.
    let ticks ≝ \snd (fetch … code_memory program_counter) in
    let instr ≝ \fst (\fst (fetch … code_memory program_counter)) in
    let newpc ≝ \snd (\fst (fetch … code_memory program_counter)) in
    match lookup_opt … newpc cost_labels return λx: option costlabel. nat with
    [ None ⇒
        let classify ≝ ASM_classify0 instr in
        match classify return λx. ∀classify_refl: x = classify. nat with
        [ cl_jump ⇒ λclassify_refl. ticks
        | cl_call ⇒ λclassify_refl. ticks + (block_cost code_memory cost_labels newpc program_size total_program_size size_invariant ? final_instr_invariant ?)
        | cl_return ⇒ λclassify_refl. ticks
        | cl_other ⇒ λclassify_refl. ticks + (block_cost code_memory cost_labels newpc program_size total_program_size size_invariant ? final_instr_invariant ?)
        ] (refl … classify)
    | Some _ ⇒ ticks
    ]
  ].
  [1:
    cases(lt_to_not_zero … absurd)
  |2,4:
    cut(nat_of_bitvector … program_counter < nat_of_bitvector … newpc)
    [1,3:
      @fetch_pc_lt_new_pc (* XXX: now @sig2 from fetch's russell type *)
      lapply(transitive_lt
        (nat_of_bitvector 16 pc)
        initial_program_size
        (size Byte 16 mem) ? ?)
      [1,2,4,5:
        assumption
      |3,6:
        #assm assumption
      ]
    |2,4:
      change with (
        S (nat_of_bitvector … 16 pc) ≤ nat_of_bitvector … newpc
      ) in ⊢ (% → ?);
      #first_le_assm
      normalize in recursive_case;
      change with (
        S (initial_program_size - nat_of_bitvector … newpc) ≤ program_size)
      lapply (le_S_S_to_le … recursive_case)
      change with (
        initial_program_size - (nat_of_bitvector … pc) ≤ program_size
      ) in ⊢ (% → ?);
      #second_le_assm
      @(transitive_le
        (S (initial_program_size-nat_of_bitvector 16 newpc))
        (initial_program_size-nat_of_bitvector_aux 16 O pc)
        program_size ? second_le_assm)
      change with (
        initial_program_size - nat_of_bitvector … pc
      ) in ⊢ (??%);
      lapply (minus_Sn_m (nat_of_bitvector … newpc) initial_program_size)
      #assm <assm in ⊢ (?%?);
      [1,3:
        @minus_m_n_le_minus_m_So_to_Sm_minus_n_le_minus_m_o
        @monotonic_le_minus_r
        assumption
      |2,4:
        @(strengthened_invariant mem initial_program_size newpc pc size_invariant final_instr_invariant)
        [1,3:
          %
        |2,4:
          assumption
        ]
      ]
    ]
  |3,5:
  ]
qed.

let rec traverse_code_internal
  (program: list Byte) (mem: BitVectorTrie Byte 16)
    (cost_labels: BitVectorTrie costlabel 16) (pc: Word) (program_size: nat)
      on program: identifier_map CostTag nat ≝
 let 〈instr,newpc,ticks〉 ≝ fetch … mem pc in
 match program with
 [ nil ⇒ empty_map …
 | cons hd tl ⇒
   match lookup_opt … pc cost_labels with
   [ None ⇒ traverse_code_internal tl mem cost_labels newpc program_size
   | Some lbl ⇒
     let cost ≝ block_cost mem cost_labels pc program_size in
     let cost_mapping ≝ traverse_code_internal tl mem cost_labels newpc program_size in
       add … cost_mapping lbl cost ]].

definition traverse_code ≝
  λprogram: list Byte.
  λmem: BitVectorTrie Byte 16.
  λcost_labels.
  λprogram_size: nat.
    traverse_code_internal program mem cost_labels (zero …) program_size.

definition compute_costs ≝
  λprogram: list Byte.
  λcost_labels: BitVectorTrie costlabel 16.
  λhas_main: bool.
  let program_size ≝ |program| + 1 in
  let memory ≝ load_code_memory program in
   traverse_code program memory cost_labels program_size.