Require Import Coq.Lists.List.
Import ListNotations.

Require Import SecureCompartmentalization.fullabst.
Require Import SecureCompartmentalization.lib.

Arguments Some {A} _.
Arguments fst {_ _} _.
Arguments snd {_ _} _.

Class interface_language (interface: Type): Type :=
  {
    compatible: list interface -> Prop;
    complete: list interface -> Prop
  }.

Let comp
    {interface: Type}
    `{interface_language interface}
    (Is : list interface) := compatible Is /\ complete Is.

Class component_language
      {interface: Type} (il: interface_language interface)
      (component program: Type): Type :=
  {
    get_interface : component -> option interface;
    link : list component -> program;
    fully_defined : list component -> list interface -> Prop;
    beh_eq : program -> program -> Prop
  }.

Definition has_interface
           {interface: Type} {il: interface_language interface}
           {component program: Type}
           {L: component_language il component program}

           (C: component) (I: interface): Prop :=
  get_interface C = Some I.

Notation "P ~~ Q" := (beh_eq P Q) (at level 60, no associativity).

Fixpoint All {A} (P : A -> Prop) (l : list A) : Prop :=
  match l with
  | [] => True
  | a :: l' => P a /\ All P l'
  end.

Lemma All_app {A} (P : A -> Prop) (l1 l2 : list A) :
  All P l1 ->
  All P l2 ->
  All P (l1 ++ l2).
Proof.
  induction l1 as [|x l1 IH]; simpl; try solve [intuition].
Qed.

Lemma All_app_l {A} (P : A -> Prop) (l1 l2 : list A) :
  All P (l1 ++ l2) ->
  All P l1.
Proof.
  induction l1 as [|x l1 IH]; simpl; try solve [intuition].
Qed.

Lemma All_app_r {A} (P : A -> Prop) (l1 l2 : list A) :
  All P (l1 ++ l2) ->
  All P l2.
Proof.
  induction l1 as [|x l1 IH]; simpl; try solve [intuition].
Qed.

Fixpoint All2 {A B} (P : A -> B -> Prop) l1 l2 : Prop :=
  match l1, l2 with
  | [], [] => True
  | a :: l1', b :: l2' => P a b /\ All2 P l1' l2'
  | _, _ => False
  end.

Lemma All2_All {A B} (P : A -> B -> Prop) l1 l2 :
  All2 P l1 l2 ->
  All (fun p => P (fst p) (snd p)) (combine l1 l2).
Proof.
  revert l2.
  induction l1 as [|x l1 IH]; intros [|y l2]; simpl; try solve [intuition].
Qed.

Lemma All_All2 {A B} (P : A -> B -> Prop) l1 l2 :
  length l1 = length l2 ->
  All (fun p => P (fst p) (snd p)) (combine l1 l2) ->
  All2 P l1 l2.
Proof.
  revert l2.
  induction l1 as [|x l1 IH]; intros [|y l2]; simpl;
  try solve [intuition|discriminate].
Qed.

Lemma All2_length {A B} (P : A -> B -> Prop) l1 l2 :
  All2 P l1 l2 ->
  length l1 = length l2.
Proof.
  revert l2; induction l1 as [|x1 l1 IH]; intros [|x2 l2]; simpl; intuition.
Qed.

Definition secure_compartmentalization
  {interface: Type} {I : interface_language interface}
  {hcomponent hprogram: Type} {lcomponent lprogram: Type}

  (H: component_language I hcomponent hprogram)
  (L: component_language I lcomponent lprogram)
  (compile : hcomponent -> lcomponent) :=

  forall PIs AIs : list interface,
  comp (PIs ++ AIs) ->
  forall Ps : list hcomponent, All2 has_interface Ps PIs -> fully_defined Ps AIs ->
  forall Qs : list hcomponent, All2 has_interface Qs PIs -> fully_defined Qs AIs ->
  ((forall As : list hcomponent, All2 has_interface As AIs -> fully_defined As PIs ->
     link (Ps ++ As) ~~ link (Qs ++ As))
   <->
   (forall az : list lcomponent, All2 has_interface az AIs ->
         link (map compile Ps ++ az)
      ~~ link (map compile Qs ++ az))).

(* Instantiating structured full abstraction to obtain something
   stronger than secure compartmentalization (we will use this to
   prove sec. comp. for our compiler) *)

Coercion is_true : bool >-> Sortclass. (* Prop *)

Definition fst_map {A B C} (f : A -> C) (p : A * B) :=
  (f (fst p), snd p).

Fixpoint somes {A} (l : list (option A)) : list A :=
  match l with
  | [] => []
  | Some a :: l' => a :: somes l'
  | None :: l' => somes l'
  end.

Lemma somes_app {A} (l1 l2 : list (option A)) :
  somes (l1 ++ l2) = somes l1 ++ somes l2.
Proof.
  induction l1 as [|[x|] l1 IH]; simpl; trivial.
  now rewrite IH.
Qed.

Lemma somes_map {A} (l : list A) :
  somes (map Some l) = l.
Proof.
  now induction l as [|x l IH]; simpl; trivial; rewrite IH.
Qed.

Lemma somes_repeat_none {T} (n : nat) :
  somes (repeat (@None T) n) = [].
Proof. now induction n. Qed.

Definition isSome {A} (oa : option A) :=
  match oa with
  | Some _ => true
  | _ => false
  end.

Definition fsts {A B} (l : list (A * B)) : list A :=
  map fst l.

Definition snds {A B} (l : list (A * B)) : list B :=
  map snd l.

Lemma fsts_combine {A B} (l1 : list A) (l2 : list B) :
  length l1 = length l2 ->
  map fst (combine l1 l2) = l1.
Proof.
  revert l2.
  induction l1 as [|x l1 IH]; simpl; intros [|y l2];
  simpl; trivial; try discriminate.
  intros H. rewrite IH; congruence.
Qed.

Lemma snds_combine {A B} (l1 : list A) (l2 : list B) :
  length l1 = length l2 ->
  map snd (combine l1 l2) = l2.
Proof.
  revert l2.
  induction l1 as [|x l1 IH]; simpl; intros [|y l2];
  simpl; trivial; try discriminate.
  intros H. rewrite IH; congruence.
Qed.

Section SFAfromSC.

Context {interface: Type} {I: interface_language interface} {component program: Type}.
Context (compl : component_language I component program).

(* We take both SFA programs and contexts to be "partial programs"
   of the following type *)
Definition pprog := list (option (component*interface)).

Fixpoint merge (p:pprog) (q:pprog) : option pprog :=
  match p, q with
  | [], [] => Some []
  | Some (c,i1) :: p', None :: q'
  | None :: p', Some (c,i1) :: q' =>
    match merge p' q' with
    | Some pq' => Some (Some (c,i1) :: pq')
    | None     => None
    end
  | None :: p', None :: q' =>
    match merge p' q' with
    | Some pq' => Some (None :: pq')
    | None     => None
    end
  | _, _ => None
  end.

Lemma merge_app p1 p2 q1 q2 :
  length p1 = length q1 ->
  merge (p1 ++ p2) (q1 ++ q2) =
  match merge p1 q1, merge p2 q2 with
  | Some m1, Some m2 => Some (m1 ++ m2)
  | _, _ => None
  end.
Proof.
  revert q1.
  induction p1 as [|[[cp ip]|] p1 IH]; intros [|[[cq iq]|] q1];
  simpl; trivial; try discriminate.
  - now destruct (merge p2 q2).
  - intros H. inversion H as [H'].
    rewrite (IH _ H').
    destruct (merge p1 q1) as [m1|]; trivial.
    destruct (merge p2 q2) as [m2|]; trivial.
  - intros H. inversion H as [H'].
    rewrite (IH _ H').
    destruct (merge p1 q1) as [m1|]; trivial.
    destruct (merge p2 q2) as [m2|]; trivial.
  - intros H. inversion H as [H'].
    rewrite (IH _ H').
    destruct (merge p1 q1) as [m1|]; trivial.
    destruct (merge p2 q2) as [m2|]; trivial.
Qed.

Lemma merge_None_Some n Ps :
  n = length Ps ->
  merge (repeat None n) (map Some Ps) = Some (map Some Ps).
Proof.
  intros ->.
  induction Ps as [|[CP IP] Ps IH]; simpl; trivial.
  now rewrite IH.
Qed.

Lemma merge_Some_None n Ps :
  n = length Ps ->
  merge (map Some Ps) (repeat None n) = Some (map Some Ps).
Proof.
  intros ->.
  induction Ps as [|[CP IP] Ps IH]; simpl; trivial.
  now rewrite IH.
Qed.

Definition interfaces_ok (p:pprog) :=
  All (fun PIP => has_interface (fst PIP) (snd PIP)) (somes p).

Definition stat_eq (p q : pprog) :=
    All2 (fun optPIP optQIQ =>
            match optPIP, optQIQ with
            | None, None => True
            | Some (P, IP), Some (Q, IQ) =>
              has_interface P IP /\ has_interface Q IQ /\ IP = IQ
            | _, _ => False
            end) p q.

Lemma merge_stat_eq As Ps Qs CP :
  merge As Ps = Some CP ->
  stat_eq Ps Qs ->
  match merge As Qs with
  | Some CQ => snds (somes CP) = snds (somes CQ)
  | None => False
  end.
Proof.
revert Ps Qs CP. induction As as [|[[A IA]|] As IH].
- unfold stat_eq.
  intros [] []; simpl; try congruence.
  intros CP H. now inversion H.
- unfold stat_eq in *.
  intros [|[[P IP]|] Ps] [|[[Q IQ]|] Qs]; simpl; try solve [congruence|intuition].
  intros CP.
  destruct (merge As Ps) as [CP'|] eqn:mergeAsPs; try discriminate.
  intros H. inversion H; subst CP; clear H.
  intros [_ HPsQs].
  specialize (IH _ _ _ mergeAsPs HPsQs).
  destruct (merge As Qs) as [CQ'|] eqn:mergeAsQs; try solve [intuition].
  now simpl; rewrite IH.
- unfold stat_eq in *.
  intros [|[[P IP]|] Ps] [|[[Q IQ]|] Qs]; simpl; try solve [congruence|intuition].
  + intros CP.
    destruct (merge As Ps) as [CP'|] eqn:mergeAsPs; try discriminate.
    intros H. inversion H; subst CP; clear H.
    intros [[PI [QI H]] HPsQs]. subst IQ. rename IP into IPQ.
    specialize (IH _ _ _ mergeAsPs HPsQs).
    destruct (merge As Qs) as [CQ'|] eqn:mergeAsQs; try solve [intuition].
    now simpl; rewrite IH.
  + intros CP.
    destruct (merge As Ps) as [CP'|] eqn:mergeAsPs; try discriminate.
    intros H. inversion H; subst CP; clear H.
    intros [_ HPsQs].
    specialize (IH _ _ _ mergeAsPs HPsQs).
    destruct (merge As Qs) as [CQ'|] eqn:mergeAsQs; solve [intuition].
Qed.

Definition shape := list (bool*interface). (* replace (bool*interface) with option interface
                                              if dropping scl_program_has_shape *)
Definition flip_shape (s : shape): shape :=
  List.map (fun bi => match bi with (b, i) => (negb b, i) end) s.

Fixpoint context_has_shape (s : shape) (c : pprog) : Prop :=
  match s, c with
  | [], [] => true
  | (true,i1)::s', (Some (p,i2))::c' =>
    has_interface p i2 /\ i1 = i2 /\ context_has_shape s' c'
  | (false,_)::s', None::c' =>
    context_has_shape s' c'
  | _, _ => false
  end.

Lemma context_has_shape_app_l s1 s2 p :
  context_has_shape (s1 ++ s2) p ->
  exists p1 p2,
    p = p1 ++ p2
    /\ context_has_shape s1 p1
    /\ context_has_shape s2 p2.
Proof.
  revert p.
  induction s1 as [|[b i] s1 IH]; intros p.
  - exists [], p; simpl.
    repeat split; trivial.
  - destruct b, p as [|[[cp ip]|] p]; simpl; try discriminate.
    + intros [H1 [H2 H3]].
      specialize (IH _ H3).
      destruct IH as [p1 [p2 [Hp1 [Hp2 Hp3]]]].
      rewrite Hp1.
      exists (Some (cp, ip) :: p1), p2; repeat split; trivial.
    + intros H. specialize (IH _ H).
      destruct IH as [p1 [p2 [Hp1 [Hp2 Hp3]]]]. rewrite Hp1.
      exists (None :: p1), p2; repeat split; trivial.
Qed.

Lemma context_has_shape_app s1 s2 p1 p2 :
  context_has_shape s1 p1 ->
  context_has_shape s2 p2 ->
  context_has_shape (s1 ++ s2) (p1 ++ p2).
Proof.
  intros H1 H2. revert p1 H1.
  induction s1 as [|[[] i] s1 IH]; intros [|[[cp ip]|] ps]; simpl;
  try discriminate; trivial.
  - intros [H11 [H12 H13]].
    subst ip. repeat split; trivial.
    now eauto.
  - now eauto.
Qed.

Lemma context_has_shape_length s p :
  context_has_shape s p ->
  length s = length p.
Proof.
  revert p. induction s as [|[b i] s IH]; intros p; simpl.
  - destruct p; simpl; trivial; discriminate.
  - destruct b, p as [|[[cp ip]|] p]; simpl; try discriminate.
    + intros [_ [_ H]]. rewrite (IH _ H). reflexivity.
    + intros H. rewrite (IH _ H). reflexivity.
Qed.

Lemma context_has_shape_false PIs AA :
  context_has_shape (map (pair false) PIs) AA ->
  AA = repeat None (length PIs).
Proof.
  revert AA; induction PIs as [|Is PIs IH]; intros [|A AA]; simpl; trivial; try discriminate.
  destruct A; trivial; try discriminate.
  intros H. rewrite (IH _ H). reflexivity.
Qed.

Lemma context_has_shape_false_conv PIs n :
  n = length PIs ->
  context_has_shape (map (pair false) PIs) (repeat None n).
Proof.
  intros ->.
  induction PIs as [|PI PIs IH]; simpl; trivial; reflexivity.
Qed.

Lemma context_has_shape_true PIs AA :
  context_has_shape (map (pair true) PIs) AA ->
  exists AA',
    AA = map Some (combine AA' PIs)
    /\ All2 has_interface AA' PIs.
Proof.
  revert AA.
  induction PIs as [|PI PIs IH]; intros AA; simpl.
  - destruct AA; try discriminate.
    intros _. exists [].
    now split; simpl; trivial.
  - destruct AA as [|[[CA IA]|] AA]; try discriminate.
    intros [H1 [H2 H3]]. simpl.
    destruct (IH _ H3) as [AA' [? H4]].
    subst AA. exists (CA :: AA'). simpl. repeat split; congruence.
Qed.

Lemma context_has_shape_true_conv PIs AA :
  All2 has_interface AA PIs ->
  context_has_shape (map (pair true) PIs) (map Some (combine AA PIs)).
Proof.
  revert AA.
  induction PIs as [|PI PIs IH]; intros [|A AA]; simpl;
  try solve [intuition discriminate|reflexivity].
  intros [H1 H2]. repeat split; auto.
Qed.

Definition context_fully_defined (p : pprog) (s : shape) : Prop :=
  fully_defined (fsts (somes p))
                (somes (map (fun i : bool * _ => if fst i then None else Some (snd i)) s)).

Definition program_has_shape (s : shape) (p : pprog) : Prop :=
  context_has_shape (flip_shape s) p.

Definition program_fully_defined (p : pprog) (s : shape) : Prop :=
  context_fully_defined p (flip_shape s).

Lemma stat_eq_shape Ps Qs s :
  stat_eq Ps Qs ->
  program_has_shape s Ps ->
  program_has_shape s Qs.
Proof.
revert Qs s. unfold stat_eq.
induction Ps as [|[[P IP]|] Ps IH];
intros [|[[Q IQ]|] Qs] [|[b Is] s]; simpl; try solve [intuition].
- intros [[HPI [HQI HI]]]. subst IQ. rename IP into IPQ.
  unfold program_has_shape; simpl.
  destruct b; simpl; trivial.
  intros H [_ [HIs HPs]]. subst Is.
  repeat split; trivial.
  exact (IH _ _ H HPs).
- unfold program_has_shape; simpl.
  destruct b; simpl; trivial.
  intros [_ H] HPs.
  exact (IH _ _ H HPs).
Qed.

Definition insert c p :=
  match merge c p with
  | Some cp => cp
  | None => [] (* shouldn't happen *)
  end.

(* FIXME: It is a bit strange that [cl_compatible] is a property of [c]
   and [p] separately when in this instance we factor it through
   [merge c p]. This makes it very similar to how we use [cl_complete]
   in e.g. [cl_stat_eq_compatible_complete].  *)

Definition comp_compatible c p :=
  match merge c p with
  | Some cp => interfaces_ok cp /\ compatible (snds (somes cp))
  | None => False
  end.

Definition comp_complete (p : pprog) :=
  All isSome p /\ complete (snds (somes p)).

Lemma merge_shape As Ps :
  (match merge As Ps with
   | Some CP => interfaces_ok CP /\ All isSome CP
   | None => False
   end)
  <-> exists s : shape,
        context_has_shape s As /\
        program_has_shape s Ps.
Proof.
revert Ps. unfold interfaces_ok, program_has_shape.
induction As as [|[[A IA]|] As IH].
- intros [|[[P IP]|] Ps]; simpl; split; try solve [intuition].
  + intros _. exists []. simpl.
    now repeat split; trivial.
  + intros [s [HAs HPs]].
    destruct s as [|[[] ?] s]; simpl in *;
    solve [inversion HPs|inversion HAs].
  + intros [s [HAs HPs]].
    destruct s as [|[[] ?] s]; simpl in *;
    solve [inversion HPs|inversion HAs].
- intros [|[[P IP]|] Ps]; simpl; split; try solve [intuition].
  + intros [s [HAs HPs]].
    destruct s as [|[[] ?] s]; simpl in *;
    solve [inversion HPs|inversion HAs].
  + intros [s [HAs HPs]].
    destruct s as [|[[] ?] s]; simpl in *;
    solve [inversion HPs|inversion HAs].
  + specialize (IH Ps).
    destruct (merge As Ps) as [CP|] eqn:mergeAsPs; try solve [intuition].
    simpl.
    rewrite and_assoc, (and_comm true), <- (and_assoc _ _ true), IH.
    intros [HIA [[s [HAs HPs]] _]]. exists ((true, IA) :: s). simpl.
    repeat split; trivial.
  + specialize (IH Ps).
    intros [s [HAs HPs]].
    destruct s as [|[[] i] s]; simpl in *;
    try solve [inversion HPs|inversion HAs].
    destruct HAs as [HIA [? HAs]]. subst i.
    assert (H : exists s', context_has_shape s' As /\ context_has_shape (flip_shape s') Ps) by eauto.
    rewrite <- IH in H.
    destruct (merge As Ps) as [CP|] eqn:mergeAsPs; try solve [intuition].
    simpl. intuition (split; eauto).
- intros [|[[P IP]|] Ps]; simpl; split; try solve [intuition].
  + intros [s [HAs HPs]].
    destruct s as [|[[] ?] s]; simpl in *;
    solve [inversion HPs|inversion HAs].
  + specialize (IH Ps).
    destruct (merge As Ps) as [CP|] eqn:mergeAsPs; try solve [intuition].
    simpl.
    rewrite and_assoc, (and_comm true), <- (and_assoc _ _ true), IH.
    intros [HIA [[s [HAs HPs]] _]]. exists ((false, IP) :: s). simpl.
    repeat split; trivial.
  + specialize (IH Ps).
    intros [s [HAs HPs]].
    destruct s as [|[[] i] s]; simpl in *;
    try solve [inversion HPs|inversion HAs].
    destruct HPs as [HIP [? HPs]]. subst i.
    assert (H : exists s', context_has_shape s' As /\ context_has_shape (flip_shape s') Ps) by eauto.
    rewrite <- IH in H.
    destruct (merge As Ps) as [CP|] eqn:mergeAsPs; try solve [intuition].
    simpl. intuition (split; eauto).
  + specialize (IH Ps).
    destruct (merge As Ps) as [CP|] eqn:mergeAsPs; try solve [intuition].
    simpl. intros [_ [contra _]]. inversion contra.
  + intros [s [HAs HPs]].
    destruct s as [|[[] ?] s]; simpl in *;
    solve [inversion HPs|inversion HAs].
Qed.

Lemma stat_eq_compatible_complete Ps Qs :
  stat_eq Ps Qs ->
  forall A,
    comp_compatible A Ps /\ comp_complete (insert A Ps) ->
    comp_compatible A Qs /\ comp_complete (insert A Qs).
Proof.
  intros HPsQs As.
  unfold comp_compatible, comp_complete, insert in *.
  destruct (merge As Ps) as [CP|] eqn:mergeAsPs; try solve [intuition].
  intros [[int_CP comp_CP] [somes_AsPs compl_CP]].
  assert (H := merge_shape As Ps).
  rewrite mergeAsPs in H.
  assert (H' := conj int_CP somes_AsPs).
  rewrite H in H'. destruct H' as [s [HAs HPs]].
  assert (H' := stat_eq_shape Ps Qs s HPsQs HPs).
  assert (H'' : exists s, context_has_shape s As /\ program_has_shape s Qs) by eauto.
  rewrite <- merge_shape in H''.
  destruct (merge As Qs) as [CQ|] eqn:mergeAsQs; try solve [intuition].
  assert (H''' := merge_stat_eq _ _ _ _ mergeAsPs HPsQs).
  rewrite mergeAsQs in H'''.
  intuition congruence.
Qed.

Instance context_lang_from_component_lang :
     context_language pprog pprog :=
{
  cl_insert := insert;
  cl_compatible := comp_compatible;
  cl_complete := comp_complete;
  cl_stat_eq := stat_eq;
  cl_beh_eq p q := beh_eq (link (fsts (somes p))) (link (fsts (somes q)));
  cl_stat_eq_compatible_complete := @stat_eq_compatible_complete

}.

Instance structured_context_lang_from_component_lang :
     structured_context_language
       context_lang_from_component_lang
       shape :=
{
  scl_context_has_shape := context_has_shape;
  scl_context_fully_defined := context_fully_defined;
  scl_program_has_shape := program_has_shape;
  scl_program_fully_defined := program_fully_defined
}.

End SFAfromSC.

Section SFAimpliesSC.

Context {interface: Type} {I : interface_language interface}
        {hcomponent hprogram: Type} {lcomponent lprogram: Type}
        (H: component_language I hcomponent hprogram)
        (L: component_language I lcomponent lprogram)
        (compile : hcomponent -> lcomponent).

Hypothesis has_interface_compile :
  forall P PI, has_interface P PI ->
               has_interface (compile P) PI.

Definition map_compile (p:@pprog interface hcomponent) :
                          @pprog interface lcomponent :=
  List.map (option_map (fst_map compile)) p.

Theorem sfa_implies_sc :
  structured_full_abstraction
    (structured_context_lang_from_component_lang H)
    (structured_context_lang_from_component_lang L)
    map_compile ->
  secure_compartmentalization H L compile.
Proof.
intros Hsfa PIs AIs Hcomp Ps HPs HPs_def Qs HQs HQs_def.
assert (HPQ : @stat_eq _ _ _ _
                       H
                       (map Some (combine Ps PIs) ++ repeat None (length AIs))
                       (map Some (combine Qs PIs) ++ repeat None (length AIs))).
{ clear Hcomp Hsfa HPs_def HQs_def. revert Ps HPs Qs HQs. unfold stat_eq.
  induction PIs as [|PI PIs IH]; intros [|P Ps] HPs [|Q Qs] HQs; simpl in *; try solve [intuition].
  generalize (length AIs). clear AIs HPs HQs.
  now induction 0 as [|n IH]; split. }
assert (HPs' : @program_has_shape _ _ _ _
                                  H
                                  (map (pair false) PIs ++ map (pair true) AIs)
                                  (map Some (combine Ps PIs) ++ repeat None (length AIs))).
{ clear HPs_def Qs HQs HQs_def HPQ Hcomp.
  revert PIs HPs. unfold program_has_shape.
  induction Ps as [|P Ps IH]; simpl; intros PIs HPs.
  - destruct PIs as [|]; simpl in *; try solve [intuition].
    clear HPs.
    induction AIs as [|AI AIs IH]; simpl; trivial.
    reflexivity.
  - destruct PIs as [|PI PIs]; simpl in *; try solve [intuition]. }
assert (HPs_def' : @program_fully_defined _ _ _ _
                                          H
                                          (map Some (combine Ps PIs) ++ repeat None (length AIs))
                                          (map (pair false) PIs ++ map (pair true) AIs)).
{ unfold program_fully_defined, context_fully_defined.
  rewrite somes_app, somes_map, somes_repeat_none, app_nil_r.
  unfold fsts. rewrite fsts_combine; try now rewrite (All2_length _ _ _ HPs).
  unfold flip_shape. rewrite !map_app, !map_map. simpl.
  now rewrite somes_app, somes_map, map_const, somes_repeat_none. }
assert (HQs' : @program_has_shape _ _ _ _
                                  H
                                  (map (pair false) PIs ++ map (pair true) AIs)
                                  (map Some (combine Qs PIs) ++ repeat None (length AIs))).
{ clear Ps HPs HPs' HPs_def HPs_def' HQs_def HPQ Hcomp.
  revert PIs HQs. unfold program_has_shape.
  induction Qs as [|Q Qs IH]; simpl; intros PIs HQs.
  - destruct PIs as [|]; simpl in *; try solve [intuition].
    clear HQs.
    induction AIs as [|AI AIs IH]; simpl; trivial.
    reflexivity.
  - destruct PIs as [|PI PIs]; simpl in *; try solve [intuition]. }
assert (HQs_def' : @program_fully_defined _ _ _ _
                                          H
                                          (map Some (combine Qs PIs) ++ repeat None (length AIs))
                                          (map (pair false) PIs ++ map (pair true) AIs)).
{ unfold program_fully_defined, context_fully_defined.
  rewrite somes_app, somes_map, somes_repeat_none, app_nil_r.
  unfold fsts. rewrite fsts_combine; try now rewrite (All2_length _ _ _ HQs).
  unfold flip_shape. rewrite !map_app, !map_map. simpl.
  now rewrite somes_app, somes_map, map_const, somes_repeat_none. }
specialize (Hsfa (map (pair false) PIs ++ map (pair true) AIs)
                 (map Some (combine Ps PIs) ++ repeat None (length AIs))
                 (map Some (combine Qs PIs) ++ repeat None (length AIs))).
specialize (Hsfa HPQ HPs' HPs_def' HQs' HQs_def'). simpl in Hsfa.
clear HPs_def HPs_def' HQs_def HQs_def'.
simpl in *. split.
- (* -> *)
  intros Heq az Haz. destruct Hsfa as [Hsfa _].
  assert (H' : (forall AA : pprog,
                 context_has_shape H (map (pair false) PIs ++ map (pair true) AIs)
                                   AA ->
                 context_fully_defined H AA
                                       (map (pair false) PIs ++ map (pair true) AIs) ->
                 comp_compatible H AA
                                 (map Some (combine Ps PIs) ++ repeat None (length AIs)) ->
                 comp_complete
                   (insert AA
                           (map Some (combine Ps PIs) ++ repeat None (length AIs))) ->
                 link
                   (fsts
                      (somes
                         (insert AA
                                 (map Some (combine Ps PIs) ++ repeat None (length AIs))))) ~~
                   link
                   (fsts
                      (somes
                         (insert AA
                                 (map Some (combine Qs PIs) ++ repeat None (length AIs))))))).
  { intros AA Hshape Hdef HcompAA Hcompl.
    (* assert (Hdef' : @fully_defined _ _ _ _ H AA PIs).*)
    apply context_has_shape_app_l in Hshape.
    destruct Hshape as [AA1 [AA2 [E1 [E2 E3]]]].
    subst AA. unfold insert in *.
    assert (Hp : length AA1 = length (map Some (combine Ps PIs))).
    { rewrite map_length, combine_length, min_l.
      - rewrite <- (context_has_shape_length _ _ _ E2).
        rewrite map_length.
        now rewrite (All2_length _ _ _ HPs).
      - now rewrite (All2_length _ _ _ HPs). }
    rewrite (merge_app AA1 AA2 _ _ Hp) in *.
    assert (Hq : length AA1 = length (map Some (combine Qs PIs))).
    { rewrite map_length, combine_length, min_l.
      - rewrite <- (context_has_shape_length _ _ _ E2).
        rewrite map_length.
        now rewrite (All2_length _ _ _ HQs).
      - now rewrite (All2_length _ _ _ HQs). }
    rewrite (merge_app AA1 AA2 _ _ Hq) in *.
    pose proof (context_has_shape_false _ _ _ E2). subst AA1.
    rewrite merge_None_Some; try rewrite combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HPs); trivial.
    rewrite merge_None_Some; try rewrite combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HQs); trivial.
    apply context_has_shape_true in E3.
    destruct E3 as [AA2' [? HAA2']]. subst AA2.
    rewrite merge_Some_None; try rewrite combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HAA2'); trivial.
    rewrite !somes_app, !somes_map.
    unfold fsts. rewrite !map_app.
    unfold context_fully_defined, fsts in Hdef.
    rewrite !somes_app, somes_repeat_none, somes_map in Hdef.
    simpl in Hdef.
    rewrite fsts_combine in Hdef; try now rewrite (All2_length _ _ _ HAA2').
    rewrite map_app, !map_map in Hdef. simpl in Hdef.
    rewrite somes_app, somes_map, map_const, somes_repeat_none, app_nil_r in Hdef.
    rewrite !fsts_combine; trivial; eauto.
    - apply (All2_length _ _ _ HQs).
    - apply (All2_length _ _ _ HAA2').
    - apply (All2_length _ _ _ HPs). }
  specialize (Hsfa H'). clear H'.
  specialize (Hsfa (repeat None (length PIs) ++ map Some (combine az AIs))).
  unfold map_compile, insert, program_has_shape, flip_shape in *.
  rewrite !map_app, !map_map in *. simpl in *.
  rewrite map_repeat in Hsfa. simpl in *.
  repeat rewrite <- (map_map (fst_map compile) Some) in Hsfa.
  rewrite merge_app in Hsfa;
  try rewrite repeat_length, !map_length, combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HPs); trivial.
  rewrite merge_None_Some in Hsfa;
  try rewrite map_length, combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HPs); trivial.
  rewrite merge_Some_None in Hsfa;
  try rewrite combine_length, min_r; trivial; try rewrite (All2_length _ _ _ Haz); trivial.
  rewrite somes_app, !somes_map in Hsfa.
  rewrite merge_app in Hsfa;
  try rewrite repeat_length, !map_length, combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HQs); trivial.
  rewrite merge_None_Some in Hsfa;
  try rewrite map_length, combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HQs); trivial.
  rewrite merge_Some_None in Hsfa;
  try rewrite map_length; try rewrite combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ Haz); trivial.
  unfold fsts in Hsfa.
  rewrite map_app, fsts_combine in Hsfa; try rewrite (All2_length _ _ _ Haz); trivial.
  rewrite somes_app, !somes_map, map_app, !(map_map _ fst) in Hsfa.
  rewrite fsts_combine in Hsfa; try rewrite (All2_length _ _ _ Haz); trivial.
  rewrite (map_ext (fun x => fst (fst_map compile x)) (fun x => compile (fst x))) in Hsfa;
  try now intros [].
  rewrite (map_ext (fun x => fst (fst_map compile x)) (fun x => compile (fst x))) in Hsfa;
  try now intros [].
  rewrite <- !(map_map fst), !fsts_combine in Hsfa;
  try rewrite (All2_length _ _ _ HPs); trivial;
  try rewrite (All2_length _ _ _ HQs); trivial.
  apply Hsfa; clear Hsfa; simpl; repeat split.
  + apply context_has_shape_app.
    * now apply context_has_shape_false_conv.
    * now apply context_has_shape_true_conv.
  + unfold comp_compatible.
    rewrite merge_app; try rewrite repeat_length, !map_length, combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HPs); trivial.
    rewrite merge_None_Some; try rewrite map_length, combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HPs); trivial.
    rewrite merge_Some_None; try rewrite combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ Haz); trivial.
    split.
    * unfold interfaces_ok.
      rewrite somes_app, !somes_map.
      apply All_app; try now apply All2_All.
      clear - Ps PIs HPs has_interface_compile.
      revert PIs HPs.
      now induction Ps as [|P Ps IH]; intros [|PI PIs]; simpl; intuition.
    * unfold snds.
      rewrite somes_app, !somes_map, map_app, snds_combine;
      try rewrite (All2_length _ _ _ Haz); trivial.
      rewrite map_map.
      rewrite (map_ext (fun x => snd (fst_map compile x)) snd); try now intros [].
      rewrite snds_combine; try rewrite (All2_length _ _ _ HPs); trivial.
      exact (proj1 Hcomp).
  + rewrite <- map_app.
    generalize (map (fst_map compile) (combine Ps PIs) ++ combine az AIs).
    clear.
    now intros PIs; induction PIs as [|PI PIs IH]; simpl; intuition.
  + unfold snds.
    rewrite somes_app, !somes_map, map_app, snds_combine;
    try rewrite (All2_length _ _ _ Haz); trivial.
    rewrite map_map.
    rewrite (map_ext (fun x => snd (fst_map compile x)) snd); try now intros [].
    rewrite snds_combine; try rewrite (All2_length _ _ _ HPs); trivial.
    exact (proj2 Hcomp).
- (* <- *)
  intros Heq As HAs HAs_def. destruct Hsfa as [_ Hsfa].
  assert (H' : forall aa : pprog,
                 context_has_shape L (map (pair false) PIs ++ map (pair true) AIs)
                                   aa ->
                 comp_compatible L aa
                                 (map_compile
                                    (map Some (combine Ps PIs) ++ repeat None (length AIs))) ->
                 comp_complete
                   (insert aa
                           (map_compile
                              (map Some (combine Ps PIs) ++ repeat None (length AIs)))) ->
                 link
                   (fsts
                      (somes
                         (insert aa
                                 (map_compile
                                    (map Some (combine Ps PIs) ++
                                         repeat None (length AIs)))))) ~~
                   link
                   (fsts
                      (somes
                         (insert aa
                                 (map_compile
                                    (map Some (combine Qs PIs) ++
                                         repeat None (length AIs))))))).
  { clear Hsfa. intros az'.
    unfold fsts, insert, map_compile in *.
    rewrite !map_app, !map_repeat; simpl.
    rewrite !map_map; simpl.
    intros H1'.
    destruct (context_has_shape_app_l _ _ _ _ H1') as [az'' [az [H1 [H2 H3]]]].
    subst az'. clear H1'.
    apply context_has_shape_false in H2. subst az''.
    apply context_has_shape_true in H3.
    destruct H3 as [aa [? Haa]]. subst az.
    unfold comp_compatible, interfaces_ok.
    rewrite merge_app;
    try rewrite repeat_length, map_length, combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HPs); trivial.
    rewrite <- !(map_map (fst_map compile) Some).
    rewrite merge_None_Some; try rewrite map_length, combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HPs); trivial.
    rewrite merge_app;
    try rewrite repeat_length, !map_length, combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HQs); trivial.
    rewrite merge_Some_None; try rewrite map_length; try rewrite combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ Haa); trivial.
    rewrite merge_None_Some; try rewrite map_length, combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HQs); trivial.
    rewrite !somes_app, !somes_map, !map_app.
    rewrite fsts_combine; try now rewrite (All2_length _ _ _ Haa).
    rewrite !(map_map _ fst).
    rewrite (map_ext (fun x => fst (fst_map compile x)) (fun x => compile (fst x))); try now intros [].
    rewrite <- (map_map fst compile), fsts_combine; try now rewrite (All2_length _ _ _ HPs).
    rewrite (map_ext (fun x => fst (fst_map compile x)) (fun x => compile (fst x))); try now intros [].
    rewrite <- (map_map fst compile), fsts_combine; try now rewrite (All2_length _ _ _ HQs).
    intros [HAll _] _. apply Heq. clear Heq.
    apply All_All2; try now rewrite (All2_length _ _ _ Haa).
    now apply All_app_r in HAll. }
  specialize (Hsfa H'). clear H'.
  specialize (Hsfa (repeat None (length PIs) ++ map Some (combine As AIs))).
  unfold insert in *.
  rewrite merge_app in Hsfa;
  try rewrite repeat_length, map_length, combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HPs); trivial.
  rewrite merge_app in Hsfa;
  try rewrite repeat_length, map_length, combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HQs); trivial.
  rewrite merge_None_Some in Hsfa;
  try rewrite combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HPs); trivial.
  rewrite merge_None_Some in Hsfa;
  try rewrite combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HQs); trivial.
  rewrite merge_Some_None in Hsfa;
  try rewrite combine_length, min_r; trivial;
  try rewrite (All2_length _ _ _ HAs); trivial.
  rewrite !somes_app, !somes_map in Hsfa.
  unfold fsts, snds in *.
  rewrite !map_app in Hsfa.
  rewrite fsts_combine in Hsfa; try rewrite (All2_length _ _ _ HPs); trivial.
  rewrite fsts_combine in Hsfa; try rewrite (All2_length _ _ _ HAs); trivial.
  rewrite fsts_combine in Hsfa; try rewrite (All2_length _ _ _ HQs); trivial.
  apply Hsfa; clear Hsfa; simpl; repeat split.
  + apply context_has_shape_app.
    * now apply context_has_shape_false_conv.
    * now apply context_has_shape_true_conv.
  + unfold context_fully_defined, fsts.
    rewrite somes_app, somes_repeat_none, somes_map; simpl.
    rewrite fsts_combine; try now rewrite (All2_length _ _ _ HAs).
    rewrite map_app, !map_map, somes_app; simpl.
    now rewrite map_const, somes_map, somes_repeat_none, app_nil_r.
  + unfold comp_compatible.
    rewrite merge_app; try rewrite repeat_length, map_length, combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HPs); trivial.
    rewrite merge_None_Some; try rewrite combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HPs); trivial.
    rewrite merge_Some_None; try rewrite combine_length, min_r; trivial;
    try rewrite (All2_length _ _ _ HAs); trivial.
    unfold interfaces_ok. rewrite !somes_app, !somes_map.
    unfold snds. rewrite map_app.
    rewrite snds_combine; try now rewrite (All2_length _ _ _ HPs).
    rewrite snds_combine; try now rewrite (All2_length _ _ _ HAs).
    split; try exact (proj1 Hcomp).
    apply All_app; try now apply All2_All.
  + rewrite <- map_app.
    generalize (combine Ps PIs ++ combine As AIs).
    clear.
    now intros Ps; induction Ps as [|P Ps IH]; simpl; intuition.
  + unfold snds.
    rewrite somes_app, !somes_map, map_app, !snds_combine;
    try rewrite (All2_length _ _ _ HPs); trivial;
    try rewrite (All2_length _ _ _ HAs); trivial.
    exact (proj2 Hcomp).
Qed.

End SFAimpliesSC.


(*
 * Copyright (c) 2009-2011, Andrew Appel, Robert Dockins and Aquinas Hobor.
 *
 *)

Require Import VST.msl.base.
Require Import VST.msl.ageable.
Require Import Coq.Logic.Eqdep_dec.
Require Import VST.msl.functors.

Import CovariantFunctor.
Import CovariantFunctorLemmas.
Import CovariantFunctorGenerator.

Local Open Scope nat_scope.

Module Type TY_FUNCTOR.
  Parameter F : functor.

  Parameter T : Type.
  Parameter T_bot : T.

  Parameter other : Type.
End TY_FUNCTOR.

Module Type KNOT.
  Declare Module TF:TY_FUNCTOR.
  Import TF.

  Parameter knot : Type.

  Parameter ag_knot : ageable knot.
  Existing Instance ag_knot.
  Existing Instance ag_prod.

  Definition predicate := (knot * other) -> T.

  Parameter squash : (nat * F predicate) -> knot.
  Parameter unsquash : knot -> (nat * F predicate).

  Definition approx (n:nat) (p:predicate) : predicate :=
     fun w => if Compare_dec.le_gt_dec n (level w) then T_bot else p w.

  Axiom squash_unsquash : forall x, squash (unsquash x) = x.
  Axiom unsquash_squash : forall n x', unsquash (squash (n,x')) = (n,fmap F (approx n) x').

  Axiom knot_level : forall k:knot,
    level k = fst (unsquash k).

  Axiom knot_age1 : forall k:knot,
    age1 k =
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.

End KNOT.

Module Knot (TF':TY_FUNCTOR) : KNOT with Module TF:=TF'.
  Module TF := TF'.
  Import TF.

  (* Put the discrete pointed order on rhs *)
  Inductive leT : T -> T -> Prop :=
   | leT_refl : forall t, leT t t
   | leT_bot: forall t, leT T_bot t.

  Lemma leT_asym: forall t t',
    leT t t' -> leT t' t -> t = t'.
  Proof.
    intros.
    inversion H; subst; auto.
    inversion H0; subst; auto.
  Qed.

  Fixpoint sinv (n: nat) : Type :=
    match n with
      | O => unit
      | S n => prodT (sinv n) ((F (sinv n) * other) -> T)
    end.

  Fixpoint floor (m:nat) (n:nat) (p:sinv (m+n)) : sinv n :=
    match m as m' return forall (p : sinv (m'+n)), sinv n with
    | O => fun p => p
    | S m' => fun p => floor m' n (fst p)
    end p.

  Definition knot := { n:nat & F (sinv n) }.

  Definition predicate := knot * other -> T.

  Fixpoint stratify (n:nat) (Q:predicate) {struct n} : sinv n :=
    match n as n' return sinv n' with
    | O => tt
    | S n' => ( stratify n' Q, fun v => Q (existT (F oo sinv) n' (fst v),snd v) )
    end.

  Lemma decompose_nat : forall (x y:nat), { m:nat & y = (m + S x) } + { ge x y }.
  Proof.
    intros x y; revert x; induction y; simpl; intros.
    right; auto with arith.
    destruct (IHy x) as [[m H]|H].
    left; exists (S m); lia.
    destruct (Peano_dec.eq_nat_dec x y).
    left; exists O; lia.
    right; lia.
  Qed.

  Definition unstratify (n:nat) (p:sinv n) : predicate := fun w =>
    match w with (existT _ nw w',e) =>
      match decompose_nat nw n with
        | inleft (existT _ m Hm) => snd (floor m (S nw) (eq_rect  n _ p (m + S nw) Hm)) (w', e)
        | inright H => T_bot
      end
    end.

  Definition proof_irr_nat := eq_proofs_unicity Peano_dec.dec_eq_nat.
  Arguments proof_irr_nat [x] [y] _ _.

  Lemma floor_shuffle:
    forall (m1 n : nat)
      (p1 : sinv (m1 + S n)) (H1 : (m1 + S n) = (S m1 + n)),
      floor (S m1) n (eq_rect (m1 + S n) sinv p1 (S m1 + n) H1) = fst (floor m1 (S n) p1).
  Proof.
    intros.
    remember (fst (floor m1 (S n) p1)) as p.
    revert n p1 H1 p Heqp.
    induction m1; simpl; intros.
    replace H1 with (refl_equal (S n)) by (apply proof_irr_nat); simpl; auto.
    assert (m1 + S n = S m1 + n) by lia.
    destruct p1 as [p1 f'].
    generalize (IHm1 n p1 H p Heqp).
    simpl.
    clear.
    revert H1; generalize H.
    revert p1 f'.
    rewrite H.
    simpl; intros.
    replace H1 with (refl_equal (S (S (m1 + n)))) by (apply proof_irr_nat).
    simpl.
    replace H0 with (refl_equal (S (m1+n))) in H2 by (apply proof_irr_nat).
    simpl in H2.
    trivial.
  Qed.

  Lemma stratify_unstratify_more : forall n m1 m2 p1 p2,
    floor m1 n p1 = floor m2 n p2 ->

    (stratify n oo unstratify (m1+n)) p1 =
    (stratify n oo unstratify (m2+n)) p2.
  Proof.
    unfold compose; induction n; simpl; intros; auto.
    apply injective_projections; simpl; trivial.

    assert ((m1 + S n) =  (S m1 + n)) by lia.
    assert ((m2 + S n) =  (S m2 + n)) by lia.
    assert (floor (S m1) n (eq_rect (m1 + S n) _ p1 _ H0) = floor (S m2) n (eq_rect (m2 + S n) _ p2 _ H1)).
    do 2 rewrite floor_shuffle.
    congruence.
    generalize (IHn (S m1) (S m2) _ _ H2).
    clear.
    generalize H0 H1.
    revert p1 p2.
    rewrite H0; clear H0.
    rewrite H1; clear H1.
    intros p1 p2 H1 H2.
    replace H1 with (refl_equal (S m1 + n)) by (apply proof_irr_nat).
    replace H2 with (refl_equal (S m2 + n)) by (apply proof_irr_nat).
    simpl; auto.

    apply extensionality; intro v.
    unfold unstratify.
    destruct (decompose_nat n (m2 + S n)) as [[r Hr]|Hr].
    2: elimtype False; lia.
    destruct (decompose_nat n (m1 + S n)) as [[s Hs]|Hs].
    2: elimtype False; lia.
    assert (m2 = r) by lia; subst r.
    assert (m1 = s) by lia; subst s.
    simpl.
    replace Hr with (refl_equal (m2 + S n)) by (apply proof_irr_nat).
    replace Hs with (refl_equal (m1 + S n)) by (apply proof_irr_nat).
    simpl.
    rewrite H; auto.
  Qed.

  Lemma stratify_unstratify : forall n,
         stratify n oo unstratify n = id (sinv n).
  Proof.
    unfold id, compose; intro n; extensionality p; revert n p.
    induction n.

    intros; destruct p; auto.

    simpl; intros [p f].
    apply injective_projections; simpl; trivial.

    replace (stratify n (unstratify (S n) (p,f))) with
            (stratify n (unstratify n p)); auto.
    replace (stratify n (unstratify n p)) with
      ((stratify n oo unstratify (0+n)) p) by trivial.
    rewrite (stratify_unstratify_more _ 0 1 p (p,f)); trivial.

    extensionality v.

    destruct (decompose_nat n (S n)) as [[r Hr]|?]; auto.
    assert (r = O) by lia; subst r.
    simpl in *.
    replace Hr with (refl_equal (S n)) by (apply proof_irr_nat); simpl; auto.
    destruct v; auto.

    elimtype False.
    lia.
  Qed.

  Lemma unstratify_stratify1 : forall n (p:predicate) w,
    leT ((unstratify n oo stratify n) p w) (p w).
  Proof.
    unfold compose; induction n; simpl; intros; unfold unstratify.

    (* 0 case *)
    destruct w as [nw rm]; simpl.
    destruct nw as [nw e].
    destruct (decompose_nat nw O) as [[r Hr]|?].
    elimtype False; lia.
    apply leT_bot.

    (* S n case *)
    case_eq w; intros nw rm Hrm.
    destruct nw as [nw e].
    destruct (decompose_nat nw (S n)) as [[r Hr]|?]; try (apply lt_rhs_top).
    destruct r; simpl.

    assert (n = nw) by lia.
    subst nw.
    simpl in Hr.
    replace Hr with (refl_equal (S n)) by apply proof_irr_nat; simpl.
    unfold compose.
    destruct w.
    apply leT_refl.

    simpl in Hr.
    assert (n = r + S nw) by lia.
    revert Hr; subst n.
    intro Hr.
    replace Hr with (refl_equal (S (r+S nw))) by apply proof_irr_nat; simpl.
    clear Hr.

    generalize (IHn p w).
    unfold unstratify.
    rewrite Hrm.
    destruct (decompose_nat nw (r + S nw)) as [[x Hx]|?].
    assert (x = r) by lia; subst x.
    replace Hx with (refl_equal (r + S nw)) by apply proof_irr_nat.
    simpl; auto.
    elimtype False; lia.
    apply leT_bot.
  Qed.

  Lemma unstratify_stratify2 : forall n p w,
     projT1 (fst w) < n ->
        leT (p w) ((unstratify n oo stratify n) p w).
  Proof.
    unfold compose.
    induction n; simpl; intros.

    (* 0 case *)
    inversion H.

    (* S n case *)
    unfold unstratify.
    case_eq w; intros [m rm] e Hw.
    assert (projT1 (fst w) = m).
    rewrite Hw; auto.

    destruct (decompose_nat m (S n)) as [[r Hr]|?].
    destruct r; simpl.

    assert (n = m) by lia.
    move H0 after H1.
    subst m. fold sinv. simpl in Hr. rewrite (proof_irr_nat Hr (refl_equal _)). clear Hr.
    simpl.
    unfold compose.
    rewrite <- Hw.
    apply leT_refl.

    simpl in Hr.
    assert (n = r + S m) by lia.
    revert Hr; subst n.
    intro Hr.
    replace Hr with (refl_equal (S (r+S m))) by apply proof_irr_nat; simpl.
    clear Hr.
    rewrite H0 in H.
    assert (m < (r +  S m)) by lia.
    specialize ( IHn p w).
    rewrite H0 in IHn.
    specialize ( IHn H1).
    revert IHn.
    unfold unstratify.
    rewrite Hw.
    destruct (decompose_nat m (r + S m)) as [[x Hx]|?].
    assert (x = r) by lia; subst x.
    replace Hx with (refl_equal (r + S m)) by apply proof_irr_nat.
    simpl; auto.
    elimtype False; lia.
    elimtype False; lia.
  Qed.

  Lemma unstratify_stratify3 : forall n (p:predicate) w,
    projT1 (fst w) >= n -> leT ((unstratify n oo stratify n) p w) T_bot.
  Proof.
    unfold compose, unstratify; intros n p w H.
    case_eq w; intros [wn rm] e.
    intro Hrm.
    rewrite Hrm in H; simpl in H.
    destruct (decompose_nat wn n) as [[r Hr]|?].
    elimtype False; lia.
    apply leT_bot.
  Qed.

  Definition squash (x:nat * F predicate) : knot :=
    match x with (n,y) => existT (F oo sinv) n (fmap F (stratify n) y) end.

  Definition unsquash (x:knot) : (nat * F predicate) :=
    match x with existT _ n y => (n, fmap F (unstratify n) y) end.

  Definition def_knot_level (k:knot) := fst (unsquash k).

  Definition def_knot_age1 (k:knot) : option knot :=
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.

  Definition def_knot_unage (k:knot) : knot :=
    match unsquash k with
    | (n,x) => squash (S n,x)
    end.

  Definition approx (n:nat) (p:predicate) : predicate :=
     fun w => if Compare_dec.le_gt_dec n (def_knot_level (fst w)) then T_bot else p w.

  Lemma squash_unsquash : forall x, squash (unsquash x) = x.
  Proof.
    intros; destruct x as [x f]; simpl.
    unfold compose.
    replace (fmap F (stratify x) (fmap F (unstratify x) f)) with
      ((fmap F (stratify x) oo fmap F (unstratify x)) f) by trivial.
    rewrite fmap_comp.
    replace (stratify x oo unstratify x) with (id (sinv x)).
    rewrite fmap_id; simpl; auto.
    unfold compose.
    extensionality z.
    replace (stratify x (unstratify x z)) with ((stratify x oo unstratify x) z) by trivial.
    rewrite stratify_unstratify; auto.
  Qed.

  Lemma unsquash_squash : forall n x', unsquash (squash (n,x')) = (n,fmap F (approx n) x').
  Proof.
    intros.
    simpl.
    replace (fmap F  (unstratify n) (fmap F (stratify n) x')) with
      ((fmap F (unstratify n) oo fmap F (stratify n)) x') by trivial.
    rewrite fmap_comp.
    apply injective_projections; simpl; trivial.
    replace (unstratify n oo stratify n) with (approx n); auto.
    extensionality p z.
    apply leT_asym.

    intuition.
    case (Compare_dec.le_gt_dec n (def_knot_level a)); intro.
    replace (approx n p (a, b)) with T_bot.
    apply leT_bot.
    unfold approx.
    simpl.
    case (Compare_dec.le_gt_dec n (def_knot_level a)); intro.
    trivial.
    elimtype False.
    lia.
    replace (approx n p (a,b)) with (p (a,b)).
    apply unstratify_stratify2.
    simpl.
    destruct a.
    unfold level in g.
    simpl in *.
    auto.
    unfold approx.
    simpl.
    case (Compare_dec.le_gt_dec n (def_knot_level a)); intro.
    elimtype False.
    lia.
    trivial.

    intuition.
    destruct (Compare_dec.le_lt_dec n (def_knot_level a)).
    replace (approx n p (a, b)) with T_bot.
    apply unstratify_stratify3; auto.
    simpl.
    destruct a.
    unfold level in l.
    simpl in *.
    auto.
    unfold approx.
    simpl.
    case (Compare_dec.le_gt_dec n (def_knot_level a)); auto.
    intro.
    elimtype False.
    lia.
    replace (approx n p (a, b)) with (p (a, b)).
    apply unstratify_stratify1; auto.
    unfold approx.
    simpl.
    case (Compare_dec.le_gt_dec n (def_knot_level a)); auto.
    intro.
    elimtype False.
    lia.
  Qed.

  Lemma unsquash_inj : forall k1 k2,
    unsquash k1 = unsquash k2 ->
    k1 = k2.
  Proof.
    intros.
    rewrite <- (squash_unsquash k1).
    rewrite <- (squash_unsquash k2).
    rewrite H.
    trivial.
  Qed.
  Arguments unsquash_inj [k1 k2] _.

  Lemma ag_knot_facts : ageable_facts knot def_knot_level def_knot_age1.
  Proof.
    constructor.

   unfold def_knot_age1; unfold def_knot_level; simpl; intros x'.
    case_eq (unsquash x'); intros.
    destruct x' as [n' xx']. simpl in *. inv H.
    exists (squash (S n, fmap F (unstratify n) xx')).
    rewrite unsquash_squash.
    f_equal.
    f_equal.
    transitivity ((fmap F (stratify n) oo fmap F (approx (S n)) oo fmap F (unstratify n)) xx'); auto.
    do 2 rewrite fmap_comp.
    replace (stratify n oo approx (S n) oo unstratify n) with (@id (sinv n)).
    rewrite fmap_id. auto.
   clear.
    rewrite <- (stratify_unstratify n).
    f_equal. extensionality  a w.
    unfold approx, compose. destruct w.
    simpl fst.
    destruct (Compare_dec.le_gt_dec (S n) (def_knot_level k)); auto.
    destruct k. simpl in *.
    unfold def_knot_level in l.
    simpl in *.
    destruct (decompose_nat x n); auto.
    destruct s. elimtype False.
    lia.

    intros.
    unfold def_knot_age1, def_knot_level.
    destruct (unsquash x); simpl.
    destruct n; intuition; try discriminate.

    unfold def_knot_age1, def_knot_level; intros.
    destruct (unsquash x).
    destruct n; inv H; simpl; auto.
  Qed.

  Definition ag_knot : ageable knot :=
    mkAgeable knot def_knot_level def_knot_age1 ag_knot_facts .
  Existing Instance ag_knot.
  Existing Instance ag_prod.

  Lemma knot_level : forall k:knot,
    level k = fst (unsquash k).
  Proof (fun k => refl_equal (def_knot_level k)).

  Lemma knot_age1 : forall k:knot,
    age1 k =
    match unsquash k with
    | (O,_) => None
    | (S n,x) => Some (squash (n,x))
    end.
  Proof (fun k => refl_equal (def_knot_age1 k)).

End Knot.


Require Import Burrow.trees.
From stdpp Require Import gmap.
From iris.prelude Require Import options.

Require Import cpdt.CpdtTactics.

Section GmapTPCM.

Context {K} `{!EqDecision K} `{!Countable K}.
Context {V} `{!EqDecision V}.

