From mathcomp Require Import all_ssreflect.

(* View level                                                           *)

Definition vlevel := unit.

Definition vL0 : vlevel. done. Qed.
Definition vOL : vlevel. done. Qed.   (* Only Logical *)
Definition vAC : vlevel. done. Qed.   (* All Core     *)
Definition vAL : vlevel. done. Qed.   (* All          *)
Definition vAB : vlevel. done. Qed.   (* All + under binder *)
Definition vL1 : vlevel. done. Qed.   (* All rules but no recursion *)

Class View (lv:vlevel) (p : Prop) (b : bool) := { r : reflect p b }.
Hint Mode View + + - : typeclass_instances.
Hint Mode View + - + : typeclass_instances.

Lemma xPm lv b p {v : View lv p b} : reflect p b.
Proof. by apply r. Qed.

Definition xP   := @xPm vL1.
Definition lxP  := @xPm vOL.
Definition rcxP := @xPm vAC.
Definition rxP  := @xPm vAL.
Definition rbxP := @xPm vAB.

(* -------------------------------------------------------------------- *)
(* Rule level                                                           *)
Definition rlevel := unit.
Definition rL0  : rlevel. done. Qed.
Definition rOL  : rlevel. done. Qed. (* Only Logical *)
Definition rOLr : rlevel. done. Qed. (* Only Logical but recursive like andP3 *)
Definition rAC  : rlevel. done. Qed.
Definition rAL  : rlevel. done. Qed.

(* -------------------------------------------------------------------- *)
(* Valid rl vl                                                          *)
(* indicate if a rule at level rl can be used at the view level vl      *)
(* -------------------------------------------------------------------- *)

Class Valid (rl:rlevel) (vl:vlevel).
  Instance V_rL0 vl : Valid rL0 vl | 0.
  (* For level 1 *)
  Instance V_OL_L1  : Valid rOL vL1 | 0.
  Instance V_AC_L1  : Valid rAC vL1 | 0.
  Instance V_AL_L1  : Valid rAL vL1 | 0.

  (* For Only Logical *)
  Instance V_OL_OL  : Valid rOL  vOL | 0.
  Instance V_OLr_OL : Valid rOLr vOL | 0.

  (* For All Core *)
  Instance V_OL_AC rl `{Valid rl vOL} : Valid rl  vAC | 1.
  Instance V_AC_AC                    : Valid rAC vAC | 0.

  (* For All  and All bind *)
  Instance V_AC_AL rl : Valid rl vAL.
  Instance V_AL_AB rl : Valid rl vAB.

(* -------------------------------------------------------------------- *)
(* ArgRec vl avl                                                        *)
(* compute the view level avl of an argument of a logic rule used with  *)
(* view level vl                                                        *)
(* -------------------------------------------------------------------- *)

Class ArgRec (vl:vlevel) (avl:vlevel).
(* Hint Mode ArgRec + -. *)

Instance ArgRec_L1     : ArgRec vL1 vL0 | 0.
Instance ArgRec_dft vl : ArgRec vl  vl  | 1.

Class ArgBind (vl:vlevel) (avl:vlevel).
(* Hint Mode ArgBind + -. *)

Instance ArgBind_AB     : ArgBind vAB vAB | 0.
Instance ArgBind_dft vl : ArgBind vl  vL0  | 1.

(* -------------------------------------------------------------------- *)
(* Level 0 rules                                                        *)
Instance idV vl (b:bool) : View vl (is_true b) b | 100.
Proof. by split;apply idP. Qed.

(* -------------------------------------------------------------------- *)
(* Logical Rules                                                        *)
Instance negV vl avl `{Valid rOL vl} p b
  `{ArgRec vl avl} `{View avl p b} : View vl (~p) (~~b) | 90.
Proof. by split; apply: (iffP negP) => HH /xPm. Qed.

Instance andV vl avl `{Valid rOL vl} p1 p2 b1 b2
  `{ArgRec vl avl} `{View avl p1 b1} `{View avl p2 b2} :
  View vl (p1 /\ p2) (b1 && b2) | 90.
Proof. by split; apply: (iffP andP) => -[] /xPm ? /xPm. Qed.

Instance orV vl avl `{Valid rOL vl} p1 p2 b1 b2
  `{ArgRec vl avl} `{View avl p1 b1} `{View avl p2 b2} :
  View vl (p1 \/ p2) (b1 || b2) | 90.
Proof. by split; apply: (iffP orP)=> -[] /xPm;auto. Qed.

Instance impV vl avl `{Valid rOL vl} p1 p2 b1 b2
  `{ArgRec vl avl} `{View avl p1 b1} `{View avl p2 b2} :
  View vl (p1 -> p2) (b1 ==> b2) | 90.
Proof. by split; apply: (iffP implyP)=> HH /xPm /HH /xPm. Qed.

(* -------------------------------------------------------------------- *)
(* Logical Recursive Rules                                              *)

Instance and3V vl avl `{Valid rOLr vl} p1 p2 p3 b1 b2 b3
  `{ArgRec vl avl} `{View avl p1 b1} `{View avl p2 b2} `{View avl p3 b3} :
  View vl [/\ p1, p2 & p3] [&& b1, b2 & b3] | 80.
Proof.
  by split; apply: (iffP and3P) => -[] /xPm ? /xPm ? /xPm;constructor.
Qed.

Instance and4V vl avl `{Valid rOLr vl} p1 p2 p3 p4 b1 b2 b3 b4
  `{ArgRec vl avl} `{View avl p1 b1} `{View avl p2 b2} `{View avl p3 b3}
  `{View avl p4 b4} : View vl [/\ p1, p2, p3 & p4] [&& b1, b2, b3 & b4] | 70.
Proof.
  by split; apply: (iffP and4P) => -[] /xPm ? /xPm ? /xPm ? /xPm;constructor.
Qed.

Instance and5V vl avl `{Valid rOLr vl} p1 p2 p3 p4 p5 b1 b2 b3 b4 b5
  `{ArgRec vl avl} `{View avl p1 b1} `{View avl p2 b2} `{View avl p3 b3}
  `{View avl p4 b4} `{View avl p5 b5} :
   View vl [/\ p1, p2, p3, p4 & p5] [&& b1, b2, b3, b4 & b5] | 60.
Proof.
  by split; apply: (iffP and5P) => -[] /xPm ? /xPm ? /xPm ? /xPm ? /xPm;constructor.
Qed.

Instance or3V vl avl `{Valid rOLr vl} p1 p2 p3 b1 b2 b3
  `{ArgRec vl avl} `{View avl p1 b1} `{View avl p2 b2} `{View avl p3 b3} :
  View vl [\/ p1, p2 | p3] [|| b1, b2 | b3] | 80.
Proof.
  split; apply: (iffP or3P)=> -[] /xPm ?;
   ((by constructor 1) || (by constructor 2) || (by constructor 3)).
Qed.

Instance or4V vl avl `{Valid rOLr vl} p1 p2 p3 p4 b1 b2 b3 b4
  `{ArgRec vl avl} `{View avl p1 b1} `{View avl p2 b2} `{View avl p3 b3}
  `{View avl p4 b4} : View vl [\/ p1, p2, p3 | p4] [|| b1, b2, b3 | b4] | 70.
Proof.
  split; apply: (iffP or4P)=> -[] /xPm ?;
   ((by constructor 1) || (by constructor 2)
   || (by constructor 3) || (by constructor 4)).
Qed.

(* -------------------------------------------------------------------- *)
(* All Core Rules                                                       *)

Instance eqV vl (T:eqType) (x y:T) `{Valid rAC vl} :
  View vl (x = y) (x == y) | 0.
