(***************************************************************)
(* Improvement of the introductory case study in               *)
(* "The Braga method", D. Larchey-Wendling and JF. Monin, 2021 *)
(***************************************************************)

(*Require Import Arith Lia Utf8 Extraction.*)
Require Import Utf8 Extraction.

Print option.
(*
Inductive option (A : Type) : Type :=
  Some : A → option A | None : option A.
 *)

(** Target : nb of g steps to get a None, with g : X -> option X

   let rec ns ox = match ox with None => 0 | Some x => S (ns (g x))

   Algorithm using an accumulator

   let rec nsa x n = match ox with None => n | Some x => nsa (g x) (S n)

   We prove that ns computes the same result as nsa
*)

Section ns_nsa.

Variable (X : Type) (g : X → option X).

Inductive 𝔻ns : option X ->  Prop :=
| D_None : 𝔻ns None
| D_Some : ∀x, 𝔻ns (g x) → 𝔻ns (Some x)
.
(** The small projection of constructor 𝔻ns_Some *)
Definition π_𝔻ns {x} (δ : 𝔻ns (Some x)) : 𝔻ns (g x) :=
  match δ in 𝔻ns ox return
        let x     := match ox with Some x => x    | _ => x     end in
        let shape := match ox with Some x => True | _ => False end in
        shape → 𝔻ns (g x)
  with
  | D_Some x δgx => λ sh, δgx
  | _            => λ sh, match sh with end
  end I.