Definition gmap_valid (p: gmap K (option V)) : Prop :=
  ∀ k , match p !! k with Some None => False | _ => True end.

Definition gmerge (a b : option (option V)) : option (option V) :=
  match a, b with
  | None, None => None
  | Some x, None => Some x
  | None, Some y => Some y
  | Some x, Some y => Some None
  end.

Definition gmap_dot (a b: gmap K (option V)) : gmap K (option V) :=
  merge gmerge a b.

Definition gmap_mov (a b: gmap K (option V)) : Prop :=
  ∀ z , gmap_valid (gmap_dot a z) -> gmap_valid (gmap_dot b z).

Lemma gmap_dot_comm x y
  : gmap_dot x y = gmap_dot y x.
Proof.
intros. unfold gmap_dot, gmerge. apply map_eq. intro. rewrite lookup_merge.
    rewrite lookup_merge. unfold diag_None. destruct (x !! i), (y !! i); trivial.
Qed.

Lemma gmap_dot_assoc x y z
  : gmap_dot x (gmap_dot y z) = gmap_dot (gmap_dot x y) z.
Proof. intros. unfold gmap_dot, gmerge. apply map_eq. intro. rewrite lookup_merge.
    rewrite lookup_merge.
    rewrite lookup_merge.
    unfold diag_None.
    rewrite lookup_merge.
    unfold diag_None.
    destruct (x !! i), (y !! i), (z !! i); trivial.
