(** Theory of finite probability distributions.  This file defines the {prob T}
type, which represents finite probability distributions with rational
coefficients over the type [T], which is required to be an ordered type.
Distributions form a monad, where the unit [dirac : T -> {prob T}] is the point
mass distribution, whereas the monadic bind [sample : {prob T} -> (T -> {prob
S}) -> {prob S}] samples from its first argument and feeds the sample to its
second argument. *)

Require Import Coq.Strings.String.
Require Import Coq.Unicode.Utf8.

From mathcomp Require Import ssreflect ssrfun ssrbool eqtype ssrnat choice seq
  ssrint rat ssralg ssrnum bigop path.

From extructures Require Import ord fset fmap ffun.
From extructures.instances Require Import instances.

From deriving Require Import deriving.

Set Implicit Arguments.
Unset Strict Implicit.
Unset Printing Implicit Defensive.

Local Open Scope ring_scope.
Local Open Scope fset_scope.

Declare Scope prob_scope.
Local Open Scope prob_scope.

Arguments bigcupP {T I x s P F}.

Section Prob.

Variable T : ordType.
Implicit Types (x : T) (X : {fset T}) (f : T -> rat).

(** We define probability distributions in two stages.  First, we define a
distribution as a map into the positive rationals that is 0 almost everywhere.
Then, we define {prob T} as the type of distributions of mass 1. *)

Record distr := Distr {
  dval :> ffun (fun _ : T => 0 : rat);
  _    :  all (fun x => 0 <= dval x) (supp dval)
}.
Definition distr_of & phant T := distr.
Notation "{ 'distr' T }" := (distr_of (Phant T))
  (at level 0, format "{ 'distr'  T }") : form_scope.
Identity Coercion distr_of_distr : distr_of >-> distr.

Canonical distr_subType  := [subType for dval].
Definition distr_eqMixin := [eqMixin of distr by <:].
Canonical distr_eqType   := EqType distr distr_eqMixin.
Definition distr_choiceMixin := [choiceMixin of distr by <:].
Canonical distr_choiceType := Eval hnf in ChoiceType distr distr_choiceMixin.
Definition distr_ordMixin := [ordMixin of distr by <:].
Canonical distr_ordType := Eval hnf in OrdType distr distr_ordMixin.

Canonical distr_of_newType := [subType of {distr T}].
Canonical distr_of_eqType  := [eqType of {distr T}].
Canonical distr_of_choiceType := [choiceType of {distr T}].
Canonical distr_of_ordType := [ordType of {distr T}].

Implicit Types (d : {distr T}).

