Require Import POCS.
Require Import OneDiskAPI.
Require Import BadBlockAPI.

Module RemappedDisk (bd : BadBlockAPI) <: OneDiskAPI.

  Import ListNotations.

  Definition read (a : addr) : proc block :=
    bs <- bd.getBadBlock;
    if a == bs then
      len <- bd.size;
      r <- bd.read (len-1);
      Ret r
    else
      r <- bd.read a;
    Ret r.

  Definition write (a : addr) (b : block) : proc unit :=
    Ret tt.

  Definition size : proc nat :=
    len <- bd.size;
    Ret (len - 1).

  Definition init' : proc InitResult :=
    len <- bd.size;
    if len == 0 then
      Ret InitFailed
    else
      bs <- bd.getBadBlock;
      if (lt_dec bs len) then
        Ret Initialized
      else
        Ret InitFailed.

  Definition init := then_init bd.init init'.

  Definition recover: proc unit :=
    bd.recover.

  Inductive remapped_abstraction (bs_state : BadBlockAPI.State) (rd_disk : OneDiskAPI.State) : Prop :=
    | RemappedAbstraction :
      let bs_disk := stateDisk bs_state in
      let bs_addr := stateBadBlock bs_state in
      forall
        (* Fill in the rest of your abstraction here. *)
        (* To refer to the contents of disk [d] at address [a], you can write [diskGet a] *)
        (* To refer to the size of disk [d], you can write [diskSize d] *)
        (* You can try to prove [read_ok] to discover what you need from this abstraction *)

        (* Hint 1: What should be true about the non-bad blocks?   Replace [True] with what needs to be true *)
        (Hgoodblocks : True)
        (* Hint 2: What should be true about the bad block? *)
        (Hbadblock : True)

        (* when writing the above,  *)
        (* Hint 3: What if the bad block address is the last address? *)
        (Hsize : diskSize bs_disk = diskSize rd_disk + 1),
      remapped_abstraction bs_state rd_disk.

  Hint Constructors remapped_abstraction : core.

  Definition abstr : Abstraction OneDiskAPI.State :=
    abstraction_compose bd.abstr {| abstraction := remapped_abstraction |}.

  (* Note that these examples aren't trivial to prove, even if you have the
     right abstraction relation. They should help you think about whether your
     abstraction relation makes sense (although you may need to modify it to
     actually prove the implementation correct). *)

  Example abst_1_ok : forall Hinbounds,
    remapped_abstraction (BadBlockAPI.mkState [block1;block0;block0] 0 Hinbounds) [block0;block0].
  Proof.
    constructor; auto.
    Admitted.

  Example abst_2_ok : forall Hinbounds,
    remapped_abstraction (BadBlockAPI.mkState [block1;block0] 0 Hinbounds) [block0].
  Proof.
    constructor; auto.
    Admitted.

  Example abst_3_ok : forall Hinbounds,
    remapped_abstraction (BadBlockAPI.mkState [block0;block0] 1 Hinbounds) [block0].
  Proof.
    constructor; auto.
    Admitted.

  Example abst_4_nok : forall Hinbounds,
    ~ remapped_abstraction (BadBlockAPI.mkState [block0;block0] 1 Hinbounds) [block1].
  Proof.
    intros.
    intro.
    inversion H; simpl in *.
    Admitted.

  Example abst_5_nok : forall Hinbounds,
    ~ remapped_abstraction (BadBlockAPI.mkState [block1;block1] 0 Hinbounds) [block0].
  Proof.
    intros.
    intro.
    inversion H; simpl in *.
    Admitted.

  (* Due to how remapped_abstraction is defined (as an inductive), it cannot be
  unfolded. This tactic identifies abstraction relations in the hypotheses and
  breaks them apart with [inversion], and also does some cleanup. *)
  Ltac invert_abstraction :=
    match goal with
    | H : remapped_abstraction _ _ |- _ => inversion H; clear H; subst; subst_var
    end.

  Theorem init_ok : init_abstraction init recover abstr inited_any.
  Proof.
    eapply then_init_compose; eauto.
    step_proc.
  Admitted.

  Theorem read_ok : forall a, proc_spec (OneDiskAPI.read_spec a) (read a) recover abstr.
  Proof.
    unfold read.
    intros.

    apply spec_abstraction_compose; simpl.
    step_proc; eauto.

    destruct (a == r).
    - admit.
    - admit.
  Admitted.

  Theorem bad_block_inbounds :
    forall state,
      stateBadBlock state < diskSize (stateDisk state).
  Proof.
    intros.
    destruct state.
    simpl.
    exact stateBadBlockInBounds.
  Qed.

  Theorem remapped_abstraction_diskUpd_remap : forall state state' s v,
    remapped_abstraction state s ->
    stateBadBlock state' = stateBadBlock state ->
    stateDisk state' = diskUpd (stateDisk state) (diskSize (stateDisk state) - 1) v ->
    remapped_abstraction state' (diskUpd s (stateBadBlock state) v).
  Proof.
    (* Hint: you may find it useful to use the [pose proof (bad_block_inbounds state)]
       if you need [lia] to make use of the fact that the bad block is in-bounds. *)
    Admitted.

  Theorem remapped_abstraction_diskUpd_noremap : forall state state' s a v,
    remapped_abstraction state s ->
    a <> diskSize (stateDisk state) - 1 ->
    a <> stateBadBlock state ->
    stateBadBlock state' = stateBadBlock state ->
    stateDisk state' = diskUpd (stateDisk state) a v ->
    remapped_abstraction state' (diskUpd s a v).
  Proof.
    Admitted.

  Hint Resolve remapped_abstraction_diskUpd_remap : core.
  Hint Resolve remapped_abstraction_diskUpd_noremap : core.

  Theorem write_ok : forall a v, proc_spec (OneDiskAPI.write_spec a v) (write a v) recover abstr.
  Proof.
    unfold write.
    intros.

    apply spec_abstraction_compose; simpl.
    Admitted.

  Theorem size_ok : proc_spec OneDiskAPI.size_spec size recover abstr.
  Proof.
    unfold diskSize.
    intros.

    apply spec_abstraction_compose; simpl.

    step_proc; eauto.
    Admitted.

  (* This proof proves that recovery corresponds to no_wipe. *)
  Theorem recover_wipe : rec_wipe recover abstr no_wipe.
  Proof.
    unfold rec_wipe.
    intros.

    apply spec_abstraction_compose; simpl.
    step_proc; eauto.
  Qed.

End RemappedDisk.

Require Import BadBlockImpl.
Module x := RemappedDisk BadBlockDisk.
Print Assumptions x.write_ok.


(**
This file is part of the Elfic library

Copyright (C) Boldo, Clément, Leclerc

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 3 of the License, or (at your option) any later version.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
COPYING file for more details.
*)

From Coq Require Import
    ssreflect
    ssrsearch
    Arith
    Lia
    Utf8

    ClassicalDescription

    Rdefinitions
    Rpower
    RIneq
.

From Coquelicot Require Import
    Hierarchy
    Rbar
.

Require Import
    measure
    sigma_algebra
    sum_Rbar_nonneg
    Rbar_compl
    measurable_fun
    LInt_p
    R_compl
.

Require Import
    square_bij
    hierarchy_notations
    Rmax_n
    topology_compl
.

Declare Scope sf_scope.
Delimit Scope sf_scope with sf.

Declare Scope fun_scope.
Delimit Scope fun_scope with fn.

Open Scope hy_scope.
Open Scope fun_scope.
Open Scope sf_scope.