Qed.

Lemma gmap_dot_empty
  : ∀ x : gmap K (option V), gmap_dot x ∅ = x.
Proof.
intros. unfold gmap_dot. apply map_eq. intro. rewrite lookup_merge. rewrite lookup_empty.
    unfold diag_None, gmerge. destruct (x !! i); trivial.
Qed.

Lemma gmap_dot_empty_left
  : ∀ x : gmap K (option V), gmap_dot ∅ x = x.
Proof.
intros. unfold gmap_dot. apply map_eq. intro. rewrite lookup_merge. rewrite lookup_empty.
    unfold diag_None, gmerge. destruct (x !! i); trivial.
Qed.

Lemma lookup_gmap_dot_left a b k
  : gmap_valid (gmap_dot a b) -> (a !! k ≠ None) -> (gmap_dot a b) !! k = a !! k.
Proof.
  unfold gmap_valid, gmap_dot. intros. have j := H k. rewrite lookup_merge.
  rewrite lookup_merge in j. unfold diag_None, gmerge in *. destruct (a !! k), (b !! k);
  trivial; contradiction.
Qed.

Lemma lookup_gmap_dot_right a b k
  : gmap_valid (gmap_dot a b) -> (b !! k ≠ None) -> (gmap_dot a b) !! k = b !! k.
