Require Import AST Linking Values Maps Globalenvs Smallstep Registers.
Require Import Coqlib Maps Events Errors Op OptionMonad.
Require Import RTL BTL BTL_SEtheory.
Require Import BTLmatchRTL BTL_Scheduler.

Definition match_prog (p tp: program) :=
  match_program (fun _ f tf => transf_fundef f = OK tf) eq p tp.

Lemma transf_program_match:
  forall prog tprog, transf_program prog = OK tprog -> match_prog prog tprog.
Proof.
  intros. eapply match_transform_partial_program_contextual; eauto.
Qed.

Section PRESERVATION.

Variable prog: program.
Variable tprog: program.

Hypothesis TRANSL: match_prog prog tprog.

Let pge := Genv.globalenv prog.
Let tpge := Genv.globalenv tprog.

Lemma symbols_preserved s: Genv.find_symbol tpge s = Genv.find_symbol pge s.
Proof.
  rewrite <- (Genv.find_symbol_match TRANSL). reflexivity.
Qed.

Lemma senv_preserved:
  Senv.equiv pge tpge.
Proof.
  eapply (Genv.senv_match TRANSL).
Qed.

Lemma functions_preserved:
  forall (v: val) (f: fundef),
  Genv.find_funct pge v = Some f ->
  exists tf cunit, transf_fundef f = OK tf /\ Genv.find_funct tpge v = Some tf /\ linkorder cunit prog.
Proof.
  intros. exploit (Genv.find_funct_match TRANSL); eauto.
  intros (cu & tf & A & B & C).
  repeat eexists; intuition eauto.
  + unfold incl; auto.
  + eapply linkorder_refl.
Qed.

Lemma transf_function_correct f f':
  transf_function f = OK f' -> match_function f f'.