Section simpl_fun_def.

    (* espace de départ *)
    Context {X : Set}.
    (* espace d'arrivé *)
    Context {A : AbsRing}.
    Context {E : ModuleSpace A}.

    Definition indic (P : X -> Prop) : X -> A :=
        fun x =>
            match (excluded_middle_informative (P x)) return A with
                | left _ => one
                | right _ => zero
            end.

    Open Scope core_scope.
    Open Scope type_scope.
    Open Scope nat_scope.

    (* Un espace mesuré *)
    Context {gen : (X -> Prop) -> Prop}.
    Context (μ : measure gen).

    (* la fonction est val ∘ which *)
    Record simpl_fun := mk_simpl_fun {
        which : X -> nat;
        val : nat -> E;
        max_which : nat;

        ax_val_max_which :
            val max_which = zero;
        ax_which_max_which :
            ∀ x : X, which x <= max_which;
        ax_measurable :
            ∀ n : nat, n <= max_which ->
                measurable gen (fun x => which x = n);
    }.

    Definition integrable_sf (sf : simpl_fun) :=
        ∀ n : nat, n < max_which sf ->
            is_finite (μ (fun x => which sf x = n)).

    Definition nth_carrier (sf : simpl_fun) (n : nat) : (X -> Prop) :=
        (fun x => sf.(which) x = n).

    Definition fun_sf (sf : simpl_fun) : X -> E :=
        fun x => sf.(val) (sf.(which) x).

    (* On se donne la possibilité de calculer la "valeur" d'une simpl_fun en l'évaluant en x : *)
    Coercion fun_sf : simpl_fun >-> Funclass.

    (* Un exemple d'utilisation est fait dans cete définition *)
    Definition is_simpl (f : X -> E) :=
        ∃ sf : simpl_fun, ∀ x : X, sf x = f x.

End simpl_fun_def.

Arguments simpl_fun {X A} E gen.
Arguments is_simpl {X A} [E] gen f.
Arguments integrable_sf {X A} [E] {gen} μ sf.

Notation "'χ(' P ')'" := (indic P) (at level 0) : fun_scope.

Section simpl_fun_indic.

    (* espace de départ *)
    Context  {X : Set}.
    (* Un espace mesuré *)
    Context (gen : (X -> Prop) -> Prop).

    Open Scope nat_scope.
    Open Scope R_scope.

    Definition sf_indic_aux (P : X -> Prop) :
        measurable gen P -> simpl_fun R_ModuleSpace gen.
    (* définition *)
        move => Pmeas.
        pose w := fun x =>
            match (excluded_middle_informative (P x)) return nat with
                | left _ => O
                | right _ => (S O)
            end.
        pose v := fun n =>
            match n return R with
                | O => 1
                | _ => 0
            end.
        pose max_w := (S O).
        apply (mk_simpl_fun w v max_w).
            unfold max_w => //.
            move => x; unfold w, max_w.
            case: (excluded_middle_informative (P x)); lia.
            unfold max_w; case.
                move => _.
                apply measurable_ext with P.
                unfold w => x.
                case: (excluded_middle_informative (P x)) => //.
                exact Pmeas.
            move => n Hn.
            assert (S n = 1%nat) by lia.
            rewrite H; clear Hn H; clear n.
            apply measurable_ext with (fun x => ¬ P x).
            unfold w => x.
            case: (excluded_middle_informative (P x)) => //.
            apply measurable_compl.
            apply measurable_ext with (fun x => P x).
            move => x; split; case: (excluded_middle_informative (P x)).
                move => _; auto.
                auto.
                auto.
                move => NPx NNPx; apply False_ind.
                exact (NNPx NPx).
            exact Pmeas.
    Defined.

    Lemma sf_indic :
        ∀ P : X -> Prop, measurable gen P
            -> is_simpl gen (χ(P): X -> R).
    Proof.
        move => P Pmeas.
        exists (sf_indic_aux P Pmeas) => x.
        unfold fun_sf, "χ( _ )" => /=.
        case: (excluded_middle_informative (P x)) => //.
    Qed.

    Context (μ : measure gen).

    Lemma integrable_sf_indic (P : X -> Prop) (π : measurable gen P) :
        is_finite (μ P) -> integrable_sf μ (sf_indic_aux P π).
    Proof.
        move => Pfin n Hn; simpl in *.
            assert (n = O) by lia.
            rewrite H; clear H Hn; clear n.
            rewrite <-(measure_ext gen _ P).
            exact Pfin.
            move => x.
            case (excluded_middle_informative (P x)) => //.
    Qed.

End simpl_fun_indic.

Arguments integrable_sf_indic {X gen μ P} π.

Section sf_zero_fun.

    (* espace de départ *)
    Context  {X : Set}.
    (* Un espace mesuré *)
    Context (gen : (X -> Prop) -> Prop).

    Context {A : AbsRing}.
    Context {E : NormedModule A}.

    Definition sf_zero : simpl_fun E gen.
        pose whichz := fun _ : X => O.
        pose valz := fun _ : nat => zero : E.
        apply (mk_simpl_fun whichz valz 0).
        unfold valz => //.
        unfold whichz => //.
        move => n /(Nat.le_0_r) ->.
        unfold whichz; apply measurable_Prop.
    Defined.

    Lemma sf_zero_zero : ∀ x : X, sf_zero x = zero.
    Proof.
       by [].
    Qed.

    Context (μ : measure gen).

    Lemma integrable_sf_zero :
        integrable_sf μ sf_zero.
    Proof.
        unfold integrable_sf => /=.
        lia.
    Qed.

End sf_zero_fun.

Section simpl_fun_norm.

    (* espace de départ *)
    Context  {X : Set}.
    (* espace d'arrivé *)
    Context {A : AbsRing}.
    Context {E : NormedModule A}.
    (* Un espace mesuré *)
    Context {gen : (X -> Prop) -> Prop}.

    Open Scope R_scope.
    Open Scope nat_scope.

    Definition fun_norm (f : X -> E) :=
        fun x => norm (f x).

    Notation "‖ f ‖" := (fun_norm f) (at level 100) : fun_scope.

    Definition sf_norm_aux (sf : simpl_fun E gen) : simpl_fun R_ModuleSpace gen.
        case: sf => which val max_which ax1 ax2 ax3.
        pose nval :=
            fun n => norm (val n).
        apply (mk_simpl_fun which nval max_which).
            (* ax_val_max_which *)
            unfold nval; rewrite ax1.
            apply norm_zero; assumption.
            (* ax_which_max_which *)
            exact ax2.
            (* ax_measurable *)
            exact ax3.
    Defined.

    Context {μ : measure gen}.

    Lemma integrable_sf_norm {sf : simpl_fun E gen} (isf : integrable_sf μ sf) :
        integrable_sf μ (sf_norm_aux sf).
    Proof.
        unfold integrable_sf in *.
        case_eq sf => wf vf mawf axf1 axf2 axf3 Eqf => /=.
        rewrite Eqf in isf; simpl in isf => //.
    Qed.

    Notation "‖ sf ‖" := (sf_norm_aux sf) (at level 100) : sf_scope.

    Lemma sf_norm :
        ∀ f : X -> E, is_simpl gen f ->
            is_simpl gen (fun_norm f).
    Proof.
        move => f.
        case => sf. case_eq sf => which val max_which ax1 ax2 ax3 Eqsf Eqf.
        exists (sf_norm_aux sf).
        rewrite Eqsf.
        move => x; unfold fun_sf, sf_norm_aux => /=.
        simpl in Eqf.
        rewrite Eqf => //.
    Qed.

    Lemma fun_sf_norm :
        ∀ sf : simpl_fun E gen,
            (∀ x : X, fun_sf (‖sf‖) x = (‖ fun_sf sf x ‖)%hy).
    Proof.
        move => sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        unfold fun_sf.
        rewrite Eqf => //.
    Qed.

End simpl_fun_norm.

Notation "‖ f ‖" := (fun_norm f) (at level 100) : fun_scope.
Notation "‖ sf ‖" := (sf_norm_aux sf) (at level 100) : sf_scope.

Section simpl_fun_power.

    (* espace de départ *)
    Context  {X : Set}.
    (* Un espace mesuré *)
    Context {gen : (X -> Prop) -> Prop}.

    Definition fun_power (f : X -> R_NormedModule) (p : posreal) :=
        (fun x => Rpow (f x) p.(pos)).

    Definition sf_power_aux (sf : simpl_fun R_NormedModule gen) (p : RIneq.posreal) : simpl_fun R_NormedModule gen.
    case: sf => which val max_which ax1 ax2 ax3.
    pose nval :=
        fun n => Rpow (val n) p.(pos).
    apply (mk_simpl_fun which nval max_which).
        (* ax_val_max_which *)
        unfold nval; rewrite ax1.
        rewrite Rpow_0' => //.
        case p => //.
        (* ax_which_max_which *)
        exact ax2.
        (* ax_measurable *)
        exact ax3.
    Defined.

    Notation "sf ^ p" := (sf_power_aux sf p).

    Lemma fun_sf_power :
        ∀ sf : simpl_fun R_NormedModule gen, ∀ p : RIneq.posreal,
            (∀ x : X, fun_sf (sf ^ p) x = (Rpow (fun_sf sf x) p)).
    Proof.
        move => sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        unfold fun_sf.
        rewrite Eqf => //.
    Qed.

    Context {μ : measure gen}.

    Lemma integrable_sf_power {sf : simpl_fun R_NormedModule gen} (p : RIneq.posreal) (isf : integrable_sf μ sf) :
        integrable_sf μ (sf ^ p).
    Proof.
        unfold integrable_sf in *.
        case_eq sf => wf vf mawf axf1 axf2 axf3 Eqf => /=.
        rewrite Eqf in isf; simpl in isf => //.
    Qed.

End simpl_fun_power.

Notation "f ^ p" := (fun_power f p) : fun_scope.
Notation "sf '^' p" := (sf_power_aux sf p) : sf_scope.

Section simpl_fun_plus.

    (* espace de départ *)
    Context  {X : Set}.
    (* espace d'arrivé *)
    Context {A : AbsRing}.
    Context {E : NormedModule A}.
    (* Un espace mesuré *)
    Context {gen : (X -> Prop) -> Prop}.

    Open Scope nat_scope.

    Definition fun_plus (f : X -> E) (g : X -> E) :=
        (fun x => plus (f x) (g x)).

    Notation "f + g" := (fun_plus f g) (left associativity, at level 50) : fun_scope.

    Definition sf_plus_aux (sf sg : simpl_fun E gen) : simpl_fun E gen.
        case: sf => wf vf maxf axf1 axf2 axf3.
        case: sg => wg vg maxg axg1 axg2 axg3.
        pose val := fun m =>
            let (nf, ng) := square_bij_inv (S maxg) m in
            plus (vf nf) (vg ng).
        pose max_which := (S maxf) * (S maxg) - 1.
        pose which := fun (x : X) =>
            let nf := wf x in
            let ng := wg x in
            square_bij (S maxg) (nf, ng).
        apply (mk_simpl_fun which val max_which).
            (* ax_val_max_which *)
            unfold val, max_which.
            rewrite (square_bij_inv_corner maxg maxf).
            rewrite axf1 axg1 plus_zero_r; reflexivity.
            (* ax_which_max_which *)
            move => x; unfold which, max_which.
            apply confined_square.
            split; apply [axf2, axg2].
            (* ax_measurable *)
            assert
                (∀ n : nat, n <= max_which ->
                ∀ c : nat * nat, c = square_bij_inv (S maxg) n ->
                ∀ nf ng : nat, (nf, ng) = c ->
                    measurable gen (λ x : X, wf x = nf ∧ wg x = ng)
                ) as measurable_inter_fg.
                move => n Hn c Eqc nf ng Hnfngc.
                pose Hnfng := confined_square_inv maxg maxf n Hn; clearbody Hnfng => /=.
                rewrite <-Eqc, <-Hnfngc in Hnfng.
                case: Hnfng => Hnf Hng.
                apply measurable_inter.
                apply axf3 => //.
                apply axg3 => //.
            (**)
            move => n Hn; unfold which.
            pose c := square_bij_inv (S maxg) n.
            case_eq c => [nf ng] Eqc.
            apply measurable_ext with (fun x => wf x = nf ∧ wg x = ng).
            move => x.
            assert (n = square_bij (S maxg) c) as Eqn.
                rewrite is_bij_square_inv => //.
            rewrite Eqn Eqc.
            split.
                case => -> -> //.
                move => Eqwn.
                assert (square_bij_inv (S maxg) (square_bij (S maxg) (wf x, wg x)) = square_bij_inv (S maxg) (square_bij (S maxg) (nf, ng)))
                    by congruence.
                rewrite is_bij_square in H.
                    2 : apply axg2.
                rewrite is_bij_square in H.
                    all : swap 1 2.
                    pose Hng := confined_snd_square_inv maxg n.
                    fold c in Hng; clearbody Hng; rewrite Eqc in Hng.
                    exact Hng.
                split; congruence.
                apply measurable_inter_fg with n c => //.
    Defined.

    Notation "sf + sg" := (sf_plus_aux sf sg) (left associativity, at level 50) : sf_scope.

    Lemma sf_plus :
        ∀ f g : X -> E,
        is_simpl gen f -> is_simpl gen g ->
        is_simpl gen (fun_plus f g).
    Proof.
        move => f g.
        case => sf Eq_sf_f; case => sg Eq_sg_g.
        exists (sf_plus_aux sf sg).
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf.
        case_eq sg => wg vg maxg axg1 axg2 axg3 Eqg.
        unfold fun_sf => /= x.
        rewrite is_bij_square.
            2 : apply axg2.
        unfold fun_plus.
        congr plus.
            unfold fun_sf in Eq_sf_f.
            rewrite Eqf in Eq_sf_f; simpl in Eq_sf_f.
            rewrite Eq_sf_f => //.
            unfold fun_sf in Eq_sg_g.
            rewrite Eqg in Eq_sg_g; simpl in Eq_sg_g.
            rewrite Eq_sg_g => //.
    Qed.

    Lemma fun_sf_plus :
        ∀ sf sg : simpl_fun E gen,
            (∀ x : X, fun_sf (sf + sg)%sf x = (fun_sf sf x + fun_sf sg x)%hy).
    Proof.
        move => sf sg.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        case_eq sg => wg vg maxg axg1 axg2 axg3 Eqg; rewrite <-Eqg => /=.
        unfold fun_sf.
        rewrite Eqf Eqg => x /=.
        rewrite is_bij_square => //.
    Qed.

    Context {μ : measure gen}.

    Lemma integrable_sf_plus {sf sg : simpl_fun E gen} :
        (integrable_sf μ sf) -> (integrable_sf μ sg) -> (integrable_sf μ (sf_plus_aux sf sg)).
    Proof.
        unfold integrable_sf.
        move => axf4 axg4.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf.
        case_eq sg => wg vg maxg axg1 axg2 axg3 Eqg => /=.
        rewrite Eqf Eqg in axf4 axg4; simpl in axf4, axg4.
        pose max_which := (S maxf) * (S maxg) - 1.
        assert
            (∀ n : nat, n <= max_which ->
            ∀ c : nat * nat, c = square_bij_inv (S maxg) n ->
            ∀ nf ng : nat, (nf, ng) = c ->
                measurable gen (λ x : X, wf x = nf ∧ wg x = ng)
            ) as measurable_inter_fg.
            move => n Hn c Eqc nf ng Hnfngc.
            pose Hnfng := confined_square_inv maxg maxf n Hn; clearbody Hnfng => /=.
            rewrite <-Eqc, <-Hnfngc in Hnfng.
            case: Hnfng => Hnf Hng.
            apply measurable_inter.
            apply axf3 => //.
            apply axg3 => //.
    (**)
    move => n Hn; unfold which.
    pose c := square_bij_inv (S maxg) n.
    case_eq c => [nf ng] Eqc.
    assert (
        ∀ x : Rbar, Rbar_le (@zero R_AbelianGroup) x -> Rbar_lt x p_infty ->
            is_finite x
    ) as Rbar_pos_lt_finite.
        move => x; case: x => //.
    apply Rbar_pos_lt_finite.
    apply meas_ge_0.
    assert (n = square_bij (S maxg) c).
        unfold c; rewrite is_bij_square_inv => //.
    rewrite H Eqc.
    replace
        (μ (λ x : X, square_bij (S maxg) (wf x, wg x) = square_bij (S maxg) (nf, ng)))
        with
        (μ (λ x : X, wf x = nf ∧ wg x = ng)).
        all : swap 1 2.
        apply measure_ext => x; split.
            move => [-> ->] => //.
            move => Eqfg.
            assert
                (square_bij_inv (S maxg) (square_bij (S maxg) (wf x, wg x)) = square_bij_inv (S maxg) (square_bij (S maxg) (nf, ng)))
                as Eqfg2 by congruence.
            rewrite is_bij_square in Eqfg2.
                2 : apply axg2.
            rewrite is_bij_square in Eqfg2.
                all : swap 1 2.
                pose Hng := confined_snd_square_inv maxg n.
                fold c in Hng; clearbody Hng; rewrite Eqc in Hng.
                exact Hng.
            split; congruence.
        (* Ici il faut distinguer le cas ou
            on est dans la composante de 0 pour f ou pour g,
            sachant que les deux ne peuvent pas se produire simultanément *)
        assert
            (n <= max_which)
            as Le_n_mw.
            apply le_Sn_le => //.
        pose Hnfng := confined_square_inv maxg maxf n Le_n_mw; clearbody Hnfng.
        fold c in Hnfng; rewrite Eqc in Hnfng; case: Hnfng => Hnf Hng.
        assert
            (Rbar_le
                (μ (λ x : X, wf x = nf ∧ wg x = ng))
                (μ (λ x : X, wf x = nf))
            ) as inter_le_f.
            apply measure_le.
            apply measurable_inter_fg with n c => //.
            apply axf3 => //.
            easy.
        assert
        (Rbar_le
            (μ (λ x : X, wf x = nf ∧ wg x = ng))
            (μ (λ x : X, wg x = ng))
        ) as inter_le_g.
            apply measure_le.
            apply measurable_inter_fg with n c => //.
            apply axg3 => //.
            easy.
        case: (le_lt_v_eq nf maxf) => // Hnf'.
            (* cas ou le dommaine pour f est mesurable *)
            assert
                (is_finite (μ (λ x : X, wf x = nf)))
                as fin_f.

                apply axf4 => //.
            assert
                (Rbar_lt (μ (λ x : X, wf x = nf)) p_infty) as fin_f'.
                unfold is_finite, real in fin_f.
                rewrite <-fin_f => //.
            apply (Rbar_le_lt_trans _ _ _ inter_le_f fin_f').
            (* cas ou nf = maxf, donc où le domaine pour g est mesurable *)
            assert (ng < maxg).
                apply not_le => Hng'.
                assert (ng = maxg)
                    as Eqgng by apply Nat.le_antisymm => //.
                rewrite Hnf' Eqgng in Eqc.
                rewrite Eqc in H; rewrite square_bij_corner in H.
                rewrite H in Hn; unfold max_which in Hn.
                exact (Nat.Private_Tac.lt_irrefl Hn).
            assert
                (is_finite (μ (λ x : X, wg x = ng)))
                as fin_g.
                apply axg4 => //.
            assert
                (Rbar_lt (μ (λ x : X, wg x = ng)) p_infty) as fin_g'.
                unfold is_finite, real in fin_g.
                rewrite <-fin_g => //.
            apply (Rbar_le_lt_trans _ _ _ inter_le_g fin_g').
    Qed.

End simpl_fun_plus.

Notation "f + g" := (fun_plus f g) (left associativity, at level 50) : fun_scope.
Notation "sf + sg" := (sf_plus_aux sf sg) (left associativity, at level 50) : sf_scope.

Section simpl_fun_scal.

    (* espace de départ *)
    Context  {X : Set}.
    (* espace d'arrivé *)
    Context {A : AbsRing}.
    Context {E : NormedModule A}.
    (* Un espace mesuré *)
    Context {gen : (X -> Prop) -> Prop}.

    Open Scope nat_scope.

    Definition fun_scal (a : A) (g : X -> E) :=
        (fun x => scal a (g x)).

    Notation "a ⋅ g" := (fun_scal a g) (left associativity, at level 45) : fun_scope.

    Definition sf_scal_aux (a : A) (sf : simpl_fun E gen) : simpl_fun E gen.
        case: sf => wf vf maxf axf1 axf2 axf3.
        pose val := fun k => scal a (vf k).
        apply (mk_simpl_fun wf val maxf).
            unfold val; rewrite axf1 scal_zero_r => //.
            exact axf2.
            exact axf3.
    Defined.

    Notation "a ⋅ sf" := (sf_scal_aux a sf) (left associativity, at level 45) : sf_scope.

    Lemma sf_scal :
        ∀ a : A, ∀ f : X -> E,
        is_simpl gen f ->
        is_simpl gen (fun_scal a f).
    Proof.
        move => a f.
        case => sf; case_eq sf => wf vf maxf axf1 axf2 axf3 Eqsf => /= Eqf.
        exists (sf_scal_aux a sf) => x.
        unfold fun_sf, val, which; rewrite Eqsf => /=.
        unfold fun_scal; rewrite Eqf => //.
    Qed.

    Lemma fun_sf_scal :
        ∀ a : A, ∀ sf : simpl_fun E gen,
            (∀ x : X, fun_sf (a ⋅ sf) x = (a ⋅ (fun_sf sf x))%hy).
    Proof.
        move => a sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        unfold fun_sf.
        rewrite Eqf => //.
    Qed.

    Context {μ : measure gen}.

    Lemma integrable_sf_scal (a : A) {sf : simpl_fun E gen} :
        integrable_sf μ sf -> integrable_sf μ (sf_scal_aux a sf).
    Proof.
        unfold integrable_sf.
        case: sf => //.
    Qed.

End simpl_fun_scal.

Notation "a ⋅ g" := (fun_scal a g) (left associativity, at level 45) : fun_scope.
Notation "a ⋅ sf" := (sf_scal_aux a sf) (left associativity, at level 45) : sf_scope.
Notation "- g" := (fun_scal (opp one) g) : fun_scope.
Notation "- sg" := (sf_scal_aux (opp one) sg) : sf_scope.
Notation "f - g" := (fun_plus f (- g))%fn : fun_scope.
Notation "sf - sg" := (sf_plus_aux sf (- sg)) : sf_scope.

Section simpl_fun_bounded.

    (* espace de départ *)
    Context  {X : Set}.
    (* espace d'arrivé *)
    Context {A : AbsRing}.
    Context {E : NormedModule A}.
    (* Un espace mesuré *)
    Context {gen : (X -> Prop) -> Prop}.
    Context {μ : measure gen}.

    Open Scope hy_scope.

    Definition sf_bounded :
        ∀ sf : simpl_fun E gen, { M : R | ∀ x : X, ‖ fun_sf sf x ‖ <= M }.
    (* definition *)
        move => sf.
        apply exist with (Rmax_n (fun k => ‖ val sf k ‖) (max_which sf)).
        move => x; unfold fun_sf.
        pose Hwx := ax_which_max_which sf x; clearbody Hwx.
        apply: fo_le_Rmax_n => //.
    Qed.

End simpl_fun_bounded.

Open Scope R_scope.
Open Scope nat_scope.

Lemma finite_sum_Rbar :
    ∀ u : nat -> Rbar, ∀ n : nat, (∀ k : nat, k ≤ n -> is_finite (u k))
    -> is_finite (sum_Rbar n u).
Proof.
    move => u; induction n => Hu.
        unfold sum_Rbar.
        assert (0 ≤ 0) by lia.
        apply Hu => //.
        simpl.
        assert (S n <= S n) by lia.
        pose HuSn := (Hu (S n) H); clearbody HuSn.
        assert (∀ k : nat, k ≤ n -> is_finite (u k)).
            move => k Hk.
            assert (k ≤ S n) by lia.
            apply Hu => //.
        pose Hsum := IHn H0; clearbody Hsum; clear H0.
        unfold is_finite in HuSn, Hsum |- *.
        rewrite <-HuSn, <-Hsum => //.
Qed.

Section simpl_fun_meas.

    (* espace de départ *)
    Context  {X : Set}.
    (* espace d'arrivé *)
    Context {A : AbsRing}.
    Context {E : NormedModule A}.
    (* Un espace mesuré *)
    Context {gen : (X -> Prop) -> Prop}.
    Context {μ : measure gen}.

    Lemma measurable_sf_preim :
        ∀ sf : simpl_fun E gen, ∀ y : E,
            measurable gen (fun x : X => sf x = y).
    Proof.
        move => sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        move => y.
        pose P := fun k => (λ x : X, vf k = y ∧ which sf x = k).
        apply measurable_ext with (λ x : X, ∃ k : nat, k ≤ maxf ∧ P k x).
            move => x; split.
                case => k; case => Hk.
                unfold P; case => <-; unfold fun_sf => ->.
                rewrite Eqf => /=; split; auto with arith.
                move => Eq_sfx_vfn.
                exists (which sf x); split.
                    replace (which sf x) with (wf x) by rewrite Eqf => //; apply axf2.
                    unfold P; split.
                    unfold fun_sf in Eq_sfx_vfn.
                    replace (vf (which sf x)) with (val sf (which sf x))
                        by rewrite Eqf => //.
                    rewrite Eq_sfx_vfn => //.
                    reflexivity.
                apply measurable_union_finite => k Hk.
                unfold P; apply measurable_inter.
        apply measurable_Prop.
        replace (which sf) with (wf) by rewrite Eqf => //.
        apply axf3 => //.
    Qed.

    Lemma sf_measurable_preim_le :
        ∀ sf : simpl_fun E gen, ∀ n : nat, n ≤ max_which sf ->
            measurable gen (λ x : X, which sf x ≤ n).
    Proof.
        move => sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        move => n Hn.
        pose B := fun k => (λ x : X, which sf x = k).
        apply measurable_ext with (fun x => ∃ k : nat, k ≤ n ∧ which sf x = k).
            move => x; split.
                case => k [Hk Eqk]; lia.
                move => Hsfx; exists (which sf x); auto.
            apply (measurable_union_finite _ ) => k Hk.
            apply ax_measurable.
            lia.
    Qed.

    Lemma sf_measurable_preim_lt :
        ∀ sf : simpl_fun E gen, ∀ n : nat, n ≤ max_which sf ->
            measurable gen (λ x : X, which sf x < n).
    Proof.
        move => sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        case.
            move => _.
            apply measurable_ext with (fun _ => False).
            move => x; split => //.
            lia.
            apply measurable_Prop.

            move => n' HSn'.
            apply measurable_ext with (fun x => which sf x ≤ n').
            move => x; split; lia.
            apply sf_measurable_preim_le; lia.
    Qed.

    Lemma is_finite_sf_preim_lt {sf : simpl_fun E gen} :
        integrable_sf μ sf -> ∀ n : nat, n ≤ max_which sf ->
            is_finite (μ (λ x : X, which sf x < n)).
    Proof.
        unfold integrable_sf => axf4.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        rewrite Eqf in axf4; simpl in axf4.
        case.
            move => _.
            rewrite (measure_ext _ _ _ (fun _ => False)).
            rewrite meas_False => //.
            lia.
        move => n Hn.
        rewrite (measure_ext _ _ _ (λ x : X, which sf x ≤ n)); swap 1 2.
            lia.
        pose B := fun k => (λ x : X, which sf x = k).
        rewrite (measure_decomp_finite _  μ _ maxf B).
        assert (n < maxf) as Ltnmaxf.
        rewrite Eqf in Hn; simpl in Hn.
        lia.
        apply finite_sum_Rbar => k Hk; unfold B.
        case_eq (k <=? n).
            move => /Nat.leb_le Hk'.
            rewrite (measure_ext _ _ _ (fun x => which sf x = k)).
            rewrite Eqf => /=; apply axf4; lia.
            move => x; split; [easy | move => -> //].
            move => /Nat.leb_gt Hk'.
            rewrite (measure_ext _ _ _ (fun _ => False)).
            rewrite meas_False => //.
            move => x; split => //.
            lia.
        apply sf_measurable_preim_le; lia.
        move => k Hk; unfold B; apply ax_measurable; rewrite Eqf => //.
        move => x; unfold B; exists (which sf x); split => //.
        rewrite Eqf => /=; apply axf2.

        move => p q x Hp Hq; unfold B; move => -> //.
    Qed.

    Lemma sf_decomp_preim_which :
        ∀ sf : simpl_fun E gen, ∀ y : E,
            (μ (fun x : X => sf x = y))
            = sum_Rbar (max_which sf) (fun k => (μ (fun x : X => val sf k = y ∧ which sf x = k))).
    Proof.
        move => sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        move => y.
        pose B := fun k => (λ x : X, which sf x = k).
            rewrite (measure_decomp_finite _ μ (fun x => sf x = y) maxf B).
            unfold B; rewrite Eqf => /=.
            apply sum_Rbar_ext => k Hk.
            apply measure_ext => x; split.
                case => Eq_sfx_y Eq_wfx_k.
                split; [rewrite <-Eq_wfx_k | rewrite Eq_wfx_k] => //.
                case => Eq_vfk_y Eq_wfx_k.
                split; [rewrite Eq_wfx_k | ] => //.
            apply measurable_sf_preim.
            move => k Hk.
            unfold B; apply ax_measurable; rewrite Eqf => //.
            move => x; unfold B; exists (which sf x).
            split; [rewrite Eqf => /=; apply axf2 | reflexivity].
            move => p q x Hp Hq; unfold B => -> //.
    Qed.

    Lemma finite_sf_preim_neq_0 {sf : simpl_fun E gen} :
        integrable_sf μ sf -> ∀ y : E, y ≠ zero ->
            is_finite (μ (fun x : X => sf x = y)).
    Proof.
        unfold integrable_sf => axf4.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        rewrite Eqf in axf4; simpl in axf4.
        move => y Hy.
        rewrite sf_decomp_preim_which.
        apply finite_sum_Rbar => k Hk.
        case: (le_lt_or_eq k (max_which sf) Hk) => Hkmaxf.
            apply Rbar_bounded_is_finite with (real 0%R) (μ (λ x : X, which sf x = k)).
            apply meas_ge_0.
            apply (measure_le _ μ (λ x : X, val sf k = y ∧ which sf x = k) (fun x => which sf x = k)).
            apply measurable_inter.
            apply measurable_Prop.
            1, 2 : rewrite Eqf => /=; apply axf3; rewrite Eqf in Hk; simpl in Hk => //.
            move => x; case; auto.
            easy.
            rewrite Eqf => /=; apply axf4.
            rewrite Eqf in Hkmaxf => //.

            rewrite (measure_ext _ _ _ (fun _ => False)).
            rewrite meas_False => //.
            move => x; split => //.
            case; move => Abs _.
            rewrite Hkmaxf in Abs.
            rewrite <-Abs in Hy.
            rewrite Eqf in Hy; simpl in Hy => //.
    Qed.

    Lemma measurable_fun_sf_aux :
        ∀ sf : simpl_fun E gen,
            (∀ P : E -> Prop, measurable open P -> measurable gen (fun x : X => P (sf x))).
    Proof.
        move => sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.

        move => P HP.
        pose B := fun (n : nat) => (fun x : X => (which sf x = n) ∧ (P (val sf n))).
        apply measurable_ext with (λ x : X, ∃ n : nat, n ≤ max_which sf ∧ B n x).
        move => x; split.
            case => n [Hn Bnx]; unfold B in Bnx.
            case: Bnx => [Eqwsfx Pvsfn]; unfold fun_sf.
            rewrite Eqwsfx => //.
            move => Psfx; exists (which sf x).
            split;
                [apply ax_which_max_which
                | unfold B; split;
                    [reflexivity | assumption]].
        apply measurable_union_finite => n Hn.
        apply measurable_inter.
        apply ax_measurable => //.
        apply measurable_Prop.
    Qed.

    Lemma measurable_fun_sf :
        ∀ sf : simpl_fun E gen,
            measurable_fun gen open sf.
    Proof.
        move => sf.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.

        move => P HP; inversion HP.
            clear A0 H0.
            apply measurable_fun_sf_aux => //.

            clear H HP P.
            constructor 2.

            clear H0 A0.
            apply measurable_fun_sf_aux => //.

            clear H0 HP.
            constructor 4.
            move => n; apply measurable_fun_sf_aux => //.
    Qed.

    Lemma finite_sf_preim_neq_0' {sf : simpl_fun E gen} :
        integrable_sf μ sf ->
            is_finite (μ (fun x : X => sf x ≠ zero)).
    Proof.
        unfold integrable_sf => axf4.
        case_eq sf => wf vf maxf axf1 axf2 axf3 Eqf; rewrite <-Eqf => /=.
        rewrite Eqf in axf4; simpl in axf4.
        apply Rbar_bounded_is_finite with 0%R (μ (fun x => ∃ n : nat, n < max_which sf ∧ which sf x = n)).
        apply meas_ge_0.
        apply measure_le.
        apply (measurable_fun_sf sf (fun x : E => x ≠ zero)).
        apply measurable_gen.
        apply NM_open_neq.
        apply measurable_union_finite' => n Hn.
        apply ax_measurable; lia.
        move => x Neq0.
        exists (which sf x); split => //.
        case: (le_lt_or_eq _ _ (ax_which_max_which sf x)) => //.
        move => Abs.
        unfold fun_sf in Neq0.
        rewrite Abs in Neq0.
        rewrite ax_val_max_which in Neq0 => //.
        easy.
        rewrite (measure_decomp_finite _ μ _ (max_which sf) (fun n => fun x => which sf x = n)).
        apply finite_sum_Rbar => k Hk.
        case: (le_lt_or_eq _ _ Hk) => Hk'.
        rewrite (measure_ext _ _ _ (λ x : X, which sf x = k)).
        rewrite Eqf in Hk' |- * => /=; simpl in Hk'.
        apply axf4 => //.
        move => x; split.
        easy.
        move => ->.
        split.
        exists k => //.
        reflexivity.
        rewrite Hk'.
        rewrite (measure_ext _ _ _ (fun _ => False)).
        rewrite meas_False => //.
        move => x; split => //.
        case; case => n; case => Hn1 ->.
        lia.
        apply measurable_union_finite'.
        move => n Hn.
        apply ax_measurable; lia.
        move => n Hn.
        apply ax_measurable => //.
        move => x.
        exists (which sf x); split.
        apply ax_which_max_which.
        reflexivity.
        move => n p x Hn Hp Eqwn <- //.
    Qed.

End simpl_fun_meas.

Require Import Lra.

Section simpl_fun_nonzero.

    (* espace de départ *)
    Context  {X : Set}.
    (* espace d'arrivé *)
    Context {A : AbsRing}.
    Context {E : NormedModule A}.
    (* Un espace mesuré *)
    Context {gen : (X -> Prop) -> Prop}.

    Lemma sf_nonzero_ind (sf : simpl_fun E gen) (n : nat) :
        { sf'   : simpl_fun E gen
                | (∀ x : X, sf x = sf' x)
                ∧ (max_which sf' ≤ max_which sf)
                ∧ ((max_which sf - max_which sf') ≤ n)
                ∧ (∀ k : nat, k < n - (max_which sf - max_which sf') -> k < max_which sf' -> val sf' k ≠ zero)
        }.
    Proof.
        induction n.
        exists sf; repeat split; lia.
        case: IHn => sf'n [Ext [Ineqmax IHn]].
        case_eq (n <? max_which sf); swap 1 2.
        move => /Nat.ltb_ge n_large.
        exists sf'n; repeat split => //.
        lia.
        move => k Hk1 Hk2;
            apply IHn; lia.
        move => /Nat.ltb_lt n_small.
        pose m := (n - (max_which sf - max_which sf'n)).
        case: (eq_dec (val sf'n m) (zero)); swap 1 2.
            move => NotEq0.
            exists sf'n; repeat split => //.
            lia.
            move => k.
            move => /le_lt_or_eq; case.
            move => H1 H2.
            apply IHn.
            1, 2 : lia.
            move => H.
            assert (k = m) as Eqkn by lia; clear H.
            rewrite Eqkn.
            move => _.
            assumption.

            move => Eq0.
            case_eq sf'n => whichn valn maxn axfn1 axfn2 axfn3 Eqsfn.
            pose valSn :=
            (
                fun (k : nat) =>
                    if (k <? m) then (valn k)
                    else valn (S k)
            ).
            pose maxSn := maxn - 1.
            pose whichSn :=
            (
                fun (x : X) =>
                    if (whichn x =? m) then maxSn
                    else if (m <? whichn x) then (whichn x - 1)
                    else whichn x
            ).
            assert (m ≤ maxn - 1) as Lemmaxn.
            replace maxn with (max_which sf'n)
                by rewrite Eqsfn => //.
            lia.
            assert (0 < maxn) as Lt0maxn.
            replace maxn with (max_which sf'n)
                by rewrite Eqsfn => //.
            lia.
            assert (valSn maxSn = zero) as axSn1.
            unfold valSn, maxSn.
            assert (maxn - 1 <? m = false).
            apply Nat.leb_gt.
            lia.
            rewrite H; clear H.
            replace (S (maxn - 1)) with maxn by lia.
            assumption.
            assert (∀ x : X, whichSn x ≤ maxSn) as axSn2.
            move => x; unfold whichSn, maxSn.
            case: (whichn x =? m) => //.
            case_eq (m <? whichn x).
            move => _.
            assert (whichn x ≤ maxn) by apply axfn2.
            lia.
            move => /Nat.ltb_ge.
            lia.
            assert (max_which sf - maxSn ≤ S n) as Ineq2.
            unfold maxSn.
            replace maxn with (max_which sf'n).
            2 : rewrite Eqsfn => //.
            lia.
            assert (∀ k : nat,
                k ≤ maxSn
                → measurable gen (λ x : X, whichSn x = k)) as axSn3.
            move => k Hk; unfold whichSn.
            case: (le_lt_or_eq _ _ Hk); swap 1 2.
            move => ->.
            apply measurable_ext with (fun (x : X) => whichn x = m ∨ whichn x = maxn).
            move => x; split.
            case.
            move => /Nat.eqb_eq -> //.
            move => Hx.
            assert (m <? whichn x = true).
            apply Nat.ltb_lt; lia.
            rewrite H; clear H.
            assert (whichn x =? m = false).
            apply Nat.eqb_neq; lia.
            rewrite H; clear H.
            unfold maxSn; lia.
            case_eq (whichn x =? m).
            move => /Nat.eqb_eq ->; left => //.
            move => /Nat.eqb_neq Hx0.
            case_eq (m <? whichn x).
            move => /Nat.ltb_lt Hx1.
            unfold maxSn => H.
            right; lia.
            move => /Nat.ltb_ge Hx1.
            lia.
            apply measurable_union.
            apply axfn3.
            lia.
            apply axfn3 => //.
            move => LtkmaxSn.
            case: (le_lt_dec m k) => Hkm.
            apply measurable_ext with (fun x : X => whichn x = S k).
            move => x; split.
            move => Hx0.
            assert (whichn x =? m = false).
            apply Nat.eqb_neq; lia.
            rewrite H; clear H.
            assert (m <? whichn x = true).
            apply Nat.ltb_lt; lia.
            rewrite H; clear H.
            lia.
            case_eq (whichn x =? m).
            move => /Nat.eqb_eq; lia.
            move => /Nat.eqb_neq Hx0.
            case_eq (m <? whichn x).
            move => /Nat.ltb_lt; lia.
            move => /Nat.ltb_ge Hx1; lia.
            apply axfn3; lia.
            apply measurable_ext with (fun x : X => whichn x = k).
            move => x; split.
            move => Hx0.
            assert (whichn x =? m = false).
            apply Nat.eqb_neq; lia.
            rewrite H; clear H.
            assert (m <? whichn x = false).
            apply Nat.ltb_ge; lia.
            rewrite H; clear H.
            lia.
            case_eq (whichn x =? m).
            move => /Nat.eqb_eq; lia.
            move => /Nat.eqb_neq Hx0.
            case_eq (m <? whichn x).
            move => /Nat.ltb_lt; lia.
            move => /Nat.ltb_ge Hx1; lia.
            apply axfn3; lia.
            assert (m = S n - (max_which sf - maxSn)) as Eqm.
            unfold maxSn.
            replace maxn with (max_which sf'n).
            2 : rewrite Eqsfn => //.
            lia.
            exists (mk_simpl_fun whichSn valSn maxSn axSn1 axSn2 axSn3) => /=.
            assert (maxSn ≤ max_which sf) as Ineq.
            unfold maxSn.
            assert (maxn ≤ max_which sf).
            rewrite Eqsfn in Ineqmax; simpl in Ineqmax.
            assumption.
            lia.
            repeat split.
            move => x; unfold valSn, whichSn.
            case: (lt_eq_lt_dec (whichn x) m).
            case.
            move => Hxm.
            assert (whichn x =? m = false).
            apply Nat.eqb_neq; lia.
            rewrite H; clear H.
            assert (m <? whichn x = false).
            apply Nat.ltb_ge; lia.
            rewrite H; clear H.
            assert (whichn x <? m = true).
            apply Nat.ltb_lt; lia.
            rewrite H; clear H.
            replace (valn (whichn x)) with (sf'n x).
                2 : unfold fun_sf.
                2 : rewrite Eqsfn => //.
            apply Ext.
            move => Hxm.
            assert (whichn x =? m = true).
            apply Nat.eqb_eq; lia.
            rewrite H; clear H.
            assert (maxSn <? m = false).
            apply Nat.ltb_ge; lia.
            rewrite H; clear H.
            replace (S maxSn) with maxn by lia.
            rewrite Ext.
            unfold fun_sf.
            rewrite Eqsfn => /=.
            rewrite Hxm.
            rewrite Eqsfn in Eq0; simpl in Eq0.
            rewrite axfn1 Eq0 => //.
            move => Hxm.
            assert (whichn x =? m = false).
            apply Nat.eqb_neq; lia.
            rewrite H; clear H.
            assert (m <? whichn x = true).
            apply Nat.ltb_lt; lia.
            rewrite H; clear H.
            assert (whichn x - 1 <? m = false).
            apply Nat.ltb_ge; lia.
            rewrite H; clear H.
            replace (S (whichn x - 1)) with (whichn x) by lia.
            replace (valn (whichn x)) with (sf'n x).
            2 : unfold fun_sf.
            2 : rewrite Eqsfn => //.
            apply Ext.
            assumption.
            assumption.
            move => k Hk1 Hk2.
            replace (match max_which sf - maxSn with
            | 0 => S n
            | S l => n - l
            end) with ((S n) - (max_which sf - maxSn)) in Hk1.
            2 : case: (max_which sf - maxSn) => //.
            case: (le_lt_or_eq _ _ Hk1).
            move => Hkn.
            case: IHn => IHn1 IHn2.
            rewrite Eqsfn in IHn2; simpl in IHn2.
            unfold valSn, m.
            rewrite Eqsfn => /=.
            assert (k < n - (max_which sf - maxn)) by lia.
            assert (k <? n - (max_which sf - maxn) = true).
                apply Nat.ltb_lt => //.
                rewrite H0; clear H0.
            apply IHn2 => //.
            lia.
            move => Eqk.
            assert (k = n - (max_which sf - maxSn)) by lia.
            unfold valSn.
            assert (k <? m = true).
            apply Nat.ltb_lt.
            unfold m; rewrite Eqsfn => /=.
            case_eq (S n - (max_which sf - maxSn)).
            lia.
            move => l Hl.
            assert (max_which sf - maxSn < S n) by lia.
            unfold maxSn in H0.
            replace (max_which sf - (maxn - 1)) with (S (max_which sf - maxn)) in H0.
            lia.
            lia.
            rewrite H0.
            case: IHn => [IHn1 IHn2].
            rewrite Eqsfn in IHn2; simpl in IHn2.
            apply IHn2.
            lia.
            lia.
    Qed.

    Lemma sf_remove_zeros (sf : simpl_fun E gen) :
        { sf'   : simpl_fun E gen
                | (∀ x : X, sf x = sf' x)
                ∧ (∀ k : nat, k < max_which sf' -> val sf' k ≠ zero)
        }.
    Proof.
        case: (sf_nonzero_ind sf (max_which sf)).
        move => sf' [Ext [Ineq [_ Hsf']]].
        exists sf'; split => //.
        move => k Hk.
        apply Hsf' => //.
        lia.
    Qed.

End simpl_fun_nonzero.


From iris.program_logic Require Export language ectx_language ectxi_language.
From iris.algebra Require Export ofe.
From stdpp Require Export strings.
From stdpp Require Import gmap.
Set Default Proof Using "Type".

Delimit Scope expr_scope with E.
Delimit Scope val_scope with V.

Module heap_lang.
Open Scope Z_scope.

(** Expressions and vals. *)
Definition loc := positive. (* Really, any countable type. *)

Inductive base_lit : Set :=
  | LitInt (n : Z) | LitBool (b : bool) | LitUnit | LitLoc (l : loc).
Inductive un_op : Set :=
  | NegOp | MinusUnOp.
Inductive bin_op : Set :=
  | PlusOp | MinusOp | MultOp | QuotOp | RemOp (* Arithmetic *)
  | AndOp | OrOp | XorOp (* Bitwise *)
  | ShiftLOp | ShiftROp (* Shifts *)
  | LeOp | LtOp | EqOp. (* Relations *)

Inductive binder := BAnon | BNamed : string → binder.
Delimit Scope binder_scope with bind.
Bind Scope binder_scope with binder.
Definition cons_binder (mx : binder) (X : list string) : list string :=
  match mx with BAnon => X | BNamed x => x :: X end.
Infix ":b:" := cons_binder (at level 60, right associativity).
Instance binder_eq_dec_eq : EqDecision binder.
Proof. solve_decision. Defined.

Instance set_unfold_cons_binder x mx X P :
  SetUnfold (x ∈ X) P → SetUnfold (x ∈ mx :b: X) (BNamed x = mx ∨ P).
Proof.
  constructor. rewrite -(set_unfold (x ∈ X) P).
  destruct mx; rewrite /= ?elem_of_cons; naive_solver.
Qed.

Inductive expr :=
  (* Base lambda calculus *)
  | Var (x : string)
  | Rec (f x : binder) (e : expr)
  | App (e1 e2 : expr)
  (* Base types and their operations *)
  | Lit (l : base_lit)
  | UnOp (op : un_op) (e : expr)
  | BinOp (op : bin_op) (e1 e2 : expr)
  | If (e0 e1 e2 : expr)
  (* Products *)
  | Pair (e1 e2 : expr)
  | Fst (e : expr)
  | Snd (e : expr)
  (* Sums *)
  | InjL (e : expr)
  | InjR (e : expr)
  | Case (e0 : expr) (e1 : expr) (e2 : expr)
  (* Concurrency *)
  | Fork (e : expr)
  (* Heap *)
  | Alloc (e : expr)
  | Load (e : expr)
  | Store (e1 : expr) (e2 : expr)
  | CAS (e0 : expr) (e1 : expr) (e2 : expr)
  | FAA (e1 : expr) (e2 : expr).

Bind Scope expr_scope with expr.

Fixpoint is_closed (X : list string) (e : expr) : bool :=
  match e with
  | Var x => bool_decide (x ∈ X)
  | Rec f x e => is_closed (f :b: x :b: X) e
  | Lit _ => true
  | UnOp _ e | Fst e | Snd e | InjL e | InjR e | Fork e | Alloc e | Load e =>
     is_closed X e
  | App e1 e2 | BinOp _ e1 e2 | Pair e1 e2 | Store e1 e2 | FAA e1 e2 =>
     is_closed X e1 && is_closed X e2
  | If e0 e1 e2 | Case e0 e1 e2 | CAS e0 e1 e2 =>
     is_closed X e0 && is_closed X e1 && is_closed X e2
  end.

Class Closed (X : list string) (e : expr) := closed : is_closed X e.
Instance closed_proof_irrel X e : ProofIrrel (Closed X e).
Proof. rewrite /Closed. apply _. Qed.
Instance closed_dec X e : Decision (Closed X e).
Proof. rewrite /Closed. apply _. Defined.

Inductive val :=
  | RecV (f x : binder) (e : expr) `{!Closed (f :b: x :b: []) e}
  | LitV (l : base_lit)
  | PairV (v1 v2 : val)
  | InjLV (v : val)
  | InjRV (v : val).

Bind Scope val_scope with val.

Fixpoint of_val (v : val) : expr :=
  match v with
  | RecV f x e => Rec f x e
  | LitV l => Lit l
  | PairV v1 v2 => Pair (of_val v1) (of_val v2)
  | InjLV v => InjL (of_val v)
  | InjRV v => InjR (of_val v)
  end.

Fixpoint to_val (e : expr) : option val :=
  match e with
  | Rec f x e =>
     if decide (Closed (f :b: x :b: []) e) then Some (RecV f x e) else None
  | Lit l => Some (LitV l)
  | Pair e1 e2 => v1 ← to_val e1; v2 ← to_val e2; Some (PairV v1 v2)
  | InjL e => InjLV <$> to_val e
  | InjR e => InjRV <$> to_val e
  | _ => None
  end.

(** We assume the following encoding of values to 64-bit words: The least 3
significant bits of every word are a "tag", and we have 61 bits of payload,
which is enough if all pointers are 8-byte-aligned (common on 64bit
architectures). The tags have the following meaning:

0: Payload is the data for a LitV (LitInt _).
1: Payload is the data for a InjLV (LitV (LitInt _)).
2: Payload is the data for a InjRV (LitV (LitInt _)).
3: Payload is the data for a LitV (LitLoc _).
4: Payload is the data for a InjLV (LitV (LitLoc _)).
4: Payload is the data for a InjRV (LitV (LitLoc _)).
6: Payload is one of the following finitely many values, which 61 bits are more
   than enough to encode:
   LitV LitUnit, InjLV (LitV LitUnit), InjRV (LitV LitUnit),
   LitV (LitBool _), InjLV (LitV (LitBool _)), InjRV (LitV (LitBool _)).
7: Value is boxed, i.e., payload is a pointer to some read-only memory area on
   the heap which stores whether this is a RecV, PairV, InjLV or InjRV and the
   relevant data for those cases. However, the boxed representation is never
   used if any of the above representations could be used.

Ignoring (as usual) the fact that we have to fit the infinite Z/loc into 61
bits, this means every value is machine-word-sized and can hence be atomically
read and written.  Also notice that the sets of boxed and unboxed values are
disjoint. *)
Definition val_is_unboxed (v : val) : Prop :=
  match v with
  | LitV _ => True
  | InjLV (LitV _) => True
  | InjRV (LitV _) => True
  | _ => False
  end.

(** The state: heaps of vals. *)
Definition state := gmap loc val.

(** Equality and other typeclass stuff *)
Lemma to_of_val v : to_val (of_val v) = Some v.
Proof.
  by induction v; simplify_option_eq; repeat f_equal; try apply (proof_irrel _).
Qed.

Lemma of_to_val e v : to_val e = Some v → of_val v = e.
Proof.
  revert v; induction e; intros v ?; simplify_option_eq; auto with f_equal.
Qed.

Instance of_val_inj : Inj (=) (=) of_val.
Proof. by intros ?? Hv; apply (inj Some); rewrite -!to_of_val Hv. Qed.

Instance base_lit_eq_dec : EqDecision base_lit.
Proof. solve_decision. Defined.
Instance un_op_eq_dec : EqDecision un_op.
Proof. solve_decision. Defined.
Instance bin_op_eq_dec : EqDecision bin_op.
Proof. solve_decision. Defined.
Instance expr_eq_dec : EqDecision expr.
Proof. solve_decision. Defined.
Instance val_eq_dec : EqDecision val.
Proof.
 refine (λ v v', cast_if (decide (of_val v = of_val v'))); abstract naive_solver.
Defined.

Instance base_lit_countable : Countable base_lit.
Proof.
 refine (inj_countable' (λ l, match l with
  | LitInt n => inl (inl n) | LitBool b => inl (inr b)
  | LitUnit => inr (inl ()) | LitLoc l => inr (inr l)
  end) (λ l, match l with
  | inl (inl n) => LitInt n | inl (inr b) => LitBool b
  | inr (inl ()) => LitUnit | inr (inr l) => LitLoc l
  end) _); by intros [].
Qed.
Instance un_op_finite : Countable un_op.
Proof.
 refine (inj_countable' (λ op, match op with NegOp => 0 | MinusUnOp => 1 end)
  (λ n, match n with 0 => NegOp | _ => MinusUnOp end) _); by intros [].
Qed.
Instance bin_op_countable : Countable bin_op.
Proof.
 refine (inj_countable' (λ op, match op with
  | PlusOp => 0 | MinusOp => 1 | MultOp => 2 | QuotOp => 3 | RemOp => 4
  | AndOp => 5 | OrOp => 6 | XorOp => 7 | ShiftLOp => 8 | ShiftROp => 9
  | LeOp => 10 | LtOp => 11 | EqOp => 12
  end) (λ n, match n with
  | 0 => PlusOp | 1 => MinusOp | 2 => MultOp | 3 => QuotOp | 4 => RemOp
  | 5 => AndOp | 6 => OrOp | 7 => XorOp | 8 => ShiftLOp | 9 => ShiftROp
  | 10 => LeOp | 11 => LtOp | _ => EqOp
  end) _); by intros [].
Qed.
Instance binder_countable : Countable binder.
Proof.
 refine (inj_countable' (λ b, match b with BNamed s => Some s | BAnon => None end)
  (λ b, match b with Some s => BNamed s | None => BAnon end) _); by intros [].
Qed.
Instance expr_countable : Countable expr.
Proof.
 set (enc := fix go e :=
  match e with
  | Var x => GenLeaf (inl (inl x))
  | Rec f x e => GenNode 0 [GenLeaf (inl (inr f)); GenLeaf (inl (inr x)); go e]
  | App e1 e2 => GenNode 1 [go e1; go e2]
  | Lit l => GenLeaf (inr (inl l))
  | UnOp op e => GenNode 2 [GenLeaf (inr (inr (inl op))); go e]
  | BinOp op e1 e2 => GenNode 3 [GenLeaf (inr (inr (inr op))); go e1; go e2]
  | If e0 e1 e2 => GenNode 4 [go e0; go e1; go e2]
  | Pair e1 e2 => GenNode 5 [go e1; go e2]
  | Fst e => GenNode 6 [go e]
  | Snd e => GenNode 7 [go e]
  | InjL e => GenNode 8 [go e]
  | InjR e => GenNode 9 [go e]
  | Case e0 e1 e2 => GenNode 10 [go e0; go e1; go e2]
  | Fork e => GenNode 11 [go e]
  | Alloc e => GenNode 12 [go e]
  | Load e => GenNode 13 [go e]
  | Store e1 e2 => GenNode 14 [go e1; go e2]
  | CAS e0 e1 e2 => GenNode 15 [go e0; go e1; go e2]
  | FAA e1 e2 => GenNode 16 [go e1; go e2]
  end).
 set (dec := fix go e :=
  match e with
  | GenLeaf (inl (inl x)) => Var x
  | GenNode 0 [GenLeaf (inl (inr f)); GenLeaf (inl (inr x)); e] => Rec f x (go e)
  | GenNode 1 [e1; e2] => App (go e1) (go e2)
  | GenLeaf (inr (inl l)) => Lit l
  | GenNode 2 [GenLeaf (inr (inr (inl op))); e] => UnOp op (go e)
  | GenNode 3 [GenLeaf (inr (inr (inr op))); e1; e2] => BinOp op (go e1) (go e2)
  | GenNode 4 [e0; e1; e2] => If (go e0) (go e1) (go e2)
  | GenNode 5 [e1; e2] => Pair (go e1) (go e2)
  | GenNode 6 [e] => Fst (go e)
  | GenNode 7 [e] => Snd (go e)
  | GenNode 8 [e] => InjL (go e)
  | GenNode 9 [e] => InjR (go e)
  | GenNode 10 [e0; e1; e2] => Case (go e0) (go e1) (go e2)
  | GenNode 11 [e] => Fork (go e)
  | GenNode 12 [e] => Alloc (go e)
  | GenNode 13 [e] => Load (go e)
  | GenNode 14 [e1; e2] => Store (go e1) (go e2)
  | GenNode 15 [e0; e1; e2] => CAS (go e0) (go e1) (go e2)
  | GenNode 16 [e1; e2] => FAA (go e1) (go e2)
  | _ => Lit LitUnit (* dummy *)
  end).
 refine (inj_countable' enc dec _). intros e. induction e; f_equal/=; auto.
Qed.
Instance val_countable : Countable val.
Proof. refine (inj_countable of_val to_val _); auto using to_of_val. Qed.

Instance expr_inhabited : Inhabited expr := populate (Lit LitUnit).
Instance val_inhabited : Inhabited val := populate (LitV LitUnit).

Canonical Structure stateC := leibnizC state.
Canonical Structure valC := leibnizC val.
Canonical Structure exprC := leibnizC expr.

(** Evaluation contexts *)
Inductive ectx_item :=
  | AppLCtx (e2 : expr)
  | AppRCtx (v1 : val)
  | UnOpCtx (op : un_op)
  | BinOpLCtx (op : bin_op) (e2 : expr)
  | BinOpRCtx (op : bin_op) (v1 : val)
  | IfCtx (e1 e2 : expr)
  | PairLCtx (e2 : expr)
  | PairRCtx (v1 : val)
  | FstCtx
  | SndCtx
  | InjLCtx
  | InjRCtx
  | CaseCtx (e1 : expr) (e2 : expr)
  | AllocCtx
  | LoadCtx
  | StoreLCtx (e2 : expr)
  | StoreRCtx (v1 : val)
  | CasLCtx (e1 : expr) (e2 : expr)
  | CasMCtx (v0 : val) (e2 : expr)
  | CasRCtx (v0 : val) (v1 : val)
  | FaaLCtx (e2 : expr)
  | FaaRCtx (v1 : val).

Definition fill_item (Ki : ectx_item) (e : expr) : expr :=
  match Ki with
  | AppLCtx e2 => App e e2
  | AppRCtx v1 => App (of_val v1) e
  | UnOpCtx op => UnOp op e
  | BinOpLCtx op e2 => BinOp op e e2
  | BinOpRCtx op v1 => BinOp op (of_val v1) e
  | IfCtx e1 e2 => If e e1 e2
  | PairLCtx e2 => Pair e e2
  | PairRCtx v1 => Pair (of_val v1) e
  | FstCtx => Fst e
  | SndCtx => Snd e
  | InjLCtx => InjL e
  | InjRCtx => InjR e
  | CaseCtx e1 e2 => Case e e1 e2
  | AllocCtx => Alloc e
  | LoadCtx => Load e
  | StoreLCtx e2 => Store e e2
  | StoreRCtx v1 => Store (of_val v1) e
  | CasLCtx e1 e2 => CAS e e1 e2
  | CasMCtx v0 e2 => CAS (of_val v0) e e2
  | CasRCtx v0 v1 => CAS (of_val v0) (of_val v1) e
  | FaaLCtx e2 => FAA e e2
  | FaaRCtx v1 => FAA (of_val v1) e
  end.

(** Substitution *)
Fixpoint subst (x : string) (es : expr) (e : expr)  : expr :=
  match e with
  | Var y => if decide (x = y) then es else Var y
  | Rec f y e =>
     Rec f y $ if decide (BNamed x ≠ f ∧ BNamed x ≠ y) then subst x es e else e
  | App e1 e2 => App (subst x es e1) (subst x es e2)
  | Lit l => Lit l
  | UnOp op e => UnOp op (subst x es e)
  | BinOp op e1 e2 => BinOp op (subst x es e1) (subst x es e2)
  | If e0 e1 e2 => If (subst x es e0) (subst x es e1) (subst x es e2)
  | Pair e1 e2 => Pair (subst x es e1) (subst x es e2)
  | Fst e => Fst (subst x es e)
  | Snd e => Snd (subst x es e)
  | InjL e => InjL (subst x es e)
  | InjR e => InjR (subst x es e)
  | Case e0 e1 e2 => Case (subst x es e0) (subst x es e1) (subst x es e2)
  | Fork e => Fork (subst x es e)
  | Alloc e => Alloc (subst x es e)
  | Load e => Load (subst x es e)
  | Store e1 e2 => Store (subst x es e1) (subst x es e2)
  | CAS e0 e1 e2 => CAS (subst x es e0) (subst x es e1) (subst x es e2)
  | FAA e1 e2 => FAA (subst x es e1) (subst x es e2)
  end.

Definition subst' (mx : binder) (es : expr) : expr → expr :=
  match mx with BNamed x => subst x es | BAnon => id end.

(** The stepping relation *)
Definition un_op_eval (op : un_op) (v : val) : option val :=
  match op, v with
  | NegOp, LitV (LitBool b) => Some $ LitV $ LitBool (negb b)
  | NegOp, LitV (LitInt n) => Some $ LitV $ LitInt (Z.lnot n)
  | MinusUnOp, LitV (LitInt n) => Some $ LitV $ LitInt (- n)
  | _, _ => None
  end.

Definition bin_op_eval_int (op : bin_op) (n1 n2 : Z) : base_lit :=
  match op with
  | PlusOp => LitInt (n1 + n2)
  | MinusOp => LitInt (n1 - n2)
  | MultOp => LitInt (n1 * n2)
  | QuotOp => LitInt (n1 `quot` n2)
  | RemOp => LitInt (n1 `rem` n2)
  | AndOp => LitInt (Z.land n1 n2)
  | OrOp => LitInt (Z.lor n1 n2)
  | XorOp => LitInt (Z.lxor n1 n2)
  | ShiftLOp => LitInt (n1 ≪ n2)
  | ShiftROp => LitInt (n1 ≫ n2)
  | LeOp => LitBool (bool_decide (n1 ≤ n2))
  | LtOp => LitBool (bool_decide (n1 < n2))
  | EqOp => LitBool (bool_decide (n1 = n2))
  end.

Definition bin_op_eval_bool (op : bin_op) (b1 b2 : bool) : option base_lit :=
  match op with
  | PlusOp | MinusOp | MultOp | QuotOp | RemOp => None (* Arithmetic *)
  | AndOp => Some (LitBool (b1 && b2))
  | OrOp => Some (LitBool (b1 || b2))
  | XorOp => Some (LitBool (xorb b1 b2))
  | ShiftLOp | ShiftROp => None (* Shifts *)
  | LeOp | LtOp => None (* InEquality *)
  | EqOp => Some (LitBool (bool_decide (b1 = b2)))
  end.

Definition bin_op_eval (op : bin_op) (v1 v2 : val) : option val :=
  if decide (op = EqOp) then Some $ LitV $ LitBool $ bool_decide (v1 = v2) else
  match v1, v2 with
  | LitV (LitInt n1), LitV (LitInt n2) => Some $ LitV $ bin_op_eval_int op n1 n2
  | LitV (LitBool b1), LitV (LitBool b2) => LitV <$> bin_op_eval_bool op b1 b2
  | _, _ => None
  end.

(** CAS just compares the word-sized representation of the two values, it cannot
look into boxed data.  This works out fine if at least one of the to-be-compared
values is unboxed (exploiting the fact that an unboxed and a boxed value can
never be equal because these are disjoint sets).  *)
Definition vals_cas_compare_safe (vl v1 : val) : Prop :=
  val_is_unboxed vl ∨ val_is_unboxed v1.
Arguments vals_cas_compare_safe !_ !_ /.

Inductive head_step : expr → state → expr → state → list (expr) → Prop :=
  | BetaS f x e1 e2 v2 e' σ :
     to_val e2 = Some v2 →
     Closed (f :b: x :b: []) e1 →
     e' = subst' x (of_val v2) (subst' f (Rec f x e1) e1) →
     head_step (App (Rec f x e1) e2) σ e' σ []
  | UnOpS op e v v' σ :
     to_val e = Some v →
     un_op_eval op v = Some v' →
     head_step (UnOp op e) σ (of_val v') σ []
  | BinOpS op e1 e2 v1 v2 v' σ :
     to_val e1 = Some v1 → to_val e2 = Some v2 →
     bin_op_eval op v1 v2 = Some v' →
     head_step (BinOp op e1 e2) σ (of_val v') σ []
  | IfTrueS e1 e2 σ :
     head_step (If (Lit $ LitBool true) e1 e2) σ e1 σ []
  | IfFalseS e1 e2 σ :
     head_step (If (Lit $ LitBool false) e1 e2) σ e2 σ []
  | FstS e1 v1 e2 v2 σ :
     to_val e1 = Some v1 → to_val e2 = Some v2 →
     head_step (Fst (Pair e1 e2)) σ e1 σ []
  | SndS e1 v1 e2 v2 σ :
     to_val e1 = Some v1 → to_val e2 = Some v2 →
     head_step (Snd (Pair e1 e2)) σ e2 σ []
  | CaseLS e0 v0 e1 e2 σ :
     to_val e0 = Some v0 →
     head_step (Case (InjL e0) e1 e2) σ (App e1 e0) σ []
  | CaseRS e0 v0 e1 e2 σ :
     to_val e0 = Some v0 →
     head_step (Case (InjR e0) e1 e2) σ (App e2 e0) σ []
  | ForkS e σ:
     head_step (Fork e) σ (Lit LitUnit) σ [e]
  | AllocS e v σ l :
     to_val e = Some v → σ !! l = None →
     head_step (Alloc e) σ (Lit $ LitLoc l) (<[l:=v]>σ) []
  | LoadS l v σ :
     σ !! l = Some v →
     head_step (Load (Lit $ LitLoc l)) σ (of_val v) σ []
  | StoreS l e v σ :
     to_val e = Some v → is_Some (σ !! l) →
     head_step (Store (Lit $ LitLoc l) e) σ (Lit LitUnit) (<[l:=v]>σ) []
  | CasFailS l e1 v1 e2 v2 vl σ :
     to_val e1 = Some v1 → to_val e2 = Some v2 →
     σ !! l = Some vl → vl ≠ v1 →
     vals_cas_compare_safe vl v1 →
     head_step (CAS (Lit $ LitLoc l) e1 e2) σ (Lit $ LitBool false) σ []
  | CasSucS l e1 v1 e2 v2 σ :
     to_val e1 = Some v1 → to_val e2 = Some v2 →
     σ !! l = Some v1 →
     vals_cas_compare_safe v1 v1 →
     head_step (CAS (Lit $ LitLoc l) e1 e2) σ (Lit $ LitBool true) (<[l:=v2]>σ) []
  | FaaS l i1 e2 i2 σ :
     to_val e2 = Some (LitV (LitInt i2)) →
     σ !! l = Some (LitV (LitInt i1)) →
     head_step (FAA (Lit $ LitLoc l) e2) σ (Lit $ LitInt i1) (<[l:=LitV (LitInt (i1 + i2))]>σ) [].

(** Basic properties about the language *)
Instance fill_item_inj Ki : Inj (=) (=) (fill_item Ki).
Proof. destruct Ki; intros ???; simplify_eq/=; auto with f_equal. Qed.

Lemma fill_item_val Ki e :
  is_Some (to_val (fill_item Ki e)) → is_Some (to_val e).
Proof. intros [v ?]. destruct Ki; simplify_option_eq; eauto. Qed.

Lemma val_head_stuck e1 σ1 e2 σ2 efs : head_step e1 σ1 e2 σ2 efs → to_val e1 = None.
Proof. destruct 1; naive_solver. Qed.

Lemma head_ctx_step_val Ki e σ1 e2 σ2 efs :
  head_step (fill_item Ki e) σ1 e2 σ2 efs → is_Some (to_val e).
Proof. destruct Ki; inversion_clear 1; simplify_option_eq; by eauto. Qed.

Lemma fill_item_no_val_inj Ki1 Ki2 e1 e2 :
  to_val e1 = None → to_val e2 = None →
  fill_item Ki1 e1 = fill_item Ki2 e2 → Ki1 = Ki2.
Proof.
  destruct Ki1, Ki2; intros; try discriminate; simplify_eq/=;
    repeat match goal with
    | H : to_val (of_val _) = None |- _ => by rewrite to_of_val in H
    end; auto.
Qed.

Lemma alloc_fresh e v σ :
  let l := fresh (dom (gset loc) σ) in
  to_val e = Some v → head_step (Alloc e) σ (Lit (LitLoc l)) (<[l:=v]>σ) [].
Proof. by intros; apply AllocS, (not_elem_of_dom (D:=gset loc)), is_fresh. Qed.

(* Misc *)
Lemma to_val_rec f x e `{!Closed (f :b: x :b: []) e} :
  to_val (Rec f x e) = Some (RecV f x e).
Proof. rewrite /to_val. case_decide=> //. do 2 f_equal; apply proof_irrel. Qed.

(** Closed expressions *)
Lemma is_closed_weaken X Y e : is_closed X e → X ⊆ Y → is_closed Y e.
Proof. revert X Y; induction e; naive_solver (eauto; set_solver). Qed.

Lemma is_closed_weaken_nil X e : is_closed [] e → is_closed X e.
Proof. intros. by apply is_closed_weaken with [], list_subseteq_nil. Qed.

Lemma is_closed_of_val X v : is_closed X (of_val v).
Proof. apply is_closed_weaken_nil. induction v; simpl; auto. Qed.

Lemma is_closed_to_val X e v : to_val e = Some v → is_closed X e.
Proof. intros <-%of_to_val. apply is_closed_of_val. Qed.

Lemma is_closed_subst X e x es :
  is_closed [] es → is_closed (x :: X) e → is_closed X (subst x es e).
Proof.
  intros ?. revert X.
  induction e=> X /= ?; destruct_and?; split_and?; simplify_option_eq;
    try match goal with
    | H : ¬(_ ∧ _) |- _ => apply not_and_l in H as [?%dec_stable|?%dec_stable]
    end; eauto using is_closed_weaken with set_solver.
Qed.
Lemma is_closed_do_subst' X e x es :
  is_closed [] es → is_closed (x :b: X) e → is_closed X (subst' x es e).
Proof. destruct x; eauto using is_closed_subst. Qed.

(* Substitution *)
Lemma subst_is_closed X e x es : is_closed X e → x ∉ X → subst x es e = e.
Proof.
  revert X. induction e=> X /=; rewrite ?bool_decide_spec ?andb_True=> ??;
    repeat case_decide; simplify_eq/=; f_equal; intuition eauto with set_solver.
Qed.

Lemma subst_is_closed_nil e x es : is_closed [] e → subst x es e = e.
Proof. intros. apply subst_is_closed with []; set_solver. Qed.

Lemma subst_subst e x es es' :
  Closed [] es' → subst x es (subst x es' e) = subst x es' e.
Proof.
  intros. induction e; simpl; try (f_equal; by auto);
    simplify_option_eq; auto using subst_is_closed_nil with f_equal.
Qed.
Lemma subst_subst' e x es es' :
  Closed [] es' → subst' x es (subst' x es' e) = subst' x es' e.
Proof. destruct x; simpl; auto using subst_subst. Qed.

Lemma subst_subst_ne e x y es es' :
  Closed [] es → Closed [] es' → x ≠ y →
  subst x es (subst y es' e) = subst y es' (subst x es e).
Proof.
  intros. induction e; simpl; try (f_equal; by auto);
    simplify_option_eq; auto using eq_sym, subst_is_closed_nil with f_equal.
Qed.
Lemma subst_subst_ne' e x y es es' :
  Closed [] es → Closed [] es' → x ≠ y →
  subst' x es (subst' y es' e) = subst' y es' (subst' x es e).
Proof. destruct x, y; simpl; auto using subst_subst_ne with congruence. Qed.

Lemma subst_rec' f y e x es :
  x = f ∨ x = y ∨ x = BAnon →
  subst' x es (Rec f y e) = Rec f y e.
Proof. intros. destruct x; simplify_option_eq; naive_solver. Qed.
Lemma subst_rec_ne' f y e x es :
  (x ≠ f ∨ f = BAnon) → (x ≠ y ∨ y = BAnon) →
  subst' x es (Rec f y e) = Rec f y (subst' x es e).
Proof. intros. destruct x; simplify_option_eq; naive_solver. Qed.

Lemma heap_lang_mixin : EctxiLanguageMixin of_val to_val fill_item head_step.
Proof.
  split; apply _ || eauto using to_of_val, of_to_val, val_head_stuck,
    fill_item_val, fill_item_no_val_inj, head_ctx_step_val.
Qed.
End heap_lang.

(** Language *)
Canonical Structure heap_ectxi_lang := EctxiLanguage heap_lang.heap_lang_mixin.
Canonical Structure heap_ectx_lang := EctxLanguageOfEctxi heap_ectxi_lang.
Canonical Structure heap_lang := LanguageOfEctx heap_ectx_lang.

(* Prefer heap_lang names over ectx_language names. *)
Export heap_lang.

(** Define some derived forms. *)
Notation Lam x e := (Rec BAnon x e) (only parsing).
Notation Let x e1 e2 := (App (Lam x e2) e1) (only parsing).
Notation Seq e1 e2 := (Let BAnon e1 e2) (only parsing).
Notation LamV x e := (RecV BAnon x e) (only parsing).
Notation LetCtx x e2 := (AppRCtx (LamV x e2)) (only parsing).
Notation SeqCtx e2 := (LetCtx BAnon e2) (only parsing).
Notation Match e0 x1 e1 x2 e2 := (Case e0 (Lam x1 e1) (Lam x2 e2)) (only parsing).

Notation Skip := (Seq (Lit LitUnit) (Lit LitUnit)).