Proof.
  unfold gmap_valid, gmap_dot. intros. have j := H k. rewrite lookup_merge.
  rewrite lookup_merge in j. unfold diag_None, gmerge in *. destruct (a !! k), (b !! k);
  trivial; contradiction.
Qed.

Lemma lookup_gmap_dot_3mid a b c k
  : gmap_valid (gmap_dot (gmap_dot a b) c) -> (b !! k ≠ None) ->
      gmap_dot (gmap_dot a b) c !! k = b !! k.
Proof.
  intros.
  rewrite gmap_dot_comm in H.
  rewrite gmap_dot_assoc in H.
  rewrite gmap_dot_comm.
  rewrite gmap_dot_assoc.
  apply lookup_gmap_dot_right; trivial.
Qed.

Lemma lookup_gmap_dot_3left a b c k
  : gmap_valid (gmap_dot (gmap_dot a b) c) -> (a !! k ≠ None) ->
      gmap_dot (gmap_dot a b) c !! k = a !! k.
Proof.
  intros.
  rewrite <- gmap_dot_assoc in H.
  rewrite <- gmap_dot_assoc.
  apply lookup_gmap_dot_left; trivial.
Qed.

Lemma lookup_gmap_dot_3right a b c k
  : gmap_valid (gmap_dot (gmap_dot a b) c) -> (c !! k ≠ None) ->
      gmap_dot (gmap_dot a b) c !! k = c !! k.
Proof.
  apply lookup_gmap_dot_right.
Qed.

Lemma gmap_valid_left
  : ∀ x y : gmap K (option V), gmap_valid (gmap_dot x y) → gmap_valid x.
Proof.
  intros. unfold gmap_valid, gmap_dot in *.
    intro. have h := H k. clear H. rewrite lookup_merge in h. unfold diag_None in h.
        unfold gmerge in h.
        destruct (x !! k); trivial.
        destruct (y !! k); trivial. contradiction.
Qed.

Lemma gmap_valid_right
  : ∀ x y : gmap K (option V), gmap_valid (gmap_dot x y) → gmap_valid y.