Proof. by split;apply eqP. Qed.
(*Typeclasses Opaque is_true. *)

(* All Rules  *)

Instance forallV vl avl `{Valid rAL vl} (T:finType) (P:T->Prop) (p:pred T)
  `{ArgBind vl avl} (V:forall x, View avl (P x) (p x)) :
  View vl (forall x : T, P x) [forall x, p x] | 90.
Proof. split; apply: forallPP=> ?;apply /xPm. Qed.

Instance allV vl avl `{Valid rAL vl} (T:eqType) (P:T->Prop) (p:pred T) l
  `{ArgBind vl avl} (V:forall x, View avl (P x) (p x)) :
  View vl (forall x : T, x \in l -> P x) (all p l) | 90.
Proof. by split; apply: (iffP allP)=> X ? /X /xPm. Qed.

Instance hasV vl avl `{Valid rAL vl} (T:eqType) (P:T->Prop) (p:pred T) l
  `{ArgBind vl avl} (V:forall x, View avl (P x) (p x)) :
  View vl (exists x, x \in l /\ P x) (has p l).
Proof.
  by (split; apply: (iffP hasP)) => [[x Hin /xPm ?]|[x [Hin /xPm ?]]]; exists x.
Qed.

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

Lemma test1 (b1 b2:bool) : b1 && b2.
Proof.
apply /xP.
Abort.

Lemma test1 (n1 n2:nat) : n1 <= n2.
Proof.
apply /xP.
Abort.

Lemma test2 (b1 b2:bool) : b1 && b2 && (b1 == b2).
Proof.
apply /rcxP.
Abort.

Lemma test3 (b1 b2 b3: bool) : [&& b1, b2 & b3].
Proof.
apply /rbxP.
Abort.

Lemma test4 (b1 b2 b3: bool) : [&& b1, b2, b3 & b3].
Proof.
apply /lxP.
Abort.

Lemma test5 (b1 b2 b3: bool) : [&& b1, b2, b3 & b3].
Proof.
apply /xP.
Abort.

Typeclasses Opaque leq.

Lemma test6 l :
  has [pred x : nat | 0 < x <= 7] l.
Proof.
 apply /rbxP.
Abort.

(*

With Canonical Structures it is harder.

Record c1P (b:bool) := mk_c1P { c1b : bool; c1p : Prop; _ : reflect c1p c1b }.
Record c2P (b:bool) := mk_c2P { c2b : bool; c2p : Prop; _ : reflect c2p c2b }.
Record c3P (b:bool) := mk_c3P { c3b : bool; c3p : Prop; _ : reflect c3p c3b }.
Record c4P (b:bool) := mk_c4P { c4b : bool; c4p : Prop; _ : reflect c4p c4b }.

Record wrapped T := Wrap { unwrap : T }.

Definition wrap4 T e := @Wrap T e.
Definition wrap3 T e := @wrap4 T e.
Definition wrap2 T e := @wrap3 T e.
Definition wrap1 T e := @wrap2 T e.
Canonical wrap1.

Record canonicalP (b:bool) := mk_cP{
  cano_bool : wrapped bool;
  cano_prop : wrapped Prop;
  _ : reflect (unwrap cano_prop) (unwrap cano_bool);
}.

Program Canonical cc1P b (c:c1P b) := @mk_cP b (wrap1 c.(c1b)) (wrap1 c.(c1p)) _.
Next Obligation. by case c. Qed.

Program Canonical cc2P b (c:c2P b) := @mk_cP b (wrap2 c.(c2b)) (wrap2 c.(c2p)) _.
Next Obligation. by case c. Qed.

Program Canonical cc3P b (c:c3P b) := @mk_cP b (wrap3 c.(c3b)) (wrap3 c.(c3p)) _.
Next Obligation. by case c. Qed.

Program Canonical cc4P b (c:c4P b) := @mk_cP b (wrap4 c.(c4b)) (wrap4 c.(c4p)) _.
Next Obligation. by case c. Qed.

Lemma xPb b (c:canonicalP b) : reflect (unwrap (cano_prop c)) (unwrap (cano_bool c)).
Proof. by case c. Qed.

Definition rxP := @xPb true.
Definition lxP := @xPb false.

Canonical cP_id mode (b:bool) := @mk_cP mode (Wrap b) (Wrap (is_true b)) (@idP b).

Notation ucb x := (unwrap (cano_bool x)).
Notation ucp x := (unwrap (cano_prop x)).

Lemma cP_and mode (b1 b2:canonicalP mode) : c4P mode.
Proof.
  refine (@mk_c4P mode (ucb b1 && ucb b2) (ucp b1 /\ ucp b2) _).
  by apply:(iffP andP) => -[] /xPb ? /xPb.
Defined.
Canonical cP_and.

(*Lemma cP_and3 (b1 b2 b3:canonicalP) : c3P.
Proof.
  refine (@mk_c3P ([&& ucb b1 , ucb b2 & ucb b3]) ([/\ ucp b1 , ucp b2 & ucp b3]) _).
  by apply:(iffP and3P) => -[] /xP ? /xP ? /xP.
Defined.
Canonical cP_and3.
*)
Lemma cP_imp mode (b1 b2:canonicalP mode) : c4P mode.
Proof.
  refine (@mk_c4P mode (ucb b1 ==> ucb b2) (ucp b1 -> ucp b2) _).
  by apply:(iffP implyP)=> H /xPb /H ?;apply /xPb.
Defined.
Canonical cP_imp.

(*Lemma cP_and4 (b1 b2 b3 b4:canonicalP) : c2P.
Proof.
  refine (@mk_c2P ([&& ucb b1 , ucb b2 & ucb b3]) ([/\ ucp b1 , ucp b2 & ucp b3]) _).
  by apply:(iffP and4P) => -[] /xP ? /xP ? /xP .
Defined.
Canonical cP_and3. *)

Program Canonical cP_has (T:eqType) (P:pred T) l := @mk_c1P true _ _ (@hasP T P l).

*)


Parameter A : Type .

Definition eq_A := @eq A.

Goal forall x, eq_A x x.
intros.
reflexivity.
Qed.


(* basic_misc library for yalla *)

(** * A few basic ingredients used in [yalla] *)

Require Import Lia.

Require Import Permutation_Type_solve.
Require Import genperm_Type.

(** ** Definitions *)
Definition option_prop {A:Type} (P:A->Prop) o :=
match o with
| Some a => P a
| None => True
end.

Definition Empty_fun {A} : Empty_set -> A := fun o => match o with end.

(** ** Tactics *)
(** Generalized version of [split] for [n] components *)

Ltac nsplit n :=
  match n with
  | 1 => idtac
  | S ?m => split ; [ | nsplit m ]
  end.