Lemma eq_distr d1 d2 : d1 =1 d2 <-> d1 = d2.
Proof.
by split=> [e|-> //]; apply/val_inj/eq_ffun.
Qed.

Lemma distr_ge0 d x : 0 <= d x.
Proof.
case: (boolP (x \in supp d))=> [|/suppPn -> //].
apply/allP: x; exact/(valP d).
Qed.

Lemma mem_suppd d x : (x \in supp d) = (0 < d x).
Proof. by rewrite mem_supp Num.Theory.lt0r distr_ge0 andbT. Qed.

Definition mass d := \sum_(x <- supp d) d x.

Lemma massE d X :
  fsubset (supp d) X -> mass d = \sum_(x <- X) d x.
Proof.
rewrite /mass => sub.
have -> : supp d = fset_filter (mem (supp d)) X.
  apply/eq_fset=> x; rewrite in_fset_filter /=.
  by case: (boolP (x \in supp d))=> // /(fsubsetP sub) ->.
rewrite val_fset_filter big_filter big_mkcond /=.
by apply/eq_big=> // x _; case: ifPn=> // /suppPn ->.
Qed.

Fact mkdistr_key : unit. Proof. exact: tt. Qed.

Lemma mkdistr_subproof (X : {fset T}) (f : T -> rat) :
  (forall x, x \in X -> 0 <= f x) ->
  let ff := @mkffun _ _ (fun _ => 0) f X in
  all (fun x => 0 <= ff x) (supp ff).
Proof.
move=> /= pos; apply/allP=> x _; rewrite mkffunE.
by case: ifP=> // /pos.
Qed.

Definition mkdistr X f pos :=
  locked_with mkdistr_key
    (Distr (@mkdistr_subproof X f pos)).

Lemma mkdistrE X f pos x :
  @mkdistr X f pos x = if (x \in X) then f x else 0%R.
Proof. by rewrite /mkdistr unlock /= mkffunE. Qed.

Lemma supp_mkdistr X f pos : supp (@mkdistr X f pos) = fset_filter (fun x => f x != 0) X.
Proof.
apply/eq_fset=> x; rewrite mem_supp mkdistrE.
by rewrite -mkffunE -mem_supp supp_mkffun in_fset mem_filter.
Qed.

Record prob :=
  Prob {pval : {distr T}; _ : mass pval == 1%R}.
Definition prob_of & phant T := prob.
Notation "{ 'prob' T }" := (prob_of (Phant T))
  (at level 0, format "{ 'prob'  T }") : form_scope.
Identity Coercion prob_of_prob : prob_of >-> prob.
Coercion pval : prob >-> distr_of.

Canonical prob_subType := [subType for pval].
Definition prob_eqMixin := [eqMixin of prob by <:].
Canonical prob_eqType := EqType prob prob_eqMixin.
Definition prob_choiceMixin := [choiceMixin of prob by <:].
Canonical prob_choiceType := Eval hnf in ChoiceType prob prob_choiceMixin.
Definition prob_ordMixin := [ordMixin of prob by <:].
Canonical prob_ordType := Eval hnf in OrdType prob prob_ordMixin.

Canonical prob_of_subType := [subType of {prob T}].
Canonical prob_of_eqType  := [eqType  of {prob T}].
Canonical prob_of_choiceType := [choiceType of {prob T}].
Canonical prob_of_ordType := [ordType of {prob T}].

Implicit Types (p : {prob T}).

Lemma mkprob_subproof X f pos :
  \sum_(x <- X) f x = 1 ->
  mass (@mkdistr X f pos) == 1.
Proof.
move=> <-; apply/eqP.
rewrite [RHS](bigID (fun x => f x == 0)) /=.
have -> : \sum_(x <- X | f x == 0) f x = 0.
  by rewrite big1 // => x /eqP.
rewrite GRing.add0r /mass supp_mkdistr val_fset_filter big_filter.
rewrite big_seq_cond [RHS]big_seq_cond.
by apply: eq_big=> // x /andP [x_X _]; rewrite mkdistrE x_X.
Qed.

Fact mkprob_key : unit. Proof. exact: tt. Qed.

Definition mkprob X f pos e : {prob T} :=
  locked_with mkprob_key (Prob (@mkprob_subproof X f pos e)).

Lemma mkprobE X f pos e x :
  @mkprob X f pos e x = if x \in X then f x else 0.
Proof. by rewrite /mkprob unlock /= mkdistrE. Qed.

Lemma suppPrN0 p : supp p != fset0.
Proof.
by apply/eqP=> e; move: (eqP (valP p)); rewrite /mass e big_nil.
Qed.

Definition pchoose p : T := xchoose (fset0Pn _ (suppPrN0 p)).

Lemma pchoose_supp p : pchoose p \in supp p.
Proof. exact: xchooseP. Qed.

Definition dirac_def x x' : rat := (x == x')%:R.

Lemma dirac_subproof1 x x' : x' \in fset1 x -> 0 <= dirac_def x x'.
Proof. by rewrite /dirac_def; case: eq_op. Qed.

Lemma dirac_subproof2 x : \sum_(x' <- fset1 x) dirac_def x x' = 1.
Proof. by rewrite /= big_seq1 /dirac_def eqxx. Qed.

(** Point mass distribution *)

Definition dirac x :=
  mkprob (@dirac_subproof1 x) (dirac_subproof2 x).

Lemma diracE x x' : dirac x x' = (x' == x)%:R.
Proof.
rewrite /dirac /dirac_def /= mkprobE in_fset1.
by rewrite eq_sym; case: (x == x').
Qed.

Lemma supp_dirac x : supp (dirac x) = fset1 x.
Proof.
apply/eq_fset=> x'.
by rewrite mem_supp in_fset1 diracE; case: (x' =P x).
Qed.

Lemma supp_diracP x x' : reflect (x' = x) (x' \in supp (dirac x)).
Proof. rewrite supp_dirac; exact: fset1P. Qed.

Lemma dirac_inj : injective dirac.
Proof.
by move=> x y e; apply/fset1_inj; rewrite -!supp_dirac e.
Qed.

Lemma eq_prob p1 p2 : p1 =1 p2 <-> p1 = p2.
Proof. by split=> [/eq_distr/val_inj|-> //]. Qed.

Lemma in_eq_probL p1 p2 : {in supp p1, p1 =1 p2} -> p1 = p2.
Proof.
move=> e; apply/eq_prob=> x.
case: (boolP (x \in supp p1))=> xP; first exact: e.
rewrite (suppPn xP).
have: \sum_(y <- supp p1) p2 y = 1.
  rewrite -(eqP (valP p1)) /mass /= big_seq [RHS]big_seq.
  by apply/eq_big=> // ? /e ->.
rewrite -(eqP (valP p2)) /mass /= [RHS](bigID (mem (supp p1))) /=.
rewrite -[in RHS]big_filter -val_fset_filter.
have ->: fset_filter (mem (supp p1)) (supp p2) = supp p1.
  apply/eq_fset=> y; rewrite in_fset_filter /=.
  case: (boolP (y \in supp p1))=> //= y_p1.
  by rewrite mem_supp -(e _ y_p1) -mem_supp.
rewrite -[LHS]GRing.addr0 => /GRing.addrI/esym/eqP.
rewrite Num.Theory.psumr_eq0; last by move=> ? _; rewrite distr_ge0.
case: (boolP (x \in supp p2)) => x_p2; last by rewrite (suppPn x_p2).
by move=> /allP/(_ _ x_p2); rewrite xP => /eqP ->.
Qed.

Lemma in_eq_probR p1 p2 : {in supp p2, p1 =1 p2} -> p1 = p2.
Proof.
by move=> e; apply/esym/in_eq_probL=> x x_p2; rewrite e.
Qed.

Lemma diracK : cancel dirac pchoose.
Proof. by move=> x; have /supp_diracP := pchoose_supp (dirac x). Qed.

Lemma eq_supp_dirac p x : (supp p == fset1 x) = (p == dirac x).
Proof.
apply/(sameP eqP)/(iffP eqP)=> [->|e]; first exact: supp_dirac.
apply/in_eq_probR => _ /supp_diracP ->.
have:= massE (fsubsetxx (supp p)); rewrite e big_seq1 (eqP (valP p)) => <-.
by rewrite diracE eqxx.
Qed.

End Prob.

Notation "{ 'distr' T }" := (distr_of (Phant T))
  (at level 0, format "{ 'distr'  T }") : form_scope.
Notation "{ 'prob' T }" := (prob_of (Phant T))
  (at level 0, format "{ 'prob'  T }") : form_scope.

Arguments dirac {_} x.
Arguments pchoose {_} p.
Arguments dirac_inj {_}.

Section Expectation.

Variable T : ordType.

Implicit Types (p : {prob T}) (x y z : T) (F : T -> rat).
Implicit Types (X Y : {fset T}).

Import GRing.

Definition expect p F : rat :=
  \sum_(x <- supp p) F x * p x.

Notation "''E_' ( x <- p ) F" := (expect p (fun x => F))
  (at level 41, F at level 41, x, p at level 50,
    format "'[' ''E_' ( x  <-  p )  '/  ' F ']'").

Lemma expectE X p F :
  fsubset (supp p) X ->
  'E_(x <- p) F x = \sum_(x <- X) F x * p x.
Proof.
rewrite /expect => p_in_X.
rewrite -(fsetID X (supp p)) (fsetIidPr p_in_X) big_fsetU /=; last first.
  by rewrite fdisjointC; apply/fsetDidPl; rewrite fsetDDl fsetUid.
rewrite [X in _ + X]big_seq /= [X in _ + X]big1 ?addr0 //.
by move=> ? /fsetDP => - [_ /suppPn ->]; rewrite mulr0.
Qed.

Lemma expect_dirac x F : 'E_(y <- dirac x) F y = F x.
Proof.
by rewrite /expect supp_dirac /= big_seq1 diracE eqxx mulr1.
Qed.

Lemma mulr_expectr r p F : r * ('E_(x <- p) F x) = 'E_(x <- p) r * F x.
Proof.
rewrite /expect mulr_sumr; apply/eq_big => [//|x _].
by rewrite mulrA.
Qed.

Lemma mulr_expectl r p F : ('E_(x <- p) F x) * r = 'E_(x <- p) F x * r.
Proof.
rewrite /expect mulr_suml; apply/eq_big => [//|x _].
by rewrite mulrAC.
Qed.

Lemma addr_expect p F1 F2 :
  'E_(x <- p) F1 x + 'E_(x <- p) F2 x =
  'E_(x <- p) (F1 x + F2 x).
Proof.
rewrite /expect -big_split; apply/eq_big => // x _.
by rewrite mulrDl.
Qed.

Lemma eq_expect_gen p1 p2 F1 F2 :
  (forall x, F1 x * p1 x = F2 x * p2 x) ->
  expect p1 F1 = expect p2 F2.
Proof.
move=> e12.
rewrite (expectE _ (fsubsetUl (supp p1) (supp p2))).
rewrite (expectE _ (fsubsetUr (supp p1) (supp p2))).
exact/eq_big.
Qed.

Lemma eq_expect_in p F1 F2 :
  {in supp p, F1 =1 F2} ->
  expect p F1 = expect p F2.
Proof.
move=> e12; apply/eq_expect_gen => x.
case: (boolP (x \in supp p)) => [x_in|x_nin]; rewrite ?(e12 _ x_in) //.
by rewrite (suppPn x_nin) !mulr0.
Qed.

Lemma eq_expect p F1 F2 : F1 =1 F2 -> expect p F1 = expect p F2.
Proof.
move=> e12; apply: eq_expect_in => ??; exact: e12.
Qed.

Lemma expect_const c p F :
  (forall x, x \in supp p -> F x = c) ->
  'E_(x <- p) F x = c.
Proof.
move=> Fc.
rewrite -[c]mul1r -(eqP (valP p)) /mass /= mulr_suml.
rewrite [LHS]big_seq [RHS]big_seq; apply/eq_big => // x /Fc ->.
by rewrite mulrC.
Qed.

Lemma sumr_expect I s (P : I -> bool) p (F : I -> T -> rat) :
  \sum_(i <- s | P i) 'E_(x <- p) F i x =
  'E_(x <- p) \sum_(i <- s | P i) F i x.
Proof.
elim: s => [|i s IH].
  by rewrite big_nil (@expect_const 0) // => ? _; rewrite big_nil.
under [in RHS]eq_expect => x do rewrite big_cons.
rewrite big_cons IH; case: (P i) => //.
by rewrite addr_expect.
Qed.

Lemma expect_ge0 p F :
  (forall x, x \in supp p -> 0 <= F x) ->
  0 <= expect p F.
Proof.
move=> F_ge0; apply: Num.Theory.sumr_ge0 => x _.
case: (boolP (x \in supp p)) => [x_in|/suppPn ->]; rewrite ?mulr0 //.
by rewrite Num.Theory.mulr_ge0 // ?F_ge0 // distr_ge0.
Qed.

End Expectation.

Notation "''E_' ( x <- p ) F" := (expect p (fun x => F))
  (at level 41, F at level 41, x, p at level 50,
    format "'[' ''E_' ( x  <-  p )  '/  ' F ']'")
  : prob_scope.

Arguments expect_const {T} c.

Lemma expectC (T S : ordType) p q (F : T -> S -> rat) :
  'E_(x <- p) 'E_(y <- q) F x y =
  'E_(y <- q) 'E_(x <- p) F x y.
Proof.
rewrite /expect.
under [LHS]eq_big => [//|x _] do rewrite GRing.mulr_suml.
under [RHS]eq_big => [//|x _] do rewrite GRing.mulr_suml.
rewrite exchange_big /=; apply/eq_big => // x _.
by under eq_big => [//|y _] do rewrite GRing.mulrAC.
Qed.

Section Sample.

Variables T S : ordType.
Variable (p : {prob T}) (f : T -> {prob S}).
Implicit Types (x : T) (y : S).

Let Y   : {fset S} := \bigcup_(x <- supp p) supp (f x).
Let P y : rat      := 'E_(x <- p) f x y.

Lemma sample_subproof1 y : y \in Y -> 0 <= P y.
Proof.
move=> _; apply: Num.Theory.sumr_ge0 => x _.
apply: Num.Theory.mulr_ge0; exact: distr_ge0.
Qed.

Lemma sample_subproof2 : \sum_(y <- Y) P y = 1.
Proof.
rewrite /P sumr_expect (expect_const 1) // => x x_in.
rewrite -(eqP (valP (f x))); apply/esym/massE.
apply/fsubsetP=> y; rewrite mem_supp => yP.
apply/bigcupP; exists x=> //.
by rewrite mem_supp.
Qed.

Definition sample :=
  mkprob sample_subproof1 sample_subproof2.

Lemma sample_defE0 y : (y \in Y) = (P y != 0).
Proof.
rewrite /P Num.Theory.psumr_eq0 -?has_predC /=; last first.
  move=> x _; apply: Num.Theory.mulr_ge0; exact: distr_ge0.
apply/(sameP bigcupP)/(iffP hasP).
- case=> /= x x_p n0; exists x=> //.
  by move: n0; rewrite GRing.mulf_eq0 negb_or mem_supp; case/andP.
- case=> /= x; rewrite !mem_supp => x_p _ y_f.
  by exists x; rewrite 1?mem_supp // GRing.mulf_neq0.
Qed.

Lemma sampleE y : sample y = P y.
Proof.
rewrite /sample mkprobE sample_defE0.
by case: (P y =P 0)=> [->|].
Qed.

Lemma supp_sample : supp sample = Y.
Proof.
apply/eq_fset=> x.
by rewrite mem_supp sample_defE0 sampleE.
Qed.

Lemma supp_sampleP y :
  reflect (exists2 x, x \in supp p & y \in supp (f x)) (y \in supp sample).
Proof.
rewrite supp_sample; apply/(iffP bigcupP).
- by case=> ????; eauto.
- by case=> ???; eexists; eauto.
Qed.

End Sample.

Arguments supp_sampleP {_ _ _ _ _}.

(** A more convenient notation for sampling. *)

Notation "'sample:' x '<-' t1 ';' t2" :=
  (sample t1 (fun x => t2))
  (at level 20, t1 at level 100, t2 at level 200,
   right associativity, format "'[' 'sample:'  x  '<-'  '[' t1 ;  ']' ']' '/' t2")
  : prob_scope.

Section SampleProperties.

Variables T S : ordType.

Import GRing.

Lemma sample_diracL (x : T) (f : T -> {prob S}) : sample (dirac x) f = f x.
Proof.
by apply/eq_prob=> y; rewrite sampleE expect_dirac.
Qed.

Lemma sample_diracR (p : {prob T}) : sample p dirac = p.
Proof.
apply/eq_prob=> x; rewrite sampleE.
under eq_expect do rewrite diracE.
case: (boolP (x \in supp p)) => [x_in|x_nin]; last first.
  rewrite (suppPn x_nin); apply/(expect_const 0) => y y_in.
  by case: eqP y_in x_nin => [->|//] ->.
rewrite /expect (big_rem _ x_in) /= eqxx mul1r big_seq big1 ?addr0 // => y.
rewrite mem_rem_uniq ?uniq_fset //= eq_sym.
by case/andP => /negbTE ->; rewrite mul0r.
Qed.

Lemma eq_sample (p : {prob T}) (f g : T -> {prob S}) :
  f =1 g -> sample p f = sample p g.
Proof.
move=> efg; apply/eq_prob=> y.
by rewrite !sampleE; apply/eq_big=> // x _; rewrite efg.
Qed.

Lemma eq_in_sample (p : {prob T}) (f g : T -> {prob S}) :
  {in supp p, f =1 g} -> sample p f = sample p g.
Proof.
move=> efg; apply/eq_prob=> y.
by rewrite !sampleE; apply/eq_expect_in => // x /efg ->.
Qed.

Lemma sample_const (px : {prob T}) (py : {prob S}) :
  (sample: _ <- px; py) = py.
Proof.
by apply/eq_prob=> y; rewrite sampleE (expect_const (py y)).
Qed.

Lemma eq_sample_dirac (p : {prob T}) (f : T -> {prob S}) y :
  sample p f = dirac y ->
  forall x, x \in supp p -> f x = dirac y.
Proof.
move=> e x x_p.
have {}e: supp (sample p f) = supp (dirac y) by rewrite e.
rewrite supp_sample supp_dirac in e.
apply/eqP; rewrite -eq_supp_dirac eqEfsubset; apply/andP; split.
  rewrite -e; exact/bigcup_sup.
rewrite fsub1set; have /fset0Pn [z zP] := suppPrN0 (f x).
suff: z \in fset1 y by move=> /fset1P => <-.
by rewrite -e; apply/bigcupP; exists x.
Qed.

End SampleProperties.

Lemma expect_sample (T S : ordType) p (f : T -> {prob S}) (G : S -> rat) :
  'E_(y <- sample p f) G y =
  'E_(x <- p) 'E_(y <- f x) G y.
Proof.
rewrite /expect.
under [LHS]eq_big => [//|x _] do rewrite sampleE mulr_expectr.
rewrite exchange_big /= [LHS]big_seq [RHS]big_seq.
apply/eq_big => // x x_in; rewrite GRing.mulr_suml.
have /fsetIidPl <-: fsubset (supp (f x)) (supp (sample: x <- p; f x)).
  apply/fsubsetP=> y fxy; rewrite supp_sample.
  by apply/bigcupP; exists x.
rewrite /fsetI val_fset_filter big_filter [RHS]big_mkcond /=.
apply/eq_big=> // y _; rewrite mem_supp.
by case: eqP=> //= ->; rewrite GRing.mulr0 GRing.mul0r.
Qed.

Lemma sampleA (T S R : ordType) p (f : T -> {prob S}) (g : S -> {prob R}) :
  (sample: y <- (sample: x <- p; f x); g y) =
  (sample: x <- p; sample: y <- f x; g y).
Proof.
apply/eq_prob=> z; rewrite !sampleE expect_sample.
by under [RHS]eq_expect => ? do rewrite sampleE.
Qed.

Lemma sampleC (T S R : ordType) (p1 : {prob T}) (p2 : {prob S}) (f : T -> S -> {prob R}) :
  (sample: x <- p1; sample: y <- p2; f x y) =
  (sample: y <- p2; sample: x <- p1; f x y).
Proof.
apply/eq_prob=> z; rewrite !sampleE.
under [LHS]eq_expect do rewrite sampleE.
under [RHS]eq_expect do rewrite sampleE.
by rewrite expectC.
Qed.

Section Uniform.

Variable T : ordType.
Variable X : {fset T}.
Hypothesis Xn0 : X != fset0.

Definition unif_def (x : T) : rat := (size X)%:R^-1.

Lemma unif_subproof1 x : x \in X -> 0 <= unif_def x.
Proof.
move=> _; by rewrite Num.Theory.invr_ge0 Num.Theory.ler0n.
Qed.

Lemma unif_subproof2 : \sum_(x <- X) unif_def x = 1.
Proof.
rewrite big_const_seq count_predT -Monoid.Theory.iteropE.
rewrite -[RHS](@GRing.mulfV _ (size X)%:R) ?GRing.mulr_natl //.
by rewrite Num.Theory.pnatr_eq0 sizes_eq0.
Qed.

Definition unif := mkprob unif_subproof1 unif_subproof2.

Lemma unifE x : unif x = if x \in X then (size X)%:R^-1 else 0.
Proof. by rewrite mkprobE. Qed.

Lemma supp_unif : supp unif = X.
Proof.
apply/eq_fset => x; rewrite mem_supp unifE.
by case: ifP => // _; rewrite GRing.invr_eq0 Num.Theory.pnatr_eq0 sizes_eq0.
Qed.

End Uniform.

(** To simplify the reasoning about our optimizations, we use probabilistic
couplings.  A coupling is a way of lifting a relation between two sets to a
relation over distributions over these sets.  (NB: We define this relation in
[Type] to avoid issues with the axiom of choice when proving [coupling_sample]
below.)  The definition is useful because (1) it has good composition
properties, as seen by [coupling_dirac] and [coupling_sample], and because (2)
it is strong enough to establish the equality of two distributions, as seen in
[coupling_same]. *)

Variant coupling (T S : ordType) (R : T -> S -> Prop) pT pS : Type :=
| Coupling p of
  pT = sample p (dirac \o fst) &
  pS = sample p (dirac \o snd) &
  (forall xy, xy \in supp p -> R xy.1 xy.2).

Lemma coupling_dirac (T S : ordType) (R : T -> S -> Prop) x y :
  R x y -> coupling R (dirac x) (dirac y).
Proof.
move=> xy; exists (dirac (x, y)); rewrite ?sample_diracL //.
by move=> [??] /supp_diracP [-> ->].
Qed.

Lemma coupling_sample (T1 S1 T2 S2 : ordType) (R1 : T1 -> S1 -> Prop) (R2 : T2 -> S2 -> Prop) pT pS f g :
  coupling R1 pT pS ->
  (forall x y, x \in supp pT -> y \in supp pS -> R1 x y ->
    coupling R2 (f x) (g y)) ->
  coupling R2 (sample pT f) (sample pS g).
Proof.
case=> /= p eT eS R1P R12.
pose def xy := sample: x' <- f xy.1; sample: y' <- g xy.2; dirac (x', y').
have WT xy : xy \in supp p -> xy.1 \in supp pT.
  move=> xyp; rewrite eT; apply/supp_sampleP.
  exists xy=> //=; exact/supp_diracP.
have WS xy : xy \in supp p -> xy.2 \in supp pS.
  move=> xyp; rewrite eS; apply/supp_sampleP.
  exists xy=> //=; exact/supp_diracP.
pose draw xy := if insub xy is Some xy then
                  let xyP := svalP xy in
                  let xP := WT _ xyP in
                  let yP := WS _ xyP in
                  let: Coupling p _ _ _ := R12 _ _ xP yP (R1P _ xyP) in p
                else def xy.
exists (sample p draw).
- rewrite eT !sampleA; apply/eq_in_sample; case=> [x y] /= xy_supp.
  by rewrite sample_diracL insubT /=; case: (R12 _ _ _ _ _).
- rewrite eS !sampleA; apply/eq_in_sample; case=> [x y] /= xy_supp.
  by rewrite sample_diracL insubT /=; by case: (R12 _ _ _ _ _).
case=> x' y' /supp_sampleP [] [x y] xy_supp.
rewrite /draw insubT /=.
case: (R12 _ _ _ _ _)=> /= pxy eT' eS' R2P; exact: R2P.
Qed.

Lemma coupling_same (T : ordType) (p : {prob T}) : coupling eq p p.
Proof.
pose pp := sample: x <- p; dirac (x, x); exists pp; rewrite ?sampleA.
- under eq_sample do rewrite sample_diracL /=.
  by rewrite sample_diracR.
- under eq_sample do rewrite sample_diracL /=.
  by rewrite sample_diracR.
by move=> xx /supp_sampleP [] x x_p /supp_diracP ->.
Qed.

Lemma coupling_eq (T : ordType) (p q : {prob T}) : coupling eq p q -> p = q.
Proof.
by case=> pq -> -> {p q} e; apply/eq_in_sample; case=> x y /e /= ->.
Qed.

Lemma coupling_sample_same (T S1 S2 : ordType) (R : S1 -> S2 -> Prop) (p : {prob T}) f g :
  (forall x, x \in supp p -> coupling R (f x) (g x)) ->
  coupling R (sample p f) (sample p g).
Proof.
move=> e; apply: coupling_sample; first exact: coupling_same.
move=> x _ x_p _ <-; exact: e.
Qed.

Lemma coupling_trivial (T S : ordType) (R : T -> S -> Prop) (pT : {prob T}) (pS : {prob S}) :
  (forall x y, x \in supp pT -> y \in supp pS -> R x y) ->
  @coupling T S R pT pS.
Proof.
move=> RP.
exists (sample: x <- pT; sample: y <- pS; dirac (x, y)).
- rewrite sampleA; under eq_sample => x.
    rewrite sampleA; under eq_sample do rewrite sample_diracL.
    rewrite /= sample_const; over.
  by rewrite sample_diracR.
- rewrite sampleA; under eq_sample=> x.
    rewrite sampleA; under eq_sample do rewrite sample_diracL.
    rewrite /= sample_diracR; over.
  by rewrite sample_const.
case=> x y /supp_sampleP [] x' xP /supp_sampleP [] y' yP /=.
move=> /supp_diracP [-> ->]; exact: RP.
Qed.

Lemma couplingT T S pT pS : @coupling T S (fun _ _ => True) pT pS.
Proof. by apply: coupling_trivial. Qed.

Lemma coupling_sampleL (T S1 S2 : ordType) (R : S1 -> S2 -> Prop) (pT : {prob T}) f pS :
  (forall x, x \in supp pT -> coupling R (f x) pS) ->
  coupling R (sample: x <- pT; f x) pS.
Proof.
move=> RP; rewrite -[pS](sample_const pS).
apply: coupling_sample (couplingT _ _) _ => x _ xP _ _.
exact: RP.
Qed.

Lemma coupling_sampleR (T S1 S2 : ordType) (R : S1 -> S2 -> Prop) (pT : {prob T}) f pS :
  (forall x, x \in supp pT -> coupling R pS (f x)) ->
  coupling R pS (sample: x <- pT; f x).
Proof.
move=> RP; rewrite -[pS](sample_const pS).
apply: coupling_sample (couplingT _ _) _ => _ x _ xP _.
exact: RP.
Qed.

Lemma couplingW (T S : ordType) (R1 R2 : T -> S -> Prop) pT pS :
  (forall x y, R1 x y -> R2 x y) ->
  @coupling T S R1 pT pS ->
  coupling R2 pT pS.
Proof. by move=> R12 [p eT eS R1P]; exists p; eauto. Qed.

Lemma in_couplingW (T S : ordType) (R1 R2 : T -> S -> Prop) (pT : {prob T}) (pS : {prob S}) :
  (forall x y, x \in supp pT -> y \in supp pS -> R1 x y -> R2 x y) ->
  coupling R1 pT pS ->
  coupling R2 pT pS.
Proof.
move=> R12 [p eT eS R1P]; exists p; eauto.
case=> x y xyP; apply: R12; last exact: R1P.
- rewrite eT; apply/supp_sampleP.
  by exists (x, y)=> //; rewrite // supp_dirac in_fset1.
- rewrite eS; apply/supp_sampleP.
  by exists (x, y)=> //; rewrite // supp_dirac in_fset1.
Qed.

(** The following definitions allow us to apply probabilistic computations to
the elements of sequences, maps and other container types.  They will be used to
define the semantics of Imp programs. *)

Definition foldrM T (S : ordType) (f : T -> S -> {prob S}) (y : S) (xs : seq T) : {prob S} :=
  foldr (fun x p => sample p (f x)) (dirac y) xs.

Fixpoint map_p T (S : ordType) (f : T -> {prob S}) (xs : seq T) : {prob seq S} :=
  match xs with
  | [::] => dirac [::]
  | x :: xs =>
    sample: y  <- f x;
    sample: ys <- map_p f xs;
    dirac (y :: ys)
  end.

Lemma eq_map_p T (S : ordType) (f g : T -> {prob S}) :
  f =1 g -> map_p f =1 map_p g.
Proof. by move=> fg; elim=> //= x xs IH; rewrite fg IH. Qed.

Lemma map_p_dirac (T : ordType) (xs : seq T) : map_p dirac xs = dirac xs.
Proof.
elim: xs=> //= x xs IH.
by rewrite sample_diracL IH sample_diracL.
Qed.

Lemma map_p_comp T S (R : ordType) (f : T -> S) (g : S -> {prob R}) xs :
  map_p g [seq f x | x <- xs] = map_p (g \o f) xs.
Proof. by elim: xs=> //= x xs ->. Qed.

Lemma map_p_sample (T S R : ordType) (g : S -> {prob R}) (f : T -> {prob S}) (xs : seq T) :
  map_p (fun x => sample: y <- f x; g y) xs =
  sample: ys <- map_p f xs; map_p g ys.
Proof.
elim: xs=> [|x xs IH] /=; first by rewrite sample_diracL.
rewrite !sampleA; apply/eq_sample=> y.
rewrite sampleA {}IH.
under eq_sample=> z do rewrite sampleA.
under [in RHS]eq_sample=> zs do rewrite sample_diracL /=.
by rewrite sampleC.
Qed.

Lemma supp_map_p T (S : ordType) (f : T -> {prob S}) xs ys :
  ys \in supp (map_p f xs) =
  all2 (fun x y => y \in supp (f x)) xs ys.
Proof.
elim: xs ys=> [|x xs IH] [|y ys] /=.
- by rewrite supp_dirac.
- by rewrite supp_dirac.
- case: supp_sampleP=> //=.
  by case=> y' y'P /supp_sampleP [ys' _ /supp_diracP].
- rewrite -IH; apply/(sameP supp_sampleP)/(iffP andP).
  + case=> [yP ysP]; exists y=> //.
    apply/supp_sampleP; exists ys=> //.
    by apply/supp_diracP.
  + by case=> [y' y'P /supp_sampleP [ys' ys'P /supp_diracP [-> ->]]].
Qed.

Section MapMapProb.

Variable T : ordType.

Definition mapim_p (S : Type) (R : ordType)
  (f : T -> S -> {prob R}) (m : {fmap T -> S}) : {prob {fmap T -> R}} :=
  let do_pair p := sample: y <- f p.1 p.2; dirac (p.1, y) in
  sample: pairs <- map_p do_pair (val m);
  dirac (mkfmap pairs).

Lemma mapim_pE S R f m :
  @mapim_p S R f m =
  foldrM (fun p m => sample: z <- f p.1 p.2; dirac (setm m p.1 z)) emptym m.
Proof.
rewrite /mapim_p /=.
elim: (val m)=> {m} [|[x y] m /= IH] //=.
- by rewrite sample_diracL //.
- rewrite !sampleA [in RHS]sampleC; apply/eq_sample=> z.
  rewrite sample_diracL sampleA -IH sampleA.
  by apply/eq_sample=> pairs; rewrite !sample_diracL.
Qed.

Lemma eq_mapim_p (S : Type) (R : ordType)
  (f g : T -> S -> {prob R}) :
  f =2 g -> mapim_p f =1 mapim_p g.
Proof.
move=> fg m; rewrite /mapim_p.
by under eq_map_p => p do rewrite fg.
Qed.

Lemma mapim_p_dirac (S : ordType) (m : {fmap T -> S}) :
  mapim_p (fun _ => dirac) m = dirac m.
Proof.
rewrite /mapim_p.
under eq_map_p => p do rewrite sample_diracL -surjective_pairing.
by rewrite map_p_dirac sample_diracL fmvalK.
Qed.

Lemma mapim_p_comp (S R U : ordType)
  (g : T -> R -> {prob U}) (f : T -> S -> R) m :
  mapim_p g (mapim f m) =
  mapim_p (fun x y => g x (f x y)) m.
Proof. by rewrite /mapim_p /= map_p_comp. Qed.

Fact mapm_p_key : unit. Proof. exact: tt. Qed.

Definition mapm_p (S : Type) (R : ordType) (f : S -> {prob R}) :=
  locked_with mapm_p_key (mapim_p (fun (x : T) => f)).

Lemma mapm_pE S R f m :
  @mapm_p S R f m =
  if splitm m is Some (x, y, m) then
    sample: z <- f y;
    sample: m <- mapm_p f m;
    dirac (setm m x z)
  else dirac emptym.
Proof.
rewrite /splitm /= /mapm_p unlock /mapim_p /=.
move: (valP m) => /=; case: (val m)=> [|[x y] ps] //=.
  by rewrite sample_diracL.
move=> /path_sorted psP.
rewrite !sampleA; apply/eq_sample=> z.
rewrite sample_diracL !sampleA mkfmapK //.
by apply/eq_sample=> /= ps'; rewrite !sample_diracL.
Qed.

Lemma eq_mapm_p (S : Type) (R : ordType) (f g : S -> {prob R}) :
  f =1 g -> mapm_p f =1 mapm_p g.
Proof.
rewrite /mapm_p !unlock.
by move=> e; apply/eq_mapim_p=> ??; eauto.
Qed.

Lemma mapm_p_dirac (S : ordType) (m : {fmap T -> S}) :
  mapm_p dirac m = dirac m.
Proof. rewrite /mapm_p !unlock; exact/mapim_p_dirac. Qed.

Lemma mapm_p_comp (S R U : ordType) (g : R -> {prob U}) (f : S -> R) m :
  mapm_p g (mapm f m) = mapm_p (g \o f) m.
Proof. rewrite /mapm_p !unlock; exact/mapim_p_comp. Qed.

Lemma supp_mapm_p (S R : ordType) (f : S -> {prob R}) m1 m2 :
  m2 \in supp (mapm_p f m1) =
  (domm m1 == domm m2) &&
  all (fun x => match m1 x, m2 x with
                | Some y, Some z => z \in supp (f y)
                | _, _ => true
                end) (domm m1).
Proof.
rewrite /mapm_p unlock /mapim_p.
apply/(sameP supp_sampleP)/(iffP andP).
- case=> /eqP edomm ecodomm; exists (val m2); last first.
    by apply/supp_diracP; rewrite fmvalK.
  have /= esize: size (val m1) = size (val m2).
    move/(congr1 (size \o val)): edomm.
    by rewrite /= !val_domm !size_map => ->.
  rewrite supp_map_p all2E esize eqxx /=.
  apply/allP=> /= - xyz /(nthP xyz) [i].
  rewrite size_zip -esize minnn=> isize.
  case: xyz=> [[x1 y] [x2 z]]; rewrite nth_zip //=; case=> e1 e2.
  have e1' : nth x1 (domm m1) i = x1.
    by rewrite val_domm (nth_map (x1, y)) // e1.
  have e2' : nth x2 (domm m2) i = x2.
    by rewrite val_domm (nth_map (x2, z)) -?esize // e2.
  have ex: x1 = x2.
    rewrite -e1' -e2' -edomm; apply/set_nth_default.
    by rewrite val_domm size_map.
  move: x1 ex e1 e2 {e1' e2'}=> {x2} x <- e1 e2.
  have {}e1: m1 x = Some y.
    by apply/getmP; rewrite -e1; apply/mem_nth.
  have {}e2: m2 x = Some z.
    by apply/getmP; rewrite -e2; apply/mem_nth; rewrite -esize.
  have xP: x \in domm m1 by rewrite mem_domm e1.
  move/allP/(_ _ xP): ecodomm; rewrite e1 e2=> yz.
  by apply/supp_sampleP; exists z; rewrite // supp_dirac in_fset1 eqxx.
- case=> {}m2 m2P /supp_diracP ->.
  move: m2P; rewrite supp_map_p all2E; case/andP=> /eqP esize ecodomm.
  have edomm : unzip1 m1 = unzip1 m2.
    move: m2 esize ecodomm; rewrite /=.
    elim: {m1} (val m1)=> [|[x1 y] m1 IH] [|[x2 z] m2] //= [esize].
    case/andP=> /supp_sampleP [{}z _ /supp_diracP [-> _]] ecodomm.
    congr cons; exact: IH.
  have m2_sorted: mkfmap m2 = m2 :> seq _.
    apply/mkfmapK; rewrite -edomm; exact: (valP m1).
  have {}edomm : domm m1 = domm (mkfmap m2).
    by apply/val_inj; rewrite /= !val_domm m2_sorted.
  split; first by rewrite edomm.
  move: esize ecodomm; rewrite -m2_sorted.
  move: (mkfmap m2) edomm=> {m2_sorted} {}m2 /= edomm esize ecodomm.
  apply/allP=> x /dommP [y yP]; rewrite yP fmvalK.
  have /dommP [z zP]: x \in domm m2 by rewrite -edomm mem_domm yP.
  rewrite zP.
  case/getmP/(nthP (x, y)): (yP)=> /= i isize ei1.
  have ei1domm : nth x (domm m1) i = x.
    by rewrite val_domm (nth_map (x, y)) // ei1.
  have ei2: nth (x, z) m2 i = (x, z).
    move: zP; rewrite -{1} ei1domm edomm.
    rewrite (getm_nth (x, z)) -?esize //; case=> {2}<-.
    move: ei1domm; rewrite edomm val_domm (nth_map (x, z)) -?esize //.
    by move=> e; rewrite -[X in _ = (X, _)]e [LHS]surjective_pairing.
  have inzip : nth ((x, y), (x, z)) (zip m1 m2) i = ((x, y), (x, z)).
    by rewrite nth_zip // ei1 ei2.
  have {}inzip: ((x, y), (x, z)) \in zip m1 m2.
    by rewrite -inzip; apply/mem_nth; rewrite size_zip -esize minnn.
  move/allP/(_ _ inzip): ecodomm; rewrite inE /=.
  by case/supp_sampleP=> ?? /supp_diracP [->].
Qed.

Lemma supp_mapm_pP (S R : ordType) (f : S -> {prob R}) m1 m2 :
  reflect (domm m1 = domm m2 /\
           forall x y z, m1 x = Some y -> m2 x = Some z -> z \in supp (f y))
          (m2 \in supp (mapm_p f m1)).
Proof.
rewrite supp_mapm_p; apply/(iffP andP).
- case=> /eqP edomm ecodomm; split=> // x y z yP zP.
  have xP: x \in domm m1 by rewrite mem_domm yP.
  by move/allP/(_ _ xP): ecodomm; rewrite yP zP.
- case=> edomm ecodomm; split; first by rewrite edomm.
  apply/allP=> x xP; case/dommP: (xP)=> y yP.
  move: (xP); rewrite edomm; case/dommP=> z zP.
  rewrite yP zP; apply: ecodomm yP zP.
Qed.

Lemma mapm_p0 (S R : ordType) f : @mapm_p S R f emptym = dirac emptym.
Proof. by rewrite mapm_pE /=. Qed.

Lemma mapm_p_setm (S R : ordType) f m x y :
  @mapm_p S R f (setm m x y) =
  sample: m <- mapm_p f m;
  sample: z  <- f y;
  dirac (setm m x z).
Proof.
rewrite /mapm_p unlock !mapim_pE [setm m x y]/setm /=.
case: m=> /=; elim=> [|[x' y'] m IH] //=.
case: Ord.ltgtP=> [//|xx'|{x'} <-] //= mP.
  rewrite IH ?(path_sorted mP) //.
  rewrite !sampleA; apply/eq_sample=> /= m'.
  rewrite sampleC sampleA; apply/eq_sample=> /= z.
  rewrite sampleA sample_diracL; apply/eq_sample=> /= z'.
  rewrite sample_diracL setmC // eq_sym.
  by case: Ord.ltgtP xx'.
rewrite sampleA; apply/eq_sample=> m'.
rewrite sampleC; apply/eq_sample=> z.
rewrite sampleA.
under eq_sample=> ? do rewrite sample_diracL setmxx.
by rewrite sample_const.
Qed.

(*
Lemma mapm_pE' (S R : ordType) f my mz :
  @mapm_p S R f my mz =
  if domm my == domm mz then
    \prod_(x <- domm my)
      (match my x, mz x with
       | Some y, Some z => f y z
       | _, _ => 0
       end)
  else 0.
Proof.
case: eqP=>[e|ne]; last first.
  apply/suppPn/supp_mapm_pP; case=> ??; congruence.
move ez: (domm mz) e=> X ey.
elim/fset_ind: X my mz ey ez=> [|x X x_X IH] my mz ey ez.
  move: ey ez=> /eqP/emptymP -> /eqP/emptymP ->.
  by rewrite domm0 big_nil mapm_pE /= diracE.
rewrite ey big_fsetU1 //=.
have /dommP [y yP]: x \in domm my by rewrite ey in_fsetU1 eqxx.
have /dommP [z zP]: x \in domm mz by rewrite ez in_fsetU1 eqxx.
set my' := remm my x; set mz' := remm mz x.
have e: (x |: X) :\ x = X.
  apply/eq_fset=> x'; rewrite in_fsetD1 in_fsetU1.
  by case: eqP=> [{x'}->|] //=; rewrite (negbTE x_X).
have eyD: domm my' = X by rewrite domm_rem ey.
have ezD: domm mz' = X by rewrite domm_rem ez.
have {}ey: my = setm my' x y.
  apply/eq_fmap=> x'; rewrite setmE remmE.
  by case: eqP yP=> [->|].
have {}ez: mz = setm mz' x z.
  apply/eq_fmap=> x'; rewrite setmE remmE.
  by case: eqP zP=> [->|].
rewrite yP zP ey ez mapm_p_setm sampleE big_seq.
under eq_big=> [m|m mP]; first over.
  rewrite sampleE.
  under eq_big=> [z'|z' z'P]; first over.
    rewrite diracE eq_setm remmI ?ezD // remmI; last first.
      by case/supp_mapm_pP: mP=> <- _; rewrite eyD.
    over.
  rewrite /=.
rewrite (IH _ _ ey ez).

have e: perm_eq (x |: X) (x :: X).
  apply: uniq_perm.
  - exact: uniq_fset.
  - by rewrite /= x_X uniq_fset.
  by move=> x'; rewrite inE in_fsetU1.
rewrite ey (perm_big _ e) /= big_cons.

move e: (size my)=> n; elim: n my mz e=> [|n IH] /= my mz.
  rewrite sizeES mapm_pE [domm my]dommES.
  case e: (splitm my)=> [[[x y] my']|] //= _.
  by rewrite diracE -eq_domm0 eq_sym big_nil.
rewrite sizeES mapm_pE [domm my]dommES.
case ey: (splitm my)=> [[[x y] my']|] //= [e].
rewrite sampleE [domm mz]dommES.
case ez: (splitm mz)=> [[[x' z] mz']|]; last first.
  rewrite eq_domm0.

rewrite mapm_pE /= -[domm _ == domm _]val_eqE /= val_domm.
(* FIXME: why doesn't under work here? *)
have e: forall x, true -> match my x, mz x with
                  | Some y, Some z => f y z
                  | _ , _ => 0
                  end =
                  match getm_def (val my) x, mz x with
                  | Some y, Some z => f y z
                  | _, _ => 0
                  end.
  by move=> x; rewrite -mkfmapE fmvalK.
rewrite (eq_big xpredT _ _ e) //= {e}.
elim: (val my) mz=> /= {my} [|[x y] my IH] /= mz.
  by rewrite big_nil diracE -eq_domm0 -val_eqE /= eq_sym.
rewrite sampleE.
*)

Lemma mapm_p_sample (S R U : ordType) (g : R -> {prob U}) (f : S -> {prob R}) (m : {fmap T -> S}) :
  mapm_p (fun x => sample: y <- f x; g y) m =
  sample: m' <- mapm_p f m; mapm_p g m'.
Proof.
elim/fmap_ind: m=> [|m IH x y x_m]; first by rewrite !(mapm_p0, sample_diracL).
rewrite !mapm_p_setm IH // !sampleA; apply/eq_sample=> m'.
rewrite sampleA; under [in LHS]eq_sample=> ? do rewrite sampleA.
rewrite sampleC; apply/eq_sample=> z.
by rewrite sample_diracL mapm_p_setm sampleC.
Qed.

End MapMapProb.

Section SampleFun.

Variables (T : ordType) (S : ordType) (def : T -> S).

Definition samplef (f : ffun (dirac \o def)) : {prob ffun def} :=
  sample: m <- mapm_p id (val f);
  dirac (mkffun (fun x => odflt (def x) (m x)) (domm m)).

Lemma samplef0 : samplef emptyf = dirac emptyf.
Proof.
rewrite /samplef /= mapm_p0 sample_diracL; congr dirac.
by apply/eq_ffun=> x; rewrite mkffunE /= domm0 /= emptyfE.
Qed.

Lemma samplef_map f : samplef (mapf dirac f) = dirac f.
Proof.
rewrite /samplef /= val_mapf; last exact: dirac_inj.
rewrite mapm_p_comp mapm_p_dirac sample_diracL; congr dirac.
apply/eq_ffun=> x; rewrite mkffunE /appf /= mem_domm.
by case: (ffval f x).
Qed.

End SampleFun.

Section MapFunProb.

Variables (T S R : ordType) (def : T -> S).

Definition mapf_p (g : S -> R) (f : ffun (dirac \o def)) : ffun (dirac \o g \o def) :=
  mkffun (fun x => sample: y <- f x; dirac (g y)) (supp f).

Lemma mapf_pE g f x :
  mapf_p g f x = sample: y <- f x; dirac (g y).
Proof.
rewrite /mapf_p mkffunE mem_supp /=.
by case: eqP=> //= ->; rewrite sample_diracL.
Qed.

End MapFunProb.

Lemma coupling_mapm_p (T S11 S12 S21 S22 : ordType)
    (R1 : T -> S11 -> S12 -> Prop) (R2 : T -> S21 -> S22 -> Prop)
    (f1 : S11 -> {prob S21}) (f2 : S12 -> {prob S22}) :
  (forall x y1 y2, R1 x y1 y2 -> coupling (R2 x) (f1 y1) (f2 y2)) ->
  forall (m1 : {fmap _}) (m2 : {fmap _}),
    (forall x y1 y2, m1 x = Some y1 -> m2 x = Some y2 -> R1 x y1 y2) ->
    coupling (fun (m1' : {fmap _}) (m2' : {fmap _}) =>
                forall x y1 y2, m1' x = Some y1 -> m2' x = Some y2 -> R2 x y1 y2)
             (mapm_p f1 m1) (mapm_p f2 m2).
Proof.
move=> RP; elim/fmap_rect.
  rewrite mapm_p0; move=> m2 em; apply: coupling_trivial.
  move=> _ m2' /supp_diracP -> //.
move=> m1 IH x y1 fresh m2 em.
have y1P y2 : m2 x = Some y2 -> R1 x y1 y2.
  by apply: em; rewrite setmE eqxx.
have /IH {}em x' y1' y2 : m1 x' = Some y1' -> m2 x' = Some y2 -> R1 x' y1' y2.
  move=> m1x' m2x'; apply: em; rewrite // setmE.
  by case: eqP m1x'=> // ->; rewrite (dommPn fresh).
case m2x: (m2 x)=> [y2|].
  rewrite -(setmI m2x) !mapm_p_setm; apply: coupling_sample em _.
  move=> /= m1' m2' _ _ em'; apply: coupling_sample.
    apply: RP; exact: y1P.
  move=> y1' y2' _ _ yP; apply: coupling_dirac.
  move=> x' y1'' y2''; rewrite !setmE; case: eqP=> [->|_].
    by move=> [<-] [<-].
  exact: em'.
rewrite mapm_p_setm -[mapm_p _ m2]sample_diracR.
apply: coupling_sample em _ => m1' m2' _ /supp_mapm_pP [ed2 _].
move=> {}em; apply: coupling_sampleL=> y1' _; apply: coupling_dirac.
move=> x' y1'' y2''; rewrite setmE; case: eqP=> [->|_]; last exact: em.
by have /dommPn -> : x \notin domm m2' by rewrite -ed2 mem_domm m2x.
Qed.

Lemma coupling_mapm_p_eq (T S1 S2 : ordType) (f : S1 -> {prob S2}) (m1 m2 : {fmap T -> S1}) (P : T -> bool) :
  {in P, m1 =1 m2} ->
  coupling (fun m1' m2' : {fmap T -> S2} => {in P, m1' =1 m2'}) (mapm_p f m1) (mapm_p f m2).
Proof.
move=> em.
pose R1 x (y1 y2 : S1) := P x -> y1 = y2.
pose R2 x (y1' y2' : S2) := P x -> y1' = y2'.
have emW x y1 y2 : m1 x = Some y1 -> m2 x = Some y2 -> R1 x y1 y2.
  move=> m1x m2x Px; apply: Some_inj; rewrite -m1x -m2x; exact: em.
have RP x y1 y2 : R1 x y1 y2 -> coupling (R2 x) (f y1) (f y2).
  move=> xP; case: (boolP (P x))=> [{}/xP ->|nPx].
    apply: couplingW; last exact: coupling_same.
    by move=> ?? ->.
  apply: coupling_trivial=> ?? _ _ contra; by rewrite contra in nPx.
apply: in_couplingW; last by apply: coupling_mapm_p emW; eauto.
move=> /= m1' m2' /supp_mapm_pP [ed1 _] /supp_mapm_pP [ed2 _].
move=> em' x Px.
case m1'x: (m1' x)=> [y1'|].
  suff /dommP [y2' m2'x] : x \in domm m2'.
    by rewrite m2'x (em' _ _ _ m1'x m2'x) // mem_domm.
  by rewrite -ed2 mem_domm -em // -mem_domm ed1 mem_domm m1'x.
by apply/esym/dommPn; rewrite -ed2 mem_domm -em // -mem_domm ed1 mem_domm m1'x.
Qed.


Require Import Axioms.
Require Import Tactics.
Require Import Ordinal.

Parameter top : ordinal.

Definition stop : ordinal := succ top.

Record page : Type :=
mk_page
  { str : ordinal;  (* stratification level *)
    cex : ordinal;  (* external candidate level *)
    cin : ordinal;  (* internal candidate level *)

    str_top : str <<= top;
    cex_top : cex <<= top;
    cin_cex : cin <<= cex;
  }.

Lemma cin_top : forall pg, cin pg <<= top.
Proof.
intros pg.
eapply le_ord_trans; eauto using cex_top, cin_cex.
Qed.

Lemma cin_stop : forall pg, cin pg <<= stop.
Proof.
intros pg.
exact (le_ord_trans _#3 (cin_top pg) (succ_nodecrease _)).
Qed.

(* Probably should change this to use w <<= cex pg instead. *)
Definition upd_cin (pg : page) {w : ordinal} (h : w <<= cin pg) : page
  :=
  mk_page (str pg) (cex pg) w (str_top pg) (cex_top pg) (le_ord_trans _#3 h (cin_cex pg)).

Lemma page_extensionality :
  forall pg pg',
    str pg = str pg'
    -> cex pg = cex pg'
    -> cin pg = cin pg'
    -> pg = pg'.
Proof.
intros pg pg' ? ? ?.
destruct pg as [wn wx wi h1 h2 h3].
destruct pg' as [wn' wx' wi' h1' h2' h3'].
cbn in *.
subst.
so (proof_irrelevance _ h1 h1'); subst h1'.
so (proof_irrelevance _ h2 h2'); subst h2'.
so (proof_irrelevance _ h3 h3'); subst h3'.
reflexivity.
Qed.

Lemma upd_cin_self :
  forall pg (h : cin pg <<= cin pg), upd_cin pg h = pg.
Proof.
intros pg h.
apply page_extensionality; auto.
Qed.

Definition le_page (pg pg' : page) : Prop
  :=
  str pg <<= str pg'
  /\ cex pg <<= cex pg'
  /\ cin pg <<= cin pg'.

Definition lt_page (pg pg' : page) : Prop
  :=
  str pg << str pg'
  /\ cex pg << cin pg'.

Definition toppg : page :=
  mk_page top top top (le_ord_refl _) (le_ord_refl _) (le_ord_refl _).

Definition succ_page (pg : page) (h : lt_page pg toppg) : page
  :=
  mk_page (succ (str pg)) (succ (cex pg)) (succ (cex pg))
    (h andel) (h ander) (le_ord_refl _).

Lemma le_page_refl :
  forall pg, le_page pg pg.
Proof.
intros pg.
do2 2 split; apply le_ord_refl.
Qed.

Lemma le_page_refl' :
  forall pg pg',
    pg = pg'
    -> le_page pg pg'.
Proof.
intros pg pg' H.
subst pg'.
apply le_page_refl.
Qed.

Lemma le_page_trans :
  forall pg1 pg2 pg3,
    le_page pg1 pg2
    -> le_page pg2 pg3
    -> le_page pg1 pg3.
Proof.
intros pg1 pg2 pg3 H12 H23.
destruct H12 as (? & ? & ?).
destruct H23 as (? & ? & ?).
do2 2 split; eauto using le_ord_trans.
Qed.

Lemma lt_le_page_trans :
  forall pg1 pg2 pg3,
    lt_page pg1 pg2
    -> le_page pg2 pg3
    -> lt_page pg1 pg3.
Proof.
intros pg1 pg2 pg3 H12 H23.
destruct H12 as (? & ?).
destruct H23 as (? & ? & ?).
split; eauto using lt_le_ord_trans.
Qed.

Lemma lt_page_succ :
  forall pg h,
    lt_page pg (succ_page pg h).
Proof.
intros pg h.
split; apply succ_increase.
Qed.

Lemma lt_page_impl_le_page :
  forall pg pg',
    lt_page pg pg'
    -> le_page pg pg'.
Proof.
intros pg pg' H.
destruct H as (Hstr & Hcex).
do2 2 split; cbn; eauto using lt_ord_impl_le_ord.
  {
  apply lt_ord_impl_le_ord.
  eapply lt_le_ord_trans; eauto.
  apply cin_cex.
  }

  {
  apply lt_ord_impl_le_ord.
  eapply lt_le_ord_trans; eauto.
  eapply le_lt_ord_trans.
    {
    apply cin_cex.
    }
  apply succ_increase.
  }
Qed.

Lemma lt_page_impl_succ_le_page :
  forall pg pg' h,
    lt_page pg pg'
    -> le_page (succ_page pg h) pg'.
Proof.
intros pg pg' h H.
destruct H as (Hstr & Hcex).
do2 2 split; cbn; auto.
eapply lt_le_ord_trans; eauto.
apply cin_cex.
Qed.

Lemma toppg_max :
  forall pg, le_page pg toppg.
Proof.
intros pg.
do2 2 split; cbn; auto using str_top, cex_top, cin_top.
Qed.

Lemma str_mono :
  forall pg pg', le_page pg pg' -> str pg <<= str pg'.
Proof.
intros pg pg' H.
destruct H as (? & ? & ?); auto.
Qed.

Lemma cex_mono :
  forall pg pg', le_page pg pg' -> cex pg <<= cex pg'.
Proof.
intros pg pg' H.
destruct H as (? & ? & ?); auto.
Qed.

Lemma cin_mono :
  forall pg pg', le_page pg pg' -> cin pg <<= cin pg'.
Proof.
intros pg pg' H.
destruct H as (? & ? & ?); auto.
Qed.

Lemma upd_cin_mono :
  forall pg pg' w w' (h : w <<= cin pg) (h' : w' <<= cin pg'),
    le_page pg pg'
    -> w <<= w'
    -> le_page (upd_cin pg h) (upd_cin pg' h').
Proof.
intros [wn wx wi] [wn' wx' wi'] w w' h h' Hle Hw.
destruct Hle as (? & ? & ?).
do2 2 split; auto.
Qed.

Definition max_page (pg pg' : page) : page :=
  mk_page
    (max_ord (str pg) (str pg'))
    (max_ord (cex pg) (cex pg'))
    (max_ord (cin pg) (cin pg'))
    (max_ord_lub _#3 (str_top pg) (str_top pg'))
    (max_ord_lub _#3 (cex_top pg) (cex_top pg'))
    (max_ord_mono _#4 (cin_cex pg) (cin_cex pg')).

Lemma le_page_max_l :
  forall pg pg', le_page pg (max_page pg pg').
Proof.
intros pg pg'.
do2 2 split; cbn; apply le_ord_max_l.
Qed.

Lemma le_page_max_r :
  forall pg pg', le_page pg' (max_page pg pg').
Proof.
intros pg pg'.
do2 2 split; cbn; apply le_ord_max_r.
Qed.

Parameter omega_top : omega <<= top.


  (***************************************************************)
  (*   This file is part of the static analyses developement -   *)
  (*   specification of kildall's algorithm, and proof it is a   *)
  (*   bytecode verifier -                                       *)
  (*							         *)
  (*   File : substitutions.v			 	         *)
  (*   Authors : S. Coupet-Grimal, W. Delobel 		         *)
  (*   Content : definition and properties of substitutions      *)
  (***************************************************************)

Section substitutions.

  Require Export machine.

  (* Name = name U {variables x_{i,j} / i,j : nat}.*)
  Inductive Name : Set :=
    | x : nat -> nat -> Name
    | symbol : name -> Name.

  Notation Get_function := (get_function_by_name functions).

  (* an expression is a tree whose nodes and leaves are marked with elements of name.
   note that in this definition, nothing prevents a node from being marked with a variable
  - functions dealing with expressions will filter these cases out
   *)
  Notation Expression := (tree Name Name).

  (* elementary substitution : xi and xj represent variable x_{xi,xj},
   the last field is the expression to be substituted to the variable;
   substituted variable may appear in that expression
   *)
  Record subst_elem : Set := mksubst {xi : nat; xj : nat; subst_expr : Expression}.

  (* substitutions : list of elementary substitutions *)
  Notation Substitution := (list subst_elem).
  Notation value := (tree name name).

  (* applying an elementary substitution to an expression;
     (apply_elem_tree {i;j;es} e) returns a copy of e where all occurences of x_{i,j} have been replaced by es;
     es itself is copied and left unchanged, that is (apply_elem_tree {i;j;es} e) may contain variable x_{i,j}
     - in case x_{i,j} appears in es -
   *)
  Fixpoint apply_elem_tree (s : subst_elem) (e : Expression) {struct e} : Expression :=
    match e with
      | Node (x i j) nil =>
        match eq_nat_dec (xi s) i with
          | left _ =>
	    match eq_nat_dec (xj s) j with
	      | left _ => (subst_expr s)
	      | right _ => e
	    end
          | right _ => e
        end
      | Node (symbol dft) forst => Node (symbol dft) (map (apply_elem_tree s) forst)
      | _ => e
    end.

  (* (apply [s1,...,sn] e) <=> (apply_elem_tree s1 (apply_elem_tree (s2 (...(apply_elem_tree sn e)...)*)
  Fixpoint apply (S : Substitution) (e : Expression) {struct S} : Expression :=
    match S with
      | nil => e
      | s :: S' => apply_elem_tree s (apply S' e)
    end.

  (* fresh p h m returns a forest whose trees are single nodes x_{p,h}, ... , x_{p,h+m-1} *)
  Fixpoint fresh (p h m : nat) {struct m} : list Expression :=
    match m with
      | 0 => nil
      | (S m') => (Node (x p h) nil) :: (fresh p (S h) m')
    end.

  Lemma eq_Name_dec : forall (n n' : Name), {n = n'} + {n <> n'}.
  Proof.
    intros n n'; destruct n as [xn yn | n]; destruct n' as [xn' yn' | n'].
    elim (eq_nat_dec xn xn'); intro case_x.
    elim (eq_nat_dec yn yn'); intro case_y; [left | right].
    subst; trivial.
    intro H; apply case_y; inversion H; trivial.
    right; intro H; apply case_x; inversion H; trivial.
    right; intro H; inversion H.
    right; intro H; inversion H.
    elim (eq_name_dec n n'); intro case_n;
      [left; subst | right; intro H; apply case_n; inversion H]; trivial.
  Qed.

  Lemma eq_Expression_dec : forall (e e' : Expression), {e=e'} + {e<>e'}.
  Proof.
    intros e' e.
    apply eq_tree_dec; apply eq_Name_dec.
  Qed.

  Lemma eq_list_Expression_dec : forall (e e' : list Expression), {e=e'} + {e<>e'}.
  Proof.
    apply list_eq_dec.
    apply eq_Expression_dec.
  Qed.

  Lemma eq_Substitution_dec : forall (s s' : Substitution), {s=s'} + {s<>s'}.
  Proof.
    apply list_eq_dec.
    intros s s'; destruct s as [i j expr]; destruct s' as [i' j' expr'].
    elim (eq_nat_dec i i'); intro casei.
    elim (eq_nat_dec j j'); intro casej.
    elim (eq_Expression_dec expr expr'); intro case; subst; [left; trivial | right].
    injection; intros; apply case; trivial.
    right; injection; intros; apply casej; trivial.
    right; injection; intros; apply casei; trivial.
  Qed.

  (* (tree_is_pattern e) means no function name appears in e *)
  Fixpoint tree_is_pattern_bool (e : Expression) : bool :=
    match e with
      | Node (x i j) nil => true
      | Node (symbol dft) forst =>
        (match Get_function dft with
           | None => true
           | _ => false
         end) &&
        (fold_left (fun (p : bool) (e : Expression) => p && (tree_is_pattern_bool e)) forst true)
      | _ => false
    end.

  Definition tree_is_pattern (e : Expression) := (tree_is_pattern_bool e = true).

  (* convertion *)
  Fixpoint tree_name_to_tree_Name (t : value) : Expression :=
    match t with
      | Leaf l => Leaf (symbol l)
      | Node n f => Node (symbol n) (map tree_name_to_tree_Name f)
    end.

  (*  *)
  Definition stack_shaping (lexpr : list Expression) (stack : list value) (rho : Substitution) : Prop :=
    map (apply rho) lexpr = map tree_name_to_tree_Name stack.

  (* initial variables of a function *)
  Definition init_vars (fcn : function) := rev_lin (fresh 0 0 (length (fcn_args fcn))).

   (* builds a substitution from the variables in le and values in lv; returns None if |le|<>|lv|,
      or le contains an element that isn't a variable
      order in the substitution is the order of variables in le
   *)
  Fixpoint make_substitution (le : list Expression) (lv : list value) {struct le} : option Substitution :=
    match le,lv with
      | nil, nil => Some nil
      | (Node (x i j) nil)::le', v :: lv' =>
        match make_substitution le' lv' with
          | Some s => Some ((mksubst i j (tree_name_to_tree_Name v))::s)
          | _ => None
        end
      | _,_ => None
    end.

  (* initial substitution of a function : let [v_1,...v_k] be arguments of f s.t. ar(f) = ar;
     (init_subst 0 ar [v_1, ..., v_k]) is Some [x_{0,1}/v_1, ..., x_{0,ar}/v_k]
     or None if k <> ar
   *)
  Definition init_subst (args : list value) :=
    make_substitution (fresh 0 0 (length args)) args.

   (* (var_in_tree i j e) means variable x_{i,j} appears in expression e *)
  Fixpoint var_in_tree_bool (e : Expression) (i j : nat) {struct e} : bool :=
    match e with
      | Node (x i' j') nil =>
        (match eq_nat_dec i i' with
           | left _ => true
           | right _ => false
         end) &&
        (match eq_nat_dec j j' with
           | left _ => true
           | right _ => false
         end)
      | Node (symbol _) forst =>
        fold_left (fun (p : bool) (e : Expression) => p || (var_in_tree_bool e i j)) forst false
      | _ => false
    end.

  Definition var_not_in_tree (e : Expression) (i j : nat) : Prop :=
    var_in_tree_bool e i j = false.

   (* (var_in_subst_elem s i j) means x_{i,j} is the substituted variable in s, or appears in the substituted expression *)
  Definition var_in_subst_elem (s : subst_elem) (i j : nat) : Prop :=
    match s with
      | mksubst i' j' e => var_in_tree_bool e i j = true \/ (i' = i /\ j' = j)
    end.

  Lemma var_not_in_subst_elem : forall (si sj : nat) (se : Expression) (i j : nat),
    ~ var_in_subst_elem (mksubst si sj se) i j -> var_in_tree_bool se i j = false /\ (si <> i \/ sj <> j).
  Proof.
    simpl; intros si sj se i j H_not_in_subst.
    elim (Decidable.not_or _ _ H_not_in_subst).
    intros H1 H2; split.
    destruct (var_in_tree_bool se i j); [elim H1 | idtac]; trivial.
    assert (Hdec : Decidable.decidable (si = i)).
    unfold Decidable.decidable.
    elim (eq_nat_dec si i); [left | right]; trivial.
    elim (Decidable.not_and _ _ Hdec H2); [left | right]; assumption.
  Qed.

  Lemma aux_init_vars_length : forall (p h m : nat),
    length (fresh p h m) = m.
  Proof.
    intros p h m; generalize h; clear h; induction m; simpl.
    trivial.
    intro h; rewrite (IHm (S h)).
    trivial.
  Qed.

  Lemma init_vars_length : forall (fcn : function), length (init_vars fcn) = length (fcn_args fcn).
  Proof.
    intro fcn; unfold init_vars.
    rewrite (rev_lin_is_rev).
    rewrite rev_length.
    rewrite (aux_init_vars_length 0 0 (length (fcn_args fcn))).
    trivial.
  Qed.

  Lemma element_at_init_vars : forall (n p h j : nat) (ej : Expression),
    (rev_lin (fresh p h n))[j] = Some ej -> ej = Node (x p (n+h-(S j))) nil.
  Proof.
    intros n p h; rewrite rev_lin_is_rev; generalize n p h; clear n p h.
    induction n; simpl.
    intros p h j ej F; inversion F.
    intros p h j ej;simpl.
    destruct n.
    simpl.
    destruct j; intro H; inversion H.
    replace (h-0) with (h); [trivial | auto with arith].
    intro H.
    compare j (S n); intro comp.
    subst.
    rewrite plus_comm; simpl.
    replace (h+S n - S n) with h; [idtac | rewrite plus_comm; auto with arith].
    cut (((rev (fresh p (S h) (S n))) ++  Node (x p h) nil :: nil)[S n] = (Node (x p h) nil :: nil)[0]).
    intro H1.
    rewrite H1 in H.
    simpl in H.
    inversion H; trivial.
    replace ((rev (fresh p (S h) (S n)) ++ Node (x p h) nil :: nil)[S n]) with
      ((rev (fresh p (S h) (S n)) ++ Node (x p h) nil :: nil)[length (rev (fresh p (S h) (S n))) + 0]).
    apply element_at_list_concat.
    rewrite rev_length.
    replace (length (fresh p (S h) (S n))) with (S n); simpl; trivial.
    rewrite plus_comm; simpl; trivial.
    cut (n = (length (fresh p (S (S h)) n))); [auto with arith | symmetry].
    apply aux_init_vars_length.
    rewrite (IHn p (S h) j ej).
    rewrite plus_comm; simpl.
    rewrite plus_comm; trivial.
    rewrite <- H.
    apply element_at_list_concat'.
    rewrite rev_length.
    cut (length (fresh p (S h) (S n)) = S n);
      [intro Hlen; rewrite Hlen; trivial | idtac].
    cut (S j < S (S n)); [auto with arith | idtac].
    cut (j < S (S n)).
    apply not_pred_lt_S_lt.
    auto with arith.
    simpl; assumption.
    replace (S (S n)) with (length (rev (fresh p (S h) (S n)) ++
      Node (x p h) nil :: nil)).
    generalize H; apply element_at_list_some.
    rewrite concat_length.
    rewrite rev_length; rewrite Hlen; simpl.
    rewrite plus_comm; simpl; trivial.
    apply aux_init_vars_length.
  Qed.

  Lemma subst_length : forall (n h : nat) (args : list value),
    length args = n -> make_substitution (fresh 0 h n) args <> None.
  Proof.
    induction n; simpl; trivial.
    intros h args Hl; destruct args.
    intro F; inversion F.
    inversion Hl.
    intros h args Hl; destruct args.
    inversion Hl.
    simpl.
    CaseEq (make_substitution (fresh 0 (S h) n) args).
    intros l H F; inversion F.
    intros H dd; apply (IHn (S h) args).
    generalize Hl; simpl; auto with arith.
    assumption.
  Qed.

  Lemma subst_unchanged : forall (l : Substitution) (i j : nat),
    (forall (s : Expression), ~ In (mksubst i j s) l) -> apply l (Node (x i j) nil) = (Node (x i j) nil).
  Proof.
    induction l; simpl; trivial.
    intros i j Hnotin.
    destruct a; simpl; trivial.
    rewrite (IHl i j); simpl.
    unfold apply_elem_tree; simpl.
    elim (eq_nat_dec xi0 i); intro case_i; simpl; trivial.
    elim (eq_nat_dec xj0 j); intro case_j; simpl; trivial.
    subst.
    elim (Hnotin subst_expr0); left; trivial.
    intros s H; apply (Hnotin s).
    right; assumption.
  Qed.

  Lemma element_at_init_subst : forall (n h j : nat) (l : Substitution) (args : list value) (v : value),
    args[j] = Some v ->
    make_substitution (fresh 0 h n) args = Some l ->
    l[j] = Some (mksubst 0 (h+j) (tree_name_to_tree_Name v)).
  Proof.
    induction n; simpl; trivial.
    intros h j l args v; destruct args; simpl.
    intro F; inversion F.
    intros dd F; inversion F.
    intros h j l args v; destruct args; simpl.
    intros dd F; inversion F.
    destruct j; simpl.
    intro H; inversion H; clear H; subst.
    rewrite plus_comm; simpl.
    CaseEq (make_substitution (fresh 0 (S h) n) args).
    intros l0 case_l0 H; inversion H.
    simpl; trivial.
    intros H F; inversion F.
    CaseEq (args[j]).
    intros v1 case H; inversion H; clear H; subst.
    CaseEq (make_substitution (fresh 0 (S h) n) args).
    intros l0 case_l0 H; inversion H; clear H.
    generalize (IHn (S h) j l0 args v case case_l0).
    replace (S h + j) with (h + S j).
    simpl; destruct (l0[j]); trivial.
    simpl.
    rewrite plus_comm; simpl.
    rewrite plus_comm; trivial.
    intros dd F; inversion F.
    intros dd F; inversion F.
  Qed.

  Lemma init_subst_member : forall (n h h' : nat) (s : Expression) (l : Substitution) (args : list value),
    make_substitution (fresh 0 h' n) args = Some l -> In (mksubst 0 h s) l -> h' <= h.
  Proof.
    induction n; simpl; intros h h' s l args.
    destruct args; intro H; inversion H.
    intro F; inversion F.
    destruct args.
    intro F; inversion F.
    CaseEq (make_substitution (fresh 0 (S h') n) args).
    intros l0 case H Hin; inversion H.
    clear H; subst.
    elim Hin; clear Hin; intro Hin.
    inversion Hin; subst; clear Hin.
    auto with arith.
    generalize (IHn h (S h') s l0 args case Hin).
    auto with arith.
    intros H F; inversion F.
  Qed.

  Lemma apply_elem_unchanged : forall (s : subst_elem) (v : value),
    apply_elem_tree s (tree_name_to_tree_Name v) = (tree_name_to_tree_Name v).
  Proof.
    intro s.
    intro v; induction v as [l | n l IHv] using tree_ind.
    simpl; trivial.
    simpl.
    assert (H : (map (apply_elem_tree s) (map tree_name_to_tree_Name l)) = (map tree_name_to_tree_Name l)).
    induction l; simpl; trivial.
    rewrite (IHv a); try left; trivial.
    rewrite IHl; trivial.
    intros t t_in_l; apply (IHv t); right; assumption.
    rewrite H; trivial.
  Qed.

  Lemma apply_elem_forest_unchanged : forall (s : subst_elem) (l : list value),
    map (apply_elem_tree s) (map tree_name_to_tree_Name l) = (map tree_name_to_tree_Name l).
  Proof.
    induction l as [ | h t IHl].
    (* case l = nil : *)
    simpl; trivial.
    (* case l = h :: t : *)
    simpl; rewrite IHl; rewrite apply_elem_unchanged; trivial.
  Qed.

  Lemma apply_unchanged : forall (S : Substitution) (v : value),
    apply S (tree_name_to_tree_Name v) = (tree_name_to_tree_Name v).
  Proof.
    induction S; simpl.
    trivial.
    intro v; rewrite (IHS v).
    apply apply_elem_unchanged.
  Qed.

  Lemma init_subst_apply : forall (n h j : nat) (l : Substitution) (args : list value) (v : value),
    args[j] = Some v -> make_substitution (fresh 0 h n) args = Some l ->
    apply l (Node (x 0 (h+j)) nil) = tree_name_to_tree_Name v.
  Proof.
    induction n; simpl; intros h j l args v.
    destruct args; simpl.
    intro F; inversion F.
    intros dd F; inversion F.
    destruct args; simpl.
    intro F; inversion F.
    destruct j; simpl.
    intro H; inversion H; clear H; subst.
    CaseEq (make_substitution (fresh 0 (S h) n) args).
    intros l0 case_l0 H; inversion H; clear H; subst.
    rewrite plus_comm; simpl.
    rewrite (subst_unchanged l0 0 h).
    unfold apply_elem_tree; simpl.
    elim (eq_nat_dec h h); [trivial | intro H; elim H; trivial].
    intros s Hin; generalize (init_subst_member n h (S h) s l0 args case_l0 Hin).
    apply not_S_le.
    intros dd F; inversion F.
    CaseEq (args[j]).
    intros v1 case_argsj H; inversion H; subst.
    CaseEq (make_substitution (fresh 0 (S h) n) args).
    clear H; intros l0 case_l0 H; inversion H; clear H.
    simpl.
    replace (h + S j) with (S h + j).
    rewrite (IHn (S h) j l0 args v case_argsj case_l0).
    apply apply_elem_unchanged.
    simpl; rewrite plus_comm; simpl; rewrite plus_comm; trivial.
    intros dd F; inversion F.
    intros dd F; inversion F.
  Qed.

  Lemma init_stack_shaping : forall (n h : nat) (args : list value) (rho : Substitution),
    make_substitution (fresh 0 h n) args = Some rho -> stack_shaping (rev_lin (fresh 0 h n)) (rev_lin args) rho.
  Proof.
    intros n h args rho Hinit.
    rewrite rev_lin_is_rev; rewrite rev_lin_is_rev.
    unfold stack_shaping.
    assert (H : map (apply rho) ((fresh 0 h n)) =
      map tree_name_to_tree_Name (args)).
    generalize n h args rho Hinit; clear Hinit rho args h n.
    induction n; intros h args rho; simpl.
    destruct args; simpl; trivial.
    intro H; inversion H; subst; unfold stack_shaping; simpl; trivial.
    destruct args; simpl; trivial.
    intro F; inversion F.
    CaseEq (make_substitution (fresh 0 (S h) n) args).
    intros rho' case_rho' H; inversion H; subst.
    generalize (IHn (S h) args rho' case_rho'); unfold stack_shaping; simpl.
    replace (apply rho' (Node (L:=Name) (x 0 h) nil)) with (Node (L := Name) (x 0 h) nil).
    simpl; elim (eq_nat_dec h h); intro case_eq; [clear case_eq | elim case_eq; trivial].
    intro Heq; rewrite <- Heq.
    cut (map (fun e : Expression => apply_elem_tree (mksubst 0 h (tree_name_to_tree_Name t)) (apply rho' e)) (fresh 0 (S h) n) =
      map (apply rho') (fresh 0 (S h) n)); [intro H'; rewrite H'; trivial | idtac].
    clear IHn case_rho' H.
    generalize args Heq; clear Heq args.
    induction (fresh 0 (S h) n); simpl; intros args Heq.
    trivial.
    destruct args; simpl; simpl in Heq.
    inversion Heq.
    simpl; rewrite (IHl args).
    inversion Heq.
    rewrite H0.
    rewrite H1.
    rewrite apply_elem_unchanged; trivial.
    inversion Heq; trivial.
    symmetry; apply subst_unchanged.
    intros s Hin; generalize (init_subst_member n h (S h) s rho' args case_rho' Hin); apply not_S_le.
    intros dd F; inversion F.
    rewrite map_rev_commut.
    rewrite  H.
    symmetry; apply map_rev_commut.
  Qed.

  Lemma sub_pattern : forall (cnm : name) (args : list Expression),
    tree_is_pattern (Node (symbol cnm) args) ->
    forall (j : nat) (t : Expression), args[j] = Some t -> tree_is_pattern t.
  Proof.
    induction args; simpl; intros Hpattern j tj Htj.
    inversion Htj.
    destruct j.
    generalize Hpattern.
    unfold tree_is_pattern; simpl.
    CaseEq (Get_function cnm).
    simpl; intros f Hget H; inversion H.
    intro Hnone; simpl; intro H.
    inversion Htj; subst.
    CaseEq (tree_is_pattern_bool tj); trivial.
    intro HF; rewrite HF in H.
    clear IHargs Hpattern; induction args; simpl in H; trivial.
    apply IHargs; trivial.
    CaseEq (element_at_list args j).
    intros tj' case_tj; rewrite case_tj in Htj.
    inversion Htj; generalize case_tj; subst; apply IHargs.
    unfold tree_is_pattern in *; simpl in Hpattern.
    CaseEq (Get_function cnm); simpl in Hpattern.
    intros f Hget; rewrite Hget in Hpattern; inversion Hpattern.
    intro Hget; rewrite Hget in Hpattern.
    unfold tree_is_pattern_bool.
    rewrite Hget; simpl.
    destruct (tree_is_pattern_bool a).
    simpl in Hpattern ;apply Hpattern.
    simpl in Hpattern.
    cut (False); [intro F; inversion F | clear IHargs case_tj].
    induction args; simpl in Hpattern.
    inversion Hpattern.
    apply IHargs; trivial.
    intro F; rewrite F in Htj; inversion Htj.
  Qed.

  Lemma apply_elem_to_forest :  forall (cnm : name) (s : subst_elem) (args : list Expression),
    apply_elem_tree s (Node (symbol cnm) args) = Node (symbol cnm) (map (apply_elem_tree s) args).
  Proof.
    induction args as [ | hargs targs]; simpl; trivial.
  Qed.

  Lemma apply_to_forest : forall (cnm : name) (rho : Substitution) (args : list Expression),
    apply rho (Node (symbol cnm) args) = Node (symbol cnm) (map (apply rho) args).
  Proof.
    induction rho; simpl.
    (* case rho nil : *)
    intro args; rewrite (map_id args); trivial.
    (* case rho not nil : *)
    intro args.
    rewrite (IHrho args).
    rewrite apply_elem_to_forest.
    replace (map (apply_elem_tree a) (map (apply rho) args)) with (map (fun e : Expression => apply_elem_tree a (apply rho e)) args).
    trivial.
    clear IHrho; induction args; simpl.
    trivial.
    rewrite IHargs.
    trivial.
  Qed.

  Lemma apply_cons : forall (cnm : name) (rho : Substitution) (args : list value) (sargs : list Expression),
    length args = length sargs ->
    (forall (vj : value) (ej : Expression) (j : nat), args[j] = Some vj -> sargs[j] = Some ej -> apply rho ej = tree_name_to_tree_Name vj) ->
    apply rho (Node (symbol cnm) sargs) = Node (symbol cnm) (map tree_name_to_tree_Name args).
  Proof.
    intros cnm rho args sargs Hlength H; rewrite (apply_to_forest cnm rho sargs).
    replace (map tree_name_to_tree_Name args) with (map (apply rho) sargs); trivial.
    generalize args Hlength H; clear H Hlength args; induction sargs.
    simpl; intros args Hlength H.
    destruct args; simpl.
    trivial.
    inversion Hlength.
    simpl; intros args Hlength H.
    destruct args; simpl.
    inversion Hlength.
    cut (length args = length sargs); [intro Hlen | generalize Hlength; simpl; auto with arith].
    cut (forall (vj : value) (ej : Expression) (j : nat),
      args[j] = Some vj -> sargs[j] = Some ej ->
      apply rho ej = tree_name_to_tree_Name vj); [intro H' | intros vj ej j Hej Hvj; apply (H vj ej (S j)); simpl; trivial].
    rewrite (IHsargs args Hlen H').
    cut (apply rho a = tree_name_to_tree_Name t); [intro Hav | apply (H t a 0); simpl; trivial].
    rewrite Hav; trivial.
  Qed.

  Lemma apply_substitution_cons : forall (s0 : subst_elem) (s : Substitution) (e : Expression),
    apply (s0::s) e = apply (s0::nil) (apply s e).
  Proof.
    intros s0 s; generalize s0; clear s0; induction s; simpl; intros s0 e; trivial.
  Qed.

  Lemma apply_substitution_concat : forall (s0 s : Substitution) (e : Expression),
    apply (s0++s) e = apply s0 (apply s e).
  Proof.
    induction s0; simpl; intros s e.
    trivial.
    rewrite (IHs0 s e); trivial.
  Qed.

  Lemma map_apply_substitution_cons : forall (s0 : subst_elem) (s : Substitution) (le : list Expression),
    map (apply (s0::s)) le = map (apply (s0::nil)) (map (apply s) le).
  Proof.
    intros s0 s le; generalize s0 s; clear s0 s; induction le; intros s0 s.
    simpl; trivial.
    generalize (IHle s0 s); simpl.
    intro h; rewrite h.
    trivial.
  Qed.

  Lemma map_apply_substitution_concat : forall (s0 s : Substitution) (le : list Expression),
    map (apply (s0++s)) le = map (apply s0) (map (apply s) le).
  Proof.
    intros s0 s le; generalize s0 s; clear s0 s; induction le; intros s0 s.
    simpl; trivial.
    generalize (IHle s0 s); simpl; intro h; rewrite h.
    rewrite (apply_substitution_concat s0 s a); trivial.
  Qed.

  Lemma make_substitution_some : forall (m p h : nat) (lv : list value),
    length lv = m -> make_substitution (fresh p h m) lv <> None.
  Proof.
    induction m; simpl; intros p h lv.
    destruct lv; simpl; intro Hlen.
    intro F; inversion F.
    inversion Hlen.
    destruct lv; simpl; intro Hlen.
    inversion Hlen.
    CaseEq (make_substitution (fresh p (S h) m) lv).
    intros; intro F; inversion F.
    intros F dd; generalize F; clear dd F.
    apply IHm.
    generalize Hlen; auto with arith.
  Qed.

  Lemma var_not_in_subst_not_in_apply_elem : forall (e : Expression) (s : subst_elem) (i j : nat),
    ~ var_in_subst_elem s i j -> var_not_in_tree e i j -> var_not_in_tree (apply_elem_tree s e) i j.
  Proof.
    intros e s i j Hnot.
    destruct s as [si sj sexpr]; simpl.
    elim (var_not_in_subst_elem si sj sexpr i j Hnot).
    generalize e si sj sexpr i j; clear Hnot e si sj sexpr i j.
    induction e as [l | n l IHe]; intros si sj se i j Hnots1 Hnots2 Hnote.
    (* case e is a leaf : *)
    simpl; trivial.
    (* case e = Node n l : *)
    simpl; destruct n as [ni nj | n].
      (* case n is a variable : *)
    destruct l; trivial.
      (* remaining case l = nil : *)
    elim (eq_nat_dec si ni); intro case_si.
        (* case si = ni : *)
    elim (eq_nat_dec sj nj); intro case_sj.
          (* case sj = nj : *)
    subst.
    unfold var_not_in_tree; trivial.
          (* case sj <> nj : *)
    assumption.
        (* case si <> ni : *)
    assumption.
      (* case n is a name : *)
    unfold var_not_in_tree; simpl.
    elim (fold_bool_or_false (f:=fun (e : Expression) => var_in_tree_bool e i j) (map (apply_elem_tree (mksubst si sj se)) l)).
    simpl; intros H1 H2.
    apply H2; clear H1 H2.
    intros a Ha ; elim (In_mapped _ _ _ _ _ Ha); intros e He.
    elim He; clear He; intros e_in_l a_eq.
    subst a; apply (IHe e e_in_l si sj se i j); trivial.
    unfold var_not_in_tree.
    elim (fold_bool_or_false (f:=fun (e : Expression) => var_in_tree_bool e i j) l).
    simpl; intros H1 H2.
    apply H1; clear H1 H2; trivial.
  Qed.

  Lemma var_not_in_subst_not_in_apply : forall (S : Substitution) (e : Expression) (i j : nat),
    (forall (s : subst_elem), In s S -> ~ var_in_subst_elem s i j) -> var_not_in_tree e i j -> var_not_in_tree (apply S e) i j.
  Proof.
    induction S; simpl; trivial.
    intros e i j HS Hnote; apply var_not_in_subst_not_in_apply_elem; trivial.
    apply (HS a).
    left; trivial.
    apply IHS; trivial.
    intros s Hin; apply (HS s); right; assumption.
  Qed.

  Lemma not_var_in_tree_value : forall (v : value) (p h : nat),
     var_not_in_tree (tree_name_to_tree_Name v) p h.
  Proof.
    induction v as [ l | n l IHv].
    (* case v leaf : *)
    intros p h; unfold var_not_in_tree; simpl; trivial.
    (* case v node : *)
    unfold var_not_in_tree in *; simpl.
    intros p h.
    elim (fold_bool_or_false (f:=fun (e : Expression) => var_in_tree_bool e p h) (map tree_name_to_tree_Name l)).
    simpl; intros H1 H2.
    apply H2; clear H1 H2.
    intros a Ha; elim (In_mapped _ _ _ _ _ Ha).
    intros t Ht; elim Ht; clear Ht; intros t_in_l Ht.
    subst a; apply (IHv t t_in_l p h).
  Qed.

  Lemma apply_elem_var_not_in_unchanged : forall (e es : Expression) (i j : nat),
    var_not_in_tree e i j -> apply_elem_tree (mksubst i j es) e = e.
  Proof.
    induction e as [ l | n l IHe]; simpl.
    (* case e leaf : *)
    trivial.
    (* case e node : *)
    intros es i j; unfold var_not_in_tree; simpl; destruct n as [ni nj | n].
      (* case n variable *)
    destruct l as [ | h t].
        (* case l empty : *)
    elim (eq_nat_dec i ni); elim (eq_nat_dec j nj); intros case_j_nj case_i_ni; try subst i j; simpl; trivial.
          (* remaining case i = ni /\ j = nj : *)
    intro H; inversion H.
        (* case l = h :: t : *)
    trivial.
      (* case n is a name : *)
    intro Hfold; cut ((map (apply_elem_tree (mksubst i j es)) l) = l).
    intro H; rewrite H; trivial.
    elim (fold_bool_or_false (f:=fun (e : Expression) => var_in_tree_bool e i j) l).
    simpl; intros H1 H2.
    assert (H : forall a : Expression, In a l -> var_in_tree_bool a i j = false).
    apply H1; trivial.
    clear Hfold H1 H2.
    elim (map_unchanged_elements (f := (apply_elem_tree (mksubst i j es))) l).
    intros H1 H2; apply H2.
    intros a a_in_l; apply IHe;trivial.
    unfold var_not_in_tree; apply H; trivial.
  Qed.

  Lemma var_not_in_subst_elem_unchanged : forall (s : subst_elem) (i j : nat),
    ~ var_in_subst_elem s i j ->
    apply_elem_tree s (Node (x i j) nil) = Node (x i j) nil.
  Proof.
    intros s i j Hnot; destruct s as [si sj se].
    elim (var_not_in_subst_elem si sj se i j Hnot); simpl.
    elim (eq_nat_dec si i); elim (eq_nat_dec sj j); intros H1 H2; subst; trivial.
    (* remaining case : si =i /\ sj =j *)
    intros H1 H2; elim H2; clear H2; intro H2; elim H2; trivial.
  Qed.

  Lemma var_not_in_subst_unchanged : forall (S : Substitution) (i j : nat),
    (forall (s : subst_elem), In s S -> ~ var_in_subst_elem s i j) ->
    apply S (Node (x i j) nil) = Node (x i j) nil.
  Proof.
    induction S as [ | hS tS IHS]; intros i j HS; simpl.
    (* case S = nil : *)
    trivial.
    (* case S = hS :: tS : *)
    rewrite IHS; trivial.
    apply var_not_in_subst_elem_unchanged.
    apply HS; left; trivial.
    intros s s_in_tS; apply HS; right; trivial.
  Qed.

  Lemma make_substitution_vars :  forall (S : Substitution) (vf : list value) (p h n : nat) ,
   make_substitution (fresh p h n) vf = Some S ->
   forall (s : subst_elem), In s S ->
     forall (i j : nat), var_in_subst_elem s i j -> i = p /\ h <= j /\ j < h + n.
  Proof.
    induction S as [ | hS tS IHS].
    (* case S = nil : *)
    intros vf p h n Hsubst s s_in_nil; inversion s_in_nil.
    (* case S = hS :: tS : *)
    intro vf; destruct vf as [ | hvf tvf].
      (* case vf = nil : *)
    destruct n as [ | n]; simpl; intro H; inversion H.
      (* case vf = hvf :: tvf : *)
    destruct n as [ | n]; simpl.
        (* case n = 0 : *)
    intro H; inversion H.
        (* case n > 0 : *)
    CaseEq (make_substitution (fresh p (S h) n) tvf).
          (* case substitution exists : *)
    intros S' Hs' H s s_in_S.
    assert (H' : mksubst p h (tree_name_to_tree_Name hvf) :: S' = hS :: tS).
    (* Proof of H' : *)
    inversion H; trivial.
    (* H' prooved *)
    elim (cons_eq _ _ _ _ _ H'); clear H H'; intros H1 H2.
    subst hS tS; elim s_in_S; clear s_in_S; intro case_s.
            (* case s = mksubst p h (tree_name_to_tree_Name hvf) : *)
    subst s; simpl; intros i j H.
    elim H; clear H; intro H.
    generalize (not_var_in_tree_value hvf i j); unfold var_not_in_tree; simpl.
    intro H'; rewrite H' in H; inversion H.
    elim H; intros; subst; split; trivial.
    split; auto with arith.
    rewrite plus_comm; simpl.
    apply lt_le_trans with (S j); auto with arith.
            (* case s in S' : *)
    assert (H : forall i j : nat, var_in_subst_elem s i j -> i = p /\ (S h) <= j < (S h) + n).
    (* Proof of H : *)
    intros i j; apply (IHS tvf p (S h) n); trivial.
    (* H prooved *)
    intros i j Hij; elim (H i j Hij).
    intros H1 H2; clear H; elim H2; clear H2; intros H2 H3.
    repeat split; auto with arith.
    rewrite plus_comm; simpl; rewrite plus_comm.
    simpl in H3; assumption.
          (* case substitution does not exist : contradiction *)
    intros dd F; inversion F.
  Qed.

  Lemma make_substitution_apply : forall (m p h : nat) (vf : list value) (svf : Substitution) (j : nat) (e : Expression) (v : value),
    make_substitution (fresh p h m) vf = Some svf ->
    (fresh p h m)[j] = Some e -> vf[j] = Some v ->
    apply svf e = tree_name_to_tree_Name v.
  Proof.
    induction m as [ | n]; simpl; intros p h vf svf j e v.
    (* case m = 0 : *)
    intros dd F; inversion F.
    (* case m = S n : *)
    destruct vf as [ | hvf tvf].
      (* case vf = nil : *)
    intro F; inversion F.
      (* case vf = hvf::tvf : *)
    CaseEq (make_substitution (fresh p (S h) n) tvf).
        (* case substitution exists : *)
    intros l Hl H; inversion H.
    destruct j as [ | j].
          (* case j = 0 : *)
    clear H; intro H; inversion H; subst; simpl.
    clear H; intro H; inversion H; clear H; subst.
    replace (apply l (Node (L := Name) (x p h) nil)) with (Node (L := Name) (x p h) nil).
    simpl.
    elim (eq_nat_dec p p); intro case_p; [clear case_p | elim case_p; trivial].
    elim (eq_nat_dec h h); intro case_h; [clear case_h | elim case_h; trivial].
    trivial.
    symmetry.
    apply var_not_in_subst_unchanged.
    (* prooving forall s : subst_elem, In s l -> ~ var_in_subst_elem s p h : *)
    intros s Hin H.
    elim (make_substitution_vars l tvf p (S h) n Hl s Hin p h H).
    intros H1 H2; elim H2; clear H1 H2; intros H1 H2.
    generalize H1; apply not_S_le.
          (* case j > 0 : *)
    CaseEq (((fresh p (S h) n))[j]).
    intros varj case_varj H2 case_v.
    inversion H2; clear H2; subst.
    simpl.
    rewrite (IHn p (S h) tvf l j e v Hl); trivial.
    apply apply_elem_unchanged.
    intros dd F; inversion F.
        (* case substitution does not exist : *)
    intros dd F; inversion F.
  Qed.

  Lemma apply_elem_in_subst : forall (e : Expression) (i j : nat) (v : value) (is js : nat) (es : Expression),
    var_not_in_tree e i j ->
    apply_elem_tree (mksubst i j (tree_name_to_tree_Name v)) (apply_elem_tree (mksubst is js es) e) =
    apply_elem_tree (mksubst is js (apply_elem_tree (mksubst i j (tree_name_to_tree_Name v)) es)) e.
  Proof.
    induction e as [n | n l IHe].
    (* case e is a leaf : *)
    simpl; trivial.
    (* case e is Node n l : *)
    intros i j v si sj sexpr.
    unfold var_not_in_tree; simpl.
    destruct n as [ni nj | n].
      (* case n is a variable : *)
    destruct l as [ | h t].
        (* case l = nil : *)
    elim (eq_nat_dec i ni); intro case_i_ni.
    elim (eq_nat_dec j nj); intro case_j_nj; subst; simpl.
          (* case ni = i /\ nj = j : *)
    intro H; inversion H.
          (* case ni = i /\ nj <> j : *)
    intro dd; clear dd.
    elim (eq_nat_dec si ni); intro case_si_ni; subst; simpl.
            (* case si = ni : *)
    elim (eq_nat_dec sj nj); intro case_sj_nj;  subst; simpl; trivial.
    elim (eq_nat_dec ni ni); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec j nj); [intro H; subst; elim case_j_nj | idtac]; trivial.
            (* case si <> ni : *)
    elim (eq_nat_dec ni ni); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec j nj); [intro H; subst; elim case_j_nj | idtac]; trivial.
          (* case ni <> i : *)
    simpl; intro dd; clear dd.
    elim (eq_nat_dec si ni); intro case_si_ni; subst; simpl.
            (* case si = ni : *)
    elim (eq_nat_dec sj nj); intro case_sj_nj;  subst; simpl; trivial.
    elim (eq_nat_dec i ni); [intro; subst; elim case_i_ni | idtac]; trivial.
            (* case si <> ni : *)
    elim (eq_nat_dec i ni); [intro; subst; elim case_i_ni | idtac]; trivial.
        (* case l = h :: t : *)
    intro dd; clear dd; simpl; trivial.
      (* case n is a name : *)
    simpl; intro Hfold.
    assert (H : (map (apply_elem_tree (mksubst i j (tree_name_to_tree_Name v)))
      (map (apply_elem_tree (mksubst si sj sexpr)) l)) = (map
        (apply_elem_tree (mksubst si sj
          (apply_elem_tree (mksubst i j (tree_name_to_tree_Name v)) sexpr))) l)).
    (* Proof of H : *)
    rewrite compose_map.
    elim (map_to_elements Expression Expression (fun a : Expression =>
      apply_elem_tree (mksubst i j (tree_name_to_tree_Name v)) (apply_elem_tree (mksubst si sj sexpr) a))
    (apply_elem_tree (mksubst si sj (apply_elem_tree (mksubst i j (tree_name_to_tree_Name v)) sexpr))) l).
    intros H1 H2; apply H1; clear H1 H2.
    intros a a_in_l; apply IHe; trivial.
    unfold var_not_in_tree; simpl.
    elim (fold_bool_or_false (f:= fun (e : Expression) => var_in_tree_bool e i j) l).
    simpl; intros H1 H2; apply H1; trivial.
    (* H prooved *)
    rewrite H; trivial.
  Qed.

  Lemma apply_applied : forall (e : Expression) (i j : nat) (v : value) (rho : Substitution),
    apply rho (Node (x i j) nil) = tree_name_to_tree_Name v ->
      apply rho (apply_elem_tree (mksubst i j (tree_name_to_tree_Name v)) e) = apply rho e.
  Proof.
    intros e i j v rho Hrho.
    generalize e; clear e.
    induction e as [l | n l IHe].
    (* case e is a leaf : *)
    simpl; trivial.
    (* case e = Node n l : *)
    simpl.
    destruct n as [ni nj | n].
      (* case n variable : *)
    destruct l as [ | h t].
        (* case l = nil : *)
    elim (eq_nat_dec i ni); elim (eq_nat_dec j nj); intros H1 H2; simpl; trivial.
          (* remaining case : i = ni, j = nj : *)
    subst; rewrite Hrho; apply apply_unchanged.
        (* case l = h :: t : *)
    trivial.
      (* case n name : *)
    rewrite apply_to_forest.
    rewrite apply_to_forest.
    assert (H : (map (apply rho) (map (apply_elem_tree (mksubst i j (tree_name_to_tree_Name v))) l)) =
      map (apply rho) l).
    (* Proof of H : *)
    rewrite compose_map.
    elim (map_to_elements Expression Expression (fun a : Expression =>
      apply rho (apply_elem_tree (mksubst i j (tree_name_to_tree_Name v)) a))
    (apply rho) l).
    simpl; intros H1 H2; apply H1; clear H1 H2.
    apply IHe.
    (* H proved *)
    rewrite H; trivial.
  Qed.

  Lemma apply_elem_commut : forall (e : Expression) (i j k l : nat) (t v : value),
    (i<>k \/ j <> l \/ t = v) ->
    apply ((mksubst i j (tree_name_to_tree_Name t))::(mksubst k l (tree_name_to_tree_Name v))::nil) e =
      apply ((mksubst k l (tree_name_to_tree_Name v))::(mksubst i j (tree_name_to_tree_Name t))::nil) e.
  Proof.
    induction e as [l | n l IHe].
    (* case e leaf : *)
    simpl; trivial.
    (* case e = Node n l : *)
    intros i j k m t v Heqneq.
    simpl.
    destruct n as [ni nj | n].
      (* case n variable : *)
    elim (eq_nat_dec k ni); intro case_k_ni.
        (* case k = ni : *)
    elim (eq_nat_dec m nj); intro case_m_nj.
          (* case m = nj : *)
    elim (eq_nat_dec i ni); intro case_i_ni.
            (* case i = ni : *)
    elim (eq_nat_dec j nj); intro case_j_nj.
              (* case j = nj *)
    subst.
    elim Heqneq; clear Heqneq; intro Heqneq; elim Heqneq; trivial.
    intro H; elim H; trivial.
    intro H; subst; trivial.
              (* case j <> nj : *)
    destruct l as [ | h tl].
                (* case l = nil : *)
    subst; simpl.
    elim (eq_nat_dec ni ni); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec nj nj); intro H; [clear H | elim H; trivial].
    apply apply_elem_unchanged.
                (* case l <> nil : *)
    subst; simpl; trivial.
            (* case i <> ni : *)
    destruct l as [ | h tl].
              (* case l = nil : *)
    subst; simpl.
    elim (eq_nat_dec ni ni); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec nj nj); intro H; [clear H | elim H; trivial].
    apply apply_elem_unchanged.
            (* case l <> nil : *)
    subst; simpl; trivial.
          (* case m <> nj : *)
    elim (eq_nat_dec i ni); intro case_i_ni.
            (* case i = ni : *)
    elim (eq_nat_dec j nj); intro case_j_nj.
              (* case j = nj *)
    subst.
    destruct l as [ | h tl]; simpl.
                (* case l = nil : *)
    elim (eq_nat_dec ni ni); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec nj nj); intro H; [clear H | elim H; trivial].
    symmetry; apply apply_elem_unchanged.
                (* case l = h ::tl : *)
    trivial.
              (* case j <> nj : *)
    subst.
    destruct l as [ | h tl]; simpl.
                (* case l = nil : *)
    elim (eq_nat_dec ni ni); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec j nj); intro H; [elim case_j_nj; subst; trivial | clear H].
    elim (eq_nat_dec m nj); intro H; [elim case_m_nj; subst; trivial | clear H].
    trivial.
                 (* case l = h ::tl : *)
    trivial.
            (* case i <> ni : *)
    subst; destruct l as [ | h tl]; simpl.
              (* case l = nil : *)
    elim (eq_nat_dec ni ni); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec i ni); intro H; [elim case_i_ni; subst; trivial | clear H].
    elim (eq_nat_dec m nj); intro H; [elim case_m_nj; subst; trivial | clear H].
    trivial.
              (* case l = h ::tl : *)
    trivial.
        (* case k <> ni : *)
    destruct l as [ | h tl]; simpl; trivial.
    elim (eq_nat_dec i ni); intro case_i_ni.
          (* case i = ni : *)
    elim (eq_nat_dec j nj); intro case_j_nj.
            (* case j = nj : *)
    symmetry; apply apply_elem_unchanged.
            (* case j <> nj : *)
    simpl; subst.
    elim (eq_nat_dec k ni); intro H; [elim case_k_ni; subst | clear H]; trivial.
          (* case i <> ni : *)
    simpl; elim (eq_nat_dec k ni); intro H; [elim case_k_ni; subst | clear H]; trivial.

      (* case n name : *)
    rewrite apply_elem_to_forest.
    rewrite apply_elem_to_forest.
    assert (H : map (apply_elem_tree (mksubst i j (tree_name_to_tree_Name t)))
      (map (apply_elem_tree (mksubst k m (tree_name_to_tree_Name v))) l) =
      map (apply_elem_tree (mksubst k m (tree_name_to_tree_Name v)))
      (map (apply_elem_tree (mksubst i j (tree_name_to_tree_Name t))) l)).
    (* Proof of H : *)
    rewrite compose_map.
    rewrite compose_map.
    elim (map_to_elements Expression Expression (fun a : Expression =>
      apply_elem_tree (mksubst i j (tree_name_to_tree_Name t))
        (apply_elem_tree (mksubst k m (tree_name_to_tree_Name v)) a))
    (fun a : Expression =>
      apply_elem_tree (mksubst k m (tree_name_to_tree_Name v))
      (apply_elem_tree (mksubst i j (tree_name_to_tree_Name t)) a)) l).
    simpl; intros H1 H2; apply H1; clear H1 H2.
    intros a a_in_l.
    generalize (IHe a a_in_l i j k m t v Heqneq); simpl; trivial.
    (* H proved *)
    rewrite H; trivial.
  Qed.

  Lemma reverse_apply : forall (S : Substitution) (n p h : nat) (e : Expression) (l : list value),
    make_substitution (fresh p h n) l = Some S ->
    apply S e = apply (rev_lin S) e.
  Proof.
    induction S as [ | s S IHS].
    (* case S nil : *)
    simpl; trivial.
    (* case S = s :: S : *)
    simpl; intros n p h e l.
    destruct n as [ | n].
    (* case n = 0 : *)
    simpl; intro H.
    destruct l; inversion H.
    (* case n>0 : *)
    destruct l as [ | hl tl]; simpl.
    intro F; inversion F.
    CaseEq (make_substitution (fresh p (Datatypes.S h) n) tl).
    intros S' case_S'.
    intro H; inversion H; subst.
    clear H.
    rewrite rev_lin_is_rev; simpl.
    rewrite apply_substitution_concat.
    simpl.
    rewrite <- rev_lin_is_rev.
    rewrite <- (IHS n p (Datatypes.S h) (apply_elem_tree (mksubst p h (tree_name_to_tree_Name hl)) e) tl); trivial.
    cut (h < Datatypes.S h); [idtac | auto with arith].
    generalize n p h e tl (Datatypes.S h) case_S'; clear IHS case_S' n p h e tl.
    induction S as [ | s S' IHS].
    simpl; trivial.
    intros n p h e tl m; simpl.
    destruct n; simpl.
    destruct tl; intro H; inversion H.
    destruct tl.
    intro F; inversion F.
    CaseEq (make_substitution (fresh p (S m) n) tl).
    intros stl case_stl H; inversion H; subst.
    clear H; generalize (apply_elem_commut (apply S' e) p h p m hl t); simpl; intros H Hpn.
    rewrite H.
    rewrite (IHS n p h e tl (S m) case_stl).
    trivial.
    generalize Hpn; auto with arith.
    right; left; intro; subst.
    apply lt_irrefl with m; trivial.
    intros dd F; inversion F.
    intros dd F; inversion F.
  Qed.

  Lemma make_substitution_invariant : forall (m p h k : nat) (t : value), k < h ->
    map (apply_elem_tree (mksubst p k (tree_name_to_tree_Name t))) (fresh p h m) = fresh p h m.
  Proof.
    induction m; simpl; trivial.
    intros p h k t lt_k_h.
    elim (eq_nat_dec p p); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec k h); intro H; [cut False; [intro F; inversion F | subst; generalize lt_k_h; apply lt_irrefl] | clear H].
    rewrite IHm; trivial.
    generalize lt_k_h; auto with arith.
  Qed.

  Lemma make_substitution_simpl_in_branch_then_elem : forall (cnm : name) (n n0 m p h : nat) (svf rho : Substitution)
    (vf : list value) (e : Expression) (v : value) (f : list value),
    (forall (j : nat), var_not_in_tree e p (h+j)) ->
    tree_name_to_tree_Name (Node cnm (f ++ vf)) = apply rho (Node (x n n0) nil) ->
      make_substitution (fresh p h m) vf = Some svf -> apply rho e = tree_name_to_tree_Name v ->
        apply rho (apply svf (apply_elem_tree (mksubst n n0 (Node (symbol cnm) ((map tree_name_to_tree_Name f)++(fresh p h m)))) e)) = tree_name_to_tree_Name v.
  Proof.
    intros cnm n n0 m p h svf rho vf e v f Hnotin Hrho Hsvf Hv.
    rewrite (reverse_apply svf m p h (apply_elem_tree (mksubst n n0 (Node (symbol cnm) ((map tree_name_to_tree_Name f)++(fresh p h m)))) e) vf).
    generalize cnm n n0 m p h rho vf e v f Hnotin Hrho Hsvf Hv; clear Hnotin Hrho Hsvf Hv cnm n n0 m p h rho vf e v f.
    induction svf; simpl; intros cnm n n0 m p h rho vf e v f Hnotin Hrho Hsvf Hv.
    cut (fresh p h m = nil /\ vf = nil).
    intro H; elim H; clear H; intros h1 h2.
    rewrite h1.
    replace (Node (symbol cnm) ((map tree_name_to_tree_Name f) ++ nil)) with (tree_name_to_tree_Name (Node cnm f)).
    rewrite apply_applied.
    assumption.
    rewrite <- Hrho; rewrite h2; simpl; trivial.
    rewrite list_concat_nil; trivial.
    simpl; rewrite list_concat_nil; trivial.
    destruct vf.
    destruct m; simpl in Hsvf.
    simpl; split; trivial.
    inversion Hsvf.
    destruct m; simpl in Hsvf.
    inversion Hsvf.
    destruct (make_substitution (fresh p (S h) m) vf); inversion Hsvf.

    rewrite rev_lin_is_rev; simpl.
    rewrite apply_substitution_concat; simpl.
    destruct vf.
    destruct m; simpl; inversion Hsvf.
    destruct m; simpl; inversion Hsvf.
    generalize H0; clear H0; CaseEq (make_substitution (fresh p (S h) m) vf); [intros l Hl H0 | intros Hl H0; inversion H0].
    inversion H0; subst; clear H0.
    rewrite apply_elem_in_subst.
    simpl.
    rewrite map_app; simpl.
    elim (eq_nat_dec p p); intro H; [clear H | elim H; trivial].
    elim (eq_nat_dec h h); intro H; [clear H | elim H; trivial].
    replace (map (apply_elem_tree (mksubst p h (tree_name_to_tree_Name t))) (fresh p (S h) m)) with (fresh p (S h) m).
    replace (map (apply_elem_tree (mksubst p h (tree_name_to_tree_Name t))) (map tree_name_to_tree_Name f)) with (map tree_name_to_tree_Name f).
    replace ((map tree_name_to_tree_Name f)++((tree_name_to_tree_Name t)::(fresh p (S h) m))) with ((map tree_name_to_tree_Name (f++(t::nil)))++(fresh p (S h) m)).
    rewrite <- rev_lin_is_rev.
    apply (IHsvf cnm n n0 m p (S h) rho vf e v); trivial.
    intro j; replace (S h + j) with (h + (S j)).
    apply (Hnotin (S j)).
    rewrite plus_comm; simpl; rewrite plus_comm; trivial.
    rewrite <- Hrho; rewrite <- list_concat_cons; simpl; trivial.
    rewrite map_app; simpl; rewrite <- list_concat_cons; trivial.
    symmetry; apply apply_elem_forest_unchanged.
    symmetry; apply make_substitution_invariant.
    apply (lt_n_Sn h).
    generalize (Hnotin 0).
    rewrite plus_comm; simpl; trivial.
    assumption.
  Qed.

  Lemma apply_elem_pattern : forall (e : Expression) (s : subst_elem), tree_is_pattern (apply_elem_tree s e) -> tree_is_pattern e.
  Proof.
    induction e as [n | n l IHe].
    (* case e leaf : *)
    simpl; trivial.
    (* case e node : *)
    simpl; intro s; destruct n as [ni nj | n].
    (* case n variable : *)
    destruct l as [ | h t].
      (* case l empty : *)
    intro H; clear H; unfold tree_is_pattern; simpl; trivial.
      (* case l not empty : *)
    trivial.
    (* case n name : *)
    unfold tree_is_pattern; simpl.
    destruct (Get_function n); simpl.
    trivial.
    intro H.
    elim (fold_bool_and_true (f := tree_is_pattern_bool) l); intros H1 H2.
    apply H2; clear H1 H2.
    intros a a_in_l; fold tree_is_pattern_bool; apply (IHe a a_in_l s).
    elim (fold_bool_and_true (f:=tree_is_pattern_bool) (map (apply_elem_tree s) l)); intros H1 H2.
    unfold tree_is_pattern; apply (H1 H (apply_elem_tree s a)).
    clear H1 H2.
    apply in_map; trivial.
  Qed.

  Lemma map_apply_to_element : forall (l : list Expression) (vf : list value) (rho : Substitution) (e : Expression) (v : value) (j : nat),
    map tree_name_to_tree_Name vf = (map (apply rho) l) -> element_at_list (rev_lin l) j = Some e ->
    (rev_lin vf)[j] = Some v -> tree_name_to_tree_Name v = apply rho e.
  Proof.
    induction l; simpl; intros vf rho e v j Heq He Hv.
    inversion He.
    cut (j < length (rev_lin l ++ a :: nil)); [intro Cj | generalize He; rewrite rev_lin_is_rev; rewrite rev_lin_is_rev; apply element_at_list_some].
    rewrite concat_length in Cj; rewrite plus_comm in Cj; simpl in Cj.
    rewrite rev_lin_is_rev in Cj; compare j (length (rev l)); intro comp.
    subst.
    replace (length (rev l)) with (length (rev l) + 0) in He.
    rewrite rev_lin_is_rev in He; simpl in He; rewrite element_at_list_concat in He.
    inversion He; subst; clear He.
    rewrite rev_lin_is_rev in Hv; destruct vf.
    simpl in Hv; inversion Hv.
    simpl in Heq.
    inversion Heq; trivial.
    simpl in Hv.
    replace (length (rev l)) with (length (rev vf) + 0) in Hv.
    rewrite element_at_list_concat in Hv.
    inversion Hv; trivial.
    rewrite plus_comm; simpl.
    rewrite rev_length; rewrite rev_length.
    rewrite (lists.map_length _ _ (apply rho) l).
    rewrite <- H1; apply lists.map_length.
    rewrite plus_comm; simpl; trivial.
    cut (j < length (rev l)).
    clear Cj; intro Cj.
    rewrite rev_lin_is_rev in He; simpl in He; rewrite <- element_at_list_concat' in He; trivial.
    destruct vf.
    simpl in Hv; inversion Hv.
    simpl in Heq; inversion Heq.
    rewrite rev_lin_is_rev in Hv; simpl in Hv; rewrite <- element_at_list_concat' in Hv.
    apply (IHl vf rho e v j H1); rewrite rev_lin_is_rev; trivial.
    replace (length (rev vf)) with (length (rev l)); trivial.
    rewrite rev_length; rewrite rev_length.
    rewrite (lists.map_length _ _ (apply rho) l).
    rewrite <- H1; symmetry; apply lists.map_length.
    generalize Cj comp; apply lt_S_neq_lt.
  Qed.

  Lemma In_fresh_not_var_in : forall (m p h h' : nat) (e : Expression),
    In e (fresh p h m) -> var_not_in_tree e p (h' + h + m).
  Proof.
    induction m; simpl; intros p h h' e Hin.
    inversion Hin.
    elim Hin; clear Hin; intro Hin.
    subst e; simpl.
    unfold var_not_in_tree; simpl.
    elim (eq_nat_dec p p); intro H; simpl; [clear H | elim H; trivial].
    elim (eq_nat_dec (h' + h + S m) h); intro case.
    cut False; [intro F; inversion F | generalize case; clear case].
    rewrite plus_comm.
    rewrite plus_assoc.
    rewrite plus_comm.
    cut (0 < S m + h').
    generalize (S m + h'); clear IHm h' p m; induction h.
    simpl; intros n H H'; subst; inversion H.
    simpl; intros n H H'.
    apply (IHh n H).
    generalize H'; auto with arith.
    (* Proof of 0 < S m + h' : *)
    simpl; auto with arith.
    trivial.
    replace  (h' + h + S m) with (h' + S h + m).
    apply IHm; trivial.
    rewrite (plus_comm h' (S h)); simpl.
    rewrite (plus_comm (h' + h) (S m)); simpl.
    replace (h + h' + m) with (m + (h' + h)); trivial.
    rewrite plus_comm; simpl.
    rewrite (plus_comm h h'); trivial.
  Qed.

  Lemma In_fresh : forall (e : Expression) (m p h h' p': nat),
    p' > p -> In e (rev_lin (fresh p h m)) -> var_not_in_tree e p' h'.
  Proof.
    induction m; intros p h h' p' Cp' Hin.
    simpl in Hin; inversion Hin.
    rewrite rev_lin_is_rev in Hin; simpl in Hin.
    cut (In e (rev (fresh p (S h) m)) \/ In e (Node (x p h) nil :: nil)).
    clear Hin; intro Hin; inversion Hin.
    (* case e in fresh p (S h) m : *)
    apply (IHm p (S h)); trivial.
    rewrite rev_lin_is_rev; trivial.
    (* case e in (Node (x p h) nil :: nil) : *)
    inversion H.
    subst e; simpl.
    unfold var_not_in_tree; simpl.
    elim (eq_nat_dec p' p); intro case_p'_p; simpl.
      (* case p = p' : *)
    subst p'; elim (gt_irrefl p); trivial.
    trivial.
    inversion H0.
    generalize Hin; apply In_concat.
  Qed.

(*
  Lemma var_in_fresh : forall (m p h p' h': nat), var_in_tree (fresh p h m) p' h' -> p' = p.
  Proof.
    induction m; simpl; intros p h p' h' Hin.
    inversion Hin.
    elim Hin; clear Hin; intro Hin.
    symmetry; elim Hin; trivial.
    generalize Hin; apply IHm; trivial.
  Qed.

  Lemma var_in_forest_concat : forall (f f' : forest Name) (p h : nat),
    var_in_forest (forest_concat f f') p h -> var_in_forest f p h \/ var_in_forest f' p h.
  Proof.
    induction f; simpl; intros f' p h.
    unfold forest_concat; simpl.
    rewrite tree_list_to_forest_forest_to_tree_list.
    right; trivial.
    intro H; elim H; clear H; intro H.
    left; left; trivial.
    cut (var_in_forest f p h \/ var_in_forest f' p h).
    intro H2; elim H2; clear H2; intro H2.
    left; right; assumption.
    right; assumption.
    apply IHf.
    apply H.
  Qed.
*)

  Lemma var_in_apply_elem_tree_aux : forall (e e' : Expression) (n n0 p h : nat),
    var_in_tree_bool (apply_elem_tree (mksubst n n0 e') e) p h = true -> var_in_tree_bool e p h = true \/ var_in_tree_bool e' p h = true.
  Proof.
    induction e as [n | n l IHe]; simpl.
    (* case e leaf : *)
    left; trivial.
    (* case e node : *)
    destruct n as [ni nj | n].
      (* case n variable : *)
    intros e' n n0 p h.
    destruct l; simpl.
    elim (eq_nat_dec n ni); intro case_n_ni.
    elim (eq_nat_dec n0 nj); intro case_n0nj.
    subst; right; assumption.
    simpl; left; trivial.
    simpl; left; trivial.
    simpl; left; trivial.
      (* case n name : *)
    simpl; intros e' n0 n1 p h H.
    elim (fold_bool_or_true (f:=fun (e : Expression) => var_in_tree_bool e p h) (map (apply_elem_tree (mksubst n0 n1 e')) l)).
    simpl; intros H1 H2; elim (H1 H).
    intros y Hy; elim Hy; clear Hy; intros Hy1 Hy2.
    elim (In_mapped _ _ _ _ _ Hy1).
    intros x' Hx'; elim Hx'; clear Hx'; intros Hx'1 Hx'2.
    assert (H3 : var_in_tree_bool (apply_elem_tree (mksubst n0 n1 e') x') p h = true).
    rewrite <- Hx'2; trivial.
    elim (IHe x' Hx'1 e' n0 n1 p h H3); intro caseIHe; [left | right; trivial].
    elim (fold_bool_or_true (f:=fun (e : Expression) => var_in_tree_bool e p h) l).
    clear H1 H2; intros H1 H2; apply H2.
    exists x'; split; trivial.
  Qed.

  Lemma var_in_apply_elem_tree_aux2 : forall (e e' : Expression) (n n0 p h : nat),
    var_not_in_tree e p h -> var_not_in_tree e' p h -> var_not_in_tree (apply_elem_tree (mksubst n n0 e') e) p h.
  Proof.
    unfold var_not_in_tree; intros e e' n n0 p h Hnot_in_e Hnot_in_e'.
    CaseEq (var_in_tree_bool (apply_elem_tree (mksubst n n0 e') e) p h); intro H; trivial.
    elim (var_in_apply_elem_tree_aux _ _ _ _ _ _ H); clear H; intro H.
    rewrite H in Hnot_in_e; inversion Hnot_in_e.
    rewrite H in Hnot_in_e'; inversion Hnot_in_e'.
  Qed.

End substitutions.