Proof.
  intro. intro. rewrite gmap_dot_comm. apply gmap_valid_left.
Qed.

Lemma gmap_valid_update_singleton j x y (m: gmap K (option V)) :
  gmap_valid (gmap_dot {[j := Some x]} m) ->
  gmap_valid (gmap_dot {[j := Some y]} m).
Proof.
  intros. unfold gmap_valid, gmap_dot in *. intro.
  have r := H k. rewrite lookup_merge. rewrite lookup_merge in r.
  unfold gmerge, diag_None in *.
  have h : Decision (j = k) by solve_decision. destruct h.
  - subst k. rewrite lookup_singleton. rewrite lookup_singleton in r.
      destruct (m !! j); trivial.
  - rewrite lookup_singleton_ne; trivial.
    rewrite lookup_singleton_ne in r; trivial.
Qed.

#[refine]
Global Instance gmap_tpcm : TPCM (gmap K (option V)) := {
  m_valid p := gmap_valid p ;
  dot a b := gmap_dot a b ;
  mov a b := gmap_mov a b ;
  unit := empty ;
}.
 - apply gmap_valid_left.
 - unfold gmap_valid. intros. rewrite lookup_empty. trivial.
 - apply gmap_dot_empty.
 - intros. apply gmap_dot_comm.
 - intros. apply gmap_dot_assoc.
  - intros. unfold gmap_mov. intro. trivial.
  - intros. unfold gmap_mov in *. intros. apply H0. apply H. trivial.
  - intros. split.
    * unfold gmap_mov in H. apply H. apply H0.
    * unfold gmap_mov in H. unfold gmap_mov. intro.
        rewrite <- gmap_dot_assoc.
        rewrite <- gmap_dot_assoc.
        apply H.
Defined.

Lemma le_of_subset (a b : gmap K (option V))
  (f: ∀ k v , a !! k = Some v -> b !! k = Some v) : tpcm_le a b.
Proof.
  unfold tpcm_le.
  assert (∀ x : K * option V, Decision (match x with (k,v) => a !! k = None end)) as X
    by solve_decision.
  exists (map_filter (λ x , match x with (k,v) => a !! k = None end) X b).
  unfold dot, gmap_tpcm, gmap_dot. apply map_eq. intro.
  have ff := f i.
  rewrite lookup_merge. unfold diag_None, gmerge.
  destruct (a !! i) eqn:ai.
  - rewrite map_filter_lookup.
    unfold "≫=", option_bind. destruct (b!!i) eqn:bi.
    + unfold mguard, option_guard. have fff := ff o. intuition. inversion H.
      destruct (X (i, o)).
      * rewrite ai in e. discriminate.
      * trivial.
    + have fff := ff o. intuition.
  - rewrite map_filter_lookup.
    unfold "≫=", option_bind. destruct (b!!i) eqn:bi; trivial.
    unfold mguard, option_guard. destruct (X (i, o)); trivial.
    contradiction.
Qed.

End GmapTPCM.


(**
Conditions to lift a biadjunction to a displayed biadjunction to a dependent product.
*)
Require Import UniMath.Foundations.All.
Require Import UniMath.MoreFoundations.All.
Require Import UniMath.CategoryTheory.Core.Categories.
Require Import UniMath.CategoryTheory.Core.Functors.
Require Import UniMath.CategoryTheory.DisplayedCats.Core.
Require Import UniMath.Bicategories.Core.Bicat. Import Bicat.Notations.
Require Import UniMath.Bicategories.Core.Invertible_2cells.
Require Import UniMath.Bicategories.Core.Univalence.
Require Import UniMath.Bicategories.Core.BicategoryLaws.
Require Import UniMath.Bicategories.Morphisms.Adjunctions.
Require Import UniMath.Bicategories.PseudoFunctors.Display.Base.
Require Import UniMath.Bicategories.PseudoFunctors.Display.Map1Cells.
Require Import UniMath.Bicategories.PseudoFunctors.Display.Map2Cells.
Require Import UniMath.Bicategories.PseudoFunctors.Display.Identitor.
Require Import UniMath.Bicategories.PseudoFunctors.Display.Compositor.
Require Import UniMath.Bicategories.PseudoFunctors.Display.PseudoFunctorBicat.
Require Import UniMath.Bicategories.PseudoFunctors.PseudoFunctor.
Import PseudoFunctor.Notations.
Require Import UniMath.Bicategories.PseudoFunctors.Biadjunction.
Require Import UniMath.Bicategories.PseudoFunctors.Examples.Identity.
Require Import UniMath.Bicategories.PseudoFunctors.Examples.Composition.
Require Import UniMath.Bicategories.Transformations.PseudoTransformation.
Require Import UniMath.Bicategories.Transformations.Examples.Whiskering.
Require Import UniMath.Bicategories.Transformations.Examples.Unitality.
Require Import UniMath.Bicategories.Transformations.Examples.Associativity.
Require Import UniMath.Bicategories.Modifications.Modification.
Require Import UniMath.Bicategories.DisplayedBicats.DispBicat.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.DisplayedCatToBicat.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.Add2Cell.
Require Import UniMath.Bicategories.DisplayedBicats.DispPseudofunctor.
Require Import UniMath.Bicategories.DisplayedBicats.DispTransformation.
Require Import UniMath.Bicategories.DisplayedBicats.DispModification.
Require Import UniMath.Bicategories.DisplayedBicats.DispBuilders.
Require Import UniMath.Bicategories.DisplayedBicats.DispBiadjunction.
Require Import UniMath.Bicategories.DisplayedBicats.Examples.DispDepProd.

Local Open Scope cat.

Section LiftPseudofunctor.
  Context {B₁ B₂ : bicat}
          (F : psfunctor B₁ B₂)
          {I : UU}
          {D₁ : I → disp_bicat B₁}
          {D₂ : I → disp_bicat B₂}
          (FF : ∏ (i : I), disp_psfunctor (D₁ i) (D₂ i) F)
          (D₂_locally_prop : ∏ (i : I), disp_2cells_isaprop (D₂ i))
          (D₂_locally_groupoid : ∏ (i : I), disp_locally_groupoid (D₂ i)).

  Definition disp_depprod_psfunctor
    : disp_psfunctor (disp_depprod_bicat I D₁) (disp_depprod_bicat I D₂) F.
  Proof.
    use make_disp_psfunctor.
    - apply disp_2cells_isaprop_depprod.
      exact D₂_locally_prop.
    - apply disp_locally_groupoid_depprod.
      exact D₂_locally_groupoid.
    - exact (λ x xx i, FF i x (xx i)).
    - exact (λ x y f xx yy ff i, pr121 (FF i) _ _ _ _ _ (ff i)).
    - exact (λ x y f g α xx yy ff gg αα i, pr1 (pr221 (FF i)) _ _ _ _ _ _ _ _ _ (αα i)).
    - refine (λ x xx, _).
      exact (λ i, pr12 (pr221 (FF i)) _ (xx i)).
    - refine (λ x y z f g xx yy zz ff gg, _).
      exact (λ i, pr22 (pr221 (FF i)) _ _ _ _ _ _ _ _ (ff i) (gg i)).
  Defined.
End LiftPseudofunctor.

Section LiftBiadjunction.
  Context {B₁ B₂ : bicat}
          {L : psfunctor B₁ B₂}
          (R : left_biadj_data L)
          {I : UU}
          {D₁ : I → disp_bicat B₁}
          {D₂ : I → disp_bicat B₂}
          (LL : ∏ (i : I), disp_psfunctor (D₁ i) (D₂ i) L)
          (RR : ∏ (i : I), disp_left_biadj_data
                             _ _
                             R
                             (LL i))
          (D₁_locally_prop : ∏ (i : I), disp_2cells_isaprop (D₁ i))
          (D₁_locally_groupoid : ∏ (i : I), disp_locally_groupoid (D₁ i))
          (D₂_locally_prop : ∏ (i : I), disp_2cells_isaprop (D₂ i))
          (D₂_locally_groupoid : ∏ (i : I), disp_locally_groupoid (D₂ i)).

  Definition disp_depprod_biadjunction
    : disp_left_biadj_data
        _ _
        R
        (disp_depprod_psfunctor L LL D₂_locally_prop D₂_locally_groupoid).
  Proof.
    use tpair.
    - use tpair.
      + exact (disp_depprod_psfunctor
                 R (λ i, pr11 (RR i))
                 D₁_locally_prop D₁_locally_groupoid).
      + split.
        * use make_disp_pstrans.
          ** apply disp_2cells_isaprop_depprod.
             exact D₁_locally_prop.
          ** apply disp_locally_groupoid_depprod.
             exact D₁_locally_groupoid.
          ** exact (λ x xx i, pr121 (RR i) _ (xx i)).
          ** exact (λ x y f xx yy ff i, pr21 (pr121 (RR i)) _ _ _ _ _ (ff i)).
        * use make_disp_pstrans.
          ** apply disp_2cells_isaprop_depprod.
             exact D₂_locally_prop.
          ** apply disp_locally_groupoid_depprod.
             exact D₂_locally_groupoid.
          ** exact (λ x xx i, pr221 (RR i) _ (xx i)).
          ** exact (λ x y f xx yy ff i, pr21 (pr221 (RR i)) _ _ _ _ _ (ff i)).
    - split.
      + use make_disp_invmodification.
        * apply disp_2cells_isaprop_depprod.
          exact D₂_locally_prop.
        * apply disp_locally_groupoid_depprod.
          exact D₂_locally_groupoid.
        * exact (λ x xx i, pr1 (pr12 (RR i)) _ (xx i)).
      + use make_disp_invmodification.
        * apply disp_2cells_isaprop_depprod.
          exact D₁_locally_prop.
        * apply disp_locally_groupoid_depprod.
          exact D₁_locally_groupoid.
        * exact (λ x xx i, pr1 (pr22 (RR i)) _ (xx i)).
  Defined.
End LiftBiadjunction.


Require Export BisimRefinement.TraceMonad.
Require Export BisimRefinement.BisimPerm.

Require Export Coq.Arith.PeanoNat.
Require Export Coq.Lists.List.

(***
 *** Machine Values
 ***)

(* Values in our simple machine are just nats *)
Inductive SMValue := | mkSMValue (n:nat).

(* The "undefined" value *)
(* Definition undefValue := mkSMValue 0. *)

(* Values are related by the built-in Coq intensional equality *)
Instance LR_SMValue : LR SMValue := LR_eq.

(* Test equality of values *)
Definition val_eqb (v1 v2 : SMValue) : bool :=
  match v1, v2 with
    mkSMValue n1, mkSMValue n2 => Nat.eqb n1 n2
  end.

(* Test if a value is "true", represented here as not equal to 0 *)
Definition val_true (v: SMValue) : bool :=
  match v with
    mkSMValue n => negb (Nat.eqb n 0)
  end.

(* The lens for viewing a value as a nat *)
Program Definition SMValue_nat_lens : Lens SMValue nat :=
  {| getL v := match v with mkSMValue n => n end;
     putL n v := mkSMValue n; |}.
Next Obligation.
  intros [ n1 ] [ n2 ] eq_v. destruct eq_v. inversion H. reflexivity.