(** The same, interactive definition *)
Definition π_𝔻ns_script {x} (δ : 𝔻ns (Some x)) : 𝔻ns (g x).
Proof.
  pose (usome ox := match ox with Some x' => x' | _ => x end).
  change x with (usome (Some x)).
  pose (shape (ox : option X) := match ox with Some x' => True | _ => False end).
  refine ((_ : shape (Some x) -> 𝔻ns (g (usome (Some x)))) I).
  destruct δ as [ | x'' δ'']; cbn.
  - intro sh. case sh.
  - intro sh. exact δ''.
Defined.

(** The same, easier to read *)
Definition π_𝔻ns' {x} (δ : 𝔻ns (Some x)) : 𝔻ns (g x) :=
  match δ in 𝔻ns ox return
        let x'     := match ox with Some x' => x'   | _ => x     end in
        let shape  := match ox with Some x' => True | _ => False end in
        shape → 𝔻ns (g x')
  with
  | D_Some x δgx => λ sh, δgx
  | _            => λ sh, match sh with end
  end I.

(** First we define ns/nsa as structural fixpoint on 
    the domain predicate *)

Fixpoint ns ox (δ : 𝔻ns ox) {struct δ} : nat :=
  match ox return 𝔻ns ox → nat with
  | None   => λ _, 0
  | Some x => λ δ, S (ns (g x) (π_𝔻ns δ))
  end δ.

Fixpoint nsa ox n (δ : 𝔻ns ox) : nat :=
  match ox with
  | None   => λ _, n
  | Some x => λ δ, nsa (g x) (S n) (π_𝔻ns δ)
  end δ.

(* ------------------------------------------------------------ *)
(** Additional (advanced) material: instructive failures.
    You can skip them. *)

Fixpoint ns ox (δ : 𝔻ns ox) {struct δ} : nat.
  refine (
  match ox return nat with
  | None   => 0
  | Some x => S (ns (g x) _)
  end).
  Fail Check (π_𝔻ns δ).
  (* stuck *)
Abort.

Fail Fixpoint ns ox (δ : 𝔻ns ox) {struct δ} : nat :=
  match ox return nat with
  | None   => 0
  | Some x => S (ns (g x) (π_𝔻ns δ))
  end.

(** Close to the solution but happens to not satisfy the "guard condition"  *)
Definition π_𝔻ns_large {x} (δ : 𝔻ns (Some x)) : 𝔻ns (g x) :=
  match δ in 𝔻ns ox return
        let x     := match ox with Some x => x    | _ => x     end in
        let shape := match ox with Some x => True | _ => False end in
        match ox with Some _ => 𝔻ns (g x) | _ =>  shape → 𝔻ns (g x) end
  with
  | D_Some x δgx => δgx
  | _            => λ sh, match sh with end
  end.

Fail Fixpoint ns_large ox (δ : 𝔻ns ox) {struct δ} : nat :=
  match ox return 𝔻ns ox → nat with
  | None   => λ _, 0
  | Some x => λ δ, S (ns_large (g x) (π_𝔻ns_large δ))
  end δ.

(** Same problem with the "guard condition" *)
Definition π_𝔻ns_large2 {x} (δ : 𝔻ns (Some x)) : 𝔻ns (g x) :=
  match δ in 𝔻ns ox return
        let x     := match ox with Some x => x    | _ => x     end in
        let shape := match ox with Some x => True | _ => False end in
        match ox with Some _ => shape → 𝔻ns (g x) | _ =>  shape → 𝔻ns (g x) end
  with
  | D_Some x δgx => λ sh, δgx
  | _            => λ sh, match sh with end
  end I.

Fail Fixpoint ns_large2 ox (δ : 𝔻ns ox) {struct δ} : nat :=
  match ox return 𝔻ns ox → nat with
  | None   => λ _, 0
  | Some x => λ δ, S (ns_large2 (g x) (π_𝔻ns_large2 δ))
  end δ.

(** The difference between π_𝔻ns_large2 and π_𝔻ns looks tiny:
    at first sight, the underlying computation is the same,
    only the typing is different.
    Maybe because that the returned type in π_𝔻ns_large2 is obtained
    using a large elimination (same observation in π_𝔻ns_large)
    Anyway, according to a private discussion with the Coq team in 2018,
    allowing this would conflict with axioms such as
    [True  =  False → False] that are fine to homotopists.
 *)
(** End of additional (advanced) material: instructive failures. *)
(* ------------------------------------------------------------- *)

(* We can show this identity by structural recursion on δ *)

Lemma ns_nsa_n : ∀ox n δ, nsa ox n δ = n + ns ox δ.
Proof.
  refine (fix loop ox n δ {struct δ} := _).
  destruct δ as [ | x δ ]; cbn.
  - apply plus_n_O. 
  - rewrite (loop (g x) (S n) δ). cbn. apply plus_n_Sm. 
Qed.

Corollary ns_nsa : ∀ox (δ: 𝔻ns ox), nsa ox 0 δ = ns ox δ.
Proof. intros. apply ns_nsa_n. Qed.

(** Remark: see "more advanced material" below
Lemma ns_nsa_n_advanced : ∀ox n δ δ', nsa ox n δ = n + ns ox δ'.
 *)

(* ====================================================================== *)
(* The relational graph of ns *)
Inductive 𝔾ns : option X → nat → Prop :=
  | G_None       :                  𝔾ns  None    0
  | G_Some {x r} : 𝔾ns (g x) r  →  𝔾ns (Some x) (S r)
.
Inductive 𝔾nsa : option X → nat → nat → Prop :=
  | Ga_None {n}     :                         𝔾nsa   None   n n
  | Ga_Some {x n r} : 𝔾nsa (g x) (S n) r  →  𝔾nsa (Some x) n r
.

Arguments exist {A P}.
Fixpoint ns_cbc ox (δ : 𝔻ns ox) {struct δ} : {r | 𝔾ns ox r} :=
  match ox return 𝔻ns ox → _ with
  | None   => λ _, exist 0                                  (* *)  G_None
  | Some x => λ δ, let (r, γ) := ns_cbc (g x) (π_𝔻ns δ) in
                   exist (S r)                              (* *)  (G_Some γ)
  end δ.

Fixpoint nsa_cbc ox n (δ : 𝔻ns ox) {struct δ} : {r | 𝔾nsa ox n r} :=
  match ox return 𝔻ns ox → _ with
  | None   => λ _, exist n                                  (* *)  Ga_None
  | Some x => λ δ, let (r, γ) := nsa_cbc (g x) (S n) (π_𝔻ns δ) in
                   exist r                                  (* *)  (Ga_Some γ)
  end δ.

(* ---------------------------------------------------------------------- *)
(** Properties of ns and nsa are studied on the relational graphs *)

(* ------------------------------------------------------ *)
(* Small inversion of 𝔾ns *)
Inductive 𝔾ns_None : nat → Prop := G_None'     : 𝔾ns_None 0.
Inductive 𝔾ns_Some x : nat → Prop := G_Some' r : 𝔾ns (g x) r → 𝔾ns_Some x (S r).
Definition 𝔾ns_dispatch (ox : option X) : nat → Prop :=
  match ox with
  | None   => 𝔾ns_None
  | Some x => 𝔾ns_Some x
  end.
Lemma 𝔾ns_inv {ox r} (γ : 𝔾ns ox r) : 𝔾ns_dispatch ox r.
Proof. destruct γ; constructor; assumption. Qed.
(* ------------------------------------------------------ *)

Lemma Gns_nsa_n : ∀ox n ra r, 𝔾nsa ox n ra → 𝔾ns ox r → ra = n + r.
Proof.
  intros ox n ra r Ga. revert r.
  induction Ga as [n | x n ra γa Hγa]; intros r G.
  - destruct (𝔾ns_inv G). apply plus_n_O.
  - destruct (𝔾ns_inv G) as [r γ]. rewrite (Hγa _ γ). apply plus_n_Sm.
Qed.

Corollary Gns_nsa : ∀ox ra r, 𝔾nsa ox 0 ra → 𝔾ns ox r → ra = r.
Proof. intros ox ra r. apply Gns_nsa_n. Qed.

(* ------------------------------------------------------ *)
(* Both graphs are deterministic/functional *)

Lemma Gns_deterministic : ∀ox r r', 𝔾ns ox r → 𝔾ns ox r' → r = r'.
  intros ox r r' G. revert r'.
  induction G as [n | x r γ Hγ]; intros r' G'.
  - destruct (𝔾ns_inv G'). reflexivity.
  - destruct (𝔾ns_inv G') as [r' γ']. rewrite (Hγ _ γ'). reflexivity.
Qed.

(** Exercise *)
Lemma Gnsa_deterministic : ∀ox n r r', 𝔾nsa ox n r → 𝔾nsa ox n r' → r = r'.
Proof.
Admitted.

(* ====================================================================== *)
(** More advanced material *)
(** Ref: "small inversions for smaller inversions", JF. Monin, TYPES'22 *)

(** Dependent small inversions of 𝔻ns *)
(** Strictly more powerful than standard Coq inversion and needed below *)
Inductive 𝔻ns_None : 𝔻ns None → Prop := D_None' : 𝔻ns_None D_None.
Inductive 𝔻ns_Some x : 𝔻ns (Some x) → Prop :=
  D_Some' : ∀ δ : 𝔻ns (g x), 𝔻ns_Some x (D_Some x δ).
Definition 𝔻ns_dispatch ox : 𝔻ns ox → Prop :=
  match ox with
  | None => 𝔻ns_None
  | Some x => 𝔻ns_Some x
  end.
Lemma 𝔻ns_inv {ox} (δ : 𝔻ns ox ) : 𝔻ns_dispatch ox δ.
Proof. destruct δ; constructor; assumption. Qed.
(* ------------------------------------------------------ *)

Lemma ns_nsa_n_advanced : ∀ox n δ δ', nsa ox n δ = n + ns ox δ'.
Proof.
  refine (fix loop ox n δ δ' {struct δ} := _).
  destruct δ as [ | x δ ]; cbn.
  - destruct (𝔻ns_inv δ'). cbn. apply plus_n_O.
  - destruct (𝔻ns_inv δ') as [δ']. cbn.
    rewrite (loop _ _ δ δ'). cbn. apply plus_n_Sm. 
Qed.

Corollary ns_nsa_advanced : ∀ox (δ δ': 𝔻ns ox), nsa ox 0 δ = ns ox δ'.
Proof. intros. apply (ns_nsa_n_advanced ox 0). Qed.


(** Exercises (domain irrelevance):
    show that the results of ns and nsa don't depend on δ *)
Lemma ns_domirr : ∀ox (δ δ': 𝔻ns ox), ns ox δ = ns ox δ'.
Admitted.

Lemma nsa_domirr : ∀ox n (δ δ': 𝔻ns ox), nsa ox n δ = nsa ox n δ'.
Admitted.

(* ====================================================================== *)

End ns_nsa.

Recursive Extraction ns ns_cbc nsa nsa_cbc.