Set Implicit Arguments.

(* To install ITree, do:
     opam install coq-itree *)
From ITree Require Import
     ITree ITreeFacts.
Import ITreeNotations.
Local Open Scope itree_scope.

From Paco Require Import paco.

Require Import Coq.Program.Basics.
Require Import Coq.micromega.Lia.
Require Import Coq.QArith.QArith.
Require Import Coq.QArith.Qabs.
Require Import Coq.micromega.Lqa.
Require Import Coq.Logic.Eqdep_dec.
Require Import RelationClasses.
Require Import Coq.Classes.Equivalence.
Local Open Scope equiv_scope.
Require Import Nat.
Local Open Scope program_scope.

Require Import ExtLib.Structures.Monad.
Require Import ExtLib.Structures.MonadState.
Require Import ExtLib.Data.Monads.StateMonad.

Import MonadNotation.
Local Open Scope monad_scope.

Require Import axioms.
Require Import borel.
Require Import cpGCL.
Require Import incomparable.
Require Import infer.
Require Import infer_regular.
Require Import itree_lang_unique.
Require Import itree_regular.
Require Import order.
Require Import misc.
Require Import nondivergent.
Require Import Q.
Require Import regular.
Require Import semiring.
Require Import tree.
Require Import unbiased_itree.

(** A biased coin event can be raised, passing a rational p to the
  environment and receiving in return a single bit with probability p
  of being 1. *)
Variant biasedE : Type -> Type :=
| GetBiased : Q -> biasedE bool.
(* | GetBiased : forall p : Q, 0 <= p <= 1 -> biasedE bool. *)

(* Variant unbiasedE : Type -> Type := *)
(* | GetUnbiased : unbiasedE bool. *)

(** Commands to itrees directly. *)
Fixpoint compile_itree (c : cpGCL) : ktree biasedE St (unit + St) :=
  match c with
  | CSkip => ret ∘ inr
  | CAbort => (const ∘ ret ∘ inl) tt
  | CAssign n e => fun st => ret (inr (upd n (e st) st))
  | CSeq c1 c2 =>
    fun st => x <- compile_itree c1 st ;;
           match x with
           | inl tt => ret (inl tt)
           | inr st' => compile_itree c2 st'
           end
  | CIte e c1 c2 =>
    fun st => if e st then compile_itree c1 st else compile_itree c1 st
  | CChoice p c1 c2 =>
    fun st => Vis (GetBiased p) (fun b => if b : bool
                                    then compile_itree c1 st
                                    else compile_itree c2 st)
  | CWhile e body =>
    fun st => ITree.iter (map (fun x =>
                              match x with
                              | inl tt => inr (inl tt)
                              | inr st' =>
                                if e st' then inl st else inr (inr st')
                              end) (compile_itree body)) st
  | CObserve e => fun st => if e st then ret (inr st) else ret (inl tt)
  end.

Fixpoint tree_to_sampler {A : Type}
         (f : nat -> bool) (t : tree A) (env : nat -> tree A) (n : nat)
  : state nat (unit + A) :=
  match n with
  | O => ret (inl tt)
  | S n' =>
    match t with
    | Leaf x => ret (inr x)
    | Fail _ lbl => tree_to_sampler f (env lbl) env n'
    | Choice _ t1 t2 =>
      i <- get ;;
      put (S i) ;;
      tree_to_sampler f (if f i then t1 else t2) env n'
    | Fix lbl t1 =>
      tree_to_sampler f t1 (fun m => if m =? lbl then t1 else env m) n'
    end
  end.

Definition tree_to_sampler' {A : Type} (f : nat -> bool) (t : tree A) (n : nat)
  : state nat (unit + A) :=
  tree_to_sampler f t (const (Fail _ O)) n.

(** Inductive trees to itrees. *)

Fixpoint tree_to_itree {A : Type} (t : tree A) : itree biasedE (nat + A) :=
  match t with
  | Leaf x => ret (inr x)
  | Fail _ lbl => ret (inl lbl)
  | Choice p t1 t2 =>
    Vis (GetBiased p) (fun b => if b : bool
                             then tree_to_itree t2
                             else tree_to_itree t1)
  | Fix lbl t1 => ITree.iter (fun _ => x <- tree_to_itree t1 ;;
                                   match x with
                                   | inl n =>
                                     if n =? lbl
                                     then ret (inl tt)
                                     else ret (inr (inl n))
                                   | inr y =>
                                     ret (inr (inr y))
                                   end) tt
  end.

(** The preimage of a set of samples [P] wrt itree [t] is the set of
    bit sequences [bs] such that [sample t bs P]. *)
Definition preimage {A : Type} (t : itree unbiasedE A) (P : A -> Prop)
  : list bool -> Prop :=
  sample t P.

Require Import List.
Import ListNotations.

Fixpoint sample_nb {A : Type} (P : A -> bool)
         (t : itree unbiasedE A) (bs : list bool) (n : nat) : bool :=
  match n with
  | O =>
    match observe t with
    | RetF x =>
      match bs with
      | [] => P x
      | _ => false
      end
    | _ => false
    end
  | S n' =>
    match observe t with
    | RetF x =>
      match bs with
      | [] => P x
      | _ => false
      end
    | TauF t' => sample_nb P t' bs n'
    | VisF e k =>
      match e in unbiasedE T return (T -> itree unbiasedE A) -> bool with
      | GetUnbiased =>
        fun f => match bs with
              | [] => false
              | b :: bs' => sample_nb P (f b) bs' n'
              end
      end k
    end
  end.

Lemma VisF_inversion {A B : Type} e (f g : bool -> itree unbiasedE B) :
  @VisF unbiasedE A (itree unbiasedE B) bool e f =
  @VisF unbiasedE A (itree unbiasedE B) bool e g ->
  f = g.
Proof.
  intro Heq; inversion Heq.
  apply Classical_Prop.EqdepTheory.inj_pair2 in H0; auto.
Qed.

Lemma sample_n_S {A : Type}
      (t : itree unbiasedE A) (P : A -> Prop) (bs : list bool) (n : nat) :
  sample_n t P bs n -> sample_n t P bs (S n).
Proof.
  intro Hsample.
  induction Hsample.
  - econstructor; eauto.
  - econstructor; eauto.
  - eapply sample_n_vis_false; eauto.
  - eapply sample_n_vis_true; eauto.
Qed.

Lemma sample_n_sumbool {A : Type}
      (t : itree unbiasedE A) (P : A -> Prop) (bs : list bool) :
  (forall x, {P x} + {~ P x}) ->
  forall n, {sample_n t P bs n} + {~ sample_n t P bs n}.
Proof.
  intros Hdec n.
  revert t bs. induction n; intros t bs.
  - destruct (observe t) eqn:Ht;
      try (right; intro HC; inversion HC; subst; rewrite Ht in H; congruence).
    destruct bs; destruct (Hdec r);
      try solve [left; eapply sample_n_ret; eauto];
      right; intro HC; inversion HC; subst.
    rewrite Ht in H; inversion H; subst; congruence.
  - destruct (IHn t bs).
    + left; apply sample_n_S; auto.
    + destruct (observe t) eqn:Ht.
      * destruct (Hdec r); destruct bs;
          try (right; intro HC; inversion HC; subst; congruence);
          left; econstructor; eauto.
      * destruct (IHn t0 bs).
        -- left; econstructor; eauto.
        -- right; intro HC; inversion HC; subst; congruence.
      * destruct e; destruct bs.
        -- right; intro HC; inversion HC; subst; congruence.
        -- destruct b.
           ++ destruct (IHn (k true) bs).
              ** left; eapply sample_n_vis_true; eauto.
              ** right; intro HC; inversion HC; subst.
                 { congruence. }
                 { rewrite Ht in H1; apply VisF_inversion in H1;
                     subst; congruence. }
           ++ destruct (IHn (k false) bs).
              ** left; eapply sample_n_vis_false; eauto.
              ** right; intro HC; inversion HC; subst.
                 { congruence. }
                 { rewrite Ht in H1; apply VisF_inversion in H1;
                     subst; congruence. }
Qed.

Fixpoint walk_path {A : Type} (t : itree unbiasedE A) (bs : list bool) (n : nat)
  : option (itree unbiasedE A) :=
  match n with
  | O => None
  | S n' =>
    match bs with
    | [] => Some t
    | b :: bs' =>
      match observe t with
      | VisF e k =>
        match e in unbiasedE T return
              (T -> itree unbiasedE A) -> option (itree unbiasedE A) with
        | GetUnbiased => fun f => walk_path (f b) bs' n'
        end k
      | TauF t' => walk_path t' bs n'
      | _ => None
      end
    end
  end.

Fixpoint exists_tau {A : Type} (t : itree unbiasedE A) (n : nat) : bool :=
  match n with
  | O => false
  | S n' =>
    match observe t with
    | TauF _ => true
    | VisF e k =>
      match e in unbiasedE T return (T -> itree unbiasedE A) -> bool with
      | GetUnbiased => fun f => exists_tau (f false) n' || exists_tau (f true) n'
      end k
    | RetF _ => false
    end
  end.

(** Will tree [t] either run out of bits in [bs] or successfully
    produce a sample within [n] steps? *)
Inductive terminates {A : Type} : itree unbiasedE A -> list bool -> nat -> Prop :=
| terminates_nil : forall t n,
    terminates t [] n
| terminates_ret : forall t x bs n,
    observe t = RetF x ->
    terminates t bs n
| terminates_tau : forall t t' bs n,
    observe t = TauF t' ->
    terminates t' bs n ->
    terminates t bs (S n)
| terminates_vis_false : forall t t1 t2 bs n,
    observe t = VisF GetUnbiased (fun b => if b : bool then t2 else t1) ->
    terminates t1 bs n ->
    terminates t (false :: bs) (S n)
| terminates_vis_true : forall t t1 t2 bs n,
    observe t = VisF GetUnbiased (fun b => if b : bool then t2 else t1) ->
    terminates t2 bs n ->
    terminates t (true :: bs) (S n).

Definition bit_sequences : nat -> list bool :=
  fun n => removelast (nat_binary (S n)).

Definition bit_sequences_inverse (l : list bool) : nat :=
  pred (binary_nat (l ++ [true])).

Inductive last_element {A : Type} : list A -> A -> Prop :=
| last_element_x : forall x,
    last_element [x] x
| last_element_cons : forall x y z l,
    last_element (z :: l) x ->
    last_element (y :: z :: l) x.

Lemma last_element_remove_last {A : Type} (l : list A) (x : A) :
  last_element l x ->
  removelast l ++ [x] = l.
Proof.
  induction l; simpl; intro Hlast.
  - inversion Hlast.
  - inversion Hlast; subst.
    + reflexivity.
    + simpl in *; rewrite IHl; auto.
Qed.

Lemma last_element_positive_binary (p : positive) :
  last_element (positive_binary p) true.
Proof.
  induction p; simpl.
  - inversion IHp; subst.
    + repeat constructor.
    + constructor; rewrite H; auto.
  - inversion IHp; subst.
    + repeat constructor.
    + constructor; rewrite H; auto.
  - constructor.
Qed.

Lemma last_element_nat_binary (n : nat) :
  n <> O ->
  last_element (nat_binary n) true.
Proof.
  intro Hneq.
  destruct n; try congruence.
  unfold nat_binary; simpl.
  apply last_element_positive_binary.
Qed.

Lemma S_pred (n : nat) :
  (0 < n)%nat ->
  S (pred n) = n.
Proof. destruct n; lia. Qed.

Lemma bit_sequences_right_inverse (bs : list bool) :
  bit_sequences (bit_sequences_inverse bs) = bs.
Proof.
  unfold bit_sequences.
  unfold bit_sequences_inverse.
  rewrite S_pred.
  rewrite nat_binary_binary_nat.
  - rewrite removelast_app.
    + simpl; rewrite app_nil_r; reflexivity.
    + intro HC; inversion HC.
  - apply no_leading_falses_app.
    + intro HC; inversion HC.
    + repeat constructor.
  - induction bs; simpl.
    + unfold binary_nat; simpl; lia.
    + unfold binary_nat, binary_N.
      rewrite no_leading_falses_strip_falses_noop; simpl.
      * destruct (bs ++ [true]) eqn:Hl; lia.
      * rewrite app_comm_cons; apply no_leading_falses_app.
        -- intro HC; inversion HC.
        -- repeat constructor.
Qed.

Lemma bit_sequences_left_inverse (n : nat) :
  bit_sequences_inverse (bit_sequences n) = n.