(** Personalised versions of [easy] *)

Ltac myeasy_pattern assump :=
  try assump ;
  try reflexivity ;
  try lia.

Ltac myeasy := myeasy_pattern assumption.
Ltac myeeasy := myeasy_pattern eassumption.

(** An [easy] tactic with automated permutation solving **)

Ltac myeasy_perm_Type :=
  myeeasy ;
  try PCperm_Type_solve ;
  try (simpl_hyp_perm_all_Type ; PCperm_Type_solve).


(*
Copyright © 2009 Valentin Blot

Permission is hereby granted, free of charge, to any person obtaining a copy of
this proof and associated documentation files (the "Proof"), to deal in
the Proof without restriction, including without limitation the rights to
use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
the Proof, and to permit persons to whom the Proof is furnished to do so,
subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Proof.

THE PROOF IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
CONNECTION WITH THE PROOF OR THE USE OR OTHER DEALINGS IN THE PROOF.
*)

(** Rings as a type class *)
Require Import Setoid Ring Morphisms.
Open Scope signature_scope.

Require Import OperationClasses.

Set Implicit Arguments.
Unset Strict Implicit.

Section Definitions.

Context {R: Type} {req: relation R}.
Class Ring {r_st : Equivalence req} (rO rI : R) (radd rmul rsub : binop R) (ropp : unop R) :=
  { r_rt : ring_theory rO rI radd rmul rsub ropp req;
    r_ree : ring_eq_ext radd rmul ropp req
}.

Class SubRing {r_st : Equivalence req} `{r_ring : Ring} (P : R -> Type) :=
  {zero_stab : P rO;
   one_stab : P rI;
   radd_int : binop_intern P radd;
   rmul_int : binop_intern P rmul;
   rsub_int : binop_intern P rsub;
   ropp_int : unop_intern P ropp}.

End Definitions.

Section Properties.

Context `{r_ring : Ring}.
Global Instance radd_morph : Proper (req==>req==>req) radd.
Proof. reduce; apply (Radd_ext r_ree); auto. Qed.
Global Instance rmul_morph : Proper (req==>req==>req) rmul.
Proof. reduce; apply (Rmul_ext r_ree); auto. Qed.
Global Instance rsub_morph : Proper (req==>req==>req) rsub.
Proof.
 reduce; rewrite -> (Rsub_def r_rt), -> (Rsub_def r_rt y y0).
 apply (Radd_ext r_ree); auto; apply (Ropp_ext r_ree); auto.
Qed.
Global Instance ropp_morph : Proper (req==>req) ropp.
Proof. reduce; apply (Ropp_ext r_ree); auto. Qed.
Global Instance radd_assoc : associative radd.
Proof. reduce; apply (Radd_assoc r_rt). Qed.
Global Instance radd_comm : commutative radd.
Proof. reduce; apply (Radd_comm r_rt). Qed.
Global Instance radd_left_comm : left_commutative radd := mulAC_comm_l.
Global Instance radd_right_comm : right_commutative radd := mulAC_comm_r.
Global Instance radd_left_unit : left_unit radd rO.
Proof. reduce; apply (Radd_0_l r_rt). Qed.
Global Instance radd_right_unit : right_unit radd rO := mulC_id_l.
Global Instance ropp_left_inverse : left_inverse radd rO ropp.
Proof. reduce; apply (Ropp_def r_rt). Qed.
Global Instance ropp_right_inverse : right_inverse radd rO ropp := mulC_inv_l.
Global Instance rmul_assoc : associative rmul.
Proof. reduce; apply (Rmul_assoc r_rt). Qed.
Global Instance rmul_comm : commutative rmul.
Proof. reduce; apply (Rmul_comm r_rt). Qed.
Global Instance rmul_left_comm : left_commutative rmul := mulAC_comm_l.
Global Instance rmul_right_comm : right_commutative rmul := mulAC_comm_r.
Global Instance rmul_left_unit : left_unit rmul rI.
Proof. reduce; apply (Rmul_1_l r_rt). Qed.
Global Instance rmul_right_unit : right_unit rmul rI := mulC_id_l.
Global Instance radd_rmul_left_distr : left_distributive radd rmul.
Proof. reduce; apply (Rdistr_l r_rt). Qed.
Global Instance radd_rmul_right_distr : right_distributive radd rmul := mulC_distr_l.
Global Instance rmul_left_zero : left_absorbing rmul rO := @opA_zero_l R req r_st radd rmul ropp rO radd_morph rmul_morph radd_assoc radd_comm ropp_left_inverse radd_left_unit radd_rmul_left_distr.
Global Instance rmul_right_zero : right_absorbing rmul rO := @mulC_zero_l R req r_st rmul rmul_comm rO rmul_left_zero.
End Properties.

Section SubRing_is_Ring.

