(** RTLpath Scheduling from an external oracle. 

This module is inspired from [Duplicate] and [Duplicateproof]

*)

Require Import AST Linking Values Maps Globalenvs Smallstep Registers.
Require Import Coqlib Maps Events Errors Op.
Require Import RTL RTLpath RTLpathLivegen RTLpathLivegenproof RTLpathSE_theory RTLpathSE_impl.
Require RTLpathWFcheck.

Notation "'ASSERT' A 'WITH' MSG 'IN' B" := (if A then B else Error (msg MSG))
         (at level 200, A at level 100, B at level 200)
         : error_monad_scope.

Local Open Scope error_monad_scope.
Local Open Scope positive_scope.

(** External oracle returning the new RTLpath function and a mapping of new path_entries to old path_entries 

NB: the new RTLpath function is generated from the fn_code, the fn_entrypoint and the fn_path
It requires to check that the path structure is wf !

*)

(* Returns: new code, new entrypoint, new pathmap, revmap
 * Indeed, the entrypoint might not be the same if the entrypoint node is moved further down
 * a path ; same reasoning for the pathmap *)
Axiom untrusted_scheduler: RTLpath.function -> code * node * path_map * (PTree.t node).

Extract Constant untrusted_scheduler => "RTLpathScheduleraux.scheduler".