Proof.
  unfold bit_sequences_inverse.
  unfold bit_sequences.
  rewrite last_element_remove_last.
  - rewrite binary_nat_nat_binary; reflexivity.
  - apply last_element_nat_binary; lia.
Qed.

Lemma bit_sequences_nil_O n :
  bit_sequences n = [] ->
  n = O.
Proof.
  intro H; eapply inj_spec.
  - apply bit_sequences_left_inverse.
  - rewrite H; reflexivity.
Qed.

Lemma is_true_dec {A : Type} (P : A -> bool) :
  forall x, is_true (P x) \/ ~ is_true (P x).
Proof. intro x; destruct (P x); intuition. Qed.

Lemma is_true_sumbool {A : Type} (P : A -> bool) :
  forall x, {is_true (P x)} + {~ is_true (P x)}.
Proof. intro x; destruct (P x); intuition. Defined.

(* (** LPO version of measure_seq (assuming decidability of termination). *) *)
Definition measure_seq {A : Type}
           (t : itree unbiasedE A) (P : A -> bool) (n : nat) : Q :=
  let bs := bit_sequences n in
  if strong_LPO (* (sample_n (fun x => is_true (P x)) t bs) *)
       (sample_n_sumbool t (fun x => is_true (P x))
                         bs (is_true_sumbool P))
  then interval_measure bs
  else 0.

(** Purely propositional version. *)

Inductive measure_seqP {A : Type}
  : itree unbiasedE A -> (A -> Prop) -> nat -> Q -> Prop :=
| measure_seq_false : forall t P n q,
    ~ preimage t P (bit_sequences n) ->
    q == 0 ->
    measure_seqP t P n q
| measure_seq_true : forall t P n q,
    preimage t P (bit_sequences n) ->
    q == interval_measure (bit_sequences n) ->
    measure_seqP t P n q.

Definition measure_chainP {A : Type}
           (t : itree unbiasedE A) (P : A -> Prop) : nat -> Q -> Prop :=
  partial_sumP (measure_seqP t P).

Lemma measure_seqP_deterministic {A : Type}
      (t : itree unbiasedE A) (P : A -> Prop) (n : nat) (p q : Q) :
  measure_seqP t P n p ->
  measure_seqP t P n q ->
  p == q.
Proof.
  intros Hp Hq; inversion Hp; inversion Hq; subst; try lra; congruence.
Qed.

Definition ubP (ub : Q) (f : nat -> Q -> Prop) : Prop :=
  forall n q, f n q -> q <= ub.

Definition supremumP (sup : Q) (f : nat -> Q -> Prop) : Prop :=
  ubP sup f /\ forall ub, ubP ub f -> sup <= ub.

Definition preimage_measure {A : Type}
           (t : itree unbiasedE A) (P : A -> Prop) (m : Q) : Prop :=
  supremumP m (measure_chainP t P).

Instance Proper_measure_seqP {A : Type} (t : itree unbiasedE A) (P : A -> Prop)
  : Proper (eq ==> Qeq ==> iff) (measure_seqP t P).
Proof.
  intros ? n ? p q Heq; subst.
  split; intro H; inversion H; subst; solve [constructor; auto; lra].
Qed.

Definition tie_itree {A : Type} {E : Type -> Type} (t : itree E (unit + A)) (lbl : nat)
  : itree E A :=
  ITree.iter (const t) tt.

Definition itree_topdown_lang {A : Type}
           (t : itree unbiasedE A) (P : A -> bool) (n : nat) : option (list bool) :=
  let bs := bit_sequences n in
  if (strong_LPO (sample_n_sumbool t (fun x => is_true (P x))
                                   bs (is_true_sumbool P)))
  then Some bs
  else None.

Definition meas_seq {A : Type}
           (t : itree unbiasedE A) (P : A -> bool) (n : nat) : meas :=
  match itree_topdown_lang t P n with
  | Some bs => meas_interval bs
  | None => meas_empty
  end.

Definition measure_seq' {A : Type}
           (t : itree unbiasedE A) (P : A -> bool) : nat -> Q :=
  option_measure ∘ itree_topdown_lang t P.

(* Weird thought: any infinite object can be represented, instead of
   by explicit coinduction, by a nat-indexed function of finite
   approximations. E.g., an infinite tree can be taken to be the
   "limit" of a function from a nat n to the nth finite approximation
   of the tree. A proof of a property of the tree could then proceed
   by 1) proving it for all finite approximations (by induction on n),
   and 2) showing that the property is continuous in that it preserves
   suprema (i.e., P(supf) iff forall n, P(f n). We are still using
   induction to prove a universal property (for all nats), it only
   amounts to an "existential" property about a particular infinite
   tree (the infinite nature of the tree "uses up" all of the
   universality of the nat induction). Is there a way to do some kind
   of transfinite induction or something similar, in a sense to gain
   "more universality", to prove universal properties of *all*
   elements of a coinductive type (or a countable subset)? *)

Lemma seq_nodup_topdown_lang {A : Type} (t : itree unbiasedE A) (P : A -> bool) :
  seq_nodup (itree_topdown_lang t P).
Proof.
  intros i j x Hnz H0 H1.
  unfold itree_topdown_lang in *.
  unfold equiv, zero in Hnz; simpl in Hnz.
  repeat destruct_LPO; try congruence.
  clear Hnz e0 e.
  rewrite <- H0 in H1; inversion H1; subst; clear H1.
  eapply inj_spec; eauto.
  apply bit_sequences_left_inverse.
Qed.

Lemma itree_lang_topdown_lang {A : Type} (t : itree unbiasedE A) (P : A -> bool) :
  itree_lang (is_true ∘ P) t (itree_topdown_lang t P).
Proof.
  revert t.
  unfold compose.
  pcofix CH; intro t.
  unfold itree_topdown_lang in *.
  pstep. unfold itree_lang_.
  destruct (observe t) eqn:Ht.
  - destruct (P r0) eqn:Hr0.
    + apply itree_langF_ret_true; auto.
      split.
      * unfold seq_injection_upto_0.
        exists id, id.
        intros i HC.
        split; auto.
        destruct_LPO.
        -- apply sample_sample_n in e.
           inversion e; subst; try solve[inversion Ht]; try congruence.
           symmetry in H. apply bit_sequences_nil_O in H; subst.
           reflexivity.
        -- exfalso; apply HC; reflexivity.
      * unfold seq_injection_upto_0.
        exists id, id. unfold id.
        intros i HC.
        split; auto.
        destruct_LPO.
        -- apply sample_sample_n in e.
           inversion e; subst; try solve[inversion Ht]; try congruence.
           symmetry in H; apply bit_sequences_nil_O in H; subst.
           reflexivity.
        -- apply singleton_seq_nonzero in HC; subst.
           exfalso; apply n.
           apply sample_sample_n.
           econstructor; eauto.
    + assert (H: (fun n : nat =>
                    if
                      strong_LPO
                        (sample_n_sumbool t (fun x : A => is_true (P x))
                                          (bit_sequences n) (is_true_sumbool P))
                    then Some (bit_sequences n)
                    else None) === seq_zero).
      { split; exists id, id; intros i Hnz; unfold id;
          split; auto; unfold equiv, seq_zero, const; destruct_LPO; auto;
            destruct e as [n HC]; inversion HC; subst; congruence. }
      rewrite H.
      apply itree_langF_ret_false; unfold is_true; congruence.
  - econstructor; eauto.
    split; exists id, id; intros i Hnz; split; auto; unfold id in *.
    + destruct_LPO; destruct_LPO.
      * reflexivity.
      * exfalso. apply n.
        destruct e as [m H0].
        exists (S m). econstructor; eauto.
      * exfalso; apply Hnz; reflexivity.
      * reflexivity.
    + destruct_LPO; destruct_LPO.
      * reflexivity.
      * exfalso. apply n.
        destruct e as [m H0].
        inversion H0; subst; try solve[inversion Ht]; try congruence.
        rewrite Ht in H. inversion H; subst.
        eexists; eauto.
      * exfalso; apply Hnz; reflexivity.
      * reflexivity.
  - destruct e.
    eapply itree_langF_vis.
    { right; apply CH. }
    { right; apply CH. }
    clear CH.
    split.
    + exists (fun n => match bit_sequences n with
               | true :: bs => (bit_sequences_inverse bs * 2)%nat
               | false :: bs => (bit_sequences_inverse bs * 2 + 1)%nat
               | nil => O
               end).
      exists (fun n => match n mod 2 with
               | O => bit_sequences_inverse (true :: bit_sequences (n / 2)%nat)
               | S _ => bit_sequences_inverse (false :: bit_sequences (n / 2)%nat)
               end).
      unfold equiv, zero.
      intros i Hnz; simpl in Hnz; split.
      * destruct_LPO; try congruence; clear Hnz.
        destruct e as [n Hsample].
        destruct (bit_sequences i) eqn:Hi.
        -- inversion Hsample; congruence.
        -- destruct b.
           ++ rewrite Nat.mod_mul; auto.
              rewrite Nat.div_mul; auto.
              rewrite bit_sequences_right_inverse.
              rewrite <- Hi, bit_sequences_left_inverse; reflexivity.
           ++ rewrite plus_comm, Nat.mod_add; auto.
              replace ((1 + bit_sequences_inverse l * 2) / 2)%nat with
                  ((Nat.b2n true + bit_sequences_inverse l * 2) / 2)%nat by reflexivity.
              rewrite mult_comm.
              rewrite (Nat.add_b2n_double_div2 true (bit_sequences_inverse l)).
              simpl.
              rewrite bit_sequences_right_inverse.
              rewrite <- Hi, bit_sequences_left_inverse; reflexivity.
      * destruct_LPO; try congruence; clear Hnz.
        destruct e as [n Hsample].
        unfold compose, seq_union.
        destruct (bit_sequences i) eqn:Hi.
        { inversion Hsample; try congruence. }
        destruct b.
        -- rewrite Nat.mod_mul; auto.
           rewrite Nat.div_mul; auto.
           rewrite bit_sequences_right_inverse.
           destruct_LPO; auto.
           exfalso; apply n0; inversion Hsample; subst; try congruence.
           rewrite Ht in H0; inversion H0.
           apply Eqdep.EqdepTheory.inj_pair2 in H1; subst.
           exists n1; auto.
        -- rewrite plus_comm, Nat.mod_add; auto.
           replace ((1 + bit_sequences_inverse l * 2) / 2)%nat with
               ((Nat.b2n true + bit_sequences_inverse l * 2) / 2)%nat by reflexivity.
           rewrite mult_comm.
           rewrite (Nat.add_b2n_double_div2 true (bit_sequences_inverse l)).
           rewrite bit_sequences_right_inverse.
           simpl.
           destruct_LPO; auto.
           exfalso; apply n0; inversion Hsample; subst; try congruence.
           rewrite Ht in H0; inversion H0.
           apply Eqdep.EqdepTheory.inj_pair2 in H1; subst.
           exists n1; auto.
    + exists (fun n => match n mod 2 with
               | O => bit_sequences_inverse (true :: bit_sequences (n / 2)%nat)
               | S _ => bit_sequences_inverse (false :: bit_sequences (n / 2)%nat)
               end).
      exists (fun n => match bit_sequences n with
               | true :: bs => (bit_sequences_inverse bs * 2)%nat
               | false :: bs => (bit_sequences_inverse bs * 2 + 1)%nat
               | nil => O
               end).
      unfold equiv, zero.
      intros i Hnz; simpl in Hnz; split.
      * destruct (mod_2_dec i) as [Hmod | Hmod]; rewrite Hmod.
        -- rewrite bit_sequences_right_inverse.
           rewrite bit_sequences_left_inverse.
           rewrite mult_comm.
           apply mod_n_div; auto.
        -- rewrite bit_sequences_right_inverse.
           rewrite bit_sequences_left_inverse.
           rewrite mult_comm.
           apply mod_n_div_plus_1; auto.
      * unfold seq_union in *.
        destruct (mod_2_dec i) as [Hmod | Hmod]; rewrite Hmod in *;
          unfold compose in *; rewrite bit_sequences_right_inverse;
            repeat destruct_LPO; auto; simpl in Hnz; try congruence;
              destruct e as [m HC]; exfalso; apply n.
        -- exists (S m); eapply sample_n_vis_true; eauto.
        -- exists (S m); eapply sample_n_vis_false; eauto.
Qed.

Definition tree_lang_meas {A : Type}
           (t : tree A) (P : A -> bool) (lbl n : nat) : meas :=
  match tree_sequence t P lbl n with
  | Some bs => meas_interval bs
  | None => meas_empty
  end.