Context `{SubRing}.
Add Ring r_r : r_rt (setoid r_st r_ree).

Let R' := sigT P.
Let proj1' := fun x : R' => projT1 x.
Coercion proj1' : R'>->R.
Let req' : relation R' := fun x y => req (projT1 x) (projT1 y).
Instance r_st' : Equivalence req'.
Proof.
 constructor; intro x; destruct x as [x Px]; try (intro y; destruct y as [y Py]);
   try (intro z; destruct z as [z Pz]);
     unfold req'; [ reflexivity | intro; symmetry; assumption | intros eqxy eqyz; rewrite eqxy; assumption ].
Qed.
Let rO' := existT P rO zero_stab.
Let rI' := existT P rI one_stab.
Let radd' : binop R' := fun x y => existT P (radd x y) (radd_int (projT2 x) (projT2 y)).
Let rmul' : binop R' := fun x y => existT P (rmul x y) (rmul_int (projT2 x) (projT2 y)).
Let rsub' : binop R' := fun x y => existT P (rsub x y) (rsub_int (projT2 x) (projT2 y)).
Let ropp' : unop R' := fun x => existT P (ropp x) (ropp_int (projT2 x)).
Global Instance sr_ring : @Ring R' req' r_st' rO' rI' radd' rmul' rsub' ropp'.
Proof.
 constructor.
  constructor; unfold R', req', radd', rmul', rsub', ropp', proj1'; intro x; destruct x as [ x Px ];
    try (intro y; destruct y as [ y Py ]); try (intro z; destruct z as [ z Pz ]); simpl; ring.
 constructor; unfold R', req', radd', rmul', rsub', ropp', proj1'; simpl;
   intros x x'; destruct x as [ x Px ]; destruct x' as [ x' Px' ]; simpl; intro eqx;
     try (intros y y'; destruct y as [ y Py ]; destruct y' as [ y' Py' ]; simpl; intro eqy; simpl);
       [ apply radd_morph | apply rmul_morph | apply ropp_morph ]; assumption.
Qed.

End SubRing_is_Ring.


(*
 * Usage: You may need "make equiv.vo" before
 * using this file.
 * In vscode, you may also need to add
 * "coqtop.args": [ "-Q", "/path/to/this/dir", "" ]
 * to settings.json.
 *)

Require Export equiv.
Require Export Lists.List.

Local Open Scope type_scope.  (* for '*' *)

(* finite control states *)

Parameter Q : Set.
Definition Q' := Q * Phi.

Inductive Com :=
  | pop
  | skip
  | push (j : nat).
Inductive Com' :=
  | pop'
  | skip'
  | push' (phi : Phi).

(* Stack and configurations *)

Definition Stack  := list (D * Theta).
Definition Stack' := list Phi.

Definition Config  := Q  * Theta * Stack.
Definition Config' := Q' * Stack'.

Definition update_stack (u : Stack) theta com :=
  match com with
  | pop => match u with
             _ :: u' => u' | nil => nil
           end
  | skip => u
  | push j => (theta j, theta) :: u
  end.

Definition update_stack' (u : Stack') com' :=
  match com' with
  | pop' => match u with
              _ :: u' => u' | nil => nil
            end
  | skip' => u
  | push' z => z :: u
  end.

Definition not_contain d (cell: D * Theta) :=
  match cell with
    (_, theta) => forall i, theta i <> d
  end.

Axiom stack_has_outside_data :
  forall (theta: Theta) (e: D) (u: Stack),
  exists d : D,
    d <> e /\ (forall i, theta i <> d) /\
    Forall (not_contain d) u.

Definition is_proper_stack (stack : Stack) :=
  let is_proper_cell cell :=
    match cell with (z, theta) => exists i, z = theta i end
  in Forall is_proper_cell stack.

(* Bisimulation relation between configs *)

(* The stack bottom of a given starting ID *)
Parameter theta_0 : Theta.
Parameter d_0 : D.
Axiom proper_bottom : exists i, d_0 = theta_0 i.

Inductive stack_R_stack'
  : Theta -> Stack -> Phi -> Stack' -> Prop :=
  | Stack_R_stack'_nil :
    forall theta phi,
      (theta_0, d_0, theta) |= phi ->
      stack_R_stack' theta nil phi nil
  | Stack_R_stack'_cons :
    forall theta ptheta d phi pphi u v,
      (ptheta, d, theta) |= phi ->
      stack_R_stack' ptheta u pphi v ->
      stack_R_stack' theta ((d, ptheta) :: u) phi (pphi :: v).

Inductive config_R_config'
  : Config -> Config' -> Prop :=
  | Config_R_config' :
    forall q theta u phi v,
      stack_R_stack' theta u phi v ->
      config_R_config' (q, theta, u) ((q, phi), v).

(* freshness_p on stack *)

Inductive Forall2 {X : Type} (P : X -> X -> Prop)
  : list X -> Prop :=
  | Forall2_nil : Forall2 P nil
  | Forall2_cons x l :
    Forall (P x) l -> Forall2 P l -> Forall2 P (x :: l).

Inductive Forall3 {X : Type} (P : X -> X -> X -> Prop)
  : list X -> Prop :=
  | Forall3_nil : Forall3 P nil
  | Forall3_cons x l :
    Forall2 (P x) l -> Forall3 P l -> Forall3 P (x :: l).

Definition freshness_p_on_triple
  (p3 p2 p1 : D * Theta) :=
  match (p3, p2, p1) with
    ((_, th3), (_, th2), (d, th1))
    => freshness_p th1 d th2 th3
  end.

Definition freshness_p_on_stack theta stack :=
  Forall3 freshness_p_on_triple ((bot, theta) :: stack).

Local Example forall2_example_1 :
  Forall2 le (0 :: 1 :: 3 :: 8 :: nil).
Proof.
  repeat apply Forall2_cons;
  repeat apply Forall_cons;
  repeat (apply le_S; try apply le_n); auto;
  try apply Forall2_nil.
Qed.

(* Lemmas *)

(* Forall, Forall2, Forall3 *)

Lemma Forall_sublist {A : Type} :
  forall p (a : A) u1 u2,
  Forall p (u1 ++ (a :: u2)) ->
  Forall p (u1 ++ u2).
Proof.
  intros p a u1 u2.
  induction u1 as [| b u1 IHu1].
  - (* u1 = nil -> ... *)
  unfold app.
  intros Hfor.
  inversion Hfor as [| x l Hpa Hfor' [EQx EQl]].
  exact Hfor'.
  - (* u1 = b :: u1 -> ... *)
  repeat rewrite<- app_comm_cons.
  intros Hfor.
  inversion Hfor as [| x l Hpb Hfor' [EQx EQl]].
  clear x EQx l EQl.
  apply Forall_cons; auto.
Qed.

Local Lemma Forall2_sublist {A : Type} :
  forall p (a : A) u1 u2,
  Forall2 p (u1 ++ (a :: u2)) ->
  Forall2 p (u1 ++ u2).
Proof.
  intros p a u1 u2.
  induction u1 as [| b u1 IHu1].
  - (* u1 = nil -> ... *)
  unfold app.
  intros Hfor2.
  inversion Hfor2 as [| x l Hfor Hfor2' [EQx EQl]].
  exact Hfor2'.
  - (* u1 = b :: u1 -> ... *)
  repeat rewrite<- app_comm_cons.
  intros Hfor2.
  inversion Hfor2 as [| x l Hfor Hfor2' [EQx EQl]].
  clear x EQx l EQl.
  apply Forall2_cons.
  + (* Forall ... *)
  now apply Forall_sublist with a.
  + (* Forall2 ... *)
  now apply IHu1.
Qed.

Lemma Forall3_sublist {A : Type} :
  forall p (a : A) u1 u2,
  Forall3 p (u1 ++ (a :: u2)) ->
  Forall3 p (u1 ++ u2).
Proof.
  intros p a u1 u2.
  induction u1 as [| b u1 IHu1].
  - (* u1 = nil -> ... *)
  unfold app.
  intros Hfor3.
  inversion Hfor3 as [| x l Hfor2 Hfor3' [EQx EQl]].
  exact Hfor3'.
  - (* u1 = b :: u1 -> ... *)
  repeat rewrite<- app_comm_cons.
  intros Hfor3.
  inversion Hfor3 as [| x l Hfor2 Hfor3' [EQx EQl]].
  clear x EQx l EQl.
  apply Forall3_cons.
  + (* Forall ... *)
  now apply Forall2_sublist with a.
  + (* Forall2 ... *)
  now apply IHu1.
Qed.

Local Lemma Forall2_hd2 {A : Type} :
  forall p (a : A) b u,
  Forall2 p (a :: b :: u) ->
  p a b.
Proof.
  intros p a b u H.
  inversion H as [| x l Hfor Hfor2 [EQx EQl]].
  clear x EQx l EQl H Hfor2.
  inversion Hfor as [| x l H Hfor' [EQx EQl]].
  exact H.
Qed.

Lemma Forall3_hd3 {A : Type} :
  forall p (a : A) b c u,
  Forall3 p (a :: b :: c :: u) ->
  p a b c.
Proof.
  intros p a b c u H.
  inversion H as [| x l Hfor2 Hfor3' [EQx EQl]].
  clear x EQx l EQl Hfor3' H.
  apply Forall2_hd2 with u.
  exact Hfor2.
Qed.

(* is_proper_stack *)

Lemma substack_is_proper_stack :
  forall a u1 u2,
  is_proper_stack (u1 ++ (a :: u2)) ->
  is_proper_stack (u1 ++ u2).
Proof.
  apply Forall_sublist.
Qed.

(* freshness_p_on_stack *)

Lemma substack_keeps_freshness_p_0 :
  forall theta d th u,
  freshness_p_on_stack theta ((d, th) :: u) ->
  freshness_p_on_stack th u.
Proof.
  intros theta d th u.
  unfold freshness_p_on_stack.
  intros Hfor3.
  inversion Hfor3 as [| x l Hfor2 Hfor3' [EQx EQl]].
  clear x l EQx EQl Hfor3.
  inversion Hfor3' as [| x l Hfor2' Hfor3 [EQx EQl]].
  clear x l EQx EQl Hfor3'.
  now apply Forall3_cons.
Qed.

Lemma substack_keeps_freshness_p :
  forall theta a u1 u2,
  freshness_p_on_stack theta (u1 ++ (a :: u2)) ->
  freshness_p_on_stack theta (u1 ++ u2).
Proof.
  intros theta a u1 u2.
  unfold freshness_p_on_stack.
  repeat rewrite app_comm_cons.
  apply Forall3_sublist.
Qed.

Lemma push_keeps_freshness_p :
  forall theta u z,
  freshness_p_on_stack theta u ->
  freshness_p_on_stack theta ((z, theta) :: u).
Proof.
  intros theta u z.
  unfold freshness_p_on_stack.
  intros H.
  apply Forall3_cons.
  - (* Forall2 ... ((z, theta) :: u) *)
  apply Forall2_cons.
  + (* Forall ... u *)
  apply Forall_forall.
  intros [d1 th1] Hth1.
  unfold freshness_p_on_triple.
  unfold freshness_p.
  split.
  * (* forall i j, th1 i = ... -> ... *)
  intros i j H1.
  inversion H as [| x l Hfor2 Hfor3 [EQx EQl]].
  clear x EQx l EQl Hfor3.
  inversion Hfor2 as [EQu | x l Hfor Hfor2' EQx].
  -- (* u = nil -> ... *)
  rewrite<- EQu in Hth1.
  apply in_nil in Hth1.
  contradiction.
  -- (* u = x :: l -> ... *)
  exists j.
  exact H1.
  * (* forall j, d1 = ... -> ... *)
  intros j H1.
  exists j.
  exact H1.
  + (* Forall2 ... u *)
  inversion H as [| x l Hfor2 Hfor3 [EQx EQl]].
  clear x EQx l EQl Hfor3.
  exact Hfor2.

  - (* Forall3 ... ((theta j, ) :: u) *)
  apply Forall3_cons.
  + (* Forall2 ... u *)
  induction u as [| [d1 th1] u IHu].
  * (* u = nil -> ... *)
  apply Forall2_nil.
  * (* u = (d1, th1) :: u -> ... *)
  apply Forall2_cons.
  -- (* Forall ... u *)
  apply Forall_forall.
  intros [d2 th2] Hth2.
  unfold freshness_p_on_triple.
  inversion H as [| x l Hfor2 Hfor3 [EQx EQl]].
  clear x EQx l EQl Hfor3.
  inversion Hfor2 as [| x l Hfor Hfor2' [EQx EQl]].
  clear x EQx l EQl Hfor2 Hfor2'.
  rewrite Forall_forall in Hfor.
  unfold freshness_p_on_triple in Hfor.
  apply (Hfor (d2, th2)).
  exact Hth2.
  -- (* Forall2 ... u *)
  apply IHu.
  apply (Forall3_sublist _ (d1, th1) ((bot, theta) :: nil)).
  unfold app.
  exact H.
  + (* Forall3 ... u *)
  apply (Forall3_sublist _ (bot, theta) nil).
  unfold app.
  exact H.
Qed.

(* config_R_config' *)

Lemma config_R_nil_nil_1 :
  forall q theta u phi,
  config_R_config' (q, theta, u) (q, phi, nil) ->
  u = nil.
Proof.
  intros q theta u phi.
  intro H.
  inversion H as [q1 theta1 u1 phi1 v1 HR].
  inversion HR.
  reflexivity.
Qed.

Lemma config_R_nil_nil_2 :
  forall q theta v phi,
  config_R_config' (q, theta, nil) (q, phi, v) ->
  v = nil.
Proof.
  intros q theta v phi.
  intro H.
  inversion H as [q1 theta1 u1 phi1 v1 HR].
  inversion HR.
  reflexivity.
Qed.


Require Import Monad.
Require Import String.

Inductive Exception : Type :=
| exception : string -> string -> Exception.

Inductive ExceptionM (A:Type) : Type :=
| exn : Exception -> ExceptionM A
| no_exn : A -> ExceptionM A.

Definition bind_exn (A : Type) (B : Type) (e : ExceptionM A) (f : A -> ExceptionM B) : ExceptionM B :=
  match e with
    | exn t => exn B t
    | no_exn a => (f a)
  end.

Definition ret_exn (A : Type) (a : A) := no_exn A a.

Program Instance Monad ExceptionM :=
  ret := ret_exn ;
  bind := bind_exn.

  Next Obligation. (* bind left ret *)
  Proof.
    unfold bind_exn ; unfold ret_exn ; destruct x ; auto.
  Qed.

  Next Obligation. (* bind right ret *)
  Proof.
    unfold bind_exn ; unfold ret_exn ; destruct x ; auto.
  Qed.

Definition raise (A : Type) (t s:string) : ExceptionM A := exn A (exception t s).
Implicit Arguments raise [A].

Definition alt (A: Type) (e1 e2 : ExceptionM A) : ExceptionM A :=
  match e1 with
    | exn t => match e2 with
                 | exn t' => exn A t'
                 | no_exn a => no_exn A a
               end
    | no_exn a => no_exn A a
  end.
Notation "A |+| B" := (alt _ A B) (at level 55).

Definition handle (A:Type) (s:string) (e: ExceptionM A) (handler : string -> ExceptionM A) : ExceptionM A :=
  match e with
    | exn t => match t with
                 | exception name msg =>
                   match string_dec s name with
                     | left _ => handler msg
                     | right _ => e
                   end
               end
    | _ => e
  end.

Notation "A |:| N , s => B" := (handle _ N A (fun s => B)) (at level 105).

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

(*
Check (raise (A:=nat) "test" "problem with zero").
Open Scope string_scope.
Definition zero_exn := "zero_minus".
Definition syntax_error := "syntax_error".
Definition P (n:nat) :=
  match n with
    | 0 => raise zero_exn "problem with zero"
    | S n' => ret n'
  end.

Eval compute in P 0.

Definition fail_test (n:nat) :=
  m <- (P n) ;; ret (m,"no prob")
  |:| zero_exn, msg => ret (n,msg)
  |:| syntax_error, msg => ret (0,msg).

Eval compute in fail_test 0.
Eval compute in P 1 |+| P 0.

*)


(******************************)
(******************************)
(****                      ****)
(****   Variable opening   ****)
(****                      ****)
(******************************)
(******************************)

Require Import Coq.Arith.Le.
Require Import Coq.Arith.Peano_dec.
Require Import Coq.Lists.List.
Require Import Coq.micromega.Lia.
Require Import Main.SystemF.FreeVar.
Require Import Main.SystemF.LocalClosure.
Require Import Main.SystemF.Syntax.
Require Import Main.Tactics.

Import Coq.Arith.PeanoNat.Nat.

Fixpoint eeOpen e1 i1 e2 :=
  match e1 with
  | eFreeVar _ => e1
  | eBoundVar i2 => if eq_nat_dec i1 i2 then e2 else e1
  | eAbs t e3 => eAbs t (eeOpen e3 (S i1) e2)
  | eApp e3 e4 => eApp (eeOpen e3 i1 e2) (eeOpen e4 i1 e2)
  | eTAbs e3 => eTAbs (eeOpen e3 i1 e2)
  | eTApp e3 t => eTApp (eeOpen e3 i1 e2) t
  end.

Fixpoint ttOpen t1 i1 t2 :=
  match t1 with
  | tFreeVar _ => t1
  | tBoundVar i2 => if eq_nat_dec i1 i2 then t2 else t1
  | tArrow t3 t4 => tArrow (ttOpen t3 i1 t2) (ttOpen t4 i1 t2)
  | tForAll t3 => tForAll (ttOpen t3 (S i1) t2)
  end.

Fixpoint etOpen e1 i1 t1 :=
  match e1 with
  | eFreeVar _ => e1
  | eBoundVar _ => e1
  | eAbs t2 e2 => eAbs (ttOpen t2 i1 t1) (etOpen e2 i1 t1)
  | eApp e2 e3 => eApp (etOpen e2 i1 t1) (etOpen e3 i1 t1)
  | eTAbs e2 => eTAbs (etOpen e2 (S i1) t1)
  | eTApp e2 t2 => eTApp (etOpen e2 i1 t1) (ttOpen t2 i1 t1)
  end.

(*****************************************************)
(* Opening a locally closed term/type has no effect. *)
(*****************************************************)

Theorem tttOpenLocallyClosed :
  forall i t1 t2,
  tLocallyClosed t1 i ->
  ttOpen t1 i t2 = t1.
Proof.
  clean. induction H; search.
Qed.

#[export] Hint Resolve tttOpenLocallyClosed : main.

Theorem eeeOpenLocallyClosed :
  forall e1 e2 ie it,
  eLocallyClosed e1 ie it ->
  eeOpen e1 ie e2 = e1.
Proof.
  clean. induction H; search.
Qed.

#[export] Hint Resolve eeeOpenLocallyClosed : main.

Theorem eetOpenLocallyClosed :
  forall e ie it t,
  eLocallyClosed e ie it ->
  etOpen e it t = e.
Proof.
  clean. induction H; search.
Qed.

#[export] Hint Resolve eetOpenLocallyClosed : main.

(***************************************************************************)
(* If the opening of a term/type is locally closed at some level, then the *)
(* term/type is locally closed at the next level.                          *)
(***************************************************************************)

Theorem ttLocallyClosedOpen :
  forall i t1 t2,
  tLocallyClosed (ttOpen t1 i t2) i ->
  tLocallyClosed t1 (S i).
Proof.
  clean. outro t1 i.
  induction t1; search; clean.
  - destruct (eq_dec i n); search.
    apply tLocalClosureMonotonic with (i1 := i); search.
  - invert H. search.
  - invert H. search.
Qed.

#[export] Hint Resolve ttLocallyClosedOpen : main.

Theorem eeLocallyClosedOpen :
  forall e1 e2 ie it,
  eLocallyClosed (eeOpen e1 ie e2) ie it ->
  eLocallyClosed e1 (S ie) it.
Proof.
  clean. outro e1 ie it.
  induction e1; search; clean.
  - destruct (eq_dec ie n); search.
    apply eLocalClosureMonotonic with (ie1 := ie) (it1 := it); search.
  - invert H. search.
  - invert H. search.
  - invert H. search.
  - invert H. search.
Qed.

#[export] Hint Resolve eeLocallyClosedOpen : main.

Theorem etLocallyClosedOpen :
  forall e ie it t,
  eLocallyClosed (etOpen e it t) ie it ->
  eLocallyClosed e ie (S it).
Proof.
  clean. outro e ie it.
  induction e; search; clean.
  - apply eLocalClosureMonotonic with (ie1 := ie) (it1 := it); search.
  - invert H. eSearch.
  - invert H. search.
  - invert H. search.
  - invert H. eSearch.
Qed.

#[export] Hint Resolve etLocallyClosedOpen : main.

(********************************)
(* Free variables of an opening *)
(********************************)

Theorem tttFreeOpen :
  forall i t1 t2,
  incl (tFreeVars (ttOpen t1 i t2)) (tFreeVars t2 ++ tFreeVars t1) /\
  incl (tFreeVars t1) (tFreeVars (ttOpen t1 i t2)).
Proof.
  clean. split; outro i.
  - induction t1; search. clean. apply incl_app.
    + apply incl_tran with (m := tFreeVars t2 ++ tFreeVars t1_1); search.
    + apply incl_tran with (m := tFreeVars t2 ++ tFreeVars t1_2); search.
  - induction t1; search. unfold incl. search.
Qed.

#[export] Hint Resolve tttFreeOpen : main.

Theorem eeeeFreeOpen :
  forall e1 e2 i,
  incl (eeFreeVars (eeOpen e1 i e2)) (eeFreeVars e2 ++ eeFreeVars e1) /\
  incl (eeFreeVars e1) (eeFreeVars (eeOpen e1 i e2)).
Proof.
  clean. split; outro i.
  - induction e1; search. clean. apply incl_app.
    + apply incl_tran with (m := eeFreeVars e2 ++ eeFreeVars e1_1); search.
    + apply incl_tran with (m := eeFreeVars e2 ++ eeFreeVars e1_2); search.
  - induction e1; search. unfold incl. search.
Qed.

#[export] Hint Resolve eeeeFreeOpen : main.

Theorem eeetFreeOpen :
  forall e i t,
  incl (eeFreeVars (etOpen e i t)) (eeFreeVars e) /\
  incl (eeFreeVars e) (eeFreeVars (etOpen e i t)).
Proof.
  clean. split; outro i; induction e; search.
Qed.

#[export] Hint Resolve eeetFreeOpen : main.

Theorem eteeFreeOpen :
  forall e1 e2 i,
  incl (etFreeVars (eeOpen e1 i e2)) (etFreeVars e2 ++ etFreeVars e1) /\
  incl (etFreeVars e1) (etFreeVars (eeOpen e1 i e2)).
Proof.
  clean. split; outro i.
  - induction e1; search; clean.
    + specialize (IHe1 (S i)). apply incl_app; search.
      apply incl_tran with (m := etFreeVars e2 ++ etFreeVars e1); search.
    + specialize (IHe1_1 i). specialize (IHe1_2 i). apply incl_app; search.
      * apply incl_tran with (m := etFreeVars e2 ++ etFreeVars e1_1); search.
      * apply incl_tran with (m := etFreeVars e2 ++ etFreeVars e1_2); search.
    + specialize (IHe1 i). apply incl_app; search.
      apply incl_tran with (m := etFreeVars e2 ++ etFreeVars e1); search.
  - induction e1; search. unfold incl. search.
Qed.

#[export] Hint Resolve eteeFreeOpen : main.

Theorem etetFreeOpen :
  forall e i t,
  incl (etFreeVars (etOpen e i t)) (tFreeVars t ++ etFreeVars e) /\
  incl (etFreeVars e) (etFreeVars (etOpen e i t)).
Proof.
  clean. split; outro i.
  - induction e; search; clean; apply incl_app.
    + apply incl_tran with (m := tFreeVars t0 ++ tFreeVars t1); search.
      apply tttFreeOpen.
    + apply incl_tran with (m := tFreeVars t0 ++ etFreeVars e); search.
    + apply incl_tran with (m := tFreeVars t0 ++ etFreeVars e1); search.
    + apply incl_tran with (m := tFreeVars t0 ++ etFreeVars e2); search.
    + apply incl_tran with (m := tFreeVars t0 ++ etFreeVars e); search.
    + apply incl_tran with (m := tFreeVars t0 ++ tFreeVars t1); search.
      apply tttFreeOpen.
  - induction e; search; clean;
      specialize (IHe i); apply incl_app; search;
      apply incl_tran with (m := tFreeVars (ttOpen t1 i t0)); search;
      apply tttFreeOpen.
Qed.

#[export] Hint Resolve etetFreeOpen : main.

(********************************************)
(* Opening binders preserves local closure. *)
(********************************************)

Theorem locallyClosedOpenForAll :
  forall i t1 t2,
  tLocallyClosed (tForAll t1) i ->
  tLocallyClosed t2 i ->
  tLocallyClosed (ttOpen t1 i t2) i.
Proof.
  clean. invert H.
  remember (S i). assert (n <= S i); search. clear Heqn. outro t2 i.
  induction H2; search; clean.
  - destruct (eq_dec i n1); search. apply tlcBoundVar. lia.
  - constructor; search.
    specialize (IHtLocallyClosed t2 (S i)).
    feed IHtLocallyClosed.
    apply tLocalClosureMonotonic with (i1 := i); search.
Qed.

#[export] Hint Resolve locallyClosedOpenForAll : main.

Theorem locallyClosedOpenAbs :
  forall e1 e2 ie it t,
  eLocallyClosed (eAbs t e1) ie it ->
  eLocallyClosed e2 ie it ->
  eLocallyClosed (eeOpen e1 ie e2) ie it.
Proof.
  clean. invert H. clear t0 H3.
  remember (S ie). assert (n <= S ie); search. clear Heqn. outro e2 ie.
  induction H6; search; clean.
  - destruct (eq_dec ie ne1); search. apply elcBoundVar. lia.
  - constructor; search.
    specialize (IHeLocallyClosed e2 (S ie)).
    apply IHeLocallyClosed; search.
    apply eLocalClosureMonotonic with (ie1 := ie) (it1 := nt); search.
  - constructor.
    specialize (IHeLocallyClosed e2 ie).
    apply IHeLocallyClosed; search.
    apply eLocalClosureMonotonic with (ie1 := ie) (it1 := nt); search.
Qed.

#[export] Hint Resolve locallyClosedOpenAbs : main.

Theorem locallyClosedOpenTAbs :
  forall e ie it t,
  eLocallyClosed (eTAbs e) ie it ->
  tLocallyClosed t it ->
  eLocallyClosed (etOpen e it t) ie it.
Proof.
  clean. invert H.
  remember (S it). assert (n <= S it); search. clear Heqn. outro t0 it.
  induction H2; search; constructor; search; clean.
  - apply locallyClosedOpenForAll; search.
    constructor.
    apply tLocalClosureMonotonic with (i1 := nt); search.
  - specialize (IHeLocallyClosed t0 (S it)).
    feed IHeLocallyClosed.
    apply tLocalClosureMonotonic with (i1 := it); search.
  - apply locallyClosedOpenForAll; search.
    constructor.
    apply tLocalClosureMonotonic with (i1 := nt); search.
Qed.

#[export] Hint Resolve locallyClosedOpenTAbs : main.


(*** Compile to finite trees first. *)
(*** NO LOOPS *)

(*To install ITree, do:
    opam install coq-itree *)
From ITree Require Import
     ITree ITreeFacts.

Import ITreeNotations.

Require Import Coq.Program.Basics.
Require Import List.
Require Import Coq.Init.Nat.
Require Import PeanoNat.
Require Import Omega.
Local Open Scope program_scope.

Import ListNotations.

Section Zar.

  Inductive val : Type -> Type :=
  | VBool : bool -> val bool.

  Definition bool_of_val (v : val bool) : bool :=
    match v with
    | VBool b => b
    end.

  Definition St : Type := nat -> val bool.
  Definition empty : St := fun _ => VBool false.
  Definition upd (x : nat) (v : val bool) (st : St) : St :=
    fun y => if Nat.eqb y x then v else st y.

  Inductive exp : Type -> Type :=
  | EVal : val bool -> exp bool
  | EVar : nat -> exp bool
  | ENot : exp bool -> exp bool
  | EAnd : exp bool -> exp bool -> exp bool
  | EOr : exp bool -> exp bool -> exp bool
  | EEq : forall {A}, exp A -> exp A -> exp bool.

  Fixpoint eval {A : Type} (e : exp A) (st : St) : val A :=
    match e with
      | EVal b => b
      | EVar x => st x
      | ENot e' =>
        match eval e' st with
        | VBool b => VBool (negb b)
        end
      | EAnd e1 e2 =>
        match (eval e1 st, eval e2 st) with
        | (VBool b1, VBool b2) => VBool (andb b1 b2)
        end
      | EOr e1 e2 =>
        match (eval e1 st, eval e2 st) with
        | (VBool b1, VBool b2) => VBool (orb b1 b2)
        end
      | EEq e1 e2 =>
        match (eval e1 st, eval e2 st) with
        | (VBool b1, VBool b2) => VBool (Bool.eqb b1 b2)
        end
    end.

  Inductive com : Type :=
  | CSkip : com
  | CAssign : forall {A}, nat -> exp A -> com
  | CSeq : com -> com -> com
  | CIte : exp bool -> com -> com -> com
  | CChoice : nat -> nat -> com -> com -> com
  | CObserve : exp bool -> com.
  (* | CWhile : Exp bool -> Com -> Com. *)

  Record prog :=
    { prog_com : com
    ; prog_query : exp bool
    }.

  Variant bitE : Type -> Type :=
  | Get : bitE bool.

  Inductive tree A : Type :=
  | Leaf : A -> tree A
  | Split : tree A -> tree A -> tree A
  | Fail : tree A.

  Fixpoint tree_fmap {A B : Type} (f : A -> B) (t : tree A) : tree B :=
    match t with
    | Leaf _ x => Leaf _ (f x)
    | Split _ t1 t2 => Split _ (tree_fmap f t1) (tree_fmap f t2)
    | Fail _ => Fail _
    end.

  Fixpoint tree_join {A : Type} (t : tree (tree A)) : tree A :=
    match t with
    | Leaf _ t' => t'
    | Split _ t1 t2 => Split _ (tree_join t1) (tree_join t2)
    | Fail _ => Fail _
    end.

  Definition tree_bind {A B : Type} (t : tree A) (k : A -> tree B) : tree B :=
    tree_join (tree_fmap k t).

  Instance Monad_tree : Monad tree := { ret := Leaf; bind := @tree_bind }.

  Definition tree_kcomp {A B C : Type} (f : A -> tree B) (g : B -> tree C) : A -> tree C :=
    fun x => bind (f x) g.

  (* (** Using option to handle failure. *) *)
  (* Fixpoint add_to_tree_aux {A : Type} (x : A) (t : tree (unit + A)) : option (tree (unit + A)) := *)
  (*   match t with *)
  (*   | Leaf _ (inl _) => Some (Leaf _ (inr x)) *)
  (*   | Leaf _ (inr _) => None *)
  (*   | Split _ t1 t2 => *)
  (*     match add_to_tree_aux x t1 with *)
  (*     | Some t1' => Some (Split _ t1' t2) *)
  (*     | None => match add_to_tree_aux x t2 with *)
  (*              | Some t2' => Some (Split _ t1 t2') *)
  (*              | None => None *)
  (*              end *)
  (*     end *)
  (*   | Fail _ => None *)
  (*   end. *)

  Definition unit_bool_eqb (x y : unit + bool) : bool :=
    match (x, y) with
    | (inl _, inl _) => true
    | (inr b1, inr b2) => Bool.eqb b1 b2 (* XNOR *)
    | _ => false
    end.

  Lemma unit_bool_eqb_spec (x y : unit + bool) : Bool.reflect (x = y) (unit_bool_eqb x y).
  Proof.
    destruct x, y.
    - destruct u, u0; left; reflexivity.
    - right; congruence.
    - right; congruence.
    - destruct b, b0; constructor; auto; congruence.
  Qed.

  Fixpoint unit_bool_tree_eqb (t1 : tree (unit + bool)) (t2 : tree (unit + bool)) : bool :=
  match (t1, t2) with
  | (Leaf _ b1, Leaf _ b2) => unit_bool_eqb b1 b2
  | ((Split _ t1 t2), (Split _ t1' t2')) =>
    andb (unit_bool_tree_eqb t1 t1') (unit_bool_tree_eqb t2 t2')
  | (Fail _, Fail _) => true
  | _ => false
  end.

  (** Using equality testing to check for success (SPECIALIZED TO BOOL) *)
  Fixpoint add_to_tree (b : bool) (t : tree (unit + bool)) : tree (unit + bool) :=
    match t with
    | Leaf _ (inl _) => Leaf _ (inr b)
    | Leaf _ (inr _) => t
    | Split _ t1 t2 =>
      let t1' := add_to_tree b t1 in
      if unit_bool_tree_eqb t1 t1'
      then Split _ t1 (add_to_tree b t2)
      else Split _ t1' t2
    | Fail _ => Fail _
    end.

  Fixpoint tree_depth {A : Type} (t : tree A) : nat :=
    match t with
    | Split _ t1 t2 => 1 + tree_depth t1 + tree_depth t2
    | _ => 0
    end.

  Fixpoint new_bool_tree (n : nat) : tree (unit + bool) :=
    match n with
    | 0 => Leaf _ (inl tt)
    | S n' => Split _ (new_bool_tree n') (new_bool_tree n')
    end.

  Fixpoint list_tree_aux (l : list bool) (t : tree (unit + bool)) : tree (unit + bool) :=
    match l with
    | [] => t
    | b :: bs =>
      let t' := add_to_tree b t in
      let t'' := if unit_bool_tree_eqb t t'
                 then Split _ t (add_to_tree b (new_bool_tree (tree_depth t)))
                 else t' in
      list_tree_aux bs t''
    end.

  Fixpoint tie_tree (t : tree (unit + bool)) : tree bool :=
    match t with
    | Leaf _ (inl _) => Fail _
    | Leaf _ (inr b) => Leaf _ b
    | Split _ t1 t2 => Split _ (tie_tree t1) (tie_tree t2)
    | Fail _ => Fail _
    end.

  Definition list_tree (l : list bool) : tree bool :=
    tie_tree (list_tree_aux l (Leaf _ (inl tt))).

  Definition bernoulli_tree (n d : nat) : tree bool :=
    list_tree (repeat true n ++ repeat false (d-n)).

  Fixpoint interp (c : com) : St -> tree St :=
    match c with
    | CSkip => ret
    | CSeq c1 c2 => tree_kcomp (interp c1) (interp c2)
    | CAssign x e => fun st => match eval e st with
                           | VBool b => ret (upd x (VBool b) st)
                           end
    | CIte e c1 c2 => fun st => match eval e st with
                            | VBool true => interp c1 st
                            | VBool false => interp c2 st
                            end
    | CChoice n d c1 c2 => fun st =>
                          bind (bernoulli_tree n d)
                               (fun b => if b:bool then interp c1 st else interp c2 st)
    | CObserve e => fun st => match eval e st with
                          | VBool true => ret st
                          | VBool false => Fail _
                          end
    end.

  Definition interp_prog (p : prog) : St -> tree bool :=
    fun st => tree_fmap (fun st' => bool_of_val (eval (prog_query p) st'))
                     (interp (prog_com p) st).

End Zar.

CoFixpoint itree_of_tree_aux {A : Type} (root : tree A) (t : tree A) : itree bitE A :=
  match t with
  | Leaf _ x => ret x
  | Split _ t1 t2 =>
    b <- trigger Get ;;
      if b:bool then Tau (itree_of_tree_aux root t1) else Tau (itree_of_tree_aux root t2)
  | Fail _ => Tau (itree_of_tree_aux root root)
  end.

Definition itree_of_tree {A : Type} (t : tree A) : itree bitE A :=
  itree_of_tree_aux t t.

Definition flip (c1 c2 : com) : com := CChoice 1 2 c1 c2.

Definition choice_bool (n d x : nat) : com :=
  CChoice n d (CAssign x (EVal (VBool true))) (CAssign x (EVal (VBool false))).

Definition flip_bool (x : nat) : com :=
  choice_bool 1 2 x.

(* Goldfish-piranha program *)
Definition pir : prog :=
  {| prog_com :=
       CSeq (flip_bool 0)
            (CSeq (CAssign 1 (EVal (VBool true)))
                  (CSeq (flip (CAssign 2 (EVar 0))
                              (CAssign 2 (EVar 1)))
                        (CObserve (EEq (EVar 2) (EVal (VBool true))))))
   ; prog_query := EVar 0 |}.

Definition test1 : prog :=
  {| prog_com :=
       CSeq (flip_bool 0)
            (CSeq (flip_bool 1)
                  (CObserve (EEq (EVar 0) (EVar 1))))
   ; prog_query := EVar 0
  |}.

Definition fair_coin : prog :=
  {| prog_com := CSeq (choice_bool 1 3 0)
                      (CSeq (choice_bool 1 3 1)
                            (CObserve (ENot (EEq (EVar 0) (EVar 1)))))
   ; prog_query := EVar 0
  |}.

(* Definition it : itree bitE bool := itree_of_tree (interp_prog pir empty). *)

Definition t : tree bool := interp_prog fair_coin empty.
Definition it : itree bitE bool := itree_of_tree t.

Extraction "extract/itree" it.