Proof.
  unfold transf_function. intros H.
  assert (OH: transf_function f = OK f') by auto. monadInv H.
  econstructor; eauto.
  eapply check_symbolic_simu_input_equiv; eauto.
  eapply check_symbolic_simu_correct; eauto.
Qed.

Lemma transf_fundef_correct f f':
  transf_fundef f = OK f' -> match_fundef f f'.
Proof.
  intros TRANSF; destruct f; simpl; inv TRANSF.
  + monadInv H0. exploit transf_function_correct; eauto.
    constructor; auto.
  + eapply match_External.
Qed.

Lemma function_ptr_preserved:
  forall v f,
  Genv.find_funct_ptr pge v = Some f ->
  exists tf,
  Genv.find_funct_ptr tpge v = Some tf /\ transf_fundef f = OK tf.
Proof.
  intros.
  exploit (Genv.find_funct_ptr_transf_partial TRANSL); eauto.
Qed.

Lemma function_sig_translated f tf: transf_fundef f = OK tf -> funsig tf = funsig f.
Proof.
  intros H; apply transf_fundef_correct in H; destruct H; simpl; eauto.
  symmetry; erewrite preserv_fnsig; eauto.
Qed.

Theorem transf_initial_states:
  forall s1, initial_state prog s1 ->
  exists s2, initial_state tprog s2 /\ match_states s1 s2.
Proof.
  intros. inv H.
  exploit function_ptr_preserved; eauto. intros (tf & FIND &  TRANSF).
  eexists. split.
  - econstructor; eauto.
    + intros; apply (Genv.init_mem_match TRANSL); eauto.
    + replace (prog_main tprog) with (prog_main prog). rewrite symbols_preserved. eauto.
      symmetry. eapply match_program_main. eauto.
    + erewrite function_sig_translated; eauto.
  - constructor; eauto.
    + constructor.
    + apply transf_fundef_correct; auto.
Qed.

Lemma transf_final_states s1 s2 r:
  match_states s1 s2 -> final_state s1 r -> final_state s2 r.
Proof.
  intros. inv H0. inv H. inv STACKS. constructor.
Qed.

Section SYMBOLIC_CTX.

Variables f1 f2: function.
Variable sp0: val.
Variable rs0: regset.
Variable m0: Memory.Mem.mem.

Lemma symbols_preserved_rev s: Genv.find_symbol pge s = Genv.find_symbol tpge s.
Proof.
  symmetry; apply symbols_preserved.
Qed.

Let ctx := Sctx f1 f2 pge tpge symbols_preserved_rev sp0 rs0 m0.

Lemma str_regs_equiv sfv1 sfv2 rs:
  sfv_simu ctx sfv1 sfv2 ->
  match_function f1 f2 ->
  str_regs (cf (bctx1 ctx)) sfv1 rs = str_regs (cf (bctx2 ctx)) sfv2 rs.
Proof.
  intros SFV MF; inv MF.
  inv SFV; simpl;
  erewrite equiv_input_regs_tr_inputs; auto.
Qed.

Lemma sfind_function_preserved ros1 ros2 fd:
  svident_simu ctx ros1 ros2 ->
  sfind_function (bctx1 ctx) ros1 = Some fd ->
  exists fd',
    sfind_function (bctx2 ctx) ros2 = Some fd'
    /\ transf_fundef fd = OK fd'.
Proof.
  unfold sfind_function. intros [sv1 sv2 SIMU|]; simpl in *.
  + rewrite !SIMU. clear SIMU.
    destruct (eval_sval _ _); try congruence.
    intros; exploit functions_preserved; eauto.
    intros (fd' & cunit & (X1 & X2 & X3)). eexists; eauto.
  + subst. rewrite symbols_preserved. destruct (Genv.find_symbol _ _); try congruence.
    intros; exploit function_ptr_preserved; eauto.
Qed.

Theorem symbolic_simu_ssem stk1 stk2 st st' s t:
  sstate_simu ctx st st' ->
  sem_sstate (bctx1 ctx) stk1 t s st ->
  list_forall2 match_stackframes stk1 stk2 ->
  match_function f1 f2 ->
  exists s',
    sem_sstate (bctx2 ctx) stk2 t s' st'
    /\ match_states s s'.
Proof.
  intros SIMU SST1 STACKS MF. inversion MF.
  exploit sem_sstate_run; eauto.
  intros (sis1 & sfv1 & rs' & m' & SOUT1 & SIS1 & SF).
  exploit sem_si_ok; eauto. intros SOK.
  exploit SIMU; simpl; eauto.
  intros (sis2 & sfv2 & SOUT2 & SSIMU & SFVSIM).
  remember SIS1 as EQSFV. clear HeqEQSFV.
  eapply SFVSIM in EQSFV.
  exploit SSIMU; eauto. intros SIS2.
  destruct SIS1 as (PRE1 & SMEM1 & SREGS1).
  try_simplify_someHyps; intros.
  set (EQSFV2:=sfv2).
  inversion EQSFV; subst; inversion SF; subst.
  - (* goto *)
    exploit (run_sem_sstate (bctx2 ctx) st' sis2 EQSFV2); eauto.
    + erewrite <- str_regs_equiv; simpl; eauto.
    + econstructor.
    + intros SST2. eexists; split; eauto.
      econstructor; simpl; eauto.
      erewrite equiv_input_regs_tr_inputs; eauto.
  - (* call *)
    exploit sfind_function_preserved; eauto. intros (fd' & SFIND & FUND).
    exploit (run_sem_sstate (bctx2 ctx) st' sis2 EQSFV2); eauto.
    + erewrite <- str_regs_equiv; eauto.
    + econstructor; eauto.
      * apply function_sig_translated; auto.
      * erewrite <- ARGS; eauto.
    + intros SST2. eexists; split; eauto.
      econstructor; simpl; eauto.
      * erewrite equiv_input_regs_tr_inputs; eauto.
        repeat (econstructor; eauto).
      * apply transf_fundef_correct; auto.
  - (* tailcall *)
    exploit sfind_function_preserved; eauto. intros (fd' & SFIND & FUND).
    exploit (run_sem_sstate (bctx2 ctx) st' sis2 EQSFV2); eauto.
    + econstructor; eauto.
      * apply function_sig_translated; auto.
      * simpl; erewrite <- preserv_fnstacksize; eauto.
      * erewrite <- ARGS; eauto.
    + intros SST2. eexists; split; eauto.
      econstructor; simpl; eauto.
      apply transf_fundef_correct; auto.
  - (* builtin *)
    exploit (run_sem_sstate (bctx2 ctx) st' sis2 EQSFV2); eauto.
    + erewrite <- str_regs_equiv; simpl; eauto.
    + econstructor.
      * eapply eval_builtin_sargs_preserved.
        eapply senv_preserved.
        eapply eval_list_builtin_sval_correct; eauto.
        inv BARGS.
        erewrite <- eval_list_builtin_sval_preserved in H0; auto.
      * simpl in *; eapply external_call_symbols_preserved; eauto.
        eapply senv_preserved.
    + intros SST2. eexists; split; eauto.
      econstructor; simpl; eauto.
      erewrite equiv_input_regs_tr_inputs; eauto.
  - (* jumptable *)
    exploit (run_sem_sstate (bctx2 ctx) st' sis2 EQSFV2); eauto.
    + erewrite <- str_regs_equiv; simpl; eauto.
    + econstructor; eauto.
      rewrite <- VAL; auto.
    + intros SST2. eexists; split; eauto.
      econstructor; simpl; eauto.
      erewrite equiv_input_regs_tr_inputs; eauto.
  - (* return *)
    exploit (run_sem_sstate (bctx2 ctx) st' sis2 EQSFV2); eauto.
    + econstructor; eauto.
      * simpl; erewrite <- preserv_fnstacksize; eauto.
      * inv OPT; eauto.
        rewrite <- SIMU0; auto.
    + intros SST2. eexists; split; eauto.
      econstructor; simpl; eauto.
  Unshelve. all: auto.
Qed.

Theorem symbolic_simu_correct stk1 stk2 ibf1 ibf2 s t:
  sstate_simu ctx (sexec f1 ibf1.(entry)) (sexec f2 ibf2.(entry)) ->
  iblock_step tr_inputs (sge1 ctx) stk1 f1 (ssp ctx) (srs0 ctx) (sm0 ctx) ibf1.(entry) t s ->
  list_forall2 match_stackframes stk1 stk2 ->
  match_function f1 f2 ->
  exists s',
    iblock_step tr_inputs (sge2 ctx) stk2 f2 (ssp ctx) (srs0 ctx) (sm0 ctx) ibf2.(entry) t s'
    /\ match_states s s'.
Proof.
  unfold sstate_simu.
  intros SIMU STEP1 STACKS MF.
  exploit (sexec_correct (bctx1 ctx)); simpl; eauto.
  intros SST1. exploit symbolic_simu_ssem; eauto.
  intros (s2 & SST2 & MS).
  exploit (sexec_exact (bctx2 ctx)); simpl; eauto.
  intros (s3 & STEP & EQUIV).
  exists s3. intuition eauto.
  eapply match_states_equiv; eauto.
Qed.

End SYMBOLIC_CTX.

Lemma tr_inputs_eqregs_list f tbl: forall rs rs' r' ores,
  (forall r : positive, rs # r = rs' # r) ->
  (tr_inputs f tbl ores rs) # r' =
  (tr_inputs f tbl ores rs') # r'.
Proof.
  induction tbl as [|pc tbl IHtbl];
  intros; destruct ores; simpl;
  rewrite !tr_inputs_get; autodestruct.
Qed.

Lemma find_function_eqregs rs rs' ros fd:
  (forall r : positive, rs # r = rs' # r) ->
  find_function tpge ros rs = Some fd ->
  find_function tpge ros rs' = Some fd.
Proof.
  intros EQREGS; destruct ros; simpl.
  rewrite EQREGS; auto.
  autodestruct.
Qed.

Lemma args_eqregs (args: list reg): forall (rs rs': regset),
  (forall r : positive, rs # r = rs' # r) ->
  rs ## args = rs' ## args.
Proof.
  induction args; simpl; trivial.
  intros rs rs' EQREGS; rewrite EQREGS.
  apply f_equal; eauto.
Qed.

Lemma f_arg_eqregs (rs rs': regset):
  (forall r : positive, rs # r = rs' # r) ->
  (fun r : positive => rs # r) = (fun r : positive => rs' # r).
Proof.
  intros EQREGS; apply functional_extensionality; eauto.
Qed.

Lemma eqregs_update (rs rs': regset) v dest:
  (forall r, rs # r = rs' # r) ->
  (forall r, (rs # dest <- v) # r = (rs' # dest <- v) # r).
Proof.
  intros EQREGS r; destruct (Pos.eq_dec dest r); subst.
    * rewrite !Regmap.gss; reflexivity.
    * rewrite !Regmap.gso; auto.
Qed.

Local Hint Resolve equiv_stack_refl tr_inputs_eqregs_list find_function_eqregs eqregs_update: core.

Lemma iblock_istep_eqregs_None sp ib: forall rs m rs' rs1 m1
  (EQREGS: forall r, rs # r = rs' # r)
  (ISTEP: iblock_istep tpge sp rs m ib rs1 m1 None),
  exists rs1',
    iblock_istep tpge sp rs' m ib rs1' m1 None
    /\ (forall r, rs1 # r = rs1' # r).
Proof.
  induction ib; intros; inv ISTEP.
  - repeat econstructor; eauto.
  - repeat econstructor; eauto.
    erewrite args_eqregs; eauto.
  - inv LOAD; repeat econstructor; eauto;
    erewrite args_eqregs; eauto.
  - repeat econstructor; eauto.
    erewrite args_eqregs; eauto.
    rewrite <- EQREGS; eauto.
  - exploit IHib1; eauto.
    intros (rs1' & ISTEP1 & EQREGS1).
    exploit IHib2; eauto.
    intros (rs2' & ISTEP2 & EQREGS2).
    eexists; split. econstructor; eauto.
    eauto.
  - destruct b.
    1: exploit IHib1; eauto.
    2: exploit IHib2; eauto.
    all:
      intros (rs1' & ISTEP & EQREGS');
      eexists; split; try eapply exec_cond;
      try erewrite args_eqregs; eauto.
Qed.

Lemma iblock_step_eqregs stk sp f ib: forall rs rs' m t s
  (EQREGS: forall r, rs # r = rs' # r)
  (STEP: iblock_step tr_inputs tpge stk f sp rs m ib t s),
  exists s',
    iblock_step tr_inputs tpge stk f sp rs' m ib t s'
    /\ equiv_state s s'.
Proof.
  induction ib; intros;
  destruct STEP as (rs2 & m2 & fin & ISTEP & FSTEP); inv ISTEP.
  - inv FSTEP;
    eexists; split; repeat econstructor; eauto.
    + destruct (or); simpl; try rewrite EQREGS; constructor; eauto.
    + erewrite args_eqregs; eauto. repeat econstructor; eauto.
    + erewrite args_eqregs; eauto. econstructor; eauto.
    + erewrite f_arg_eqregs; eauto.
    + destruct res; simpl; eauto.
    + rewrite <- EQREGS; auto.
  - exploit IHib1; eauto.
    + econstructor. do 2 eexists; split; eauto.
    + intros (s' & STEP & EQUIV). clear FSTEP.
      destruct STEP as (rs3 & m3 & fin2 & ISTEP & FSTEP).
      repeat (econstructor; eauto).
  - exploit iblock_istep_eqregs_None; eauto.
    intros (rs1' & ISTEP1 & EQREGS1).
    exploit IHib2; eauto.
    + econstructor. do 2 eexists; split; eauto.
    + intros (s' & STEP & EQUIV). clear FSTEP.
      destruct STEP as (rs3 & m3 & fin2 & ISTEP & FSTEP).
      eexists; split; eauto.
      econstructor; do 2 eexists; split; eauto.
      eapply exec_seq_continue; eauto.
  - destruct b.
    1: exploit IHib1; eauto.
    3: exploit IHib2; eauto.
    1,3: econstructor; do 2 eexists; split; eauto.
    all:
      intros (s' & STEP & EQUIV); clear FSTEP;
      destruct STEP as (rs3 & m3 & fin2 & ISTEP & FSTEP);
      eexists; split; eauto; econstructor;
      do 2 eexists; split; eauto;
      eapply exec_cond; try erewrite args_eqregs; eauto.
Qed.

Lemma iblock_step_simulation stk1 stk2 f1 f2 sp rs rs' m ibf t s1 pc
  (STEP: iblock_step tr_inputs pge stk1 f1 sp rs m ibf.(entry) t s1)
  (PC: (fn_code f1) ! pc = Some ibf)
  (EQREGS: forall r : positive, rs # r = rs' # r)
  (STACKS: list_forall2 match_stackframes stk1 stk2)
  (TRANSF: match_function f1 f2)
  :exists ibf' s2,
    (fn_code f2) ! pc = Some ibf'
    /\ iblock_step tr_inputs tpge stk2 f2 sp rs' m ibf'.(entry) t s2
    /\ match_states s1 s2.
Proof.
  exploit symbolic_simu_ok; eauto. intros (ib2 & PC' & SYMBS).
  exploit symbolic_simu_correct; repeat (eauto; simpl).
  intros (s2 & STEP2 & MS).
  exploit iblock_step_eqregs; eauto. intros (s3 & STEP3 & EQUIV).
  clear STEP2. exists ib2; exists s3. repeat split; auto.
  eapply match_states_equiv; eauto.
Qed.

Local Hint Constructors step: core.

Theorem step_simulation s1 s1' t s2
  (STEP: step tr_inputs pge s1 t s1')
  (MS: match_states s1 s2)
  :exists s2',
    step tr_inputs tpge s2 t s2'
  /\ match_states s1' s2'.
Proof.
  destruct STEP as [stack ibf f sp n rs m t s PC STEP | | | ]; inv MS.
  - (* iblock *)
    simplify_someHyps. intros PC.
    exploit iblock_step_simulation; eauto.
    intros (ibf' & s2 & PC2 & STEP2 & MS2). 
    eexists; split; eauto.
  - (* function internal *)
    inversion TRANSF as [xf tf MF |]; subst.
    eexists; split.
    + econstructor. erewrite <- preserv_fnstacksize; eauto.
    + erewrite preserv_fnparams; eauto.
      erewrite preserv_entrypoint; eauto.
      econstructor; eauto. 
  - (* function external *)
    inv TRANSF. eexists; split; econstructor; eauto.
    eapply external_call_symbols_preserved; eauto. apply senv_preserved.
  - (* return *)
    inv STACKS. destruct b1 as [res' f' sp' pc' rs'].
    eexists; split.
    + econstructor.
    + inv H1. econstructor; eauto.
Qed.

Theorem transf_program_correct:
  forward_simulation (fsem prog) (fsem tprog).
Proof.
  eapply forward_simulation_step with match_states; simpl; eauto. (* lock-step with respect to match_states *)
  - eapply senv_preserved.
  - eapply transf_initial_states.
  - eapply transf_final_states.
  - intros; eapply step_simulation; eauto.
Qed.

End PRESERVATION.