Require Import Coq.Bool.Sumbool.
Require Import Coq.Logic.Eqdep_dec.
Require Import Crypto.Compilers.SmartMap.
Require Import Crypto.Compilers.Relations.
Require Import Crypto.Compilers.Syntax.
Require Import Crypto.Compilers.Named.Context.
Require Import Crypto.Compilers.Named.Syntax.
Require Import Crypto.Compilers.Named.ContextDefinitions.
Require Import Crypto.Compilers.Named.ContextProperties.
Require Import Crypto.Compilers.Named.ContextProperties.SmartMap.
Require Import Crypto.Compilers.Named.InterpSideConditions.
Require Import Crypto.Compilers.Named.InterpSideConditionsInterp.
Require Import Crypto.Compilers.Named.MapCast.
Require Import Crypto.Util.ZUtil.
Require Import Crypto.Util.Bool.
Require Import Crypto.Util.Option.
Require Import Crypto.Util.Sigma.
Require Import Crypto.Util.Decidable.
Require Import Crypto.Util.PointedProp.
Require Import Crypto.Util.Tactics.BreakMatch.
Require Import Crypto.Util.Tactics.SpecializeBy.
Require Import Crypto.Util.Tactics.DestructHead.
Require Import Crypto.Util.Tactics.RewriteHyp.