Program Definition function_builder (tfr: RTL.function) (tpm: path_map) : 
  { r : res RTLpath.function | forall f', r = OK f' -> fn_RTL f' = tfr} :=
  match RTLpathWFcheck.function_checker tfr tpm with
  | false => Error (msg "In function_builder: (tfr, tpm) is not wellformed")
  | true => OK {| fn_RTL := tfr; fn_path := tpm |}
  end.
Next Obligation.
  apply RTLpathWFcheck.function_checker_path_entry. auto.
Defined. Next Obligation.
  apply RTLpathWFcheck.function_checker_wellformed_path_map. auto.
Defined.

Definition entrypoint_check (dm: PTree.t node) (fr tfr: RTL.function) : res unit :=
  match dm ! (fn_entrypoint tfr) with
  | None => Error (msg "No mapping for (entrypoint tfr)")
  | Some etp => if (Pos.eq_dec (fn_entrypoint fr) etp) then OK tt
                else Error (msg "Entrypoints do not match")
  end.

Lemma entrypoint_check_correct fr tfr dm:
  entrypoint_check dm fr tfr = OK tt ->
  dm ! (fn_entrypoint tfr) = Some (fn_entrypoint fr).
Proof.
  unfold entrypoint_check. explore; try discriminate. congruence.
Qed.

Definition path_entry_check_single (pm tpm: path_map) (m: node * node) :=
  let (pc2, pc1) := m in
  match (tpm ! pc2) with
  | None => Error (msg "pc2 isn't an entry of tpm")
  | Some _ =>
      match (pm ! pc1) with
      | None => Error (msg "pc1 isn't an entry of pm")
      | Some _ => OK tt
      end
  end.

Lemma path_entry_check_single_correct pm tpm pc1 pc2:
  path_entry_check_single pm tpm (pc2, pc1) = OK tt ->
  path_entry tpm pc2 /\ path_entry pm pc1.
Proof.
  unfold path_entry_check_single. intro. explore.
  constructor; congruence.
Qed.

(* Inspired from Duplicate.verify_mapping_rec *)
Fixpoint path_entry_check_rec (pm tpm: path_map) lm :=
  match lm with
  | nil => OK tt
  | m :: lm => do u1 <- path_entry_check_single pm tpm m;
               do u2 <- path_entry_check_rec pm tpm lm;
               OK tt
  end.

Lemma path_entry_check_rec_correct pm tpm pc1 pc2: forall lm,
  path_entry_check_rec pm tpm lm = OK tt ->
  In (pc2, pc1) lm ->
  path_entry tpm pc2 /\ path_entry pm pc1.
Proof.
  induction lm.
  - simpl. intuition.
  - simpl. intros. explore. destruct H0.
    + subst. eapply path_entry_check_single_correct; eauto.
    + eapply IHlm; assumption.
Qed.

Definition path_entry_check (dm: PTree.t node) (pm tpm: path_map) := path_entry_check_rec pm tpm (PTree.elements dm).

Lemma path_entry_check_correct dm pm tpm:
  path_entry_check dm pm tpm = OK tt ->
  forall pc1 pc2, dm ! pc2 = Some pc1 ->
  path_entry tpm pc2 /\ path_entry pm pc1.
Proof.
  unfold path_entry_check. intros. eapply PTree.elements_correct in H0.
  eapply path_entry_check_rec_correct; eassumption.
Qed.

Definition function_equiv_checker (dm: PTree.t node) (f: RTLpath.function) (tf: RTLpath.function) : res unit :=
  let pm := fn_path f in
  let fr := fn_RTL f in
  let tpm := fn_path tf in
  let tfr := fn_RTL tf in
  do _ <- entrypoint_check dm fr tfr;
  do _ <- path_entry_check dm pm tpm;
  do _ <- simu_check dm f tf;
  OK tt.

Lemma function_equiv_checker_entrypoint f tf dm:
  function_equiv_checker dm f tf = OK tt ->
  dm ! (fn_entrypoint tf) = Some (fn_entrypoint f).
Proof.
  unfold function_equiv_checker. intros. explore.
  eapply entrypoint_check_correct; eauto.
Qed.

Lemma function_equiv_checker_pathentry1 f tf dm:
  function_equiv_checker dm f tf = OK tt ->
  forall pc1 pc2, dm ! pc2 = Some pc1 ->
  path_entry (fn_path tf) pc2.
Proof.
  unfold function_equiv_checker. intros. explore.
  exploit path_entry_check_correct. eassumption. all: eauto. intuition.
Qed.

Lemma function_equiv_checker_pathentry2 f tf dm:
  function_equiv_checker dm f tf = OK tt ->
  forall pc1 pc2, dm ! pc2 = Some pc1 ->
  path_entry (fn_path f) pc1.
Proof.
  unfold function_equiv_checker. intros. explore.
  exploit path_entry_check_correct. eassumption. all: eauto. intuition.
Qed.

Lemma function_equiv_checker_correct f tf dm:
  function_equiv_checker dm f tf = OK tt ->
  forall pc1 pc2, dm ! pc2 = Some pc1 ->
  sexec_simu dm f tf pc1 pc2.
Proof.
  unfold function_equiv_checker. intros. explore.
  eapply simu_check_correct; eauto.
Qed.

Definition verified_scheduler (f: RTLpath.function) : res (RTLpath.function * (PTree.t node)) :=
  let (tctetpm, dm) := untrusted_scheduler f in
  let (tcte, tpm) := tctetpm in
  let (tc, te) := tcte in
  let tfr := mkfunction (fn_sig f) (fn_params f) (fn_stacksize f) tc te in
  do tf <- proj1_sig (function_builder tfr tpm);
  do tt <- function_equiv_checker dm f tf; 
  OK (tf, dm).

Theorem verified_scheduler_correct f tf dm:
  verified_scheduler f = OK (tf, dm) ->
  fn_sig f = fn_sig tf
  /\ fn_params f = fn_params tf
  /\ fn_stacksize f = fn_stacksize tf
  /\ dm ! (fn_entrypoint tf) = Some (fn_entrypoint f)
  /\ (forall pc1 pc2, dm ! pc2 = Some pc1 -> path_entry (fn_path f) pc1)
  /\ (forall pc1 pc2, dm ! pc2 = Some pc1 -> path_entry (fn_path tf) pc2)
  /\ (forall pc1 pc2, dm ! pc2 = Some pc1 -> sexec_simu dm f tf pc1 pc2)
.
Proof.
  intros VERIF. unfold verified_scheduler in VERIF. explore.
  Local Hint Resolve function_equiv_checker_entrypoint
    function_equiv_checker_pathentry1 function_equiv_checker_pathentry2
    function_equiv_checker_correct: core.
  destruct (function_builder _ _) as [res H]; simpl in * |- *; auto.
    apply H in EQ2. rewrite EQ2. simpl.
  repeat (constructor; eauto).
  exploit function_equiv_checker_entrypoint. eapply EQ4. rewrite EQ2. intuition.
Qed.

Record match_function (dupmap: PTree.t node) (f1 f2: RTLpath.function): Prop := {
  preserv_fnsig: fn_sig f1 = fn_sig f2;
  preserv_fnparams: fn_params f1 = fn_params f2;
  preserv_fnstacksize: fn_stacksize f1 = fn_stacksize f2;
  preserv_entrypoint: dupmap!(f2.(fn_entrypoint)) = Some f1.(fn_entrypoint);
  dupmap_path_entry1: forall pc1 pc2, dupmap!pc2 = Some pc1 -> path_entry (fn_path f1) pc1;
  dupmap_path_entry2: forall pc1 pc2, dupmap!pc2 = Some pc1 -> path_entry (fn_path f2) pc2;
  dupmap_correct: forall pc1 pc2, dupmap!pc2 = Some pc1 -> sexec_simu dupmap f1 f2 pc1 pc2;
}.

Program Definition transf_function (f: RTLpath.function):
  { r : res RTLpath.function | forall f', r = OK f' -> exists dm, match_function dm f f'} :=
  match (verified_scheduler f) with
  | Error e => Error e
  | OK (tf, dm) => OK tf
  end.
Next Obligation.
  exploit verified_scheduler_correct; eauto.
  intros (A & B & C & D & E & F & G (* & H *)).
  exists dm. econstructor; eauto.
Defined.

Theorem match_function_preserves f f' dm:
  match_function dm f f' ->
  fn_sig f = fn_sig f' /\ fn_params f = fn_params f' /\ fn_stacksize f = fn_stacksize f'.
Proof.
  intros.
  destruct H as [SIG PARAM SIZE ENTRY CORRECT].
  intuition.
Qed.

Definition transf_fundef (f: fundef) : res fundef :=
  transf_partial_fundef (fun f => proj1_sig (transf_function f)) f.

Definition transf_program (p: program) : res program :=
  transform_partial_program transf_fundef p.

(** * Preservation proof *)

Local Notation ext alive := (fun r => Regset.In r alive).

Inductive match_fundef: RTLpath.fundef -> RTLpath.fundef -> Prop :=
  | match_Internal dupmap f f': match_function dupmap f f' -> match_fundef (Internal f) (Internal f')
  | match_External ef: match_fundef (External ef) (External ef).

Inductive match_stackframes: stackframe -> stackframe -> Prop :=
  | match_stackframe_intro dupmap res f sp pc rs1 rs2 f' pc' path
      (TRANSF: match_function dupmap f f')
      (DUPLIC: dupmap!pc' = Some pc)
      (LIVE: liveness_ok_function f)
      (PATH: f.(fn_path)!pc = Some path)
      (EQUIV: forall v, eqlive_reg (ext path.(input_regs)) (rs1 # res <- v) (rs2 # res <- v)):
      match_stackframes (Stackframe res f sp pc rs1) (Stackframe res f' sp pc' rs2).

Inductive match_states: state -> state -> Prop :=
  | match_states_intro dupmap st f sp pc rs1 rs2 m st' f' pc' path
      (STACKS: list_forall2 match_stackframes st st')
      (TRANSF: match_function dupmap f f')
      (DUPLIC: dupmap!pc' = Some pc)
      (LIVE: liveness_ok_function f)
      (PATH: f.(fn_path)!pc = Some path)
      (EQUIV: eqlive_reg (ext path.(input_regs)) rs1 rs2):
      match_states (State st f sp pc rs1 m) (State st' f' sp pc' rs2 m)
  | match_states_call st st' f f' args m
      (STACKS: list_forall2 match_stackframes st st')
      (TRANSF: match_fundef f f')
      (LIVE: liveness_ok_fundef f):
      match_states (Callstate st f args m) (Callstate st' f' args m)
  | match_states_return st st' v m
      (STACKS: list_forall2 match_stackframes st st'):
      match_states (Returnstate st v m) (Returnstate st' v m).

Lemma match_stackframes_equiv stf1 stf2 stf3:
  match_stackframes stf1 stf2 -> equiv_stackframe stf2 stf3 -> match_stackframes stf1 stf3.
Proof.
  destruct 1; intros EQ; inv EQ; try econstructor; eauto.
  intros; eapply eqlive_reg_trans; eauto.
  rewrite eqlive_reg_triv in * |-.
  eapply eqlive_reg_update.
  eapply eqlive_reg_monotonic; eauto.
  simpl; auto.
Qed.

Lemma match_stack_equiv stk1 stk2:
  list_forall2 match_stackframes stk1 stk2 -> 
  forall stk3, list_forall2 equiv_stackframe stk2 stk3 -> 
  list_forall2 match_stackframes stk1 stk3.
Proof.
  Local Hint Resolve match_stackframes_equiv: core.
  induction 1; intros stk3 EQ; inv EQ; econstructor; eauto.
Qed.

Lemma match_states_equiv s1 s2 s3: match_states s1 s2 -> equiv_state s2 s3 -> match_states s1 s3.
Proof.
  Local Hint Resolve match_stack_equiv: core.
  destruct 1; intros EQ; inv EQ; econstructor; eauto.
  intros; eapply eqlive_reg_triv_trans; eauto.
Qed.

Lemma eqlive_match_stackframes stf1 stf2 stf3:
  eqlive_stackframes stf1 stf2 -> match_stackframes stf2 stf3 -> match_stackframes stf1 stf3.
Proof.
  destruct 1; intros MS; inv MS; try econstructor; eauto.
  try_simplify_someHyps. intros; eapply eqlive_reg_trans; eauto.
Qed.

Lemma eqlive_match_stack stk1 stk2:
  list_forall2 eqlive_stackframes stk1 stk2 -> 
  forall stk3, list_forall2 match_stackframes stk2 stk3 -> 
  list_forall2 match_stackframes stk1 stk3.
Proof.
  induction 1; intros stk3 MS; inv MS; econstructor; eauto.
  eapply eqlive_match_stackframes; eauto.
Qed.

Lemma eqlive_match_states s1 s2 s3: eqlive_states s1 s2 -> match_states s2 s3 -> match_states s1 s3.
Proof.
  Local Hint Resolve eqlive_match_stack: core.
  destruct 1; intros MS; inv MS; try_simplify_someHyps; econstructor; eauto.
  eapply eqlive_reg_trans; eauto.
Qed.

Lemma eqlive_stackframes_refl stf1 stf2: match_stackframes stf1 stf2 -> eqlive_stackframes stf1 stf1.
Proof.
  destruct 1; econstructor; eauto.
  intros; eapply eqlive_reg_refl; eauto.
Qed.

Lemma eqlive_stacks_refl stk1 stk2:
  list_forall2 match_stackframes stk1 stk2 -> list_forall2 eqlive_stackframes stk1 stk1.
Proof.
  induction 1; simpl; econstructor; eauto.
  eapply eqlive_stackframes_refl; eauto.
Qed.

Lemma transf_fundef_correct f f':
  transf_fundef f = OK f' -> match_fundef f f'.
Proof.
  intros TRANSF; destruct f; simpl; monadInv TRANSF.
  + destruct (transf_function f) as [res H]; simpl in * |- *; auto.
    destruct (H _ EQ).
    intuition subst; auto.
    eapply match_Internal; eauto.
  + eapply match_External.
Qed.