Defined.
Next Obligation.
  intros n1 n2 [ eq_n eq_n' ] v1 v2 eq_v. rewrite eq_n. reflexivity.
Defined.
Next Obligation.
  destruct a. reflexivity.
Defined.
Next Obligation.
  reflexivity.
Defined.
Next Obligation.
  reflexivity.
Defined.

(* A tuple of 0 or more values *)
Definition SMValues n := NTuple SMValue n.

Instance LR_SMValues n : LR (SMValues n) := Build_LR _ _ _.

(***
 *** Register Sets
 ***)

(* A set of SSA register values *)
Definition SMRegs := list SMValue.

Instance LR_SMRegs : LR SMRegs := LR_eq.

(* Read a register value *)
Definition readReg (regs:SMRegs) n : option SMValue := nth_error regs n.

(* Push a new register value *)
Definition pushReg (regs:SMRegs) val : SMRegs := val :: regs.

(* Push a list of new register values *)
Fixpoint pushRegs (regs:SMRegs) n : SMValues n -> SMRegs :=
  match n return SMValues n -> SMRegs with
  | 0 => fun _ => regs
  | S n' => fun vals => pushRegs (pushReg regs (fst vals)) n' (snd vals)
  end.

(***
 *** Data Memories
 ***)

(* A data memory = a map from "pointers" (represented by nats) to values (which
are also nats) *)
Inductive DataMem := mkDataMem (vmap: nat -> SMValue).

(* Data memories are related by extensional functional equivalence *)
Instance LR_DataMem : LR DataMem :=
  { lr_leq := fun dmem1 dmem2 =>
                match dmem1, dmem2 with
                | mkDataMem vmap1, mkDataMem vmap2 => vmap1 <lr= vmap2
                end;
  }.
Proof.
  constructor.
  - intros [ vmap ]; reflexivity.
  - intros [ vmap1 ] [ vmap2 ] [ vmap3 ] R12 R23; etransitivity; eassumption.
Defined.

(* Read a pointer value *)
Definition readPtr (mem:DataMem) (ptr:SMValue) : SMValue :=
  match mem, ptr with
  | mkDataMem vmap, mkSMValue n => vmap n
  end.

Instance Proper_readPtr : Proper (lr_leq ==> lr_leq ==> lr_leq) readPtr.
Proof.
  intros [ vmap1 ] [ vmap2 ] R_vmap [ n1 ] [ n2 ] Rn. simpl in Rn.
  rewrite Rn. apply R_vmap.
Qed.

Instance Proper_readPtr_equiv : Proper (lr_equiv ==> lr_equiv ==> lr_equiv) readPtr.
Proof.
  intros mem1 mem2 [ Rmem Rmem' ] ptr1 ptr2 [ Rptr Rptr' ].
  split; apply Proper_readPtr; assumption.
Qed.

(* Write a pointer value *)
Definition writePtr (mem:DataMem) ptr val : DataMem :=
  match mem, ptr with
  | mkDataMem vmap, mkSMValue n =>
    mkDataMem (fun n' => if Nat.eqb n n' then val else vmap n')
  end.

Instance Proper_writePtr :
  Proper (lr_leq ==> lr_leq ==> lr_leq ==> lr_leq) writePtr.
Proof.
  intros [ vmap1 ] [ vmap2 ] R_vmap [ ptr1 ] [ ptr2 ] Rptr
         [ v1 ] [ v2 ] Rv. simpl in Rptr, Rv. rewrite Rptr. rewrite Rv.
  intro n. destruct (Nat.eqb ptr2 n); try reflexivity; apply R_vmap.
Qed.

Instance Proper_writePtr_equiv :
  Proper (lr_equiv ==> lr_equiv ==> lr_equiv ==> lr_equiv) writePtr.
Proof.
  intros mem1 mem2 [ Rmem Rmem' ] ptr1 ptr2 [ Rptr Rptr' ] v1 v2 [ Rv Rv' ].
  split; apply Proper_writePtr; assumption.
Qed.

(* readPtr and writePtr form a lens *)
Program Definition ptr_dmem_lens (ptr: SMValue) : Lens DataMem SMValue :=
  {| getL mem := readPtr mem ptr;
     putL val mem := writePtr mem ptr val; |}.
Next Obligation.
  intros mem1 mem2 eq_mem. rewrite eq_mem. reflexivity.
Defined.
Next Obligation.
  intros v1 v2 eq_v mem1 mem2 eq_mem. rewrite eq_mem. rewrite eq_v. reflexivity.
Defined.
Next Obligation.
  destruct a as [ vmap ]. destruct ptr as [ ptr ].
  split; intro n; case_eq (Nat.eqb ptr n); intro; try reflexivity.
  - rewrite Nat.eqb_eq in H. rewrite H. reflexivity.
  - rewrite Nat.eqb_eq in H. rewrite H. reflexivity.
Defined.
Next Obligation.
  destruct a as [ vmap ]. destruct ptr as [ ptr ]. destruct b as [ v ].
  simpl. rewrite (proj2 (Nat.eqb_eq _ _) eq_refl). reflexivity.
Defined.
Next Obligation.
  destruct a as [ vmap ]. destruct ptr as [ ptr ].
  destruct b1 as [ v1 ]. destruct b2 as [ v2 ].
  split; intro n; case_eq (Nat.eqb ptr n); intro; reflexivity.
Defined.

(***
 *** Machine Expressions
 ***)

(* A "machine expression" is any value you can build from the registers,
possibly with an error *)
Definition SMExpr := SMRegs -> option SMValue.

Instance LR_SMExpr : LR SMExpr := Build_LR _ _ _.

(* The expression that returns an undefined value *)
(* Definition undefSMExpr : SMExpr := fun _ => undefValue. *)

(* A tuple of expressions with a fixed size *)
Definition SMExprs (n:nat) := NTuple SMExpr n.

Instance LR_SMExprs n : LR (SMExprs n) := Build_LR _ _ _.

(***
 *** Machine Instructions
 ***)

Inductive SMInstr : nat -> Type :=

(* Assign values to the next registers *)
| SMAssign n (exprs: SMExprs n) : SMInstr n

(* Read a value from memory into the next register *)
| SMRead (ptr: SMExpr) : SMInstr 1

(* Write a value to memory *)
| SMWrite (ptr val: SMExpr) : SMInstr 0

(* Call a function pointer; calls always have 1 return value *)
| SMCall (fptr: SMExpr) : SMInstr 1

(* If-then-else, where the SSA phi node with N variables is represented as
having the two bodies return N values *)
| SMIf n (expr: SMExpr) (body1 body2: SMBlock n) : SMInstr n

(* While loop, where the SSA phi node with N variables is represented as having
the body return N values and having N default return values for the case that no
iterations happen *)
| SMWhile n (expr: SMExpr) (init:SMExprs n) (body:SMBlock n) : SMInstr n

with
SMBlock : nat -> Type :=

(* Return n values as the result of the block *)
| SMReturn n (exprs: SMExprs n) : SMBlock n

(* Run an instruction with n return values, push them into the next registers,
and then execute the remainder of the block in that extended context *)
| SMLet n (inst: SMInstr n) m (rest: SMBlock m) : SMBlock m
.

(* FIXME: using Coq intensional equality here adds a dependency on functional
extensionality to express equality of SMExprs, which is not really
necessary. Instead, these relations should be defined inductively on
instructions on blocks, using LR_SMExprs where needed. *)
Instance LR_SMInstr n : LR (SMInstr n) := LR_eq.
Instance LR_SMBlock n : LR (SMBlock n) := LR_eq.

(***
 *** Instruction Memories
 ***)

(* An instruction memory = a partial map from "pointers" (represented by nats)
to function bodies with 1 return value *)
Inductive InstMem : Type := mkInstMem (imap:nat -> option (SMBlock 1)).

(* Instruction memories are related by the built-in Coq intensional equality *)
(* FIXME HERE: this should be an extensional equality *)
Instance LR_InstMem : LR InstMem := LR_eq.

(* Read a function pointer to get a function body, defaulting to "return undef"
if the function pointer does not have a binding *)
Definition readFunPtr (mem:InstMem) (fptr:SMValue) : option (SMBlock 1) :=
  match mem, fptr with
  | mkInstMem imap, mkSMValue n => imap n
  end.

(***
 *** Machine States
 ***)

(* A machine state = registers and two types of memory *)
Definition SMState : Type := (SMRegs * (DataMem * InstMem)).

Instance LR_SMState : LR SMState := LR_product.

(* Lenses for the top-level components of SMState *)
Definition regs_lens : Lens SMState SMRegs := fst_lens _ _.
Definition dmem_lens : Lens SMState DataMem :=
  compose_lens (snd_lens _ _) (fst_lens _ _).
Definition imem_lens : Lens SMState InstMem :=
  compose_lens (snd_lens _ _) (snd_lens _ _).

(* The lens for a pointer into the DataMem of an SMState *)
Definition ptr_lens (ptr: SMValue) : Lens SMState SMValue :=
  compose_lens dmem_lens (ptr_dmem_lens ptr).

(***
 *** Machine Semantics
 ***)

(* The monad for executing simple machines *)
Definition SMMonad A `{LR A} := TraceM SMState A.

(* Evaluate an expression in a monad *)
Definition evalExpr (expr: SMExpr) : SMMonad SMValue :=
  bindM getM (fun smst =>
                match expr (getL regs_lens smst) with
                | Some v => returnM v
                | None => errorM
                end).

(* Evaluate a tuple of expressions in a monad *)
Fixpoint evalExprs {n} : SMExprs n -> SMMonad (SMValues n) :=
  match n return SMExprs n -> SMMonad (SMValues n) with
  | 0 => fun _ => returnM tt
  | S n' =>
    fun exprs =>
      bindM (evalExpr (fst exprs))
            (fun v =>
               bindM (evalExprs (snd exprs))
                     (fun vs => returnM (v, vs)))
  end.

(* Evaluate a machine instruction in a monad, assuming a recursive helper
function to "tie the knot" when evaluating function pointers *)
Fixpoint evalInstr' (recEval: SMBlock 1 -> SMMonad (SMValues 1))
         {n} (inst: SMInstr n) : SMMonad (SMValues n) :=
  match inst with
  | SMAssign n exprs =>
    evalExprs exprs
  | SMRead ptr_expr =>
    bindM (evalExpr ptr_expr)
          (fun ptr =>
             bindM getM (fun smst =>
                           returnM (getL (ptr_lens ptr) smst, tt)))
  | SMWrite ptr_expr val_expr =>
    bindM (evalExpr ptr_expr)
          (fun ptr =>
             bindM (evalExpr val_expr)
                   (fun val =>
                      bindM getM
                            (fun smst =>
                               putM (putL (ptr_lens ptr) val smst))))
  | SMCall fptr_expr =>
    bindM (evalExpr fptr_expr)
          (fun fptr =>
             bindM getM
                   (fun smst =>
                      match readFunPtr (getL imem_lens smst) fptr with
                      | Some block => recEval block
                      | None => errorM
                      end))
  | SMIf n expr body1 body2 =>
    bindM (evalExpr expr)
          (fun v =>
             if val_true v then
               evalBlock' recEval body1
             else
               evalBlock' recEval body2)
  | SMWhile n expr init_exprs body =>
    bindM (evalExprs init_exprs)
          (fun init_vs =>
             fixM (fun f prev_vs =>
                     bindM (evalExpr expr)
                           (fun v =>
                              if val_true v then
                                bindM (evalBlock' recEval body)
                                      (fun new_vs =>
                                         f new_vs)
                              else
                                returnM prev_vs))
                  init_vs)
  end

with
evalBlock' (recEval: SMBlock 1 -> SMMonad (SMValues 1))
          {n} (block: SMBlock n) : SMMonad (SMValues n) :=
  match block with
  | SMReturn n exprs =>
    evalExprs exprs
  | SMLet n inst m rest =>
    bindM (evalInstr' recEval inst)
          (fun vs =>
             bindM getM
                   (fun smst =>
                      bindM (putM (modifyL regs_lens
                                           (fun regs => pushRegs regs _ vs)
                                           smst))
                            (fun _ => evalBlock' recEval rest)))
  end.

(* Tie the knot for evalBlock on a function body (with 1 return value) *)
Definition evalFunBody : SMBlock 1 -> SMMonad (SMValues 1) :=
  fixM (fun recEval body => evalBlock' recEval body).

(* Tie the knot for evalInstr *)
Definition evalInstr {n} (inst: SMInstr n) : SMMonad (SMValues n) :=
  evalInstr' evalFunBody inst.

(* Tie the knot for evalBlock *)
Definition evalBlock {n} (block: SMBlock n) : SMMonad (SMValues n) :=
  evalBlock' evalFunBody block.

(***
 *** Simple Machine Permissions
 ***)

(* A state permission is a permission to read the state as some type A *)
Definition StPerm A `{LR A} : Type := BPerm SMState A.

(* A value permission is a permission to read a value-in-state as some type A *)
Definition ValPerm A `{LR A} : Type := BPerm (SMValue * SMState) A.

(* A permission for 0 or more values *)
Definition ValsPerm n A `{LR A} : Type := BPerm (SMValues n * SMState) A.

(* A permission to view a value as a natural number *)
Definition nat_perm : ValPerm nat :=
  lens_bperm (compose_lens (fst_lens _ _) SMValue_nat_lens).

(* A permission to view a value as a pointer to a (non-cyclic) linked list *)
(* FIXME HERE: define this! *)
Definition linked_list_perm A `{LR A} (bperm: ValPerm A) : ValPerm (list A).
  admit.
Admitted.

(***
 *** Bisimulation Abstraction
 ***)


(* --------------------------------------------------------------------
 * Copyright (c) - 2006--2012 - IMDEA Software Institute
 * Copyright (c) - 2006--2012 - Inria
 * Copyright (c) - 2006--2012 - Microsoft Coprporation
 *
 * Distributed under the terms of the CeCILL-B-V1 license
 * -------------------------------------------------------------------- *)

Set Implicit Arguments.

Require Export BoolEquality.
Require Export CCMisc.
Require Export BaseDef.
Require Import Arith Even Div2.

Set Implicit Arguments.

Open Scope U_scope.
Open Scope O_scope.

Hint Resolve Uabs_diff_compat_eq.

(* ***********  Boolean comparison in [U] *********** *)

 Variable Uleb : U -> U -> bool.
 Hypothesis Uleb_spec : forall (x y: U),
    if Uleb x y then (x <= y) else ~(x <= y).

 Lemma Uleb_correct : forall (x y: U),
    Uleb x y -> (x <= y).
 Proof.
  intros; generalize (Uleb_spec x y).
  rewrite H; trivial.
 Qed.

 Lemma Uleb_complete : forall (x y: U),
   Uleb x y = false -> (y < x).
 Proof.
  intros x y.
  generalize (Uleb_spec x y).
  case (Uleb x y).
    intros _ H; discriminate H.
    intros H _; auto.
 Qed.

(* ************************************************** *)

(* *************** MOVE THIS SOMEWHERE ELSE *************** *)
 Lemma Uabs_diff_zero : forall x y,  Uabs_diff x y == 0 <-> x==y.
 Proof.
  unfold Uabs_diff; split; intros H.

  apply (Ueq_orc x y); [auto|auto|].
  intro H'.
  apply (Ule_total x y); [auto|intro H''|intro H''].
    absurd (0 == x - y + (y - x))%U.
      rewrite Uplus_sym.
      apply Uplus_neq_zero_left.
      apply Uminus_lt_non_zero.
      apply Ule_diff_lt; assumption.
      apply Oeq_sym in H; assumption.
    absurd (0 == x - y + (y - x))%U.
      apply Uplus_neq_zero_left.
      apply Uminus_lt_non_zero.
      apply Ule_diff_lt; [assumption|auto].
      apply Oeq_sym in H; assumption.

  rewrite H, Uminus_eq; auto.
 Qed.

 Lemma Uminus_plus_aux: forall a b c d,
  d <= c ->
  a <= [1-] c ->
  a - b <= (a+c) - (b+d).
 Proof.
  intros.
  rewrite (Uplus_sym b d), <-Uminus_assoc_left.
  apply Uminus_le_compat; [|trivial].
  rewrite <-Uplus_minus_assoc_right;
  [|apply Uinv_le_perm_right|]; trivial.
 Qed.

 Lemma  fminus_fplus_simpl : forall (A:Type) (f h:A -o> U),
   h <= f ->
   f == fplus h (fminus f h).
 Proof.
  intros.
  apply ford_eq_intro; intro.
  unfold fplus, fminus.
  rewrite Uplus_sym, Uminus_plus_simpl; trivial.
 Qed.

 Lemma Uabs_diff_plus_bound: forall a b c d E,
   c <= a ->
   b <= d ->
   Uabs_diff a c <= E ->
   Uabs_diff d b <= E ->
   Uabs_diff (a+b) (c+d) <= E.
 Proof.
   intros a b c d E Hca Hbd HacE HdbE.
   apply (Ule_total (a+b) (c+d)); trivial; intro H.
    (* case [a+b <= c+d] *)
    rewrite Uabs_diff_compat_le; trivial.
    rewrite <- Uminus_assoc_left, Uplus_minus_perm_le.
    rewrite (Uminus_le_zero _ _ Hca), Uplus_zero_left.
    rewrite <- HdbE; apply Ule_plus_right.
    (* case [a+b >= c+d] *)
    rewrite Uabs_diff_sym, Uabs_diff_compat_le; trivial.
    rewrite <- Uminus_assoc_left, Uplus_minus_perm_le, Uplus_minus_le.
    rewrite (Uminus_le_zero _ _ Hbd),Uplus_zero_right.
    rewrite <- HacE; apply Ule_plus_right.
 Qed.

 Definition fabs_diff (A:Type) (f h : MF A) : MF A :=
               fun x =>  Uabs_diff (f x) (h x).
 Implicit Arguments fabs_diff [A].

 Lemma fabs_diff_eq_compat: forall (A:Type) (f1 h1 f2 h2: MF A),
   f1 == f2 -> h1 == h2 ->
   fabs_diff f1 h1 == fabs_diff f2 h2.
 Proof.
   unfold fabs_diff; intros.
   refine (ford_eq_intro _); intro a.
   rewrite (ford_eq_elim H a), (ford_eq_elim H0 a); trivial.
 Qed.

 Lemma Uabs_diff_mu_compat : forall (A:Type) (f' g':A-o>U) (d: Distr A),
   Uabs_diff (mu d f') (mu d g') <= mu d (fabs_diff f' g').
 Proof.
  intros.
  rewrite (@mu_stable_plus _ d (fminus f' g') (fminus g' f')).
  apply Uplus_le_compat; apply mu_stable_le_minus.
    intros x.
    apply (Ule_total (f' x) (g' x)); [auto| |].
      intro H'; rewrite (Uminus_le_zero _ _ H'); trivial.
      intro H'; apply Uinv_le_perm_right; rewrite (Uminus_le_zero _ _ H'); trivial.
 Qed.

 Definition comp (A:Type) (P Q : relation A) : relation A :=
  fun a1 a2 => exists a, P a1 a /\ Q a a2.
 Notation "P 'o' Q" := (comp P Q) (at level 40).

 Lemma PER_r : forall (A: Type) (R:relation A) x1 x2,
   PER _ R ->
   R x1 x2 ->
   R x2 x2.
 Proof.
  intros A R x1 x2  [Hsym Htrans] H.
  exact (Htrans _ _ _ (Hsym _ _ H) H).
 Qed.

 Lemma comp_PER : forall  (A:Type) (R : relation A),
  PER _ R -> same_relation _ R (R o R).
 Proof.
  unfold comp; split; intros.
   intros x1 x2 Hx; exists x2; split.
     assumption.
     apply (PER_r _ _ H Hx).
   intros x1 x2 [x [H1 H2] ].
   destruct H; apply (per_trans _ _ _ H1 H2).
 Qed.

 Lemma PER_Oeq : forall (A:Type), PER _ (@Oeq (A -o> U)).
 Proof.
  intros; constructor.
    refine (@Oeq_sym (A-o>U)).
    refine (@Oeq_trans (A-o>U)).
 Qed.

 Hint Unfold same_relation inclusion.

 Lemma same_relation_refl : forall (A:Type) (R: relation A), same_relation _ R R.
 Proof.
  split; intros; auto.
 Qed.

 Lemma same_relation_sym : forall (A:Type) (R S: relation A),
   same_relation _ R S -> same_relation _ S R.
 Proof.
  intros A R S [? ?]; split; assumption.
 Qed.

 Add Parametric Relation (A:Type) : (relation A) (@same_relation A)
   reflexivity proved by (@same_relation_refl A)
   symmetry proved by (@same_relation_sym A)
 as RSD_rel.

(*
 Definition Mlincomb (A:Type) (m1 m2 : M A) (x1 x2 :U) : M A.
  intros.
  exists (fun f => x1 * m1 f + x2 * m2 f).
  intros f g H.
  apply Uplus_le_compat; auto.
 Defined.

 Lemma stable_inv_lincomb : forall (A:Type) (m1 m2 : M A) (x1 x2 :U),
  stable_inv m1 ->
  stable_inv m2 ->
  x1 <= [1-] x2 ->
  stable_inv (Mlincomb m1 m2 x1 x2).
 Proof.
  intros  A m1 m2 x1 x2 H1 H2 Hx f; simpl.
  rewrite (H1 f), (H2 f).
  apply (Uinv_bary_le _ _ Hx).
 Qed.

 Lemma stable_plus_lincomb : forall (A:Type) (m1 m2 : M A) (x1 x2 :U),
  stable_plus m1 ->
  stable_plus m2 ->
  stable_inv m1 ->
  stable_inv m2 ->
  x1 <= [1-] x2 ->
  stable_plus (Mlincomb m1 m2 x1 x2).
 Proof.
  intros  A m1 m2 x1 x2 H1 H2 H1' H2' Hx f g H; simpl.
  rewrite (H1 _ _ H), (H2 _ _ H).
  repeat rewrite Udistr_plus_left; [ |
     rewrite <-(H2' g); auto | rewrite <-(H1' g); auto ].
  repeat rewrite Uplus_assoc; apply Uplus_eq_compat_left;
    repeat rewrite <-Uplus_assoc; apply Uplus_eq_compat_right;
      apply Uplus_sym.
 Qed.

 Lemma stable_mult_lincomb : forall (A:Type) (m1 m2 : M A) (x1 x2 :U),
  stable_mult m1 ->
  stable_mult m2 ->
  x1 <= [1-] x2 ->
  stable_mult (Mlincomb m1 m2 x1 x2).
 Proof.
  intros A m1 m2 x1 x2 H1 H2 Hx x f; simpl.
  rewrite (H1 x), (H2 x).
  rewrite Udistr_plus_left; [ | apply Uinv_le_trans with x1 x2; trivial ].
  repeat rewrite Umult_assoc; rewrite (Umult_sym x1 x), (Umult_sym x2 x); trivial.
 Qed.

 Lemma continuous_lincomb : forall (A:Type) (m1 m2 : M A) (x1 x2 :U),
  continuous m1 ->
  continuous m2 ->
  stable_mult m1 ->
  stable_mult m2 ->
  x1 <= [1-] x2 ->
  continuous (Mlincomb m1 m2 x1 x2).
 Proof.
  intros A m1 m2 x1 x2 H1 H2 H1' H2' Hx f; simpl.
  rewrite <-(@lub_le_compat _ ((UPlus @2 Mmult (fcte A x1) m1 @ f) (Mmult (fcte A x2) m2 @ f))
    (Mlincomb m1 m2 x1 x2 @ f)).
  rewrite lub_eq_plus.
  rewrite (H1 f), (H2 f).
  rewrite <-(lub_eq_mult x1), <-(lub_eq_mult x2).
  apply Uplus_le_compat; apply lub_le_compat.
    simpl; intro x; rewrite (H1' x1 (f x)); trivial.
    simpl; intro x; rewrite (H2' x2 (f x)); trivial.
  simpl; intro n.
  rewrite (H1' x1 (f n)), (H2' x2 (f n)); trivial.
 Qed.

 Definition distr_lincomb : forall (A:Type)  (d1 d2 : Distr A) (x1 x2 : U) (Hok:x1 <= [1-] x2), Distr A.
  intros.
  exists (Mlincomb (mu d1) (mu d2) x1 x2).
    apply stable_inv_lincomb; auto.
    apply stable_plus_lincomb; auto.
    apply stable_mult_lincomb; auto.
    apply continuous_lincomb; auto.
 Defined.

 Lemma distr_lincomb_simpl : forall (A:Type) (d1 d2 : Distr A) x1 x2 (Hok:x1 <= [1-] x2) f ,
    mu (distr_lincomb d1 d2 x1 x2 Hok) f == x1 * (mu d1 f) + x2 * (mu d2 f).
 Proof.
  trivial.
 Qed.

 Lemma dist_lincomb_comm: forall (A:Type)  (d1 d2 : Distr A) (x1 x2 : U) (Hok:x1 <= [1-] x2),
   distr_lincomb d1 d2 x1 x2 Hok == distr_lincomb d2 d1 x2 x1 (Uinv_le_perm_right _ _ Hok).
 Proof.
  intros.
  apply eq_distr_intro; intros.
  repeat rewrite distr_lincomb_simpl.
  apply Uplus_sym.
 Qed.

 (* Avarage distribution *)
 Definition Dmean := fun (A:Type) (d1 d2 : Distr A) =>
   (distr_lincomb d1 d2 [1/2] [1/2] (proj1 Unth_one)).

 Lemma Dmean_simpl: forall (A:Type) (d1 d2 : Distr A) f,
  mu (Dmean d1 d2) f == [1/2] * (mu d1 f) + [1/2] * (mu d2 f).
 Proof.
  unfold Dmean; intros.
  apply distr_lincomb_simpl.
 Qed.

 Lemma Dmean_comm : forall (A:Type) (d1 d2 : Distr A),
   Dmean d1 d2 == Dmean d2 d1.
 Proof.
  intros.
  apply eq_distr_intro; intros.
  repeat rewrite Dmean_simpl; trivial.
 Qed.
*)

 Definition relTr (A B :Type) (F : A -o> Distr B) (R : relation (A -o> U)) : relation (B -o> U) :=
   fun f g => R (fun x : A => (mu (F x)) f) (fun x : A => (mu (F x)) g).

 Lemma even_2n : forall n, even (2*n).
 Proof.
  intro n.
  apply even_mult_l.
  repeat constructor.
 Qed.

 Lemma odd_S2n : forall n, odd (S(2*n)).
 Proof.
  intro n.
  constructor.
  apply even_2n.
 Qed.

(*
  parity_split f g (2m) = f m
  parity_split f g (2m+1) = g m
*)
 Definition parity_split := fun (A:Type) (f g : nat -> A) n  =>
   match (even_odd_dec n) with
   | left x => f (div2 n)
   | right x => g (div2 (pred n))
   end.

(* Given two discrete disctribution over [A], restate them in terms
   of the same enumearion of [A]'s elements *)
Section SamePoints.

 Variable A : Type.
 Variable AeqU : A -> A -O-> U.
 Hypothesis cover_uR: forall a : A, cover (eq a) (AeqU a).

 Variable d1 d2 : Distr A.
 Hypothesis is_Discr_d1 : is_Discrete d1.
 Hypothesis is_Discr_d2 : is_Discrete d2.

 Definition p1 := D_points is_Discr_d1.
 Definition p2 := D_points is_Discr_d2.
 Definition D1 := D_Discr is_Discr_d1.
 Definition D2 := D_Discr is_Discr_d2.

 Definition c1 : nat -> U := parity_split (coeff AeqU p1 d1) (fzero _).
 Definition c2 : nat -> U := parity_split (fzero _) (coeff AeqU p2 d2).
 Definition p  : nat -> A := parity_split p1 p2.

 Lemma disc_eqpoint_l : Discrete (@eq A) p d1.
 Proof.
  apply range_weaken with (2:=D1); fold p1.
  unfold In_classes; intros x Hx.
  unfold p, parity_split.
  unfold exc in *.
  intros P HP H.
  apply (Hx _ HP).
  intros n Hn.
  generalize (H (2*n)%nat).
  elim (even_odd_dec (2*n)%nat); intro Hn'.
     rewrite div2_double.
     exact (fun n => n Hn).
     elimtype False; refine (not_even_and_odd _ (even_2n _) Hn').
 Qed.

 Lemma disc_eqpoint_r : Discrete (@eq A) p d2.
 Proof.
  apply range_weaken with (2:=D2); fold p2.
  unfold In_classes; intros x Hx.
  unfold p, parity_split.
  unfold exc in *.
  intros P HP H.
  apply (Hx _ HP).
  intros n Hn.
  generalize (H (S(2*n))).
  elim (even_odd_dec (S(2*n))); intro Hn'.
     elimtype False; refine (not_even_and_odd _ Hn' (odd_S2n _)).
     rewrite <-pred_Sn, div2_double; exact (fun n => n Hn).
 Qed.

 End SamePoints.

(* ******************************************************** *)

(********************************************************************)
(* *************** RELATIONAL STATISTICAL DISTANCE  *************** *)
(********************************************************************)

 Definition RSD (A:Type) (R : relation _) (d1 d2:Distr A) (ep:U) :=
   forall (f g: A -o> U),
   R f g ->
   Uabs_diff (mu d1 f) (mu d2 g) <= ep.

 Lemma RSD_eq_distr_compat: forall (A:Type) R (d1 d2 d3 d4:Distr A) (ep:U),
   d1 == d3 ->
   d2 == d4 ->
   RSD R d3 d4 ep ->
   RSD R d1 d2 ep.
 Proof.
  unfold RSD; intros.
  rewrite (eq_distr_elim H), (eq_distr_elim H0).
  apply (H1 _ _ H2).
 Qed.

 Lemma RSD_weaken: forall (A:Type) (R R' : relation _) (d1 d2:Distr A) (ep ep':U),
   inclusion _ R R' ->
   ep' <= ep ->
   RSD R' d1 d2 ep' ->
   RSD R d1 d2 ep.
 Proof.
  unfold RSD; intros.
  rewrite <-H0.
  apply (H1 _ _ (H _ _ H2)).
 Qed.

 Add Parametric Morphism (A:Type) : (@RSD A)
   with signature (@inclusion _) ==> (@Oeq (Distr A)) ==> (@Oeq (Distr A))
    ==> (@Ole (tcpo U)) --> inverse impl as RSD_impl_le_compat_mor.
 Proof.
  intros R R' HR d1 d3 H31 d2 d4 H24 ep ep' Hep H.
  apply RSD_eq_distr_compat with (1:=H31) (2:=H24).
  refine (RSD_weaken _ HR Hep H).
 Qed.

 Add Parametric Morphism (A:Type) : (@RSD A)
   with signature (@same_relation _) ==> (@Oeq (Distr A)) ==> (@Oeq (Distr A))
    ==> (@Oeq (tcpo U)) ==> inverse impl as RSD_eq_compat_mor.
 Proof.
  intros R R' [HR _] d1 d3 H31 d2 d4 H24 ep ep' [_ Hep] H.
  apply RSD_eq_distr_compat with (1:=H31) (2:=H24).
  refine (RSD_weaken _ HR Hep H).
 Qed.

 Lemma RSD_one : forall (A:Type) (d1 d2:Distr A) (R : relation _),
  RSD R d1 d2 1.
 Proof.
  unfold RSD; intros; trivial.
 Qed.

 Lemma RSD_tranp: forall (A:Type) (d1 d2:Distr A) (R : relation _) (ep:U),
   RSD R d1 d2 ep -> RSD (transp _ R) d2 d1 ep.
 Proof.
  unfold RSD; intros.
  rewrite Uabs_diff_sym.
  apply (H _ _ H0).
 Qed.

 Lemma RSD_sym: forall (A:Type) (d1 d2:Distr A) (R : relation _) (ep:U),
   symmetric _ R ->
   RSD R d2 d1 ep ->
   RSD R d1 d2 ep.
 Proof.
  intros.
  apply RSD_tranp.
  apply RSD_weaken with R ep; trivial.
  intros f g; apply H.
 Qed.

 Lemma RSD_null: forall (A:Type) (d1 d2:Distr A) (R : relation _),
   RSD R d1 d2 0 <->
   (forall f g, R f g  -> mu d1 f == mu d2 g).
 Proof.
  split; [ | unfold RSD ]; intros.
    rewrite <-Uabs_diff_zero; apply Ule_zero_eq; auto.
    rewrite (H _ _ H0), Uabs_diff_compat_eq; trivial.
 Qed.

 Lemma RSD_trans : forall (A:Type) R1 R2 (d1 d2 d3:Distr A) (ep1 ep2:U),
   RSD R1 d1 d2 ep1 ->
   RSD R2 d2 d3 ep2 ->
   RSD (R1 o R2) d1 d3 (ep1+ep2) .
 Proof.
  unfold RSD; intros.
  destruct H1 as [h [Hh1 Hh2] ].
  rewrite (Uabs_diff_triangle_ineq _ _  (mu d2 h)).
  apply Uplus_le_compat; [ apply (H _ _ Hh1) |  apply (H0 _ _ Hh2) ].
 Qed.

 Lemma RSD_Mlet : forall (A:Type) R (d1 d2 :Distr A) (B: Type) (F:A -> Distr B) (ep:U),
   RSD R d1 d2 ep  ->
   RSD (relTr F R) (Mlet d1 F) (Mlet d2 F) ep.
 Proof.
  unfold RSD; intros.
  repeat rewrite Mlet_simpl.
  apply (H _ _ H0).
 Qed.

 Lemma RSD_Mlet' : forall (A:Type) (d :Distr A) (B: Type) (F G : A -> Distr B) (S : relation (B -o> U)) (ep:U),
 (forall f g, S f g ->
   mu d (fabs_diff (fun a => mu (F a) f) (fun a => mu (G a) g)) <= ep) ->
   RSD S (Mlet d F) (Mlet d G) ep.
 Proof.
  unfold RSD; intros.
  repeat rewrite Mlet_simpl.
  rewrite Uabs_diff_mu_compat.
  auto.
 Qed.

(********************************************************************)
(* *************** CLASSICAL STATISTICAL DISTANCE  *************** *)
(********************************************************************)

Hint Unfold symmetric reflexive inclusion.

(* Statistical distance w.r.t {0,1}-valued functions *)
 Definition SD (A:Type) (d1 d2:Distr A) (ep:U)  :=
   forall (P : A -o> boolO),
     (Uabs_diff (mu d1 (charfun P)) (mu d2 (charfun P))) <= ep.

(* Statistical distance w.r.t [0,1]-valued functions *)
 Definition GSD (A:Type) (d1 d2:Distr A) (ep:U)  :=
  RSD (@Oeq _) d1 d2 ep.

 Lemma GSD_sym : forall (ep:U) (A:Type) (d1 d2:Distr A) ,
   GSD d1 d2 ep -> GSD  d2 d1 ep.
 Proof.
  intros.
  refine (RSD_sym _ H); auto.
 Qed.

 Lemma GSD_trans : forall (A:Type) (d1 d2 d3:Distr A) (ep ep':U),
   GSD  d1 d2 ep -> GSD d2 d3 ep' -> GSD d1 d3 (ep+ep') .
 Proof.
  intros;  unfold GSD.
  rewrite (comp_PER (PER_Oeq A)).
  apply RSD_trans with (1:=H) (2:=H0).
 Qed.

 Lemma GSD_eq_distr_compat : forall (A:Type) (d1 d2 d3 d4:Distr A) (ep:U),
   GSD d1 d2 ep -> d3 == d1 -> d4==d2 -> GSD d3 d4 ep.
 Proof.
  intros.
  refine (RSD_eq_distr_compat H0 H1 H).
 Qed.

 Lemma GSD_weaken : forall (A:Type) (d1 d2:Distr A) (ep ep':U),
   GSD d1 d2 ep' -> ep' <= ep -> GSD d1 d2 ep.
 Proof.
  intros.
  refine (RSD_weaken _ _ H0 H); auto.
 Qed.

 Add Parametric Morphism (A:Type) : (@GSD A)
   with signature (@Oeq (Distr A)) ==> (@Oeq (Distr A)) ==> (@Ole (tcpo U)) --> inverse impl
 as GSD_le_eq_compat_mor.
 Proof.
  intros  d3 d1 H31 d4 d2 H24 ep' ep Hep H.
  refine (GSD_weaken _ _ Hep).
  refine (GSD_eq_distr_compat H H31 H24).
 Qed.

 Add Parametric Morphism (A:Type) : (@GSD A)
   with signature (@eq (Distr A)) ==> (@eq (Distr A)) ==> ((Oeq (O:=U))) --> impl
  as GSD_eq_compat_mor.
 Proof.
  unfold impl; intros.
  refine (GSD_weaken _ H0 (Oeq_le_sym H)).
 Qed.

 Lemma GSD_NonDeg : forall (A:Type) (d1 d2:Distr A),
   GSD d1 d2 0 <-> d1 == d2.
 Proof.
  split; intro H.
    apply eq_distr_intro; intro f.
    apply (proj1 (RSD_null _ _ _) H); trivial.
    intros a1 a2 Ha.
    rewrite (mu_stable_eq d1 _ _ Ha), (eq_distr_elim H).
    rewrite Uabs_diff_compat_eq; trivial.
 Qed.

 Lemma GSD_refl_0 : forall (A:Type) (d:Distr A),
   GSD d d 0.
 Proof.
  intros; rewrite GSD_NonDeg; trivial.
 Qed.

 Lemma GSD_refl : forall (ep:U) (A:Type) (d:Distr A),
   GSD d d ep.
 Proof.
  intros.
  rewrite <-(Upos ep).
  apply GSD_refl_0.
 Qed.

 Lemma GSD_Mlet : forall (A:Type) (d1 d2 :Distr A) (B: Type) (F:A -> Distr B) (ep:U),
   GSD d1 d2 ep  ->
   GSD (Mlet d1 F) (Mlet d2 F) ep.
 Proof.
  unfold GSD; intros.
  apply RSD_weaken with (relTr F (@Oeq _)) ep; trivial.
    unfold relTr, inclusion; intros;
    refine (ford_eq_intro (fun a => mu_stable_eq (F a) _ _ H0)).
    apply (RSD_Mlet _ H).
 Qed.

(*
 Lemma Dmean_GSD : forall (A:Type) (d1 d2 : Distr A) ep,
   GSD d1 d2 ep ->
   GSD d1 (Dmean d1 d2) ([1/2] * ep).
 Proof.
  intros; apply (Dmean_RSD H).
 Qed.

 Lemma Uabs_diff_distr : forall (A:Type) (d1 d2 :Distr A) ep f g,
   GSD d1 d2 ep  ->
   Uabs_diff (mu d1 f) (mu d2 g) <= ep + mu (Dmean d1 d2) (fabs_diff f g).
 Proof.
  intros.
  rewrite (Uabs_diff_triangle_ineq _ _  (mu (Dmean d1 d2) f)),
    (Uabs_diff_sym _ (mu d2 g)),
    (Uabs_diff_triangle_ineq (mu d2 g) _ (mu (Dmean d1 d2) g)),
    Uplus_assoc.
  rewrite <-(Umult_one_left ep), <-Unth_one_plus,
    Udistr_plus_right; [ | rewrite <-Unth_one; trivial ].
  repeat apply Uplus_le_compat.
    apply Dmean_GSD; [ | intros a1 a2 Ha; rewrite Ha ]; trivial.
    rewrite (eq_distr_elim (Dmean_comm d1 d2) g).
    apply Dmean_GSD; [ apply GSD_sym | intros a1 a2 Ha; rewrite Ha ]; trivial.
    rewrite Uabs_diff_sym; apply  Uabs_diff_mu_compat.
 Qed.

 (* REMARK: it is enough to bound [|f-g|] only in the
   union of the ranges of [d1] and [d2] *)
 Lemma Uabs_diff_distr_bound : forall (A:Type) (d1 d2 :Distr A) ep f g ep',
   GSD d1 d2 ep  ->
   (forall a, Uabs_diff (f a) (g a) <= ep') ->
   Uabs_diff (mu d1 f) (mu d2 g) <= ep + ep'.
 Proof.
  intros.
  rewrite (@Uabs_diff_distr _ _ _ _ _ _ H).
  Usimpl.
  rewrite <-(mu_cte_le (Dmean d1 d2) ep').
  apply mu_le_compat; auto.
 Qed.
*)

 Theorem SD_le_GSD : forall (A:Type) (d1 d2 :Distr A) (ep:U),
   GSD d1 d2 ep -> SD d1 d2 ep.
 Proof.
  intros A d1 d2 ep H f.
  apply H; trivial.
 Qed.

(* For discrete distributions, it holds that [GSD <= SD]. The
   converse (trivially) holds for any pair of distributions *)
Section Discrete_distr.

 Variable A : Type.
 Variable AeqU : A -> A -O-> U.
 Hypothesis cover_uR: forall a : A, cover (eq a) (AeqU a).

 Variable d1 d2 : Distr A.
 Hypothesis d1_discr : is_Discrete d1.
 Hypothesis d2_discr : is_Discrete d2.

 Definition  R : A -> boolO :=
   fun a => Uleb (mu d2 (AeqU a)) (mu d1 (AeqU a)).

 Definition p' := p d1_discr d2_discr.

 Lemma discrete1 : Discrete (@eq A) p' d1.
 Proof.  apply disc_eqpoint_l. Qed.

 Lemma discrete2 : Discrete (@eq A) p' d2.
 Proof. apply disc_eqpoint_r. Qed.

 Lemma nul_in_R : forall (f:A-o>U),
   (forall a, R a -> f a == 0) -> mu d1 f <= mu d2 f.
 Proof.
  intros.
  rewrite (mu_restr_split d1 R f).
  transitivity  ((mu d1 (fzero _)) + mu d1 (restr (negP R) f)).
    apply Uplus_le_compat_left.
    apply mu_monotonic.
      refine (ford_le_intro _); intro a.
      unfold restr, fzero; generalize (H a).
      case (R a); auto.
  rewrite mu_zero, Uplus_zero_left.
  transitivity (mu d2 (restr (negP R) f)).
    rewrite (mu_discrete AeqU cover_uR (@eq_Transitive A)(@eq_Symmetric A) discrete1);
      [ | intros x y Heq; rewrite Heq; trivial ].
    rewrite (mu_discrete AeqU cover_uR (@eq_Transitive A)(@eq_Symmetric A) discrete2);
      [ | intros x y Heq; rewrite Heq; trivial ].
    unfold coeff; simpl.
    apply serie_le_compat; intro k.
    repeat rewrite <-Umult_assoc; apply Umult_le_compat_right.
    unfold restr, negP, negb, R.
    generalize (@Uleb_complete (mu d2 (AeqU (p' k)))(mu d1 (AeqU (p' k)))).
    case_eq (Uleb (mu d2 (AeqU (p' k))) (mu d1 (AeqU (p' k)))).
      intros; Usimpl; trivial.
      intros _ H'; Usimpl.
      apply Ult_le; apply H'; trivial.
  apply mu_monotonic.
    refine (ford_le_intro _); intro a.
    unfold restr; case ((negP R) a); auto.
  Qed.

 Lemma nul_in_notR : forall (f:A-o>U),
   (forall a, (negP R) a -> f a == 0) -> mu d2 f <= mu d1 f.
 Proof.
  intros.
  rewrite (mu_restr_split d2 R f).
  transitivity  ((mu d2 (restr R f)) + (mu d2 (fzero _))).
    apply Uplus_le_compat_right.
    apply mu_monotonic.
      refine (ford_le_intro _); intro a.
      unfold restr, fzero; generalize (H a).
      case ((negP R) a); auto.
  rewrite mu_zero, Uplus_zero_right.
  transitivity (mu d1 (restr R f)).
    rewrite (mu_discrete AeqU cover_uR (@eq_Transitive A)(@eq_Symmetric A) discrete1);
      [ | intros x y Heq; rewrite Heq; trivial ].
    rewrite (mu_discrete AeqU cover_uR (@eq_Transitive A)(@eq_Symmetric A) discrete2);
      [ | intros x y Heq; rewrite Heq; trivial ].
    unfold coeff; simpl.
    apply serie_le_compat; intro k.
    repeat rewrite <-Umult_assoc; apply Umult_le_compat_right.
    unfold restr, R.
    case_eq (Uleb (mu d2 (AeqU (p' k))) (mu d1 (AeqU (p' k)))); intro H'; Usimpl.
      apply Uleb_correct; apply H'; trivial.
      trivial.
  apply mu_monotonic.
    refine (ford_le_intro _); intro a.
    unfold restr; case (R a); auto.
  Qed.

 Lemma monot_Uabsdiff_notR: forall (f h :A-o>U),
   h <= f ->
   (forall a, (negP R) a -> f a == 0 /\ h a == 0) ->
   Uabs_diff (mu d1 h)  (mu d2 h) <= Uabs_diff (mu d1 f) (mu d2 f).
 Proof.
  intros f h H1 H2.
  assert (H' : fplusok h (fminus f h)) by
    (apply (@ford_le_intro _ _ h);unfold fplusok, fminus, finv; auto).
  repeat rewrite (range_eq (range_True _) _ _
    (fun a _ => (ford_eq_elim  (fminus_fplus_simpl H1)) a)).
  rewrite (@mu_stable_plus _ d2 h (fminus f h)); [|exact H'].
  rewrite (@mu_stable_plus _ d1 h (fminus f h)); [|exact H'].

  assert (H_1 : mu d2 (fminus f h) <= mu d1 (fminus f h)).
   apply nul_in_notR.
   intros a Ha; destruct (H2 a Ha) as [Hfa Hha]; clear H2.
   unfold fminus; rewrite Hfa, Hha; trivial.
  assert (H_2 : mu d2 h <= mu d1 h).
   apply nul_in_notR; intros a Ha; destruct (H2 a Ha) as [_ ?]; assumption.
  rewrite Uabs_diff_sym, (Uabs_diff_compat_le (mu d2 h) (mu d1 h));
  [| exact H_2].
  rewrite  Uabs_diff_sym, (Uabs_diff_compat_le
    (mu d2 h + mu d2 (fminus f h)) (mu d1 h + mu d1 (fminus f h)));
    [ | apply Uplus_le_compat; [exact H_2|exact H_1] ].

  apply  Uminus_plus_aux.
   apply nul_in_notR.
   intros a Ha; destruct (H2 a Ha) as [Hfa Hha]; clear H2.
   unfold fminus; rewrite Hfa, Hha; trivial.
   apply mu_fplusok; exact H'.
 Qed.

 Lemma monot_Uabsdiff_R: forall (f h :A-o>U),
   h <= f ->
   (forall a, R a -> f a == 0 /\ h a == 0) ->
    Uabs_diff (mu d2 h)  (mu d1 h) <= Uabs_diff (mu d1 f) (mu d2 f).
 Proof.
  intros f h H1 H2.
  assert (H' : fplusok h (fminus f h)) by
    (apply (@ford_le_intro _ _ h);unfold fplusok, fminus, finv; auto).
  repeat rewrite (range_eq (range_True _) _ _
    (fun a _ => (ford_eq_elim  (fminus_fplus_simpl H1)) a)).
  rewrite (@mu_stable_plus _ d2 h (fminus f h)); [|exact H'].
  rewrite (@mu_stable_plus _ d1 h (fminus f h)); [|exact H'].

  assert (H_1 : mu d1 (fminus f h) <= mu d2 (fminus f h)).
   apply nul_in_R.
   intros a Ha; destruct (H2 a Ha) as [Hfa Hha]; clear H2.
   unfold fminus; rewrite Hfa, Hha; trivial.
  assert (H_2 : mu d1 h <= mu d2 h).
   apply nul_in_R; intros a Ha; destruct (H2 a Ha) as [_ ?]; assumption.
  rewrite Uabs_diff_sym, (Uabs_diff_compat_le (mu d1 h) (mu d2 h)); [| exact H_2].
  rewrite (Uabs_diff_compat_le
    (mu d1 h + mu d1 (fminus f h)) (mu d2 h + mu d2 (fminus f h)));
    [ | apply Uplus_le_compat; [exact H_2|exact H_1] ].

  apply  Uminus_plus_aux.
   apply nul_in_R.
   intros a Ha; destruct (H2 a Ha) as [Hfa Hha]; clear H2.
   unfold fminus; rewrite Hfa, Hha; trivial.
   apply mu_fplusok; exact H'.
 Qed.

 Lemma GSD_le_SD: forall (ep:U),
   SD d1 d2 ep -> GSD d1 d2 ep.
 Proof.
  intros ep HSD f g H.
  rewrite <-(mu_stable_eq d2 _ _ H).
  rewrite (mu_restr_split d1 R f), (mu_restr_split d2 R f).
  apply Uabs_diff_plus_bound.
   apply nul_in_notR; unfold restr, negP; intro a.
   case (R a); intros; [discriminate H0 | trivial].
   apply nul_in_R; unfold restr, negP, negb; intro a.
   case (R a); intros; [trivial | discriminate H0 ].
   transitivity (Uabs_diff (mu d1 (restr R (@fone A))) (mu d2 (restr R (@fone A)))).
    apply monot_Uabsdiff_notR.
     apply restr_le_compat; auto.
     unfold negP, negb, restr; intros a.
     case (R a); intros; [discriminate H0 | split; trivial].
    apply HSD.
   transitivity (Uabs_diff (mu d1 (restr (negP R) (@fone A))) (mu d2 (restr (negP R) (@fone A)))).
    apply monot_Uabsdiff_R.
     apply restr_le_compat; auto.
     unfold negP, negb, restr; intros a.
     case (R a); intros; [split; trivial |discriminate H0].
    apply HSD.
 Qed.

End Discrete_distr.