Local Open Scope nexpr_scope.
Section language.
  Context {base_type_code : Type}
          {op : flat_type base_type_code -> flat_type base_type_code -> Type}
          {Name : Type}
          {interp_base_type_bounds : base_type_code -> Type}
          (interp_op_bounds : forall src dst, op src dst -> interp_flat_type interp_base_type_bounds src -> interp_flat_type interp_base_type_bounds dst)
          (pick_typeb : forall t, interp_base_type_bounds t -> base_type_code).
  Local Notation pick_type t v := (SmartFlatTypeMap pick_typeb (t:=t) v).
  Context (cast_op : forall t tR (opc : op t tR) args_bs,
              op (pick_type _ args_bs) (pick_type _ (interp_op_bounds t tR opc args_bs)))
          {BoundsContext : Context Name interp_base_type_bounds}
          (BoundsContextOk : ContextOk BoundsContext)
          {interp_base_type : base_type_code -> Type}
          (interp_op : forall src dst,
              op src dst -> interp_flat_type interp_base_type src -> interp_flat_type interp_base_type dst)
          (interped_op_side_conditions : forall s d, op s d -> interp_flat_type interp_base_type s -> pointed_Prop)
          (cast_backb: forall t b, interp_base_type (pick_typeb t b) -> interp_base_type t).
  Let cast_back : forall t b, interp_flat_type interp_base_type (@pick_type t b) -> interp_flat_type interp_base_type t
    := fun t b => SmartFlatTypeMapUnInterp cast_backb.
  Context {Context : Context Name interp_base_type}
          (ContextOk : ContextOk Context)
          (inboundsb : forall t, interp_base_type_bounds t -> interp_base_type t -> Prop).
  Let inbounds : forall t, interp_flat_type interp_base_type_bounds t -> interp_flat_type interp_base_type t -> Prop
    := fun t => interp_flat_type_rel_pointwise inboundsb (t:=t).
  Context (interp_op_bounds_correct:
             forall t tR opc bs
                    (v : interp_flat_type interp_base_type t)
                    (H : inbounds t bs v)
                    (Hside : to_prop (interped_op_side_conditions _ _ opc v)),
               inbounds tR (interp_op_bounds t tR opc bs) (interp_op t tR opc v))
          (pull_cast_back:
             forall t tR opc bs
                    (v : interp_flat_type interp_base_type (pick_type t bs))
                    (H : inbounds t bs (cast_back t bs v))
                    (Hside : to_prop (interped_op_side_conditions _ _ opc (cast_back t bs v))),
               interp_op t tR opc (cast_back t bs v)
               =
               cast_back _ _ (interp_op _ _ (cast_op _ _ opc bs) v))
          (base_type_dec : DecidableRel (@eq base_type_code))
          (Name_dec : DecidableRel (@eq Name)).

  Local Notation mapf_cast := (@mapf_cast _ op Name _ interp_op_bounds pick_typeb cast_op BoundsContext).
  Local Notation map_cast := (@map_cast _ op Name _ interp_op_bounds pick_typeb cast_op BoundsContext).

  Local Ltac handle_options_step :=
    match goal with
    | _ => progress inversion_option
    | [ H : ?x = Some _ |- context[?x] ] => rewrite H
    | [ H : ?x = None |- context[?x] ] => rewrite H
    | [ H : ?x = Some _, H' : context[?x] |- _ ] => rewrite H in H'
    | [ H : ?x = None, H' : context[?x] |- _ ] => rewrite H in H'
    | [ H : Some _ <> None \/ _ |- _ ] => clear H
    | [ H : Some ?x <> Some ?y |- _ ] => assert (x <> y) by congruence; clear H
    | [ H : None <> Some _ |- _ ] => clear H
    | [ H : Some _ <> None |- _ ] => clear H
    | [ H : ?x <> ?x \/ _ |- _ ] => destruct H; [ exfalso; apply H; reflexivity | ]
    | [ H : _ \/ None = Some _ |- _ ] => destruct H; [ | exfalso; clear -H; congruence ]
    | [ H : _ \/ Some _ = None |- _ ] => destruct H; [ | exfalso; clear -H; congruence ]
    | [ H : ?x = Some ?y, H' : ?x = Some ?y' |- _ ]
      => assert (y = y') by congruence; (subst y' || subst y)
    | _ => progress simpl @option_map
    end.

  Local Ltac handle_lookupb_step :=
    let do_eq_dec dec t t' :=
        first [ constr_eq t t'; fail 1
              | lazymatch goal with
                | [ H : t = t' |- _ ] => fail 1
                | [ H : t <> t' |- _ ] => fail 1
                | [ H : t = t' -> False |- _ ] => fail 1
                | _ => destruct (dec t t')
                end ] in
    let do_type_dec := do_eq_dec base_type_dec in
    match goal with
    | _ => progress unfold dec in *
    | _ => handle_options_step
    (* preprocess *)
    | [ H : context[lookupb (extend _ _ _) _] |- _ ]
      => first [ rewrite (lookupb_extend base_type_dec Name_dec) in H by assumption
               | setoid_rewrite (lookupb_extend base_type_dec Name_dec) in H; [ | assumption.. ] ]
    | [ |- context[lookupb (extend _ _ _) _] ]
      => first [ rewrite (fun C => lookupb_extend C base_type_dec Name_dec) by assumption
               | setoid_rewrite (lookupb_extend base_type_dec Name_dec); [ | assumption.. ] ]
    | _ => progress subst
    (* handle multiple hypotheses *)
    | [ H : find_Name _ ?n ?N = Some ?t', H'' : context[find_Name_and_val _ _ ?t ?n ?N ?x ?default] |- _ ]
      => do_type_dec t t'
    (* clear the default value *)
    | [ H : context[find_Name_and_val ?tdec ?ndec ?t ?n (T:=?T) ?N ?V ?default] |- _ ]
      => lazymatch default with None => fail | _ => idtac end;
         rewrite find_Name_and_val_split in H
    (* generic handlers *)
    | [ H : find_Name _ ?n ?N = Some ?t', H' : ?t <> ?t', H'' : context[find_Name_and_val _ _ ?t ?n ?N ?x ?default] |- _ ]
      => erewrite find_Name_and_val_wrong_type in H'' by eassumption
    | [ H : context[find_Name _ _ (SmartFlatTypeMapInterp2 _ _ _)] |- _ ]
      => rewrite find_Name_SmartFlatTypeMapInterp2 with (base_type_code_dec:=base_type_dec) in H
    | [ H : find_Name_and_val _ _ _ _ _ _ _ = None |- _ ]
      => apply find_Name_and_val_None_iff in H
    (* destructers *)
    | [ |- context[find_Name_and_val ?tdec ?ndec ?t ?n ?N ?V ?default] ]
      => destruct (find_Name_and_val tdec ndec t n N V default) eqn:?
    | [ H : context[match find_Name_and_val ?tdec ?ndec ?t ?n ?N ?V ?default with _ => _ end] |- _ ]
      => destruct (find_Name_and_val tdec ndec t n N V default) eqn:?
    | [ H : context[match find_Name ?ndec ?n ?N with _ => _ end] |- _ ]
      => destruct (find_Name ndec n N) eqn:?
    | [ H : context[match base_type_dec ?x ?y with _ => _ end] |- _ ]
      => destruct (base_type_dec x y)
    | [ H : context[match Name_dec ?x ?y with _ => _ end] |- _ ]
      => destruct (Name_dec x y)
    end.

  Local Ltac handle_exists_in_goal :=
    lazymatch goal with
    | [ |- exists v, Some ?k = Some v /\ @?B v ]
      => exists k; split; [ reflexivity | ]
    | [ |- (exists v, None = Some v /\ @?B v) ]
      => exfalso
    | [ |- ?A /\ (exists v, Some ?k = Some v /\ @?B v) ]
      => cut (A /\ B k); [ clear; solve [ intuition eauto ] | cbv beta ]
    | [ |- ?A /\ (exists v, None = Some v /\ @?B v) ]
      => exfalso
    end.
  Local Ltac handle_bounds_side_conditions_step :=
    match goal with
    | [ H : interpf ?e = Some ?v, H' : interpf_side_conditions_gen _ _ _ ?e = Some (_, ?v')%core |- _ ]
      => first [ constr_eq v v'; fail 1
               | assert (Some v = Some v')
                 by (erewrite <- H, snd_interpf_side_conditions_gen_eq, H'; reflexivity);
                 inversion_option; (subst v || subst v') ]
    end.
  Local Ltac fin_inbounds_cast_back_t_step :=
    match goal with
    | [ |- inboundsb _ _ _ /\ _ ]
      => split; [ eapply interp_flat_type_rel_pointwise__find_Name_and_val; eassumption | ]
    | [ |- cast_backb _ _ _ = _ ]
      => eapply find_Name_and_val__SmartFlatTypeMapInterp2__SmartFlatTypeMapUnInterp__Some_Some; [ | eassumption.. ]
    end.
  Local Ltac specializer_t_step :=
    match goal with
    | [ H : ?T, H' : ?T |- _ ] => clear H
    | [ H : forall x, Some _ = Some x -> _ |- _ ] => specialize (H _ eq_refl)
    | [ H : ?x = Some _, IH : forall a b, ?x = Some _ -> _ |- _ ]
      => specialize (IH _ _ H)
    | [ H : ?x = Some _, IH : forall a, ?x = Some _ -> _ |- _ ]
      => specialize (IH _ H)
    | [ H : forall t n v, lookupb ?ctx n = _ -> _, H' : lookupb ?ctx ?n' = _ |- _ ]
      => specialize (H _ _ _ H')
    | _ => progress specialize_by auto
    end.

  Local Ltac break_t_step :=
    first [ progress destruct_head'_ex
          | progress destruct_head'_and ].

  Local Ltac t_step :=
    first [ progress intros
          | break_t_step
          | handle_lookupb_step
          | handle_exists_in_goal
          | solve [ auto ]
          | specializer_t_step
          | fin_inbounds_cast_back_t_step
          | handle_options_step
          | handle_bounds_side_conditions_step ].
  Local Ltac t := repeat t_step.

  Local Ltac do_specialize_IHe :=
    repeat match goal with
           | [ IH : context[interpf ?e], H' : interpf (ctx:=?ctx) ?e = _ |- _ ]
             => let check_tac _ := (rewrite H' in IH) in
                first [ specialize (IH ctx); check_tac ()
                      | specialize (fun a => IH a ctx); check_tac ()
                      | specialize (fun a b => IH a b ctx); check_tac () ]
           | [ IH : context[mapf_cast _ ?e], H' : mapf_cast ?ctx ?e = _ |- _ ]
             => let check_tac _ := (rewrite H' in IH) in
                first [ specialize (IH ctx); check_tac ()
                      | specialize (fun a => IH a ctx); check_tac ()
                      | specialize (fun a b => IH a b ctx); check_tac () ]
           | [ H : forall x y z, Some _ = Some _ -> _ |- _ ]
             => first [ specialize (H _ _ _ eq_refl)
                      | specialize (fun x => H x _ _ eq_refl) ]
           | [ H : forall x y, Some _ = Some _ -> _ |- _ ]
             => first [ specialize (H _ _ eq_refl)
                      | specialize (fun x => H x _ eq_refl) ]
           | _ => progress specialize_by_assumption
           end.

  Lemma mapf_cast_correct
        {t} (e:exprf base_type_code op Name t)
    : forall
      (oldValues:Context)
      (newValues:Context)
      (varBounds:BoundsContext)
      {b} e' (He':mapf_cast varBounds e = Some (existT _ b e'))
      (Hctx:forall {t} n v,
          lookupb (t:=t) oldValues n = Some v
          -> exists b, lookupb (t:=t) varBounds n = Some b
                       /\ @inboundsb _ b v
                       /\ exists v', lookupb (t:=pick_typeb t b) newValues n = Some v'
                                     /\ cast_backb t b v' = v)
      r (Hr:interpf (interp_op:=interp_op) (ctx:=oldValues) e = Some r)
      r' (Hr':interpf (interp_op:=interp_op) (ctx:=newValues) e' = Some r')
      (Hside : prop_of_option (interpf_side_conditions interp_op interped_op_side_conditions oldValues e))
    , interpf (interp_op:=interp_op_bounds) (ctx:=varBounds) e = Some b
      /\ @inbounds _ b r /\ cast_back _ _ r' = r.
  Proof using Type*.
    induction e; simpl interpf; simpl mapf_cast; unfold option_map, cast_back in *; intros;
      unfold interpf_side_conditions in *; simpl in Hside;
      repeat (repeat handle_options_step; break_match_hyps; inversion_option; inversion_sigma; autorewrite with push_to_prop in *; simpl in *; unfold option_map in *; subst; try tauto).
    { destruct (Hctx _ _ _ Hr) as [b' [Hb'[Hb'v[v'[Hv' Hv'v]]]]]; clear Hctx Hr; subst.
      repeat match goal with
               [H: ?e = Some ?x, G:?e = Some ?x' |- _] =>
               pose proof (eq_trans (eq_sym G) H); clear G; inversion_option; subst
             end.
      auto. }
    { do_specialize_IHe.
      repeat (handle_options_step || destruct_head_and || specialize_by_assumption).
      subst; intuition eauto; try (symmetry; rewrite_hyp ?*; eauto);
        repeat handle_bounds_side_conditions_step; auto. }
    { cbv [LetIn.Let_In] in *.
      do_specialize_IHe.
      destruct_head'_and.
      destruct IHe1 as [IHe1_eq IHe1]; rewrite_hyp *; try assumption.
      { apply IHe2; clear IHe2; try reflexivity; [ | t ].
        intros ??? Hlookup.
        let b := fresh "b" in
        let H' := fresh "H'" in
        match goal with |- exists b0, ?v = Some b0 /\ _ => destruct v as [b|] eqn:H' end;
          [ exists b; split; [ reflexivity | ] | exfalso ];
          revert Hlookup H'; t. } }
    { do_specialize_IHe.
      t. }
  Qed.

  Lemma map_cast_correct
        {t} (e:expr base_type_code op Name t)
        (input_bounds : interp_flat_type interp_base_type_bounds (domain t))
    : forall
        (oldValues:Context)
        (newValues:Context)
        (varBounds:BoundsContext)
        {b} e' (He':map_cast varBounds e input_bounds = Some (existT _ b e'))
        (Hctx:forall {t} n v,
            lookupb (t:=t) oldValues n = Some v
            -> exists b, lookupb (t:=t) varBounds n = Some b
                         /\ @inboundsb _ b v
                         /\ exists v', lookupb (t:=pick_typeb t b) newValues n = Some v'
                                       /\ cast_backb t b v' = v)
        v v' (Hv : @inbounds _ input_bounds v /\ cast_back _ _ v' = v)
        r (Hr:interp (interp_op:=interp_op) (ctx:=oldValues) e v = Some r)
        r' (Hr':interp (interp_op:=interp_op) (ctx:=newValues) e' v' = Some r')
        (Hside : prop_of_option (interp_side_conditions interp_op interped_op_side_conditions oldValues e v))
        , interp (interp_op:=interp_op_bounds) (ctx:=varBounds) e input_bounds = Some b
          /\ @inbounds _ b r /\ cast_back _ _ r' = r.
  Proof using Type*.
    unfold map_cast, option_map, interp; simpl; intros.
    repeat first [ progress subst
                 | progress inversion_option
                 | progress inversion_sigma
                 | progress break_match_hyps
                 | progress destruct_head' sigT
                 | progress simpl in * ].
    eapply mapf_cast_correct; try eassumption.
    t.
  Qed.
End language.


(** * Push-Button Synthesis of Bernstein-Yang Inversion: Reification Cache *)
Require Import Coq.ZArith.ZArith.
Require Import Coq.derive.Derive.
Require Import Crypto.Util.Tactics.Head.
Require Import Crypto.Util.ZUtil.Pow.
Require Import Crypto.Arithmetic.Core.
Require Import Crypto.Arithmetic.Partition.
Require Import Crypto.Arithmetic.Freeze.
Require Import Crypto.Arithmetic.ModOps.
Require Import Crypto.Arithmetic.WordByWordMontgomery.
Require Import Crypto.Arithmetic.BYInv.
Require Import Rewriter.Language.Language.
Require Import Crypto.Language.API.
Require Import Crypto.PushButtonSynthesis.ReificationCache.
Local Open Scope Z_scope.

Require Import Crypto.Arithmetic.UniformWeight.
Require Import Lists.List.

Import
  Language.Wf.Compilers
  Language.Compilers
  Language.API.Compilers.
Import Compilers.API.

Import Associational Positional.

Local Set Keyed Unification. (* needed for making [autorewrite] fast, c.f. COQBUG(https://github.com/coq/coq/issues/9283) *)

Module Export WordByWordMontgomeryInversion.
  Import WordByWordMontgomery.WordByWordMontgomery.

  Definition msat bitwidth n m := Partition.partition (uweight bitwidth) n m. (* m in saturated representation *)

  Derive reified_divstep_gen
         SuchThat (is_reification_of reified_divstep_gen divstep)
         As reified_divstep_gen_correct.
  Proof. Time cache_reify (). Time Qed.
  Hint Extern 1 (_ = _) => apply_cached_reification divstep (proj1 reified_divstep_gen_correct) : reify_cache_gen.
  Hint Immediate (proj2 reified_divstep_gen_correct) : wf_gen_cache.
  Hint Rewrite (proj1 reified_divstep_gen_correct) : interp_gen_cache.
  Local Opaque reified_divstep_gen. (* needed for making [autorewrite] not take a very long time *)

  Derive reified_msat_gen
         SuchThat (is_reification_of reified_msat_gen msat)
         As reified_msat_gen_correct.
  Proof.
    Time cache_reify ().
    Time Qed.
  Hint Extern 1 (_ = _) => apply_cached_reification msat (proj1 reified_msat_gen_correct) : reify_cache_gen.
  Hint Immediate (proj2 reified_msat_gen_correct) : wf_gen_cache.
  Hint Rewrite (proj1 reified_msat_gen_correct) : interp_gen_cache.
  Local Opaque reified_msat_gen. (* needed for making [autorewrite] not take a very long time *)

  Derive reified_eval_twos_complement_gen
         SuchThat (is_reification_of reified_eval_twos_complement_gen eval_twos_complement)
         As reified_eval_twos_complement_gen_correct.
  Proof. Time cache_reify (). Time Qed.
  Hint Extern 1 (_ = _) => apply_cached_reification eval_twos_complement (proj1 reified_eval_twos_complement_gen_correct) : reify_cache_gen.
  Hint Immediate (proj2 reified_eval_twos_complement_gen_correct) : wf_gen_cache.
  Hint Rewrite (proj1 reified_eval_twos_complement_gen_correct) : interp_gen_cache.
  Local Opaque reified_eval_twos_complement_gen. (* needed for making [autorewrite] not take a very long time *)
End WordByWordMontgomeryInversion.


(*! Building simple pipelines !*)
Require Import Koika.Frontend.

Inductive reg_t := r0 | outputReg | inputReg | invalid | correct.
Inductive ext_fn_t := Stream | F | G.
Inductive rule_name_t := doF | doG.

Definition sz := (pow2 5).

Definition R r :=
  match r with
  | r0 => bits_t sz
  | inputReg => bits_t sz
  | outputReg => bits_t sz
  | invalid | correct => bits_t 1
  end.

Definition r reg : R reg :=
  match reg with
  | r0 => Bits.zero
  | outputReg => Bits.zero
  | inputReg => Bits.zero
  | invalid => Ob~1
  | correct => Ob~1
  end.

Definition Sigma (fn: ext_fn_t) : ExternalSignature :=
  match fn with
  | Stream => {$ bits_t sz ~> bits_t sz $}
  | F => {$ bits_t sz ~> bits_t sz $}
  | G => {$ bits_t sz ~> bits_t sz $}
  end.

Definition _doF : uaction _ _ :=
  {{
     let v := read0(inputReg) in
     write0(inputReg, extcall Stream(v));
     let invalid := read1(invalid) in
     if invalid then
       write1(invalid, Ob~0);
       write0(r0,extcall F(v))
     else
       fail
  }}.

Definition _doG : uaction _ _ :=
  {{
      let invalid := read0(invalid) in
      if !invalid then
        let data := read0(r0) in
        let v := read0(outputReg) in
        write0(outputReg, extcall Stream(v));
        write0(invalid, Ob~1);
        if extcall G(data) == extcall G(extcall F(v)) then
          pass
        else
          write0(correct, Ob~0)
      else
        fail
  }}.

Definition rules :=
  tc_rules R Sigma
           (fun rl => match rl with
                   | doF => _doF
                   | doG => _doG
                   end).

Definition pipeline : scheduler :=
  doG |> doF |> done.

Definition external (r: rule_name_t) := false.

Definition circuits :=
  compile_scheduler rules external pipeline.

Definition circuits_result sigma :=
  interp_circuits sigma circuits (lower_r (ContextEnv.(create) r)).

Definition cpp_extfuns := "class extfuns {
public:
  static bits<32> stream(bits<32> lfsr) {
    return lfsr + bits<32>{1};
  }

  static bits<32> f(bits<32> x) {
    return ~(x << bits<32>{2}) - bits<32>{1};
  }

  static bits<32> g(bits<32> x) {
    return bits<32>{5} + ((x + bits<32>{1}) >> bits<32>{1});
  }
};".

Definition ext_fn_names fn :=
  match fn with
  | Stream => "stream"
  | F => "f"
  | G => "g"
  end.

Definition package :=
  {| ip_koika := {| koika_reg_types := R;
                   koika_reg_init reg := r reg;
                   koika_ext_fn_types := Sigma;
                   koika_rules := rules;
                   koika_rule_external := external;
                   koika_scheduler := pipeline;
                   koika_module_name := "pipeline" |};

     ip_sim := {| sp_ext_fn_specs fn :=
                   {| efs_name := ext_fn_names fn;
                      efs_method := false |};
                 sp_prelude := Some cpp_extfuns |};

     ip_verilog := {| vp_ext_fn_specs fn :=
                       {| efr_name := ext_fn_names fn;
                          efr_internal := true |} |} |}.

Definition prog := Interop.Backends.register package.
Extraction "pipeline.ml" prog.


Require Import SpecDeps.
Require Import RData.
Require Import EventReplay.
Require Import MoverTypes.
Require Import Constants.
Require Import CommonLib.
Require Import TableDataOpsIntro.Spec.
Require Import TableDataOpsRef1.Spec.

Local Open Scope Z_scope.

Section SpecLow.

  Definition table_destroy2_spec0 (g_rd: Pointer) (map_addr: Z64) (rtt_addr: Z64) (level: Z64) (adt: RData) : option (RData * Z64) :=
    match g_rd, map_addr, rtt_addr, level with
    | (_g_rd_base, _g_rd_ofst), VZ64 _map_addr, VZ64 _rtt_addr, VZ64 _level =>
      rely is_int64 _map_addr;
      rely is_int64 _rtt_addr;
      rely is_int64 _level;
      when' _t'1, adt == table_destroy1_spec (_g_rd_base, _g_rd_ofst) (VZ64 _map_addr) (VZ64 _rtt_addr) (VZ64 _level) adt;
      rely is_int64 _t'1;
      Some (adt, (VZ64 _t'1))
     end
    .

End SpecLow.


(* *********************************************************************)
(*                                                                     *)
(*              The Compcert verified compiler                         *)
(*                                                                     *)
(*          Xavier Leroy, INRIA Paris-Rocquencourt                     *)
(*                                                                     *)
(*  Copyright Institut National de Recherche en Informatique et en     *)
(*  Automatique.  All rights reserved.  This file is distributed       *)
(*  under the terms of the INRIA Non-Commercial License Agreement.     *)
(*                                                                     *)
(* *********************************************************************)

(** Locations are a refinement of RTL pseudo-registers, used to reflect
  the results of register allocation (file [Allocation]). *)

Require Import OrderedType.
Require Import Coqlib.
Require Import Maps.
Require Import Ordered.
Require Import AST.
Require Import Values.
Require Export Machregs.

(** * Representation of locations *)

(** A location is either a processor register or (an abstract designation of)
  a slot in the activation record of the current function. *)

(** ** Processor registers *)

(** Processor registers usable for register allocation are defined
  in module [Machregs]. *)

(** ** Slots in activation records *)

(** A slot in an activation record is designated abstractly by a kind,
  a type and an integer offset.  Three kinds are considered:
- [Local]: these are the slots used by register allocation for
  pseudo-registers that cannot be assigned a hardware register.
- [Incoming]: used to store the parameters of the current function
  that cannot reside in hardware registers, as determined by the
  calling conventions.
- [Outgoing]: used to store arguments to called functions that
  cannot reside in hardware registers, as determined by the
  calling conventions. *)

Inductive slot: Type :=
  | Local
  | Incoming
  | Outgoing.

(** Morally, the [Incoming] slots of a function are the [Outgoing]
slots of its caller function.

The type of a slot indicates how it will be accessed later once mapped to
actual memory locations inside a memory-allocated activation record:
as 32-bit integers/pointers (type [Tint]) or as 64-bit floats (type [Tfloat]).

The offset of a slot, combined with its type and its kind, identifies
uniquely the slot and will determine later where it resides within the
memory-allocated activation record.  Offsets are always positive.
*)

Lemma slot_eq: forall (p q: slot), {p = q} + {p <> q}.
Proof.
  decide equality.
Defined.

Open Scope Z_scope.

Definition typesize (ty: typ) : Z :=
  match ty with
  | Tint => 1
  | Tlong => 2
  | Tfloat => 2
  | Tsingle => 1
  | Tany32 => 1
  | Tany64 => 2
  end.

Lemma typesize_pos:
  forall (ty: typ), typesize ty > 0.
Proof.
  destruct ty; compute; auto.
Qed.

(** ** Locations *)

(** Locations are just the disjoint union of machine registers and
  activation record slots. *)

Inductive loc : Type :=
  | R (r: mreg)
  | S (sl: slot) (pos: Z) (ty: typ).

Module Loc.

  Definition type (l: loc) : typ :=
    match l with
    | R r => mreg_type r
    | S sl pos ty => ty
    end.

  Lemma eq: forall (p q: loc), {p = q} + {p <> q}.
  Proof.
    decide equality.
    apply mreg_eq.
    apply typ_eq.
    apply zeq.
    apply slot_eq.
  Defined.

(** As mentioned previously, two locations can be different (in the sense
  of the [<>] mathematical disequality), yet denote
  overlapping memory chunks within the activation record.
  Given two locations, three cases are possible:
- They are equal (in the sense of the [=] equality)
- They are different and non-overlapping.
- They are different but overlapping.

  The second case (different and non-overlapping) is characterized
  by the following [Loc.diff] predicate.
*)
  Definition diff (l1 l2: loc) : Prop :=
    match l1, l2 with
    | R r1, R r2 =>
        r1 <> r2
    | S s1 d1 t1, S s2 d2 t2 =>
        s1 <> s2 \/ d1 + typesize t1 <= d2 \/ d2 + typesize t2 <= d1
    | _, _ =>
        True
    end.

  Lemma same_not_diff:
    forall l, ~(diff l l).
  Proof.
    destruct l; unfold diff; auto.
    red; intros. destruct H; auto. generalize (typesize_pos ty); omega.
  Qed.

  Lemma diff_not_eq:
    forall l1 l2, diff l1 l2 -> l1 <> l2.
  Proof.
    unfold not; intros. subst l2. elim (same_not_diff l1 H).
  Qed.

  Lemma diff_sym:
    forall l1 l2, diff l1 l2 -> diff l2 l1.
  Proof.
    destruct l1; destruct l2; unfold diff; auto.
    intuition.
  Qed.

  Definition diff_dec (l1 l2: loc) : { Loc.diff l1 l2 } + { ~Loc.diff l1 l2 }.
  Proof.
    intros. destruct l1; destruct l2; simpl.
  - destruct (mreg_eq r r0). right; tauto. left; auto.
  - left; auto.
  - left; auto.
  - destruct (slot_eq sl sl0).
    destruct (zle (pos + typesize ty) pos0).
    left; auto.
    destruct (zle (pos0 + typesize ty0) pos).
    left; auto.
    right; red; intros [P | [P | P]]. congruence. omega. omega.
    left; auto.
  Defined.

(** We now redefine some standard notions over lists, using the [Loc.diff]
  predicate instead of standard disequality [<>].

  [Loc.notin l ll] holds if the location [l] is different from all locations
  in the list [ll]. *)

  Fixpoint notin (l: loc) (ll: list loc) {struct ll} : Prop :=
    match ll with
    | nil => True
    | l1 :: ls => diff l l1 /\ notin l ls
    end.

  Lemma notin_iff:
    forall l ll, notin l ll <-> (forall l', In l' ll -> Loc.diff l l').
  Proof.
    induction ll; simpl.
    tauto.
    rewrite IHll. intuition. subst a. auto.
  Qed.

  Lemma notin_not_in:
    forall l ll, notin l ll -> ~(In l ll).
  Proof.
    intros; red; intros. rewrite notin_iff in H.
    elim (diff_not_eq l l); auto.
  Qed.

  Lemma notin_dec (l: loc) (ll: list loc) : {notin l ll} + {~notin l ll}.
  Proof.
    induction ll; simpl.
    left; auto.
    destruct (diff_dec l a).
    destruct IHll.
    left; auto.
    right; tauto.
    right; tauto.
  Defined.

(** [Loc.disjoint l1 l2] is true if the locations in list [l1]
  are different from all locations in list [l2]. *)

  Definition disjoint (l1 l2: list loc) : Prop :=
    forall x1 x2, In x1 l1 -> In x2 l2 -> diff x1 x2.

  Lemma disjoint_cons_left:
    forall a l1 l2,
    disjoint (a :: l1) l2 -> disjoint l1 l2.
  Proof.
    unfold disjoint; intros. auto with coqlib.
  Qed.
  Lemma disjoint_cons_right:
    forall a l1 l2,
    disjoint l1 (a :: l2) -> disjoint l1 l2.
  Proof.
    unfold disjoint; intros. auto with coqlib.
  Qed.

  Lemma disjoint_sym:
    forall l1 l2, disjoint l1 l2 -> disjoint l2 l1.
  Proof.
    unfold disjoint; intros. apply diff_sym; auto.
  Qed.

  Lemma in_notin_diff:
    forall l1 l2 ll, notin l1 ll -> In l2 ll -> diff l1 l2.
  Proof.
    intros. rewrite notin_iff in H. auto.
  Qed.

  Lemma notin_disjoint:
    forall l1 l2,
    (forall x, In x l1 -> notin x l2) -> disjoint l1 l2.
  Proof.
    intros; red; intros. exploit H; eauto. rewrite notin_iff; intros. auto.
  Qed.

  Lemma disjoint_notin:
    forall l1 l2 x, disjoint l1 l2 -> In x l1 -> notin x l2.
  Proof.
    intros; rewrite notin_iff; intros. red in H. auto.
  Qed.

(** [Loc.norepet ll] holds if the locations in list [ll] are pairwise
  different. *)

  Inductive norepet : list loc -> Prop :=
  | norepet_nil:
      norepet nil
  | norepet_cons:
      forall hd tl, notin hd tl -> norepet tl -> norepet (hd :: tl).

  Lemma norepet_dec (ll: list loc) : {norepet ll} + {~norepet ll}.
  Proof.
    induction ll.
    left; constructor.
    destruct (notin_dec a ll).
    destruct IHll.
    left; constructor; auto.
    right; red; intros P; inv P; contradiction.
    right; red; intros P; inv P; contradiction.
  Defined.

(** [Loc.no_overlap l1 l2] holds if elements of [l1] never overlap partially
  with elements of [l2]. *)

  Definition no_overlap (l1 l2 : list loc) :=
   forall r, In r l1 -> forall s, In s l2 ->  r = s \/ Loc.diff r s.

End Loc.

(** * Mappings from locations to values *)

(** The [Locmap] module defines mappings from locations to values,
  used as evaluation environments for the semantics of the [LTL]
  and [Linear] intermediate languages.  *)

Set Implicit Arguments.

Module Locmap.

  Definition t := loc -> val.

  Definition init (x: val) : t := fun (_: loc) => x.

  Definition get (l: loc) (m: t) : val := m l.

  (** The [set] operation over location mappings reflects the overlapping
      properties of locations: changing the value of a location [l]
      invalidates (sets to [Vundef]) the locations that partially overlap
      with [l].  In other terms, the result of [set l v m]
      maps location [l] to value [v], locations that overlap with [l]
      to [Vundef], and locations that are different (and non-overlapping)
      from [l] to their previous values in [m].  This is apparent in the
      ``good variables'' properties [Locmap.gss] and [Locmap.gso].

      Additionally, the [set] operation also anticipates the fact that
      abstract stack slots are mapped to concrete memory locations
      in the [Stacking] phase.  Hence, values stored in stack slots
      are normalized according to the type of the slot. *)

  Definition set (l: loc) (v: val) (m: t) : t :=
    fun (p: loc) =>
      if Loc.eq l p then
        match l with R r => v | S sl ofs ty => Val.load_result (chunk_of_type ty) v end
      else if Loc.diff_dec l p then
        m p
      else Vundef.

  Lemma gss: forall l v m,
    (set l v m) l =
    match l with R r => v | S sl ofs ty => Val.load_result (chunk_of_type ty) v end.
  Proof.
    intros. unfold set. apply dec_eq_true.
  Qed.

  Lemma gss_reg: forall r v m, (set (R r) v m) (R r) = v.
  Proof.
    intros. unfold set. rewrite dec_eq_true. auto.
  Qed.

  Lemma gss_typed: forall l v m, Val.has_type v (Loc.type l) -> (set l v m) l = v.
  Proof.
    intros. rewrite gss. destruct l. auto. apply Val.load_result_same; auto.
  Qed.

  Lemma gso: forall l v m p, Loc.diff l p -> (set l v m) p = m p.
  Proof.
    intros. unfold set. destruct (Loc.eq l p).
    subst p. elim (Loc.same_not_diff _ H).
    destruct (Loc.diff_dec l p).
    auto.
    contradiction.
  Qed.

  Fixpoint undef (ll: list loc) (m: t) {struct ll} : t :=
    match ll with
    | nil => m
    | l1 :: ll' => undef ll' (set l1 Vundef m)
    end.

  Lemma guo: forall ll l m, Loc.notin l ll -> (undef ll m) l = m l.
  Proof.
    induction ll; simpl; intros. auto.
    destruct H. rewrite IHll; auto. apply gso. apply Loc.diff_sym; auto.
  Qed.

  Lemma gus: forall ll l m, In l ll -> (undef ll m) l = Vundef.
  Proof.
    assert (P: forall ll l m, m l = Vundef -> (undef ll m) l = Vundef).
      induction ll; simpl; intros. auto. apply IHll.
      unfold set. destruct (Loc.eq a l).
      destruct a. auto. destruct ty; reflexivity.
      destruct (Loc.diff_dec a l); auto.
    induction ll; simpl; intros. contradiction.
    destruct H. apply P. subst a. apply gss_typed. exact I.
    auto.
  Qed.

  Fixpoint setlist (ll: list loc) (vl: list val) (m: t) {struct ll} : t :=
    match ll, vl with
    | l1 :: ls, v1 :: vs => setlist ls vs (set l1 v1 m)
    | _, _ => m
    end.

  Lemma gsetlisto: forall l ll vl m, Loc.notin l ll -> (setlist ll vl m) l = m l.
  Proof.
    induction ll; simpl; intros.
    auto.
    destruct vl; auto. destruct H. rewrite IHll; auto. apply gso; auto. apply Loc.diff_sym; auto.
  Qed.

End Locmap.

(** * Total ordering over locations *)

Module IndexedTyp <: INDEXED_TYPE.
  Definition t := typ.
  Definition index (x: t) :=
    match x with
    | Tany32 => 1%positive
    | Tint => 2%positive
    | Tsingle => 3%positive
    | Tany64 => 4%positive
    | Tfloat => 5%positive
    | Tlong => 6%positive
    end.
  Lemma index_inj: forall x y, index x = index y -> x = y.
  Proof. destruct x; destruct y; simpl; congruence. Qed.
  Definition eq := typ_eq.
End IndexedTyp.

Module OrderedTyp := OrderedIndexed(IndexedTyp).

Module IndexedSlot <: INDEXED_TYPE.
  Definition t := slot.
  Definition index (x: t) :=
    match x with Local => 1%positive | Incoming => 2%positive | Outgoing => 3%positive end.
  Lemma index_inj: forall x y, index x = index y -> x = y.
  Proof. destruct x; destruct y; simpl; congruence. Qed.
  Definition eq := slot_eq.
End IndexedSlot.

Module OrderedSlot := OrderedIndexed(IndexedSlot).

Module OrderedLoc <: OrderedType.
  Definition t := loc.
  Definition eq (x y: t) := x = y.
  Definition lt (x y: t) :=
    match x, y with
    | R r1, R r2 => Plt (IndexedMreg.index r1) (IndexedMreg.index r2)
    | R _, S _ _ _ => True
    | S _ _ _, R _ => False
    | S sl1 ofs1 ty1, S sl2 ofs2 ty2 =>
        OrderedSlot.lt sl1 sl2 \/ (sl1 = sl2 /\
        (ofs1 < ofs2 \/ (ofs1 = ofs2 /\ OrderedTyp.lt ty1 ty2)))
    end.
  Lemma eq_refl : forall x : t, eq x x.
  Proof (@refl_equal t).
  Lemma eq_sym : forall x y : t, eq x y -> eq y x.
  Proof (@sym_equal t).
  Lemma eq_trans : forall x y z : t, eq x y -> eq y z -> eq x z.
  Proof (@trans_equal t).
  Lemma lt_trans : forall x y z : t, lt x y -> lt y z -> lt x z.
  Proof.
    unfold lt; intros.
    destruct x; destruct y; destruct z; try tauto.
    eapply Plt_trans; eauto.
    destruct H.
    destruct H0. left; eapply OrderedSlot.lt_trans; eauto.
    destruct H0. subst sl0. auto.
    destruct H. subst sl.
    destruct H0. auto.
    destruct H.
    right.  split. auto.
    intuition.
    right; split. congruence. eapply OrderedTyp.lt_trans; eauto.
  Qed.
  Lemma lt_not_eq : forall x y : t, lt x y -> ~ eq x y.
  Proof.
    unfold lt, eq; intros; red; intros. subst y.
    destruct x.
    eelim Plt_strict; eauto.
    destruct H. eelim OrderedSlot.lt_not_eq; eauto. red; auto.
    destruct H. destruct H0. omega.
    destruct H0. eelim OrderedTyp.lt_not_eq; eauto. red; auto.
  Qed.
  Definition compare : forall x y : t, Compare lt eq x y.
  Proof.
    intros. destruct x; destruct y.
  - destruct (OrderedPositive.compare (IndexedMreg.index r) (IndexedMreg.index r0)).
    + apply LT. red. auto.
    + apply EQ. red. f_equal. apply IndexedMreg.index_inj. auto.
    + apply GT. red. auto.
  - apply LT. red; auto.
  - apply GT. red; auto.
  - destruct (OrderedSlot.compare sl sl0).
    + apply LT. red; auto.
    + destruct (OrderedZ.compare pos pos0).
      * apply LT. red. auto.
      * destruct (OrderedTyp.compare ty ty0).
        apply LT. red; auto.
        apply EQ. red; red in e; red in e0; red in e1. congruence.
        apply GT. red; auto.
      * apply GT. red. auto.
    + apply GT. red; auto.
  Defined.
  Definition eq_dec := Loc.eq.

(** Connection between the ordering defined here and the [Loc.diff] predicate. *)

  Definition diff_low_bound (l: loc) : loc :=
    match l with
    | R mr => l
    | S sl ofs ty => S sl (ofs - 1) Tany64
    end.

  Definition diff_high_bound (l: loc) : loc :=
    match l with
    | R mr => l
    | S sl ofs ty => S sl (ofs + typesize ty - 1) Tlong
    end.

  Lemma outside_interval_diff:
    forall l l', lt l' (diff_low_bound l) \/ lt (diff_high_bound l) l' -> Loc.diff l l'.
  Proof.
    intros.
    destruct l as [mr | sl ofs ty]; destruct l' as [mr' | sl' ofs' ty']; simpl in *; auto.
    - assert (IndexedMreg.index mr <> IndexedMreg.index mr').
      { destruct H. apply sym_not_equal. apply Plt_ne; auto. apply Plt_ne; auto. }
      congruence.
    - assert (RANGE: forall ty, 1 <= typesize ty <= 2).
      { intros; unfold typesize. destruct ty0; omega.  }
      destruct H.
      + destruct H. left. apply sym_not_equal. apply OrderedSlot.lt_not_eq; auto.
        destruct H. right.
        destruct H0. right. generalize (RANGE ty'); omega.
        destruct H0.
        assert (ty' = Tint \/ ty' = Tsingle \/ ty' = Tany32).
        { unfold OrderedTyp.lt in H1. destruct ty'; auto; compute in H1; congruence. }
        right. destruct H2 as [E|[E|E]]; subst ty'; simpl typesize; omega.
      + destruct H. left. apply OrderedSlot.lt_not_eq; auto.
        destruct H. right.
        destruct H0. left; omega.
        destruct H0. exfalso. destruct ty'; compute in H1; congruence.
  Qed.

  Lemma diff_outside_interval:
    forall l l', Loc.diff l l' -> lt l' (diff_low_bound l) \/ lt (diff_high_bound l) l'.
  Proof.
    intros.
    destruct l as [mr | sl ofs ty]; destruct l' as [mr' | sl' ofs' ty']; simpl in *; auto.
    - unfold Plt, Pos.lt. destruct (Pos.compare (IndexedMreg.index mr) (IndexedMreg.index mr')) eqn:C.
      elim H. apply IndexedMreg.index_inj. apply Pos.compare_eq_iff. auto.
      auto.
      rewrite Pos.compare_antisym. rewrite C. auto.
    - destruct (OrderedSlot.compare sl sl'); auto.
      destruct H. contradiction.
      destruct H.
      right; right; split; auto. left; omega.
      left; right; split; auto.
      assert (EITHER: typesize ty' = 1 /\ OrderedTyp.lt ty' Tany64 \/ typesize ty' = 2).
      { destruct ty'; compute; auto. }
      destruct (zlt ofs' (ofs - 1)). left; auto.
      destruct EITHER as [[P Q] | P].
      right; split; auto. omega.
      left; omega.
  Qed.

End OrderedLoc.


Require Import FcEtt.sigs.
Require Export FcEtt.imports.
Require Export FcEtt.ett_inf.
Require Export FcEtt.ett_ind.

Require Import FcEtt.tactics.
Require Import FcEtt.ett_par.

Require Import FcEtt.erase_syntax.

Require Export FcEtt.fc_wf.

(* can remove this parameter *)
Module fc_weak (wf : fc_wf_sig) <: fc_weak_sig.

Set Bullet Behavior "Strict Subproofs".
Set Implicit Arguments.

(* Weakening *)

(* ------------------------------------------------------------------- *)

Lemma ann_respects_atoms_eq_mutual :
  (forall G a A,       AnnTyping  G a A       -> True) /\
  (forall G phi,       AnnPropWff G phi       -> True) /\
  (forall G D g p1 p2, AnnIso     G D g p1 p2 -> forall D', D [=] D' -> AnnIso   G D' g p1 p2) /\
  (forall G D g A B,   AnnDefEq   G D g A B   -> forall D', D [=] D' -> AnnDefEq G D' g A B) /\
  (forall G,           AnnCtx     G           -> True).
Proof.
  eapply ann_typing_wff_iso_defeq_mutual.
  all: pre; subst; eauto 3.
  all: econstructor; eauto.
  fsetdec.
Qed.

Definition AnnIso_respects_atoms_eq   := third  ann_respects_atoms_eq_mutual.
Definition AnnDefEq_respects_atoms_eq := fourth ann_respects_atoms_eq_mutual.

Instance AnnIso_atoms_eq_mor : Morphisms.Proper
                                 (eq ==> AtomSetImpl.Equal ==> eq ==> eq ==> eq ==> iff)
                                 AnnIso.
Proof.
  simpl_relation; split=> ?;
  eauto using AnnIso_respects_atoms_eq, AtomSetProperties.equal_sym.
Qed.

Instance AnnDefEq_atoms_eq_mor : Morphisms.Proper
                                   (eq ==> AtomSetImpl.Equal ==> eq ==> eq ==> eq ==> iff)
                                   AnnDefEq.
Proof.
  simpl_relation; split=> ?;
  eauto using AnnDefEq_respects_atoms_eq, AtomSetProperties.equal_sym.
Qed.

Lemma ann_strengthen_noncovar:
  (forall G1  a A,   AnnTyping G1 a A -> True) /\
  (forall G1  phi,   AnnPropWff G1 phi -> True) /\
  (forall G1 D g p1 p2, AnnIso G1 D g p1 p2 -> forall x, not (exists phi, binds x (Co phi) G1) ->
                     AnnIso G1 (remove x D) g p1 p2) /\
  (forall G1 D g A B,   AnnDefEq G1 D g A B ->  forall x, not (exists phi, binds x (Co phi) G1) ->
                    AnnDefEq G1 (remove x D) g A B) /\
  (forall G1 ,       AnnCtx G1 -> True).
Proof.
  apply ann_typing_wff_iso_defeq_mutual; eauto 3; try done.
  - econstructor; eauto.
  - intros. destruct (x == c).
    + subst. assert False. apply H0. eexists. eauto. done.
    + eapply An_Assn; eauto.
  - econstructor; eauto.
  - intros.
    eapply (An_PiCong (L \u singleton x \u dom G)); eauto.
    intros. eapply H0; auto.
    unfold not in *. intros. destruct H6. apply H4.
    simpl in H6.
    destruct (binds_cons_1 _ x x0 _ (Tm A1) G H6). destruct H7. inversion H8.
    exists x1. auto.
  - intros.
    eapply (An_AbsCong (L \u singleton x)); eauto.
    intros. eapply H0; auto.
    unfold not in *. intros. destruct H6. apply H4.
    simpl in H6.
    destruct (binds_cons_1 _ x x0 _ (Tm A1) G H6). destruct H7. inversion H8.
    exists x1. auto.
  - eauto.
  - eauto.
  - intros.
    eapply (An_CPiCong (L \u singleton x)); eauto.
    intros.
    eapply H0; auto.
    unfold not in *. intros. destruct H6. apply H4.
    simpl in H6.
    destruct (binds_cons_1 _ x c _ (Co phi1) G H6). destruct H7.
    subst. fsetdec.
    exists x0. auto.
  - intros.
    eapply (An_CAbsCong (L \u singleton x)); eauto.
    move=> c Fr.
    eapply H0; first fsetdec.
    move=> [phi b].
    move: b => /binds_cons_iff [[? [?]] | /= b]; first (subst; fsetdec).
      by apply H5; exists phi.
  - eauto.
Qed. (* strengthen_nocovar *)

Lemma AnnDefEq_strengthen_available_tm :
  forall G D g A B, AnnDefEq G D g A B ->  forall x A', binds x (Tm A') G ->
                    forall D', D' [=] remove x D ->
                    AnnDefEq G D' g A B.
Proof.
  intros. eapply ann_respects_atoms_eq_mutual.
  eapply (fourth ann_strengthen_noncovar). eauto.
  unfold not.
  intros b. destruct b as [phi b].
  assert (Tm A' = Co phi). eapply binds_unique; eauto with ctx_wff.
  inversion H2.
  fsetdec.
Qed.

Lemma ann_weaken_available_mutual:
  (forall G1  a A,   AnnTyping G1 a A -> True) /\
  (forall G1  phi,   AnnPropWff G1 phi -> True) /\
  (forall G1 D g p1 p2, AnnIso G1 D g p1 p2 -> forall D', D [<=] D' -> AnnIso G1 D' g p1 p2) /\
  (forall G1 D g A B,   AnnDefEq G1 D g A B -> forall D', D [<=] D' -> AnnDefEq G1 D' g A B) /\
  (forall G1 ,       AnnCtx G1 -> True).
Proof.
  apply ann_typing_wff_iso_defeq_mutual; eauto 3; try done.
  all: econstructor; eauto.
Qed.

Lemma ann_remove_available_mutual:
  (forall G1  a A,   AnnTyping G1 a A -> True) /\
  (forall G1  phi,   AnnPropWff G1 phi -> True) /\
  (forall G1 D g p1 p2, AnnIso G1 D g p1 p2 ->
                   AnnIso G1 (AtomSetImpl.inter D (dom G1)) g p1 p2) /\
  (forall G1 D g A B,   AnnDefEq G1 D g A B ->
                   AnnDefEq G1 (AtomSetImpl.inter D (dom G1)) g A B) /\
  (forall G1 ,       AnnCtx G1 -> True).
Proof.
  apply ann_typing_wff_iso_defeq_mutual; eauto 3; try done.
  - intros L G D. intros.
    eapply (An_PiCong (L \u dom G \u D)); eauto.
    intros.
    eapply (fourth ann_respects_atoms_eq_mutual). eapply H0. eauto.
    simpl. fsetdec.
  - intros L G D. intros.
    eapply (An_AbsCong (L \u dom G \u D)); eauto.
    intros.
    eapply (fourth ann_respects_atoms_eq_mutual). eapply H0. eauto.
    simpl. fsetdec.
  - intros L G D. intros.
    eapply (An_CPiCong (L \u dom G \u D)); eauto.
    intros.
    eapply (fourth ann_respects_atoms_eq_mutual). eapply H0. eauto.
    simpl. fsetdec.
  - intros L G D. intros.
    eapply (An_CAbsCong (L \u dom G \u D)); eauto 1.
    intros.
    eapply (fourth ann_respects_atoms_eq_mutual). eapply H0. eauto.
    simpl. fsetdec.
    eauto.
Qed.

Lemma AnnDefEq_weaken_available :
  forall G D g A B, AnnDefEq G D g A B -> AnnDefEq G (dom G) g A B.
Proof.
  intros.
  remember (AtomSetImpl.inter D (dom G)) as D'.
  eapply (fourth ann_weaken_available_mutual).
  eapply (fourth ann_remove_available_mutual).
  eauto. subst. fsetdec.
Qed.

Lemma AnnIso_weaken_available :
  forall G D g A B, AnnIso G D g A B -> AnnIso G (dom G) g A B.
Proof.
  intros G D. intros.
  remember (AtomSetImpl.inter D (dom G)) as D'.
  eapply (third ann_weaken_available_mutual).
  eapply (third ann_remove_available_mutual).
  eauto. subst. fsetdec.
Qed.

Instance AnnIso_atoms_sub_mor : Morphisms.Proper
                                    (eq ==> AtomSetImpl.Subset ==> eq ==> eq ==> eq ==> impl)
                                    AnnIso.
Proof.
  simpl_relation; eapply (third ann_weaken_available_mutual); eassumption.
Qed.

Instance AnnDefEq_atoms_sub_mor : Morphisms.Proper
                                    (eq ==> AtomSetImpl.Subset ==> eq ==> eq ==> eq ==> impl)
                                    AnnDefEq.
Proof.
  simpl_relation; eapply (fourth ann_weaken_available_mutual); eassumption.
Qed.

(* FIXME: temporary hack *)
Ltac ann_weak_speedup :=
  first [eapply An_AppCong | eapply An_PiSnd].

(* ------------------------------------------------------------------- *)

Lemma ann_typing_weakening_mutual:
  (forall G0 a A,       AnnTyping  G0 a A       ->
     forall E F G, (G0 = F ++ G) -> AnnCtx (F ++ E ++ G) -> AnnTyping (F ++ E ++ G) a A) /\
  (forall G0 phi,       AnnPropWff G0 phi       ->
     forall E F G, (G0 = F ++ G) ->
        AnnCtx (F ++ E ++ G) -> AnnPropWff (F ++ E ++ G) phi) /\
  (forall G0 D g p1 p2, AnnIso     G0 D g p1 p2 ->
     forall E F G, (G0 = F ++ G) ->
        AnnCtx (F ++ E ++ G) -> AnnIso (F ++ E ++ G) D g p1 p2) /\
  (forall G0 D g A B,   AnnDefEq   G0 D g A B   ->
     forall E F G, (G0 = F ++ G) ->
        AnnCtx (F ++ E ++ G) -> AnnDefEq (F ++ E ++ G) D g A B) /\
  (forall G0,           AnnCtx     G0           ->
     forall E F G, (G0 = F ++ G) ->
        AnnCtx (F ++ E ++ G) -> AnnCtx (F ++ E ++ G)).
Proof.
  eapply ann_typing_wff_iso_defeq_mutual.
  all: pre; subst.
  all: eauto 3.
  all: try first [ ann_weak_speedup
                 | An_pick_fresh x;
                   try auto_rew_env;
                   try apply_first_hyp;
                   try simpl_env];
                   eauto 3.
  all: try solve [econstructor; eauto 2].
  all: try solve [eapply AnnDefEq_weaken_available; eauto 2].
  all: try solve [try rewrite <- dom_app; try rewrite <- dom_app;
                  eapply AnnDefEq_weaken_available;  eauto].
  all: try solve [econstructor; eauto 2;
                  eapply AnnDefEq_weaken_available; eauto 2].
  all: try solve [
    (* These are all AnnCtx goals. Need to show the new assumption
       is well-formed by using induction on a term that mentions it.
     *)
    econstructor; eauto 2;
      by move: (H1 E F G0 eq_refl ltac:(auto)); inversion 1].

  (* Left/Right
  eapply An_Right with (a:=a)(a':=a'); eauto 2;
    eapply AnnDefEq_weaken_available; eauto 2. *)
Qed.

Definition AnnTyping_weakening  := first  ann_typing_weakening_mutual.
Definition AnnPropWff_weakening := second ann_typing_weakening_mutual.
Definition AnnIso_weakening     := third  ann_typing_weakening_mutual.
Definition AnnDefEq_weakening   := fourth ann_typing_weakening_mutual.
Definition AnnCtx_weakening     := fifth  ann_typing_weakening_mutual.

End fc_weak.


Set Implicit Arguments.

Require Import GLabel GLabelMap GLabelMapFacts ConvertLabel GoodModule GoodFunction NameDecoration.
Export GLabel GLabelMap GLabelMapFacts ConvertLabel GoodModule GoodFunction NameDecoration.
Import GLabelMap.

Definition name_marker (id : glabel) : PropX W (settings * state) := (Ex s, [| s = id |])%PropX.

Require Import ADT.

Module Make (Import E : ADT).

  Require Import Semantics.
  Module Import SemanticsMake := Make E.
  Export Semantics SemanticsMake.

  Section TopSection.

    Variable modules : list GoodModule.

    Variable imports : t ForeignFuncSpec.

    Notation FName := SyntaxFunc.Name.
    Notation MName := GoodModule.Name.

    Definition label_in (lbl : glabel) :=
      (exists m f,
         List.In m modules /\
         List.In f (Functions m) /\
         lbl = (MName m, FName f)) \/
      In lbl imports.

    Definition label_mapsto lbl spec :=
      (exists ispec m f,
         spec = Internal ispec /\
         List.In m modules /\
         List.In f (Functions m) /\
         ispec = f /\
         lbl = (MName m, FName f)) \/
      (exists fspec,
         spec = Foreign fspec /\
         find lbl imports = Some fspec).

    Definition stn_good_to_use (stn : settings) :=
      forall lbl : glabel,
        label_in lbl ->
        Labels stn lbl <> None.

    Definition stn_injective (stn : settings) :=
      forall lbl1 lbl2 p,
        label_in lbl1 ->
        label_in lbl2 ->
        Labels stn lbl1 = Some p ->
        Labels stn lbl2 = Some p ->
        lbl1 = lbl2.

    Definition fs_good_to_use (fs : settings -> W -> option Callee) (stn : settings) :=
      forall p spec,
        fs stn p = Some spec <->
        exists lbl : glabel,
          Labels stn lbl = Some p /\
          label_mapsto lbl spec.

    Definition env_good_to_use stn fs :=
      stn_good_to_use stn /\
      stn_injective stn /\
      fs_good_to_use fs stn.

    Definition func_export_IFS m (f : GoodFunction) := ((MName m, FName f), f : InternalFuncSpec).

    Definition module_exports_IFS m :=
      List.map (func_export_IFS m) (Functions m).

    Require Import ListFacts1.

    Definition exports_IFS :=
      to_map
        (app_all
           (List.map module_exports_IFS modules)).

    Section fs.

      Variable stn : settings.

      Definition labels (lbl : glabel) : option W := Labels stn lbl.

      Definition is_label_map_to_word lbl p :=
        match labels lbl with
          | Some p' =>
            if weq p p' then
              true
            else
              false
          | None => false
        end.

      Definition is_label_map_to_word' A p (x : glabel * A) := is_label_map_to_word (fst x) p.

      Definition find_by_word A m (p : W) : option A :=
        match List.find (is_label_map_to_word' p) m with
          | Some (_, a) => Some a
          | None => None
        end.

      Definition is_export := find_by_word (elements exports_IFS).

      Definition is_import := find_by_word (elements imports).

      Definition fs (p : W) : option Callee :=
        match is_export p with
          | Some spec => Some (Internal spec)
          | None =>
            match is_import p with
              | Some spec => Some (Foreign spec)
              | None => None
            end
        end.

    End fs.

  End TopSection.

  Require Import RepInv.

  Module Make (Import M : RepInv E).

    Require Import CompileFuncSpec.
    Module Import CompileFuncSpecMake := Make E M.
    Import InvMake2.
    Export CompileFuncSpec CompileFuncSpecMake InvMake2.

    Section TopSection.

      Variable modules : list GoodModule.

      Variable imps : t ForeignFuncSpec.

      Notation fs := (fs modules imps).

      Definition func_spec (id : glabel) f : assert := (st ~> name_marker id /\ [| env_good_to_use modules imps (fst st) fs |] ---> spec_without_funcs_ok f fs st)%PropX.

      Definition foreign_func_spec id spec : assert :=
        st ~> name_marker id /\ ExX, foreign_spec _ spec st.

      Definition imports := mapi (foreign_func_spec) imps.

      Notation FName := SyntaxFunc.Name.
      Notation MName := GoodModule.Name.

      Definition func_export module (f : GoodFunction) :=
        let lbl := (MName module, FName f) in
        (lbl, func_spec lbl f).

      Definition module_exports m :=
        of_list
          (List.map
             (func_export m)
             (Functions m)).

      Definition exports := update_all (List.map module_exports modules).

      Definition impl_label mod_name f_name : glabel := (impl_module_name mod_name, f_name).

      Definition func_impl_export m (f : GoodFunction) := (impl_label (MName m) (FName f), spec f).

      Definition module_impl_exports m :=
        of_list
          (List.map
             (func_impl_export m)
             (Functions m)).

      Definition impl_exports := update_all (List.map module_impl_exports modules).

      Definition all_exports := update exports impl_exports.

    End TopSection.

  End Make.

End Make.


Require Import String List Bool.
Require Import ExtLib.Structures.Maps.
Require Import ExtLib.Structures.Monads.
Require Import ExtLib.Structures.Reducible.
Require Import ExtLib.Data.Option.
Require Import ExtLib.Data.Lists.
Require Import ExtLib.Core.RelDec.
Require Import ExtLib.Data.Set.ListSet.
Require Import ExtLib.Data.Map.FMapAList.
Require Import CoqCompile.CpsK.
Require Import CoqCompile.Analyze.CFA0.
Require Import CoqCompile.CpsKUtil.

Import CpsK.CPSK.

Section REACHABLE.

  Definition useful_ptv (v:PtValue unit) : lset (@eq (var + cont)) :=
    match v with
      | Ctor c => Sets.empty
      | Int i => Sets.empty
      | Ptr l => Sets.singleton (inl l)
      | Clo c pp ks xs e =>
        match pp with
          | inl v => free_vars_decl true (Fn_d v ks xs e)
          | inr k => free_vars_cont k xs e
        end
    end.

  Definition useful (vs:lset (@eq (PtValue unit))) : lset (@eq (var + cont)) :=
    fold (fun v acc => Sets.union acc (useful_ptv v)) Sets.empty vs.

  Definition D : Type := alist (var + cont) (lset (@eq (var + cont))).

  Definition reachable_1 (env:alist (var + cont) (Value unit)) (dom:D) : D.
  refine (
    fold (fun kv acc =>
      let '(k, v) := kv in
      let res := match v with
                   | Any => keys _ _ env
                   | Values vs => _
                 end in
      Maps.add k res acc) dom env).
  refine (match Maps.lookup k acc with
            | Some old_set => Sets.union old_set (useful vs)
            | None => useful vs
          end).
  Defined.

  Definition reachable_2 (heap:alist var (list (Value unit))) (dom : D) : D :=
    fold (fun kv acc =>
      let '(k,v) := kv in
      let res := fold (fun v' acc' =>
        match v' with
          | inl v' =>
            match Maps.lookup v' heap with
              | None => acc'
              | Some vs => List.fold_left (fun acc x =>
                match x with
                  | Any => map (fun x => inl x) (keys _ (s := lset (@eq var)) _ heap)
                  | Values vs => Sets.union (useful vs) acc
                end ) vs acc'
            end
          | inr _ => acc'
        end) v v in
      Maps.add k res acc) dom dom.

  Require Import BinNums.
  Require Import ExtLib.Data.Monads.FuelMonad.
  Require Import ExtLib.Data.Monads.IdentityMonad.

  Definition least_fixpoint {A} (leq : A -> A -> bool) (f : A -> A) (init : A) (fuel : N) : option A :=
    unIdent (runGFixT (mfix (fun recur x =>
      let x' := f x in
      if leq x' x then ret x' else recur x') init) fuel).

  Definition dom_eq (d1 d2 : D) : bool :=
    let f := fun s1 s2 => (Sets.subset s1 s2) && Sets.subset s2 s1 in
    submap_with _ f d1 d2 &&
    submap_with _ f d2 d1.

  Definition reachable (env:alist (var + cont) (Value unit)) (heap:alist var (list (Value unit))) (fuel : N) : option D :=
    least_fixpoint dom_eq (reachable_2 heap) (reachable_1 env Maps.empty) fuel.

End REACHABLE.

Section LIVENESS.

  Variable reach : alist (var + cont) (lset (@eq (var + cont))).

  (* let x = <1, 2>
          .
          .
          .
     let y = <3, 4>
          .
          .
          .
     in k z

   *)

  (* Do we need this??? *)
  Fixpoint escape_cont_exp (e:exp) (acc:lset (@eq (var + cont))) : (lset (@eq (var + cont))) :=
    match e with
      | App_e o ks os => acc
      | Let_e d e' => escape_cont_exp e' (Sets.union (escape_cont_decl d acc) acc)
      | Letrec_e ds e' =>
        escape_cont_exp e' (List.fold_left (fun acc' x => Sets.union (escape_cont_decl x acc) acc') ds acc)
      | Switch_e o arms def =>
        let arms_escape := List.fold_left (fun acc' x =>
          let '(p, e) := x in Sets.union (escape_cont_exp e acc) acc') arms acc in
        match def with
          | Some e => Sets.union arms_escape (escape_cont_exp e acc)
          | None => arms_escape
        end
      | Halt_e o1 o2 => acc
      | AppK_e k os => List.fold_left (fun acc' x => match x with
                                                      | Var_o x => Sets.add (inl x) acc'
                                                      | _ => acc'
                                                    end) os acc
      | LetK_e kves e' =>
        List.fold_left (fun acc' x => let '(k,x,e) := x in
          Sets.union (escape_cont_exp e acc) acc') kves acc
    end
  with escape_cont_decl (d:decl) (acc:lset (@eq (var + cont))) : (lset (@eq (var + cont))) :=
    match d with
      | Op_d x os => acc
      | Prim_d x p os => acc
      | Fn_d x ks xs e => escape_cont_exp e acc
      | Bind_d x w m os => acc
    end.

  Definition live_set (e:exp) : (lset (@eq (var + cont))) :=
    (* let escape := Sets.union (free_vars_exp e) (escape_cont_exp e Sets.empty) in *)
    let escape := free_vars_exp e in
    fold (fun x acc => match Maps.lookup x reach with
                         | Some s => Sets.union s acc
                         | None => acc
                       end) Sets.empty escape.

  Definition var_of_decl (d:decl) : var :=
    match d with
      | Op_d x _ => x
      | Prim_d x _ _ => x
      | Fn_d x _ _ _ => x
      | Bind_d x _ _ _ => x
    end.

  Fixpoint live_exp (e:exp) (dom:alist (var + cont) (lset (@eq (var + cont)))) :
    alist (var + cont) (lset (@eq (var + cont))) :=
    match e with
      | App_e o ks os => dom
      | Let_e d e' =>
        let live_stuff := live_set e' in
        let x := var_of_decl d in
        let live_stuff := match Maps.lookup (inl x) dom with
                            | Some old_set => Sets.union old_set live_stuff
                            | None => live_stuff
                          end in
        let dom' := live_decl d (Maps.add (inl x) live_stuff dom) in
        live_exp e' dom'
      | Letrec_e ds e' =>
        List.fold_left (fun acc d =>
          let live_stuff := live_set e' in
          let x := var_of_decl d in
          let live_stuff := match Maps.lookup (inl x) acc with
                              | Some old_set => Sets.union old_set live_stuff
                              | None => live_stuff
                            end in
          let dom' := Maps.add (inl x) live_stuff acc in
          let dom' := live_decl d dom' in
          live_exp e' dom') ds dom
      | Switch_e o arms def =>
        let dom' := List.fold_left (fun acc x => let '(p, e') := x in
          Maps.combine (fun k v1 v2 => Sets.union v1 v2) (live_exp e' acc) acc) arms dom in
        match def with
          | Some e' => Maps.combine (fun k v1 v2 => Sets.union v1 v2) (live_exp e' dom') dom'
          | None => dom'
        end
      | Halt_e o1 o2 => dom
      | AppK_e k os => dom
      | LetK_e kxse e' =>
        let dom' := List.fold_left (fun acc x => let '(k, xs, e') := x in
          live_exp e' acc) kxse dom in
        live_exp e' dom'
    end
  with live_decl (d:decl) (dom:alist (var + cont) (lset (@eq (var + cont)))) :
    alist (var + cont) (lset (@eq (var + cont))) :=
    match d with
      | Op_d x os => dom
      | Prim_d x p os => dom
      | Fn_d x ks xs e => live_exp e dom
      | Bind_d x w m os => dom
    end.

End LIVENESS.

Section monadic.
  Require Import CoqCompile.TraceMonad.
  Import MonadNotation.
  Local Open Scope monad_scope.

  Variable m : Type -> Type.
  Context {Monad_m : Monad m}.
  Context {MonadTrace_m : MonadTrace string m}.
  Context {MonadExc_m : MonadExc string m}.

  Definition sanitize {A B} {_ : RelDec (@eq B)} (env:alist (unit * B) A) : (alist B A) :=
    fold (fun kv acc =>
      let '(k, v) := kv in
      let '(_, k) := k in
      Maps.add k v acc) Maps.empty env.

  Require Import CoqCompile.Opt.CopyPropCpsK.

  Definition construct_live_map (e:exp) (fuel:N) : m (option (alist (var + cont) (lset (@eq (var + cont))))) :=
    catch (
      let e := CopyProp.copyprop e in
      domain <- cfa_n _ 0 e fuel ;;
      let sanitized := sanitize (env _ domain) in
      let sanitized_heap := sanitize (heap _ domain) in
      match reachable sanitized sanitized_heap fuel with
        | None => raise "reachable ran out of fuel"%string
        | Some dom' =>
        let live := live_exp dom' e Maps.empty in
          ret (Some live)
      end
    ) (fun err =>
      mlog ("construct_live_map failed: " ++ err)%string ;;
      ret None
    ).
End monadic.