(* Definition tree_lang_measures {A : Type} *)
(*            (t : tree A) (P : A -> bool) (lbl n : nat) : Q := *)
(*   match tree_sequence t P lbl n with *)
(*   | Some bs => interval_measure bs *)
(*   | None => 0 *)
(*   end. *)

(** The preimage of a set P under itree t is a measurable set. *)
Definition preimage_meas {A : Type}
           (t : itree unbiasedE A) (P : A -> bool) : meas :=
  meas_union (meas_seq t P).

Definition preimage_tree_lang_meas {A : Type}
           (t : tree A) (P : A -> bool) (lbl : nat) : meas :=
  meas_union (tree_lang_meas t P lbl).

Lemma sample_n_is_prefix_eq {A : Type}
      (t : itree unbiasedE A) (P Q : A -> Prop) (l1 l2 : list bool) (n m : nat) :
  is_prefix l1 l2 ->
  sample_n t P l1 n ->
  sample_n t Q l2 m ->
  l1 = l2.
Proof.
  revert t l1 l2 m.
  induction n; intros t l1 l2 m Hpre H0 H1.
  - inversion H0; subst; inversion H1; subst; congruence.
  - inversion H0; subst; inversion H1; subst; try congruence.
    + rewrite H2 in H; inversion H; subst.
      eapply IHn; eauto.
    + f_equal.
      rewrite H2 in H; apply VisF_inversion in H; subst.
      inversion Hpre; subst.
      eapply IHn; eauto.
    + inversion Hpre.
    + inversion Hpre.
    + f_equal.
      rewrite H2 in H; apply VisF_inversion in H; subst.
      inversion Hpre; subst.
      eapply IHn; eauto.
Qed.

Lemma sample_n_comparable_eq {A : Type}
      (t : itree unbiasedE A) (P Q : A -> Prop) (l1 l2 : list bool) (n m : nat) :
  comparable l1 l2 ->
  sample_n t P l1 n ->
  sample_n t Q l2 m ->
  l1 = l2.
Proof.
  intros Hcomp H0 H1.
  apply is_prefix_comparable in Hcomp.
  destruct Hcomp.
  - eapply sample_n_is_prefix_eq; eauto.
  - symmetry; eapply sample_n_is_prefix_eq; eauto.
Qed.

(* Lemma itree_lang_comparable_eq {A : Type} *)
(*       (t : itree unbiasedE A) (P Q : A -> Prop) (l1 l2 : list bool)  : *)
(*   comparable l1 l2 -> *)
(*   in_itree_lang P t l1 -> *)
(*   in_itree_lang Q t l2 -> *)
(*   l1 = l2. *)
(* Proof. *)
(*   intros Hcomp H0 H1. *)
(*   apply is_prefix_comparable in Hcomp. *)
(*   destruct Hcomp. *)
(*   - eapply in_itree_lang_is_prefix_eq; eauto. *)
(*   - symmetry; eapply in_itree_lang_is_prefix_eq; eauto. *)
(* Qed. *)

Inductive last_element_same {A : Type} : list A -> list A -> Prop :=
| last_element_same_x : forall x,
    last_element_same [x] [x]
| last_element_same_cons1 : forall x xs ys,
    last_element_same xs ys ->
    last_element_same (x :: xs) ys
| last_element_same_cons2 : forall xs y ys,
    last_element_same xs ys ->
    last_element_same xs (y :: ys).

(** The last elements may be the same but don't have to be. *)
Inductive all_but_last_same {A : Type} : list A -> list A -> Prop :=
(* | all_but_last_same_nil : *)
(*     all_but_last_same [] [] *)
| all_but_last_same_xy : forall x y,
    all_but_last_same [x] [y]
| all_but_last_same_cons : forall x xs ys,
    all_but_last_same xs ys ->
    all_but_last_same (x :: xs) (x :: ys).

Lemma not_last_element_same_nil {A : Type} (l : list A) :
  ~ last_element_same [] l.
Proof.
  induction l; simpl; intro HC; inversion HC; subst; congruence.
Qed.

Instance Symmetric_last_element_same {A : Type}
  : Symmetric (@last_element_same A).
Proof. induction 1; constructor; auto. Qed.

Lemma last_element_same_singleton {A : Type} (x y : A) :
  last_element_same [x] [y] ->
  x = y.
Proof.
  intro Hlast.
  inversion Hlast; subst; auto.
  - apply not_last_element_same_nil in H2; contradiction.
  - exfalso. eapply not_last_element_same_nil; symmetry; eauto.
Qed.

Lemma last_element_same_last_element {A : Type} (l : list A) (x : A) :
  last_element_same l [x] ->
  last_element l x.
Proof.
  induction l; intro Hlast.
  - apply not_last_element_same_nil in Hlast; contradiction.
  - inversion Hlast; subst.
    + constructor.
    + destruct l.
      * apply not_last_element_same_nil in H2; contradiction.
      * constructor; auto.
    + symmetry in H1.
      apply not_last_element_same_nil in H1; contradiction.
Qed.

Lemma removelast_all_but_last_same {A : Type} (l1 l2 : list A) :
  l1 <> nil -> l2 <> nil ->
  removelast l1 = removelast l2 <-> all_but_last_same l1 l2.
Proof.
  intros H0 H1.
  split.
  - revert H0 H1. revert l2.
    induction l1; intros l2 H0 h1 Heq; try congruence.
    destruct l2; try congruence.
    destruct l1, l2; simpl in *; try congruence.
    + constructor.
    + inversion Heq; subst.
      constructor.
      apply IHl1; auto; intros HC; inversion HC.
  - induction 1.
    + reflexivity.
    + simpl.
      destruct xs, ys; auto.
      * inversion H.
      * inversion H.
      * f_equal; apply IHall_but_last_same; intros HC; inversion HC.
Qed.

Lemma last_element_same_cons {A : Type} (x y : A) (l1 l2 : list A) :
  last_element_same l1 (x :: y :: l2) ->
  last_element_same l1 (y :: l2).
Proof.
  revert x y l2.
  induction l1; intros x y l2 Hlast.
  - apply not_last_element_same_nil in Hlast; contradiction.
  - inversion Hlast; subst; auto.
    apply IHl1 in H2; constructor; auto.
Qed.

Lemma last_element_same_cons' {A : Type} (x y : A) (l1 l2 : list A) :
  last_element_same (x :: l1) (y :: l2) ->
  (l1 = nil /\ l2 = nil /\ x = y) \/
  (l1 = nil /\ last_element l2 x) \/
  (l2 = nil /\ last_element l1 y) \/
  last_element_same l1 l2.
Proof.
  intros Hlast.
  inversion Hlast; subst; auto.
  - inversion H2; subst; auto.
    + right. right. left. split; auto; constructor.
    + right. right.
      destruct l2.
      * left; split; auto.
        apply last_element_same_last_element in H.
        destruct xs.
        inversion H. constructor; auto.
      * right; constructor.
        eapply last_element_same_cons; eauto.
  - inversion H1; subst; auto.
    + right. left. split; auto.
      apply last_element_same_last_element; auto.
    + right.
      destruct l1.
      * left; split; auto.
        apply last_element_same_last_element.
        constructor.
        symmetry; auto.
      * right. right.
        symmetry.
        eapply last_element_same_cons.
        constructor.
        symmetry; eauto.
Qed.

Lemma all_but_last_same_length {A : Type} (l1 l2 : list A) :
  all_but_last_same l1 l2 ->
  length l1 = length l2.
Proof. induction 1; simpl; auto. Qed.

Lemma last_element_all_but_last_eq {A : Type} (l1 l2 : list A) :
  last_element_same l1 l2 ->
  all_but_last_same l1 l2 ->
  l1 = l2.
Proof.
  revert l2.
  induction l1; intros l2 Hlast Hall.
  - apply not_last_element_same_nil in Hlast; contradiction.
  - destruct l2.
    + symmetry in Hlast.
      apply not_last_element_same_nil in Hlast; contradiction.
    + apply last_element_same_cons' in Hlast.
      destruct Hlast as [[? [? ?]] | [[? ?] | [[? ?] | ?]]]; subst.
      * reflexivity.
      * apply all_but_last_same_length in Hall.
        simpl in *. destruct l2; simpl in *; try lia.
        inversion H0.
      * apply all_but_last_same_length in Hall.
        simpl in *. destruct l1; simpl in *; try lia.
        inversion H0.
      * inversion Hall; subst.
        -- inversion H.
        -- erewrite IHl1; eauto.
Qed.

Lemma removelast_inj {A : Type} (l1 l2 : list A) :
  last_element_same l1 l2 ->
  removelast l1 = removelast l2 ->
  l1 = l2.
Proof.
  destruct l1, l2; auto; intros Hlast Heq.
  - apply not_last_element_same_nil in Hlast; contradiction.
  - symmetry in Hlast.
    apply not_last_element_same_nil in Hlast; contradiction.
  - apply removelast_all_but_last_same in Heq;
      try solve[intro HC; inversion HC].
    apply last_element_all_but_last_eq; auto.
Qed.

Lemma last_element_singleton_same {A : Type} (l : list A) (x : A) :
  last_element l x ->
  last_element_same [x] l.
Proof.
  induction l; simpl; intro Hlast; inversion Hlast; subst.
  - constructor.
  - apply last_element_same_cons2; auto.
Qed.

Lemma last_element_last_element_same {A : Type} (l1 l2 : list A) (x : A) :
  last_element l1 x ->
  last_element l2 x ->
  last_element_same l1 l2.
Proof.
  revert l2.
  induction l1; intros l2 H0 H1.
  - inversion H0.
  - inversion H0; subst.
    + apply last_element_singleton_same; auto.
    + apply last_element_same_cons1; auto.
Qed.

Lemma bit_sequences_inj (i j : nat) :
  bit_sequences i = bit_sequences j ->
  i = j.
Proof.
  unfold bit_sequences.
  intro H.
  apply removelast_inj in H.
  - apply nat_binary_injective in H; lia.
  - clear H.
    eapply last_element_last_element_same;
      apply last_element_nat_binary; lia.
Qed.

Lemma meas_seq_disjoint {A : Type}
      (t : itree unbiasedE A) (P : A -> bool) (i j : nat) :
  i <> j ->
  borel.disjoint (meas_seq t P i) (meas_seq t P j).
Proof.
  intro Hneq; unfold meas_seq, itree_topdown_lang.
  repeat destruct_LPO; try constructor.
  destruct e as [n H0].
  destruct e0 as [m H1].
  unfold interval_disjoint.
  apply incomparable_not_comparable.
  intro HC.
  eapply sample_n_comparable_eq in HC; eauto.
  apply bit_sequences_inj in HC; congruence.
Qed.

Lemma measure_seq_measure_seq' {A : Type} (t : itree unbiasedE A) (P : A -> bool) :
  forall i, measure_seq t P i = measure_seq' t P i.
Proof.
  unfold measure_seq, measure_seq', compose, itree_topdown_lang.
  intro i; destruct_LPO; auto.
Qed.

Instance Proper_option_measure : Proper (equiv ==> equiv) option_measure.
Proof. intros [] [] Heq; congruence. Qed.

Lemma seq_bijection_upto_0_tree_sequence_measure_seq
      {A : Type} (P : A -> bool) (t : tree A) (n : nat) :
  wf_tree t ->
  nondivergent'' n t ->
  not_bound_in n t ->
  (forall l : nat, free_in l t -> l = n) ->
  option_measure ∘ tree_sequence t P n === measure_seq (unbiased_tree_to_itree' t) P.
Proof.
  intros Hwf Hnd Hnotbound Hfree.
  transitivity (measure_seq' (unbiased_tree_to_itree' t) P).
  2: { apply ext_eq_seq_bijection_upto_0.
       intro i; rewrite measure_seq_measure_seq'; reflexivity. }
  unfold measure_seq'.
  apply compose_seq_bijection_upto_0; try reflexivity.
  { apply Proper_option_measure. }
  eapply itree_fintau_itree_lang_unique with (P0 := is_true ∘ P).
  - eapply nondivergent''_itree_fintau; eauto.
  - apply itree_lang_tree'; auto.
  - apply itree_lang_topdown_lang.
  - apply (@seq_nodup_tree_sequence A t [n] P); auto.
    + inversion Hnd; auto.
    + intros m Hbound [? | []]; subst.
      apply bound_in_not_bound_in in Hnotbound; congruence.
  - apply seq_nodup_topdown_lang.
Qed.

Lemma infer_supremum_measure {A : Type} (P : A -> bool) (t : tree A) (n : nat) :
  wf_tree t ->
  unbiased t ->
  (forall m, free_in m t -> m = n) ->
  not_bound_in n t ->
  nondivergent'' n t ->
  supremum (infer (fun x : A => if P x then 1 else 0) t)
           (partial_sum (measure_seq (unbiased_tree_to_itree' t) P)).
Proof.
  intros Hwf Hunbiased Hfree Hnotbound Hnd.
  apply seq_bijection_upto_0_supremum with (s1:=option_measure ∘ tree_sequence t P n).
  - intro i; apply option_measure_nonnegative.
  - intro i; unfold measure_seq; destruct_LPO; try lra;
      apply interval_measure_nonnegative.
  - apply seq_bijection_upto_0_tree_sequence_measure_seq; auto.
  - apply infer_supremum; auto.
Qed.

(* (* TODO: replace measure_seq with the analogue using tree_lang_meas *)
(*    and see how far we can go in the sampling theorem using this and *)
(*    preimage_tree_lang_meas_infer instead of the topdown lang stuff. *) *)
(* Lemma infer_supremum_measure' {A : Type} (P : A -> bool) (t : tree A) (n : nat) : *)
(*   wf_tree t -> *)
(*   unbiased t -> *)
(*   (forall m, free_in m t -> m = n) -> *)
(*   not_bound_in n t -> *)
(*   nondivergent'' n t -> *)
(*   supremum (infer (fun x : A => if P x then 1 else 0) t) *)
(*            (partial_sum (option_measure ∘ tree_sequence t P n)). *)
(* Proof. *)
(*   intros Hwf Hunbiased Hfree Hnotbound Hnd. *)
(*   apply infer_supremum; auto. *)
(* Qed. *)
(*   apply seq_bijection_upto_0_supremum with (s1:=option_measure ∘ tree_sequence t P n). *)
(*   - intro i; apply option_measure_nonnegative. *)
(*   - intro i; unfold measure_seq; destruct_LPO; try lra; *)
(*       apply interval_measure_nonnegative. *)
(*   - apply seq_bijection_upto_0_tree_sequence_measure_seq; auto. *)
(*   - apply infer_supremum; auto. *)
(* Qed. *)

Lemma preimage_tree_lang_meas_infer {A : Type} (t : tree A) (P : A -> bool) (n : nat) :
  wf_tree t ->
  unbiased t ->
  (forall m, free_in m t -> m = n) ->
  not_bound_in n t ->
  nondivergent'' n t ->
  measure (preimage_tree_lang_meas t P n)
          (infer (fun x => if P x then 1 else 0) t).
Proof.
  intros Hwf Hunbiased Hfree Hnotbound Hnd.
  unfold preimage_tree_lang_meas.
  econstructor.
  3: { eapply infer_supremum; eauto. }
  unfold tree_lang_meas.
  - intros i j Hneq.
    destruct (tree_sequence t P n i) eqn:Hi, (tree_sequence t P n j) eqn:Hj;
      constructor.
    eapply seq_incomparable_tree_sequence; eauto.
    + inversion Hnd; subst; eauto.
    + intros m Hbound [Hin | []]; subst.
      apply bound_in_not_bound_in in Hnotbound; congruence.
  - unfold compose, option_measure.
    intro i. unfold tree_lang_meas.
    destruct (tree_sequence t P n i); constructor; reflexivity.
Qed.

Lemma preimage_meas_infer {A : Type} (t : tree A) (P : A -> bool) (n : nat) :
  wf_tree t ->
  unbiased t ->
  (forall m, free_in m t -> m = n) ->
  not_bound_in n t ->
  nondivergent'' n t ->
  measure (preimage_meas (unbiased_tree_to_itree' t) P)
          (infer (fun x => if P x then 1 else 0) t).
Proof.
  intros Hwf Hunbiased Hfree Hnotbound Hnd.
  apply measure_union with (g := measure_seq (unbiased_tree_to_itree' t) P).
  - intros i j Hneq; apply meas_seq_disjoint; auto.
  - intro i; unfold meas_seq, measure_seq, itree_topdown_lang.
    destruct_LPO; constructor; reflexivity.
  - eapply infer_supremum_measure; eauto.
Qed.

Lemma sample_n_singleton_P {A : Type}
      (t : itree unbiasedE A) (P : A -> Prop) (bs : list bool) (n m : nat) (x : A) :
  sample_n t (fun y => y = x) bs n ->
  sample_n t P bs m ->
  P x.
Proof.
  revert t bs m.
  induction n; intros t bs m H0 H1;
    inversion H0; subst; inversion H1; subst; try congruence.
  - rewrite H2 in H; inversion H; subst; eapply IHn; eauto.
  - rewrite H2 in H3; apply VisF_inversion in H3; subst.
    eapply IHn; eauto.
  - rewrite H2 in H3; apply VisF_inversion in H3; subst.
    eapply IHn; eauto.
Qed.

Lemma sample_n_singleton_P' {A : Type}
      (t : itree unbiasedE A) (P : A -> Prop) (bs : list bool) (n : nat) (x : A) :
  sample_n t (fun y => y = x) bs n ->
  P x ->
  sample_n t P bs n.
Proof.
  revert t bs.
  induction n; intros t bs H0 H1;
    inversion H0; subst; solve[econstructor; eauto].
Qed.

Lemma sample_n_real_in_measb {A : Type}
      (t : itree unbiasedE A) (r : real) (x : A) (P : A -> bool) (n m : nat) :
  sample_n t (fun y : A => y = x) (prefix r m) n ->
  (exists n0, real_in_measb (meas_seq t P n0) r = true) <-> P x = true.
Proof.
  intro Hsample. split.
  - intros [n0 Hin].
    unfold meas_seq, itree_topdown_lang in Hin.
    destruct_LPO.
    + destruct e as [n' H].
      unfold real_in_measb in Hin.
      unfold real_in_intervalb in Hin.
      destruct (real_in_intervalb_aux_spec r (bit_sequences n0) 0);
        try congruence; clear Hin.
      apply real_in_interval_exists_prefix in r0.
      rewrite <- r0 in H; clear r0.
      assert (Heq: prefix r m = prefix r (length (bit_sequences n0))).
      { eapply sample_n_comparable_eq; eauto.
        apply comparable_prefix. }
      rewrite <- Heq in H; clear Heq.
      set (P' := fun x => P x = true).
      assert (P' x).
      { eapply sample_n_singleton_P; eauto. }
      apply H0.
    + simpl in Hin; congruence.
  - intro HP.
    unfold meas_seq, itree_topdown_lang.
    exists (bit_sequences_inverse (prefix r m)).
    destruct_LPO.
    + simpl.
      destruct e as [n0 H].
      rewrite bit_sequences_right_inverse.
      destruct (real_in_intervalb_spec r (prefix r m)); auto.
      exfalso; apply n1.
      apply real_in_interval_prefix.
    + exfalso.
      apply n0. exists n. rewrite bit_sequences_right_inverse; auto.
      eapply sample_n_singleton_P'; eauto.
Qed.

(* idea: use the top-down construction of the sequence using
   sample_n and strong LPO as before, and show that it's a language of
   the itree. Then, equivalence to the inductive construction follows
   from uniqueness of the itree language (which only holds under
   nondivergence). *)

(* TODO: sampling fixpoint on inductive trees? using a map from labels to subtrees *)

Lemma uniform_samples_nondivergent {A : Type} `{EqType A} (t : tree A) (lbl : nat)
      (reals : nat -> real) (samples : nat -> A) :
  uniform reals ->
  (forall i, exists m n, sample_n (unbiased_tree_to_itree' t)
                        (fun x => x = samples i) (prefix (reals i) m) n) ->
  nondivergent'' lbl t.
Proof.
  intros Huniform Hsample.
  unfold unbiased_tree_to_itree', tie_itree' in Hsample.

  assert (H0: forall i, exists n, sample (unbiased_tree_to_itree' t)
                           (eq (samples i)) (prefix (reals i) n)).
  { intros i; specialize (Hsample i).
    destruct Hsample as [m [n Hsample]].
    exists m; apply sample_sample_n; exists n.
    eapply sample_n_singleton_P'; eauto. }
  clear Hsample.
  unfold nondivergent''.
  destruct (nondivergent''_dec lbl t) as [? | Hdiv]; auto.
  exfalso.
  unfold divergent' in Hdiv.

  (* may need induction on Hdiv? *)
  inversion Hdiv; subst.
  destruct H3.
  - destruct H1 as [H1 H2].
    admit.
  -

Admitted.

(** Sampling theorem. *)
Section sampling_theorem.
  Context (A : Type) `{EqType A} (t : tree A) (lbl : nat)
          (reals : nat -> real) (samples : nat -> A) (P : A -> bool)
          (Hwf : wf_tree t) (Hunbiased : unbiased t)
          (Hfree : forall m, free_in m t -> m = lbl)
          (Hnotbound : not_bound_in lbl t)
          (Hnd : nondivergent'' lbl t).

  Definition it : itree unbiasedE A := unbiased_tree_to_itree' t.

  Hypothesis reals_uniform : uniform reals.
  Hypothesis reals_samples
    : forall i, exists m n, sample_n it (fun x => x = samples i) (prefix (reals i) m) n.

  Definition P_prob : Q := infer (fun x => if P x then 1 else 0) t.

  (* Context (A : Type) `{EqType A}. *)
  (* Variable t : tree A. *)
  (* Variable lbl : nat. *)
  (* Hypothesis Hwf : wf_tree t. *)
  (* Hypothesis Hunbiased : unbiased t. *)
  (* Hypothesis Hfree : forall m, free_in m t -> m = lbl. *)
  (* Hypothesis Hnotbound : not_bound_in lbl t. *)
  (* Hypothesis Hnd : nondivergent'' lbl t. (* eliminate somehow? should *)
  (*                                           be implied by reals_samples below. *) *)

  (* Definition it : itree unbiasedE A := unbiased_tree_to_itree' t. *)

  (* Variable reals : nat -> real. *)
  (* Hypothesis reals_uniform : uniform reals. *)

  (* Variable samples : nat -> A. *)
  (* Hypothesis reals_samples *)
    (* : forall i, exists m n, sample_n it (fun x => x = samples i) (prefix (reals i) m) n. *)

  (* Variable P : A -> bool. *)
  (* Definition P_prob : Q := infer (fun x => if P x then 1 else 0) t. *)

  Theorem samples_converge :
    forall eps : Q,
      0 < eps ->
      exists n0 : nat, forall n, (n0 <= n)%nat -> Qabs (P_prob - rel_freq P (prefix samples n)) < eps.
  Proof.
    intros eps Heps.

    specialize (reals_uniform (preimage_meas_infer P Hwf Hunbiased Hfree Hnotbound Hnd) Heps).
    destruct reals_uniform as [n0 Habs].
    exists n0. unfold rel_freq' in Habs.
    intros n Hle.
    assert (Heq: rel_freq (real_in_measb (preimage_meas it P))
                          (prefix reals n) ==
                 rel_freq P (prefix samples n)).
    { rewrite rel_freq_eq with (Q:=P) (g:=samples).
      - reflexivity.
      - intro i; destruct (reals_samples i) as [m [fuel reals_sample]].
        apply sample_n_real_in_measb with (P0:=P) in reals_sample.
        simpl; destruct_LPO.
        + apply reals_sample in e; auto.
        + destruct (P (samples i)); auto.
          exfalso; apply n1; apply reals_sample; auto. }
    rewrite <- Heq; auto.

    (* specialize (reals_uniform (preimage_tree_lang_meas_infer P Hwf Hunbiased Hfree Hnotbound Hnd) Heps). *)
    (* destruct reals_uniform as [n Habs]. *)
    (* exists n. unfold rel_freq' in Habs. *)
    (* assert (Heq: rel_freq (real_in_measb (preimage_tree_lang_meas t P lbl)) *)
    (*                       (prefix reals n) == *)
    (*              rel_freq P (prefix samples n)). *)
    (* { rewrite rel_freq_eq with (Q:=P) (g:=samples). *)
    (*   - reflexivity. *)
    (*   - intro i; destruct (reals_samples i) as [m [fuel reals_sample]]. *)
    (*     apply sample_n_real_in_measb_tree with (P0:=P) (lbl0:=lbl) in reals_sample. *)
    (*     simpl; destruct_LPO. *)
    (*     + apply reals_sample in e; auto. *)
    (*     + destruct (P (samples i)); auto. *)
    (*       exfalso; apply n0; apply reals_sample; auto. } *)
    (* rewrite <- Heq; auto. *)

  Qed.
End sampling_theorem.

(** TODO: translating biased choice events into bernoulli trees with
  unbiased choice events. *)
Definition handle_biased {A : Type} (e : biasedE A) : itree unbiasedE A :=
  match e with
  | GetBiased p => diverge
  end.

(** TODO:

 * Two versions of inference? "high" and "low" level inference
     (parameterized probabilistic choice and unbiased (uniform)
     choice).

   Compiling high to low level (replacing choice nodes with bernoulli
     subtrees) and correctness theorem?

 * Fueled inference function on itrees.

 * NTS either:

   - given an inductive tree, fueled inference on the itree compiled
     from it converges in the limit to inference on the inductive
     tree.

   - given a command, fueled inference on the itree compiled from it
     converges in the limit to cwp of the command (wrt some initial
     state).

   The first implies the second due to the already established
   connection between inductive abstract trees and cwp.

 * Need a bernoulli tree construction, either directly as an itree or
   first as an inductive tree that can then be converted to an itree.

 ONE EXTREME: compile commands directly to abstract itrees and then
 translate those to ky itrees, reasoning directly about fueled
 inference (wrt cwp first, then wrt the two versions of fueled
 inference). It could be tricky connecting the two versions of fueled
 inference because the low level version will require more fuel in
 general.

 The theorem could say something like: forall t : itree, forall n :
 nat, exists n' : nat, infer n t == infer' n' (translate t)

 The inductive case for choice nodes would be the meat of the proof,
 requiring a proof of correctness of the bernoulli tree construction.

 OTHER EXTREME: compile commands to inductive abstract trees, then to
 inductive ky trees, then to ky itrees, with a theorem similar to
 above connecting inference on the inductive tree types, and then
 connecting to itrees only at the lowest level.

 IN BETWEEN: commands to abstract trees, then to abstract itrees, then
 translate to ky itrees.

 DIAGRAM:

 cpGCL -> abstract tree -> ky tree

      \ | | \ v v | \ v -> abstract itree -> ky itree

 cpGCL' -> ky tree or ky itree directly

* cpGCL -> abstract tree : done. state monad compiler.

* abstract tree -> ky tree : replace choice nodes with bernoulli trees
                             (with unbiased flip nodes)

* cpGCL -> abstract itree : done. relatively straightforward.

* abstract itree -> ky itree : event handler that constructs ky itrees
                               given a rational. Should be similar to
                               the bernoulli tree construction for
                               inductive trees. Could even just use
                               the inductive tree construction and
                               convert it.

* abstract tree -> abstract itree : done. relatively straightforward.

* ky tree -> ky itree : should be almost the exact same as for
                        abstract trees.

* cpGCL -> cpGCL' : desugaring probabilistic choices to loops and
                    unbiased choices.

Three main features of cpGCL:
* loops
* probabilistic choice
* conditioning

Compiling from cpGCL to trees can be seen as eliminating conditioning
(replacing with a big outer loop). Compilation is only correct when
all loops are iid, and the outer loop is always iid.

Translating from "abstract" to "ky" is reduction from parametric
probabilistic choice to more primitive uniform (unbiased) coin
flips.

Translating from inductive trees to coinductive trees can be seen
as eliminating loops by taking their infinite expansions.

In the end we arrive at "ky itrees" in which conditioning and loops
have been eliminated, and probabilistic choice is reduced to its most
primitive form.

*)

(** Abstract trees -> abstract itrees

- Theorem: converges (infer_n f t) (infer f t)

  How to prove? Possibly just by induction and using a geometric
  series argument in the fix case...

*)


From Test Require Import tactic.

Section FOFProblem.

Variable Universe : Set.
Variable UniverseElement : Universe.

Variable pl_ : Universe -> Universe -> Prop.
Variable or6_ : Universe -> Universe -> Universe -> Universe -> Universe -> Universe -> Universe -> Universe -> Prop.
Variable or3_ : Universe -> Universe -> Universe -> Universe -> Prop.
Variable or2_ : Universe -> Universe -> Universe -> Prop.
Variable goal_ : Prop.
Variable ep_ : Universe -> Universe -> Prop.
Variable el_ : Universe -> Universe -> Prop.
Variable dom_ : Universe -> Prop.
Variable col_ : Universe -> Universe -> Universe -> Universe -> Prop.

Variable s_ : Universe.
Variable r_ : Universe.
Variable q_ : Universe.
Variable p_ : Universe.
Variable o_ : Universe.
Variable n_ : Universe.
Variable m_ : Universe.
Variable l_ : Universe.
Variable i_ : Universe.
Variable h_ : Universe.
Variable g_ : Universe.
Variable f_ : Universe.
Variable e_ : Universe.
Variable d_ : Universe.
Variable c_ : Universe.
Variable b_ : Universe.
Variable a_ : Universe.

Variable initial_model_1 : (dom_ a_ /\ (dom_ b_ /\ (dom_ c_ /\ (dom_ d_ /\ (dom_ e_ /\ (dom_ f_ /\ (dom_ g_ /\ (dom_ h_ /\ (dom_ i_ /\ (dom_ l_ /\ (dom_ m_ /\ (dom_ n_ /\ (dom_ o_ /\ (dom_ p_ /\ (dom_ q_ /\ (dom_ r_ /\ (dom_ s_ /\ (col_ a_ b_ c_ l_ /\ (col_ d_ e_ f_ m_ /\ (col_ b_ f_ g_ n_ /\ (col_ c_ e_ g_ o_ /\ (col_ b_ d_ h_ p_ /\ (col_ a_ e_ h_ q_ /\ (col_ c_ d_ i_ r_ /\ col_ a_ f_ i_ s_)))))))))))))))))))))))).
Variable goal1_2 : (el_ n_ o_ -> goal_).
Variable goal2_3 : (el_ p_ q_ -> goal_).
Variable goal3_4 : (el_ s_ r_ -> goal_).
Variable goal4_5 : (forall U : Universe, ((el_ U U /\ (pl_ g_ U /\ (pl_ h_ U /\ pl_ i_ U))) -> goal_)).
Variable col_elim1_6 : (forall P Q R L : Universe, (col_ P Q R L -> pl_ P L)).
Variable col_elim2_7 : (forall P Q R L : Universe, (col_ P Q R L -> pl_ Q L)).
Variable col_elim3_8 : (forall P Q R L : Universe, (col_ P Q R L -> pl_ R L)).
Variable or_6_9 : (forall A B C M D E F L : Universe, (or6_ A B C M D E F L -> (or3_ A B C M \/ or3_ D E F L))).
Variable or_3_10 : (forall P Q R L : Universe, (or3_ P Q R L -> (pl_ P L \/ or2_ Q R L))).
Variable or_2_11 : (forall Q R L : Universe, (or2_ Q R L -> (pl_ Q L \/ pl_ R L))).
Variable pref_12 : (forall P L : Universe, (pl_ P L -> ep_ P P)).
Variable psym_13 : (forall P Q : Universe, (ep_ P Q -> ep_ Q P)).
Variable ptra_14 : (forall P Q R : Universe, ((ep_ P Q /\ ep_ Q R) -> ep_ P R)).
Variable lref_15 : (forall P L : Universe, (pl_ P L -> el_ L L)).
Variable lsym_16 : (forall L M : Universe, (el_ L M -> el_ M L)).
Variable ltra_17 : (forall L M N : Universe, ((el_ L M /\ el_ M N) -> el_ L N)).
Variable pcon_18 : (forall P Q L : Universe, ((ep_ P Q /\ pl_ Q L) -> pl_ P L)).
Variable lcon_19 : (forall P L M : Universe, ((pl_ P L /\ el_ L M) -> pl_ P M)).
Variable unique_20 : (forall P L M Q : Universe, ((pl_ P L /\ (pl_ P M /\ (pl_ Q L /\ pl_ Q M))) -> (ep_ P Q \/ el_ L M))).
Variable line_21 : (forall P Q : Universe, ((ep_ P P /\ ep_ Q Q) -> (exists L : Universe, (dom_ L /\ (pl_ P L /\ pl_ Q L))))).
Variable papp1_22 : (forall A B C L D E F M G N O H P Q I R S : Universe, ((col_ A B C L /\ (col_ D E F M /\ (col_ B F G N /\ (col_ C E G O /\ (col_ B D H P /\ (col_ A E H Q /\ (col_ C D I R /\ col_ A F I S))))))) -> ((exists T : Universe, (dom_ T /\ col_ G H I T)) \/ or6_ A B C M D E F L))).
Variable point_23 : (forall L M : Universe, ((el_ L L /\ el_ M M) -> (exists P : Universe, (dom_ P /\ (pl_ P L /\ pl_ P M))))).

Theorem pp_24 : goal_.
Proof.
  time tac.
Qed.

End FOFProblem.


From FourCerty Require Import Utility StackLang.

Import Utility StackLang.

Module StackLangRel.

Open Scope Z_scope.

Inductive do_cmp_R : ins_cmp -> Z -> Z -> bool -> Prop :=
  | DC_Eq : forall i1 i2,
      do_cmp_R C_Eq i1 i2 (i1 =? i2)
  | DC_Ne : forall i1 i2,
      do_cmp_R C_Ne i1 i2 (negb (i1 =? i2))
  | DC_Lt : forall i1 i2,
      do_cmp_R C_Lt i1 i2 (i1 <? i2)
  | DC_Le : forall i1 i2,
      do_cmp_R C_Le i1 i2 (i1 <=? i2)
  | DC_Gt : forall i1 i2,
      do_cmp_R C_Gt i1 i2 (i1 >? i2)
  | DC_Ge : forall i1 i2,
      do_cmp_R C_Ge i1 i2 (i1 >=? i2).

Inductive do_uop_R : ins_uop -> ins_val -> ins_val -> Prop :=
  | DU_Add1 : forall i,
      do_uop_R U_Add1 (V_Int i) (V_Int (i + 1))
  | DU_Sub1 : forall i,
      do_uop_R U_Sub1 (V_Int i) (V_Int (i - 1))
  | DU_NotFalse : do_uop_R U_Not (V_Bool false) (V_Bool true)
  | DU_NotTrue : do_uop_R U_Not (V_Bool true) (V_Bool false)
  | DU_NotInt : forall i, do_uop_R U_Not (V_Int i) (V_Bool false).

Inductive do_bop_R : ins_bop -> ins_val -> ins_val -> ins_val -> Prop :=
  | DB_Add : forall i1 i2,
      do_bop_R B_Add (V_Int i1) (V_Int i2) (V_Int (i1 + i2))
  | DB_Sub : forall i1 i2,
      do_bop_R B_Sub (V_Int i1) (V_Int i2) (V_Int (i1 - i2))
  | DB_And : forall b1 b2,
      do_bop_R B_And (V_Bool b1) (V_Bool b2) (V_Bool (b1 && b2))
  | DB_Or : forall b1 b2,
      do_bop_R B_Or (V_Bool b1) (V_Bool b2) (V_Bool (b1 || b2)).

Inductive eval_tm_R : nat -> partial_map stk_fun -> stk_tm -> list ins_val -> list ins_val -> Prop :=
  | E_End : forall f funs vs,
      eval_tm_R f funs End vs vs
  | E_InsCall : forall f f' funs inss' inss'' args rst fn n vs1 vs2 vs3,
      f = S f' ->
      args = firstn n vs1 ->
      rst = skipn n vs1 ->
      List.length args = n ->
      Ok (Fun fn n inss'') = lookup funs fn ->
      eval_tm_R f' funs inss'' args vs2 ->
      eval_tm_R f funs inss' (vs2 ++ rst) vs3 ->
      eval_tm_R f funs (Ins (Call fn n) inss') vs1 vs3
  | E_InsPush : forall f funs v inss' vs1 vs2,
      eval_tm_R f funs inss' (v :: vs1) vs2 ->
      eval_tm_R f funs (Ins (Push v) inss') vs1 vs2
  | E_InsPop : forall f funs inss' v rst vs,
      eval_tm_R f funs inss' rst vs ->
      eval_tm_R f funs (Ins Pop inss') (v :: rst) vs
  | E_InsSwap : forall f funs inss' v1 v2 rst vs,
      eval_tm_R f funs inss' (v2 :: v1 :: rst) vs ->
      eval_tm_R f funs (Ins Swap inss') (v1 :: v2 :: rst) vs
  | E_InsStkRef : forall f funs inss' n v vs1 vs2,
      Some v = nth_error vs1 n ->
      eval_tm_R f funs inss' (v :: vs1) vs2 ->
      eval_tm_R f funs (Ins (StkRef n) inss') vs1 vs2
  | E_InsUop : forall f funs inss' op v1 v2 rst vs,
      do_uop_R op v1 v2 ->
      eval_tm_R f funs inss' (v2 :: rst) vs ->
      eval_tm_R f funs (Ins (Uop op) inss') (v1 :: rst) vs
  | E_InsBop : forall f funs inss' op v1 v2 v3 rst vs,
      do_bop_R op v1 v2 v3 ->
      eval_tm_R f funs inss' (v3 :: rst) vs ->
      eval_tm_R f funs (Ins (Bop op) inss') (v2 :: v1 :: rst) vs
  | E_InsCmp : forall f funs inss' op i1 i2 b rst vs,
      do_cmp_R op i1 i2 b ->
      eval_tm_R f funs inss' ((V_Bool b) :: rst) vs ->
      eval_tm_R f funs (Ins (Cmp op) inss') ((V_Int i2) :: (V_Int i1) :: rst) vs
  | E_IfFalse : forall f funs thn els nxt rst vs1 vs2,
      eval_tm_R f funs els rst vs1 ->
      eval_tm_R f funs nxt vs1 vs2 ->
      eval_tm_R f funs (If thn els nxt) ((V_Bool false) :: rst) vs2
  | E_IfTrue : forall f funs thn els nxt v rst vs1 vs2,
      eval_tm_R f funs thn rst vs1 ->
      eval_tm_R f funs nxt vs1 vs2 ->
      eval_tm_R f funs (If thn els nxt) (v :: rst) vs2.

Inductive extract_funs_R : list stk_fun -> partial_map stk_fun -> Prop :=
  | EF_Nil : extract_funs_R [] empty
  | EF_Cons : forall fs es l n ins,
      extract_funs_R fs es ->
      extract_funs_R ((Fun l n ins) :: fs) (update es l (Fun l n ins)).

Inductive eval_R : nat -> stk_prg -> list ins_val -> Prop :=
  | E_Prg : forall f fs es inss v,
      extract_funs_R fs es ->
      eval_tm_R f es inss [] v ->
      eval_R f (Prg fs inss) v.

End StackLangRel.


Require Import ZArith.
Require Import Lia.
Require Import List.
Import ListNotations.
Require Import LibTactics.
Require Import Utils.

Require Import Instr.
Require Import Lattices.
Require Import Concrete.
Require Import CodeGen.
Require Import CodeTriples.
Require Import CodeSpecs.
Require Import Arrays.
Require Import Memory.
Require Import CLattices.

Local Open Scope Z_scope.

Section SOP.

Variable cblock : block privilege.
Hypothesis stamp_cblock : Mem.stamp cblock = Kernel.
Variable ctable : CSysTable.

Definition labToVal (l: Zset.t) (t: val privilege) (m: memory (val privilege) privilege) : Prop :=
  exists c b,
    b <> cblock /\
    Mem.stamp b = Kernel /\
    t = Vptr (b, 0) /\
    memarr m b (map Vint c) /\
    (forall a, In a c <-> Zset.In a l).

Lemma mem_eq_except_cache_memseq :
  forall m m' b z vs
         (NCACHE : b <> cblock)
         (KERNEL : Mem.stamp b = Kernel)
         (EQ : mem_eq_except_cache cblock m m')
         (SEQ : memseq m b z vs),
    memseq m' b z vs.
Proof.
  intros.
  destruct EQ as [H1 H2].

  gdep z.
  induction vs as [|v vs IH]; intros; try solve [constructor].
  inv SEQ; econstructor; eauto.
  unfold load in *. simpl in *.
  destruct (Mem.get_frame m b) as [fr|] eqn:E; try congruence.
  apply H2 in E; eauto.
  rewrite E.
  eassumption.
Qed.

Lemma labToVal_cache_SOP :
  forall l v m m'
         (LAB : labToVal l v m)
         (EQ : mem_eq_except_cache cblock m m'),
    labToVal l v m'.
Proof.
  intros.
  unfold labToVal in *.
  destruct LAB as (c & b & NEQ & KERNEL & ? & ARR & IN). subst.
  eexists c, b.
  repeat (split; eauto).
  destruct ARR.
  econstructor; eauto using mem_eq_except_cache_memseq.
  unfold load in *. simpl in *.
  destruct (Mem.get_frame m b) as [fr|] eqn:FRAME; try congruence.
  destruct EQ as [_ EQ].
  eapply EQ in FRAME; eauto.
  rewrite FRAME. eassumption.
Qed.

Lemma memseq_extension_comp : forall m1 m2 l z off,
  memseq m1 z off l ->
  extends m1 m2 ->
  memseq m2 z off l.
Proof.
  intros m1 m2 l z off SEQ EXT.
  gdep off.
  induction l ; intros.
  constructor.
  inv SEQ.
  econstructor; eauto using extends_load.
Qed.

Lemma labToVal_extension_comp_SOP :
  forall m1 m2 l v,
    labToVal l v m1 -> extends m1 m2 -> mem_def_on_cache cblock m1 -> labToVal l v m2.
Proof.
  unfold labToVal.
  intros m1 m2 l v (c & b & NEQ & KERNEL & ? & ARR & IN) EXT DEF. subst.
  eexists c, b. repeat (split; eauto).
  destruct ARR.
  econstructor; eauto using extends_load, memseq_extension_comp.
Qed.

Definition add_principal_to_lab (ps : Zset.t) (v : val privilege) : Zset.t :=
  match v with
    | Vint p => Zset.add p ps
    | _ => ps
  end.

Definition valToLab (v : val privilege) (m: memory (val privilege) privilege) : Zset.t :=
  match v with
    | Vint _ => Zset.empty
    | Vptr (b, _) =>
      match Mem.get_frame m b with
        | Some ((Vint len, _) :: ps) =>
          fold_left add_principal_to_lab (map fst (take (Z.to_nat len) ps)) Zset.empty
        | _ => Zset.empty
      end
  end.

Definition genFlows : list Instr := subset_arrays.

Definition genBot : list Instr := push 0 ++ alloc_array.
Definition genJoin : list Instr := concat_arrays.

Global Instance SOPClatt : ConcreteLattice Zset.t :=
{ labToVal := labToVal
; valToLab := valToLab
; genBot := genBot
; genJoin:= genJoin
; genFlows:= genFlows
}.

Lemma memseq_incl : forall m l1 l2 b off,
 memseq m b off l1 ->
 length l1 = length l2 ->
 memseq m b off l2 -> incl l1 l2.
Proof.
  induction l1 ; intros.
  - inv H. inv H1. unfold incl. auto. inv H0.
  - inv H. destruct l2. inv H0.  inv H0.
    unfold incl. intros. inv H.
    inv H1. allinv'. constructor. auto.
    right.
    eapply IHl1; eauto.
    inv H1; eauto.
Qed.

Lemma memseq_get_frame :
  forall m b off vs fr
         (POS : off >= 0)
         (FRAME : Mem.get_frame m b = Some fr)
         (SEQ : memseq m b off vs),
    vs = map fst (take (length vs) (dropZ off fr)).
Proof.
  intros.
  induction SEQ as [off|off v t vs LOAD SEQ IH]; intros; eauto.
  unfold load in LOAD. simpl in LOAD.
  rewrite FRAME in LOAD.
  assert (H : dropZ off fr = (v, t) :: dropZ (off + 1) fr).
  { clear - LOAD POS.
    unfold read_m, dropZ in *.
    destruct (off <? 0); try congruence.
    destruct (off + 1 <? 0) eqn:LT.
    - rewrite Z.ltb_lt in LT. lia.
    - rewrite Z2Nat.inj_add; try lia.
      gdep (Z.to_nat off). clear.
      intros off LOAD.
      change (off + Z.to_nat 1)%nat with (off + 1)%nat.
      rewrite plus_comm.
      gdep off.
      induction fr as [|v' fr IH]; intros [|off] H; simpl in *; eauto; try congruence.
      inv H. reflexivity. }
  rewrite H.
  simpl.
  f_equal.
  apply IH.
  lia.
Qed.

Lemma labToVal_inj_SOP : forall l1 l2 t m, labToVal l1 t m -> labToVal l2 t m -> l1 = l2.
Proof.
  unfold labToVal.
  intros l1 l2 t m
         (c1 & b & NEQ1 & KERNEL1 & E1 & ARR1 & EQ1)
         (c2 & b' & _ & _ & E2 & ARR2 & EQ2).
  assert (b' = b) by congruence. subst. clear E2.
  destruct ARR1 as [? t1 LOAD1 SEQ1 ?].
  destruct ARR2 as [? t2 LOAD2 SEQ2 ?]. subst.
  assert (LENGTHS : Z.of_nat (length (map Vint c1 : list (val privilege))) =
                    Z.of_nat (length (map Vint c2 : list (val privilege)))) by congruence.
  clear LOAD1 LOAD2.
  rewrite Nat2Z.inj_iff in LENGTHS.

  eapply Zset.antisym; eapply Zset.incl_spec.

  - assert (INCL := memseq_incl _ _ _ _ _ SEQ1 LENGTHS SEQ2).
    unfold incl. intros.
    unfold Zset.In in *.
    rewrite <- EQ2. rewrite <- EQ1 in H.
    eapply in_map in H. apply INCL in H.
    rewrite in_map_iff in H. destruct H as (? & ? & ?). congruence.

  - assert (INCL := memseq_incl _ _ _ _ _ SEQ2 (eq_sym LENGTHS) SEQ1).
    unfold incl. intros.
    unfold Zset.In in *.
    rewrite <- EQ1. rewrite <- EQ2 in H.
    eapply in_map in H. apply INCL in H.
    rewrite in_map_iff in H. destruct H as (? & ? & ?). congruence.
Qed.

Lemma Zset_union_empty : forall S, Zset.union S Zset.empty = S.
Proof.
  intros.
  apply Zset.antisym;  eapply Zset.incl_spec.
  - intros x Hin.
    eapply Zset.In_union in Hin.
    inv Hin; eauto.
    unfold Zset.In in H. rewrite Zset.elements_empty in H. inv H.
  - intros x Hin.
    eapply Zset.In_union.
    left; eauto.
Qed.

Lemma Zset_incl_spec : forall s1 s2, Zset.incl s1 s2 = true <->
                                     (forall p, Zset.In p s1 -> Zset.In p s2).
Proof.
  intros.
  split; rewrite Zset.incl_spec in *; unfold Zset.In; intros; eauto.
Qed.

Lemma Zset_incl_empty : forall s, Zset.incl Zset.empty s = true.
Proof.
  intros.
  rewrite Zset.incl_spec, Zset.elements_empty.
  intros z H. inv H.
Qed.

Lemma Zset_incl_trans : forall s1 s2 s3, Zset.incl s1 s2 = true ->
                                         Zset.incl s2 s3 = true ->
                                         Zset.incl s1 s3 = true.
Proof.
  intros.
  rewrite Zset.incl_spec in *.
  eauto using incl_tran.
Qed.

Lemma Zset_incl_union : forall s1 s2 s3, Zset.incl (Zset.union s1 s2) s3 = true <->
                                         (Zset.incl s1 s3 = true) /\ (Zset.incl s2 s3 = true).
Proof.
  intros.
  split; intros; repeat rewrite Zset_incl_spec in *.
  - split;
    intros p IN;
    apply H;
    rewrite Zset.In_union; eauto.
  - destruct H as [H1 H2].
    intros p IN.
    rewrite Zset.In_union in IN.
    destruct IN; eauto.
Qed.

Lemma Zset_incl_add : forall p s1 s2, Zset.incl (Zset.add p s1) s2 = true <->
                                      Zset.In p s2 /\ Zset.incl s1 s2 = true.
Proof.
  intros.
  split; intros; repeat rewrite Zset_incl_spec in *.
  - split.
    + apply H.
      rewrite Zset.In_add. eauto.
    + intros p' IN;
      apply H;
      rewrite Zset.In_add; eauto.
  - destruct H as [H1 H2].
    intros p' IN.
    rewrite Zset.In_add in IN.
    destruct IN; subst; eauto.
Qed.

Lemma labToVal_valToLab_id_SOP : forall l t m, labToVal l t m -> valToLab t m = l.
Proof.
  unfold labToVal, valToLab ; intros.
  destruct H as (vs & b & NEQ & KERNEL & E & ARR & IN). subst.
  destruct ARR as [? t LOAD SEQ ?]. subst.
  unfold load in *. simpl in LOAD.
  destruct (Mem.get_frame m b) as [fr|] eqn:FRAME; try congruence.
  destruct fr as [|len fr]; try solve [inv LOAD].
  unfold read_m in LOAD. simpl in LOAD. inv LOAD.
  rewrite Nat2Z.id.
  change (take (length (map Vint vs : list (val privilege))) fr)
    with (take (length (map Vint vs : list (val privilege)))
               (dropZ 1 ((Vint (Z.of_nat (length (map Vint vs : list (val privilege)))), t) :: fr))).
  exploit memseq_get_frame; eauto; try lia.
  intros H.
  unfold Atom in H. (* Invisible, grr... *)
  rewrite <- H.
  clear - IN.
  assert (IN' : forall a, (In a vs \/ Zset.In a Zset.empty) <-> Zset.In a l).
  { intuition.
    - rewrite <- IN. eauto.
    - unfold Zset.In in *.
      rewrite Zset.elements_empty in *.
      simpl in *. intuition.
    - left. rewrite IN. eauto. }
  clear IN. rename IN' into IN.
  gdep Zset.empty. gdep l.
  induction vs as [|p vs IH]; simpl; eauto.
  - intros.
    apply Zset.antisym; rewrite Zset_incl_spec; intros p H; firstorder.
  - intros.
    apply IH. intros a.
    rewrite <- IN.
    intuition; rewrite Zset.In_add in *; eauto; intuition.
Qed.

Lemma genBot_spec_SOP : forall m0 (Hm0: mem_def_on_cache cblock m0) (Q:HProp),
   HT cblock ctable genBot
      (fun m s => extends m0 m /\
                  (forall m1 z t,
                     extends m m1 -> labToVal bot z m1 -> Q m1 ((z,t):::s)))
      Q.
Proof.
  intros.
  eapply HT_compose_flip.
  eapply alloc_array_spec; eauto.
  eapply HT_strengthen_premise.
  eapply push_spec; eauto.
  intros m s (EXT & POST).
  do 3 eexists.
  repeat (split; eauto; try lia).
  intros b m1 t EXT' VALID [t' LOAD] FRESH KERNEL.
  apply POST; eauto.
  eapply extends_mem_def_on_cache with (m' := m) in Hm0; eauto using extends_trans.
  destruct Hm0.
  eexists [], b.
  repeat (split; eauto); try congruence.
  - econstructor; try constructor.
    rewrite LOAD; eauto.
  - intros [].
  - intros IN.
    simpl in IN.
    unfold Zset.In in *.
    rewrite Zset.elements_empty in IN.
    inversion IN.
Qed.

Lemma genJoin_spec_SOP: forall m0 (Hm0: mem_def_on_cache cblock m0) (Q: HProp),
       HT cblock ctable genJoin
         (fun m s =>
          exists s0 l z t l' z' t',
             s = (z, t) ::: (z', t') ::: s0 /\
             extends m0 m /\
             labToVal l z m /\ labToVal l' z' m /\
             (forall m1 zz' t, extends m m1 -> labToVal (l \_/ l') zz' m1 ->
                         Q m1 ((zz', t) ::: s0)))
         Q.
Proof.
  intros.
  eapply HT_strengthen_premise.
  eapply concat_arrays_spec ; eauto.
  go_match.
  destruct H as (s0 & l & v & t & l' & v' & t' & ? & EXT & Hl & Hl' & POST). subst.
  unfold labToVal in *.
  destruct Hl as (c & b & Nb & Kb & ? & ARR & IN).
  destruct Hl' as (c' & b' & Nb' & Kb' & ? & ARR' & IN'). subst.
  do 7 eexists. repeat (split; try solve [intuition eauto]).
  intros r m1 t'' EXT' ARR'' FRESH Kr.
  apply POST; eauto.
  eexists (c' ++ c), r.
  eapply extends_mem_def_on_cache in Hm0; eauto. destruct Hm0.
  rewrite map_app.
  repeat (split; [solve [eauto | congruence]|]).
  intros a.
  rewrite in_app_iff, Zset.In_union, IN, IN'. intuition.
Qed.

Lemma incl_val_list_subset_b : forall l1 l2,
                                 incl l1 l2 <->
                                 val_list_subset_b l1 l2 = true.
Proof.
  intros. split.
  - intros. eapply forallb_forall. intros.
    eapply H in H0.  eapply existsb_exists.
    exists x; intuition.
    destruct (EquivDec.equiv_dec x x); try congruence.
  - intros.
    unfold incl. intros.
    unfold val_list_subset_b in H.
    eapply forallb_forall in H; eauto.
    unfold val_list_in_b in H.
    eapply existsb_exists in H.
    destruct H. intuition.
    destruct (EquivDec.equiv_dec a x) in H2; try congruence.
Qed.

Lemma Zsetincl_val_list_subset : forall l l' x x',
                                 forall (Hxl : forall a, In a x <-> Zset.In a l)
                                        (Hx'l' : forall a, In a x' <-> Zset.In a l'),
                                   Zset.incl l l' =
                                   val_list_subset_b (map Vint x) (map Vint x').
Proof.
  intros.
  rewrite Bool.eq_iff_eq_true.
  split; intros.
  - apply Zset.incl_spec in H.
    eapply incl_val_list_subset_b; eauto.
    unfold incl, Zset.In in *.
    intros a IN.
    rewrite in_map_iff in *.
    destruct IN as (i & ? & IN). subst.
    eexists. split; eauto.
    rewrite Hxl in IN.
    rewrite Hx'l'. eauto.

  - apply Zset.incl_spec.
    intros a IN.
    unfold Zset.In in *.
    eapply Hx'l'.
    rewrite <- Hxl in IN.
    apply incl_val_list_subset_b in H.
    eapply in_map in IN. apply H in IN.
    rewrite in_map_iff in IN. destruct IN. intuition congruence.
Qed.

Lemma genFlows_spec_SOP: forall m0 (Hm0: mem_def_on_cache cblock m0) (Q: HProp),
                   HT cblock ctable genFlows
                       (fun m s =>
                          exists l l' z z' t t' s0,
                            extends m0 m /\
                            labToVal l z m /\ labToVal l' z' m /\
                            s = (z,t):::(z',t'):::s0 /\
                            forall t'',
                            Q m ((boolToVal(l <: l'),t''):::s0))
                       Q.
Proof.
  intros.
  unfold genFlows; intros.
  eapply HT_strengthen_premise.
  eapply subset_arrays_spec; eauto.
  intros m s (l & l' & z & z' & t & t' & s0 & EXT & Hl & Hl' & ? & POST). subst.
  destruct Hl as (c & b & Nb & Kb & ? & ARR & IN).
  destruct Hl' as (c' & b' & Nb' & Kb' & ? & ARR' & IN'). subst.
  eexists b, b', (map Vint c), (map Vint c'). do 4 eexists; repeat split; eauto.
  intros t''.
  erewrite <- Zsetincl_val_list_subset; eauto.
Qed.

Global Instance SOPwf (fh: list Instr) : WfConcreteLattice cblock ctable Zset.t ZsetLat SOPClatt :=
{ labToVal_cache := labToVal_cache_SOP
; labToVal_valToLab_id := labToVal_valToLab_id_SOP
; labToVal_inj := labToVal_inj_SOP
; labToVal_extension_comp := labToVal_extension_comp_SOP
; genBot_spec := genBot_spec_SOP
; genJoin_spec := genJoin_spec_SOP
; genFlows_spec := genFlows_spec_SOP
}.

End SOP.


Require Import String.
Require Import Functors.
Require Import MonadLib.
Require Import Names.
Require Import EffPure.
Require Import EffState.
Require Import Arith.
Require Import Ref.
Require Import ESoundS.

Open Scope string_scope.

Section Test_Section.

  Definition D := AType :+: UnitType :+: RefType.

  Definition E := Arith :+: RefE.

  Definition V := StuckValue :+: NatValue :+: LocValue :+: UnitValue.

  Variable MT : Set -> Set.
  Context `{Fail_MT : FailMonad MT}.
  Context {Inj_MT : InjMonad MT}.
  Context {Reasonable_MT : Reasonable_Monad MT}.
  Context {MT_eq_dec : forall (A : Set) (mta : MT A),
    {exists a, mta = return_ a} + {mta = fail}}.

  Variable ME : Set -> Set.
  Context `{State_ME : StateM ME (list (Value V))}.
  Context {Fail_ME : FailMonad ME}.

  Definition WFV := (WFValue_VI V D (list (Names.DType D))) ::+::
    (WFValue_Unit D V _) ::+:: (WFValue_Loc D V _).
  Definition WFVM := (WFValueM_base D V MT ME _ WFV) ::+::
    (WFValueM_State D V MT ME _
      (TypContextCE := DType_Env_CE _)
      (TypContext_WFE := DType_Env_WFE _ _ WFV)).

  Instance typeof_alg : forall T, FAlgebra TypeofName T (typeofR D MT) (E).
  Proof.
    intros; eauto 150 with typeclass_instances.
  Defined.

  Lemma eval_Sound :
    forall (Sigma : Env (DType D)) (e : Exp E) (T : DType D) (env : Env (Value V)),
      WF_Environment D V _ WFV Sigma env Sigma ->
      typeof D E MT (proj1_sig e) = return_ T ->
      exists v : Value V,
        exists env' : Env (Value V),
          exists Sigma' : Env (DType D),
            (put env) >> evalM V E ME (proj1_sig e) =
            put env' >> return_ (M := ME) v /\
            WFValueC D V _ WFV Sigma' v T.
  Proof.
    intros; eapply eval_State_Sound with (WFVM' := WFVM);
      eauto 10 with typeclass_instances.
    eauto 15 with typeclass_instances.
    intros; repeat apply @P2Algebra_Plus.
    eapply Arith_eval_soundness'' with (WFV := WFV) (WFVM := WFVM);
      eauto 250 with typeclass_instances.
    eapply Ref_eval_soundness'' with (WFV := WFV) (WFVM := WFVM)
      (TypContextCE :=  _) (TypContext_S := _)
      (TypContext_WFE := DType_Env_WFE _ _ WFV);
      eauto 250 with typeclass_instances.
    simpl; eauto.
  Qed.

  Eval compute in ("Soundness for 'Arith :+: References' Proven!").

End Test_Section.

(*
*** Local Variables: ***
*** coq-prog-args: ("-emacs-U" "-impredicative-set") ***
*** End: ***
*)


Require Import Coq.Lists.List. Import ListNotations.
Require Import Coq.ZArith.ZArith. Open Scope Z_scope.
Require Import Coq.Strings.String. Open Scope string_scope.
Require Import coqutil.Byte.
Require Import coqutil.Word.Interface coqutil.Word.Properties.
Require Import coqutil.Tactics.fwd coqutil.Tactics.Tactics.
Require Import bedrock2.ZnWords.
Require Import compiler.SeparationLogic.
Require Import compiler.LowerPipeline.
Require Import Bedrock2Experiments.RiscvMachineWithCavaDevice.ExtraRiscvMachine.
Require Import Bedrock2Experiments.RiscvMachineWithCavaDevice.InternalMMIOMachine.
Require Import HmacEnd2End.CavaHmacDevice.
Require Import HmacSoftware.Sha256ToRiscV.
Require Import HmacSpec.SHA256.
Require Import Bedrock2Experiments.RiscvMachineWithCavaDevice.Bedrock2ToCava.

Definition binary: list byte := Eval compute in Pipeline.instrencode sha256_asm.

Definition mach_valid(imem: mem -> Prop)(mach: ExtraRiscvMachine hmac_device): Prop :=
  subset (footpr imem) (of_list (getXAddrs mach)) /\
  getNextPc mach = word.add (getPc mach) (word.of_Z 4) /\
  regs_initialized (getRegs mach) /\
  map.undef_on (getMem mach) device.isMMIOAddr /\
  disjoint (of_list (getXAddrs mach)) device.isMMIOAddr.

Local Notation "a * b" := (sep a b) : type_scope.
Local Notation "a * b" := (sep a b).

Add Ring wring : (Properties.word.ring_theory (word := word))
                   (preprocess [autorewrite with rew_word_morphism],
                    morphism (Properties.word.ring_morph (word := word)),
                    constants [Properties.word_cst]).

Declare Scope word_scope.
Infix "+" := word.add : word_scope.
Infix "-" := word.sub : word_scope.
Delimit Scope word_scope with word.
Local Open Scope word_scope.

Notation ShaDevice := hmac_device.

(* Trying to keep toplevel theorem at 53 chars wide:
01234567890123456789012345678901234567890123456789012 *)

Theorem sha256_end2end_correct:
forall a_outp a_inp len sp_val a_ret a_code stack_lo
  (inp: list byte) (Rdata Rexec: mem->Prop)
  (sched: nat->nat) (m: ExtraRiscvMachine ShaDevice),
map.get m.(getRegs) RegisterNames.a0 = Some a_outp ->
map.get m.(getRegs) RegisterNames.a1 = Some a_inp ->
map.get m.(getRegs) RegisterNames.a2 = Some len ->
map.get m.(getRegs) RegisterNames.sp = Some sp_val ->
map.get m.(getRegs) RegisterNames.ra = Some a_ret ->
m.(getPc) = a_code ->
Z.of_nat (length inp) = word.unsigned len ->
a_inp <> word.of_Z 0 ->
a_outp <> word.of_Z 0 ->
256 <= word.unsigned (sp_val - stack_lo) ->
word.unsigned (sp_val - stack_lo) mod 4 = 0 ->
word.unsigned a_code mod 4 = 0 ->
word.unsigned a_ret mod 4 = 0 ->
device.is_ready_state m.(getExtraState) ->
mach_valid (ptsto_bytes a_code binary * Rexec) m ->
(bytearray a_inp inp *
 mem_available a_outp (a_outp + (word.of_Z 32)) *
 mem_available stack_lo sp_val * Rdata *
 ptsto_bytes a_code binary * Rexec) m.(getMem) ->
exists n m',
run sched n m = Some m' /\
map.get m'.(getRegs) RegisterNames.sp = Some sp_val /\
map.agree_on callee_saved m.(getRegs) m'.(getRegs) /\
m'.(getPc) = a_ret /\
device.is_ready_state m'.(getExtraState) /\
mach_valid (ptsto_bytes a_code binary * Rexec) m' /\
(bytearray a_inp inp *
 bytearray a_outp (sha256 inp)  *
 mem_available stack_lo sp_val * Rdata *
 ptsto_bytes a_code binary * Rexec) m'.(getMem).
Proof.
  intros.
  change binary with (Pipeline.instrencode sha256_asm).
  assert (((mem_available a_outp (a_outp + (word.of_Z 32)) * bytearray a_inp inp) *
           (mem_available stack_lo sp_val * ptsto_bytes a_code binary * Rexec * Rdata))
            m.(getMem)) as M by ecancel_assumption.
  unfold sep in M at 1. destruct M as (mH & mL & Sp & MH & ML).
  eapply mem_available_to_exists in MH.
  destruct MH as (output_placeholder & HLen & MH).
  ring_simplify (a_outp + word.of_Z 32 - a_outp) in HLen.
  edestruct bedrock2_and_cava_system_correct with
      (f_entry_name := "b2_sha256") (p_functions := a_code)
      (stack_start := stack_lo) (stack_pastend := sp_val)
      (postH := fun mH' (retvals: list word) =>
                  (bytearray a_inp inp * bytearray a_outp (sha256 inp)) mH')
    as (steps_remaining & finalL & mH' & retvals & Rn & GM & A & Eq & M & HP);
    lazymatch goal with
    | |- _ mod _ = _ => idtac
    | |- _ => try eassumption
    end.
  { exact funcs_valid. }
  { apply List.dedup_NoDup_iff. reflexivity. }
  { exact sha256_compile_result_eq. }
  { exact sha256_asm_valid_instructions. }
  { reflexivity. }
  { change sha256_req_stack with (sha256_req_stack * 4 / 4).
    eapply Z.div_le_mono. 1: reflexivity.
    etransitivity. 2: eassumption.
    vm_compute. discriminate. }
  { assumption. }
  { assumption. }
  { assumption. }
  { (* Note: Conveniently, since "b2_sha256" is last in the alphabet among all used functional
       names, sha256_relative_pos equals 0, so the theorem becomes simpler (just jump to a_code
       instead of (a_code+sha256_relative_pos)).
       Compute sha256_finfo. *)
    replace (getPc m) with a_code. symmetry. apply add_0_r. }
  { eapply WeakestPreconditionProperties.Proper_call.
    2: eapply link_sha256.
    6: { instantiate (4 := a_outp). instantiate (5 := a_inp). instantiate (1 := mH).
         unfold ptsto_bytes in *. ecancel_assumption. }
    6: { cbn. instantiate (1 := List.repeat Byte.x00 32). reflexivity. }
    all: try eassumption.
    unfold Morphisms.pointwise_relation, Basics.impl, Sha256ExampleProperties.sha256_post.
    intros. fwd.
    split. 1: ecancel_assumption.
    eexists. split. 1: eassumption. cbn.
    apply sha256_len. }
  { cbn -[map.get].
    repeat match goal with
           | |- context[match ?x with _ => _ end] =>
             match goal with
             | H: ?x' = Some ?y |- _ => replace x with (Some y) by (symmetry; exact H)
             end
           end.
    reflexivity. }
  { unfold machine_ok, mach_valid in *. fwd.
    ssplit; try eassumption.
    change (Pipeline.instrencode sha256_asm) with binary.
    apply sep_comm.
    unfold sep at 1. exists mH, mL. split. 1: exact Sp. split. 1: reflexivity.
    ecancel_assumption. }
  { unfold machine_ok in M. unfold mach_valid. fwd.
    do 2 eexists. ssplit.
    - unfold run. rewrite Rn. reflexivity.
    - assumption.
    - assumption.
    - reflexivity.
    - assumption.
    - assumption.
    - assumption.
    - assumption.
    - assumption.
    - assumption.
    - unfold sep at 1 in Mp0. destruct Mp0 as (mL' & mH'' & Sp' & ML' & ?).  subst mH''.
      enough (((bytearray a_inp inp * bytearray a_outp (sha256 inp)) *
              (mem_available stack_lo sp_val * Rdata *
               ptsto_bytes (getPc m) (Pipeline.instrencode sha256_asm) * Rexec)) (getMem finalL))
        by ecancel_assumption.
      apply sep_comm.
      unfold sep at 1.
      exists mL', mH'. split. 1: exact Sp'. split. 2: exact HP.
      ecancel_assumption. }
  Unshelve.
  all: repeat constructor.
  all: try exact (word.of_Z 0).
Qed.

(* Goal: bring this list down to only standard axioms like functional and propositional extensionality
Print Assumptions sha256_end2end_correct.
 *)


Require Import Crypto.Specific.Framework.RawCurveParameters.
Require Import Crypto.Util.LetIn.

(***
Modulus : 2^130 - 5
Base: 43 + 1/3
***)

Definition curve : CurveParameters :=
  {|
    sz := 3%nat;
    base := 43 + 1/3;
    bitwidth := 64;
    s := 2^130;
    c := [(1, 5)];
    carry_chains := Some [seq 0 (pred 3); [0; 1]]%nat;

    a24 := None;
    coef_div_modulus := Some 2%nat;

    goldilocks := None;
    karatsuba := None;
    montgomery := false;
    freeze := Some true;
    ladderstep := false;

    mul_code := None;

    square_code := None;

    upper_bound_of_exponent_loose := None;
    upper_bound_of_exponent_tight := None;
    allowable_bit_widths := None;
    freeze_extra_allowable_bit_widths := None;
    modinv_fuel := None
  |}.

Ltac extra_prove_mul_eq _ := idtac.
Ltac extra_prove_square_eq _ := idtac.


Require Import Crypto.Specific.Framework.RawCurveParameters.
Require Import Crypto.Util.LetIn.

(***
Modulus : 2^379 - 19
Base: 54 + 1/7
***)

Definition curve : CurveParameters :=
  {|
    sz := 7%nat;
    base := 54 + 1/7;
    bitwidth := 64;
    s := 2^379;
    c := [(1, 19)];
    carry_chains := Some [seq 0 (pred 7); [0; 1]]%nat;

    a24 := None;
    coef_div_modulus := Some 2%nat;

    goldilocks := None;
    karatsuba := None;
    montgomery := false;
    freeze := Some true;
    ladderstep := false;

    mul_code := None;

    square_code := None;

    upper_bound_of_exponent_loose := None;
    upper_bound_of_exponent_tight := None;
    allowable_bit_widths := None;
    freeze_extra_allowable_bit_widths := None;
    modinv_fuel := None
  |}.

Ltac extra_prove_mul_eq _ := idtac.
Ltac extra_prove_square_eq _ := idtac.


Require Import Crypto.Specific.Framework.RawCurveParameters.
Require Import Crypto.Util.LetIn.

(***
Modulus : 2^379 - 19
Base: 47.375
***)

Definition curve : CurveParameters :=
  {|
    sz := 8%nat;
    base := 47 + 3/8;
    bitwidth := 64;
    s := 2^379;
    c := [(1, 19)];
    carry_chains := Some [seq 0 (pred 8); [0; 1]]%nat;

    a24 := None;
    coef_div_modulus := Some 2%nat;

    goldilocks := None;
    karatsuba := None;
    montgomery := false;
    freeze := Some true;
    ladderstep := false;

    mul_code := None;

    square_code := None;

    upper_bound_of_exponent_loose := None;
    upper_bound_of_exponent_tight := None;
    allowable_bit_widths := None;
    freeze_extra_allowable_bit_widths := None;
    modinv_fuel := None
  |}.

Ltac extra_prove_mul_eq _ := idtac.
Ltac extra_prove_square_eq _ := idtac.


Require Import Events.
Require Import TraceModel.
Require Import Properties.
Require Import CommonST.
Require Import Robustdef.
Require Import Criteria.
Require Import Setoid.
Require Import ClassicalExtras.
Require Import Logic.ClassicalFacts.
Require Import TechnicalLemmas.

(** This file proves that R2RTP implies RTEP *)

(** *our assumptions *)
(**********************************************************)
Hypothesis input_totality_tgt : @input_totality tgt.
Hypothesis determinacy_src    : @determinacy src.
Hypothesis tgt_sem            : @semantics_safety_like tgt.
(**********************************************************)

Theorem R2rTP_RTIP : R2rTP -> RTIP.
Proof.
  intros r2rpp. unfold RTIP. apply NNPP.
  intros hf. rewrite not_forall_ex_not in hf. destruct hf as [P1 hf].
  rewrite not_forall_ex_not in hf. destruct hf as [P2 hf].
  rewrite not_imp in hf. destruct hf as [H1 hf].
  unfold "⊆" in *.
  rewrite not_forall_ex_not in hf. destruct hf as [Ct hf].
  rewrite not_forall_ex_not in hf. destruct hf as [t hf].
  rewrite not_imp in hf. destruct hf as [k1 k2].
  specialize (r2rpp P1 P2 traces_match).
  assert (Hsat : rsat2 P1 P2 traces_match).
  { intros C b1 b2 Hb1 Hb2. apply (H1 C b1) in Hb1.
    unfold beh in Hb1.
    now apply (determinacy_src (C [P2])). }
  specialize (r2rpp Hsat Ct t); clear Hsat.
  assert (H : forall t2, sem tgt (Ct [P2 ↓]) t2 -> traces_match t t2).
  { intros t2 H. now apply r2rpp. } clear r2rpp H1.
  destruct (longest_in_psem tgt_sem (Ct [P2 ↓]) t k2) as [l [Hmt [Hpsem_m Hmax_m]]].
  destruct Hpsem_m as [t2 [Hmt2 Hsem_t2]].
  specialize (H t2 Hsem_t2). destruct H; try now subst.
  destruct H as [ll [i1 [i2 [Hi1 [Hi2 [Hdiffi [Hpref1 Hpref2]]]]]]].
  assert (psem (Ct [P2 ↓]) (ftbd (snoc ll i1))).
  { apply (input_totality_tgt (Ct [P2 ↓]) ll i2 i1); auto. now exists t2. }
  assert (fpr (ftbd (snoc ll i1)) l) by now apply Hmax_m.
  apply (fpr_pref_pref _ _ t2) in H0; auto.
  destruct (same_ext (ftbd (snoc ll i1)) (ftbd (snoc ll i2)) t2); auto;
    simpl in H1.
  ++ apply (list_snoc_pointwise ll ll i1 i2 i2 i2) in H1; auto.
     inversion H1. congruence. now apply list_list_prefix_ref.
  ++ apply (list_snoc_pointwise ll ll i2 i1 i1 i1) in H1; auto.
     inversion H1. congruence. now apply list_list_prefix_ref.
Qed.

Theorem R2rTP_RTEP : R2rTP -> RTEP.
Proof.  intros H. apply RTIP_RTEP. exact (R2rTP_RTIP H). Qed.