From Stdlib Require Import Program.Basics Program.Tactics.
From Equations.Prop Require Import Equations.
From Stdlib Require Import Vectors.VectorDef.
From Stdlib Require Import List.
Import ListNotations.
Set Equations Transparent.

Derive Signature NoConfusion NoConfusionHom for t.

Inductive Ty : Set :=
| unit : Ty
| bool : Ty
| arrow (t u : Ty) : Ty
| ref : Ty -> Ty.

Derive NoConfusion for Ty.

Infix "⇒" := arrow (at level 80).


Reserved Notation " x ∈ s " (at level 70, s at level 10).

#[universes(template)]
Inductive In {A} (x : A) : list A -> Type :=
| here {xs} : x ∈ (x :: xs)
| there {y xs} : x ∈ xs -> x ∈ (y :: xs)
where " x ∈ s " := (In x s).
Derive Signature NoConfusion for In.
Arguments here {A x xs}.
Arguments there {A x y xs} _.

From Equations Require Import Fin.

Section ScopeGraphs.
  Context (k : nat).

  Definition Scope := fin k.

  Definition Graph := Scope -> (list Ty * list Scope).

  Context (g : Graph).

  Definition declsOf (s : Scope) := fst (g s).
  Definition edgesOf (s : Scope) := snd (g s).

  Inductive sedge (s : Scope) : Scope -> Set :=
  | nilp : sedge s s
  | consp {s' s''} : s' ∈ edgesOf s -> sedge s' s'' -> sedge s s''.

  Notation "s ⟶ s'" := (sedge s s') (at level 20).

  Inductive resolve (s : Scope) (t : Ty) :=
  | path {s'} : s ⟶ s' -> t ∈ declsOf s' -> resolve s t.

  Notation "s ↦ t" := (resolve s t) (at level 20).
End ScopeGraphs.

Arguments resolve {k g}.
Notation "s ↦ t" := (resolve s t) (at level 20).

Section Expr.
  Context (k : nat).
  Context (g : Graph k).

  Inductive Expr (Γ : Scope k) : Ty -> Set :=
  | tt : Expr Γ unit
  | true : Expr Γ bool
  | false : Expr Γ bool
  | ite {t} : Expr Γ bool -> Expr Γ t -> Expr Γ t -> Expr Γ t
  | var {t} : Γ ↦ t -> Expr Γ t
  | abs {t u} : Expr (t :: Γ) u -> Expr Γ (t ⇒ u)
  | app {t u} : Expr Γ (t ⇒ u) -> Expr Γ t -> Expr Γ u
  | new {t} : Expr Γ t -> Expr Γ (ref t)
  | deref {t} : Expr Γ (ref t) -> Expr Γ t
  | assign {t} : Expr Γ (ref t) -> Expr Γ t -> Expr Γ unit.

Derive Signature NoConfusion NoConfusionHom for Expr.

#[universes(template)]
Inductive All {A} (P : A -> Type) : list A -> Type :=
| all_nil : All P []
| all_cons {x xs} : P x -> All P xs -> All P (x :: xs).
Arguments all_nil {A} {P}.
Arguments all_cons {A P x xs} _ _.
Derive Signature NoConfusion NoConfusionHom for All.

Section MapAll.
  Context  {A} {P Q : A -> Type} (f : forall x, P x -> Q x).

  Equations map_all {l : list A} : All P l -> All Q l :=
  map_all all_nil := all_nil;
  map_all (all_cons p ps) := all_cons (f _ p) (map_all ps).
End MapAll.

Definition StoreTy := list Ty.

Inductive Val : Ty -> StoreTy -> Set :=
| val_unit {Σ} : Val unit Σ
| val_true {Σ} : Val bool Σ
| val_false {Σ} : Val bool Σ
| val_closure {Σ Γ t u} : Expr (t :: Γ) u -> All (fun t => Val t Σ) Γ -> Val (t ⇒ u) Σ
| val_loc {Σ t} : t ∈ Σ -> Val (ref t) Σ.

Derive Signature NoConfusion NoConfusionHom for Val.

Definition Env (Γ : Ctx) (Σ : StoreTy) : Set := All (fun t => Val t Σ) Γ.

Definition Store (Σ : StoreTy) := All (fun t => Val t Σ) Σ.

Equations lookup : forall {A P xs} {x : A}, All P xs -> x ∈ xs -> P x :=
  lookup (all_cons p _) here := p;
  lookup (all_cons _ ps) (there ins) := lookup ps ins.

Equations update : forall {A P xs} {x : A}, All P xs -> x ∈ xs -> P x -> All P xs :=
  update (all_cons p ps) here        p' := all_cons p' ps;
  update (all_cons p ps) (there ins) p' := all_cons p (update ps ins p').

Equations lookup_store {Σ t} : t ∈ Σ -> Store Σ -> Val t Σ :=
  lookup_store l σ := lookup σ l.

Equations update_store {Σ t} : t ∈ Σ -> Val t Σ -> Store Σ -> Store Σ :=
  update_store l v σ := update σ l v.
Import Sigma_Notations.

Definition store_incl (Σ Σ' : StoreTy) := &{ Σ'' : _ & Σ' = Σ'' ++ Σ }.
Infix "⊑" := store_incl (at level 10).

Lemma app_assoc {A} (x y z : list A) : x ++ y ++ z = (x ++ y) ++ z.
Proof. induction x; simpl; auto.
       now rewrite IHx.
Defined.

Section StoreIncl.
  Context {Σ Σ' : StoreTy} (incl : Σ ⊑ Σ').

  Lemma pres_in t : t ∈ Σ -> t ∈ Σ'.
  Proof. destruct incl. subst. clear. induction pr1. intros. exact H.
         intros H. specialize (IHpr1 H). constructor 2. apply IHpr1.
  Defined.

  Equations(noind) weaken_val {t} (v : Val t Σ) : Val t Σ' := {
   weaken_val val_unit := val_unit;
   weaken_val val_true := val_true;
   weaken_val val_false := val_false;
   weaken_val (val_closure b e) := val_closure b (map_all (fun t v => weaken_val v) e); (* (weaken_vals e);  *)
   weaken_val (val_loc H) := val_loc (pres_in _ H) }.
  (* where weaken_vals {l} (a : All (fun t => Val t Σ) l) : All (fun t => Val t Σ') l by struct a := *)
  (* weaken_vals all_nil := all_nil; *)
  (* weaken_vals (all_cons p ps) := all_cons (weaken_val p) (weaken_vals ps). *)


  Definition weaken_env {Γ} (v : Env Γ Σ) : Env Γ Σ' :=
    map_all (@weaken_val) v.

  Lemma refl_incl : Σ ⊑ Σ.
  Proof. exists []. reflexivity. Defined.

  Lemma trans_incl {Σ''} (incl' : Σ' ⊑ Σ'') : Σ ⊑ Σ''.
  Proof.
    destruct incl as [? ->], incl' as [? ->].
    exists (pr0 ++ pr1). now rewrite app_assoc.
  Defined.

  Lemma store_ext_incl {t} : Σ ⊑ (t :: Σ).
  Proof. now exists [t]. Defined.

End StoreIncl.

Infix "⊚" := trans_incl (at level 10).

Equations M : forall (Γ : Ctx) (P : StoreTy -> Type) (Σ : StoreTy), Type :=
  M Γ P Σ := forall (E : Env Γ Σ) (μ : Store Σ), option &{ Σ' : _ & &{ _ : Store Σ' & &{ _ : P Σ' & Σ ⊑ Σ'}}}.

From Stdlib Require Import Utf8.
Notation "( x , .. , y , z )" := (sigmaI _ x .. (sigmaI _ y z) ..) : core_scope.

Equations bind {Σ Γ} {P Q : StoreTy -> Type} (f : M Γ P Σ) (g : ∀ {Σ'}, P Σ' -> M Γ Q Σ') : M Γ Q Σ :=
  bind f g E μ with f E μ :=
    { | None := None;
      | Some (Σ', μ', x, ext) with g _ x (weaken_env ext E) μ' :=
          { | None := None;
            | Some (_, μ'', y, ext') := Some (_, μ'', y, ext ⊚ ext') } }.

Infix ">>=" := bind (at level 20, left associativity).

Definition transp_op {Γ Σ P} (x : Store Σ -> P Σ) : M Γ P Σ :=
  fun E μ => Some (Σ, μ, x μ, refl_incl).

Equations ret : ∀ {Γ Σ P}, P Σ → M Γ P Σ :=
  ret (Σ:=Σ) a E μ := Some (Σ, μ, a, refl_incl).

Equations getEnv : ∀ {Γ Σ}, M Γ (Env Γ) Σ :=
  getEnv (Σ:=Σ) E μ := Some (Σ, μ, E, refl_incl).

Equations usingEnv {Γ Γ' Σ P} (E : Env Γ Σ) (m : M Γ P Σ) : M Γ' P Σ :=
  usingEnv E m E' μ := m E μ.

Equations timeout : ∀ {Γ Σ P}, M Γ P Σ :=
  timeout _ _ := None.

Section StoreOps.
  Context {Σ : StoreTy} {Γ : Ctx} {t : Ty}.

  Equations storeM (v : Val t Σ) : M Γ (Val (ref t)) Σ :=
    storeM v E μ :=
      let v : Val t (t :: Σ) := weaken_val store_ext_incl v in
      let μ' := map_all (fun t' => weaken_val store_ext_incl) μ in
      Some (t :: Σ, all_cons v μ', val_loc here, store_ext_incl).

  Equations derefM (l : t ∈ Σ) : M Γ (Val t) Σ :=
    derefM l := transp_op (lookup_store l).

  Equations updateM (l : t ∈ Σ) (v : Val t Σ) : M Γ (Val unit) Σ :=
    updateM l v E μ := Some (Σ, update_store l v μ, val_unit, refl_incl).
End StoreOps.

Reserved Notation "P ⊛ Q" (at level 10).

Inductive storepred_prod (P Q : StoreTy -> Type) : StoreTy -> Type :=
  | storepred_pair {Σ} : P Σ -> Q Σ -> (P ⊛ Q) Σ
where "P ⊛ Q" := (storepred_prod P Q).
Arguments storepred_pair {P Q Σ}.

Class Weakenable (P : StoreTy -> Type) : Type :=
  weaken : forall {Σ Σ'}, Σ ⊑ Σ' -> P Σ -> P Σ'.

Instance val_weaken {t} : Weakenable (Val t).
Proof. intros Σ Σ' incl. apply (weaken_val incl). Defined.

Instance env_weaken {Γ} : Weakenable (Env Γ).
Proof. intros Σ Σ' incl. apply (weaken_env incl). Defined.

Instance loc_weaken (t : Ty) : Weakenable (In t).
Proof. intros Σ Σ' incl. apply (pres_in incl). Defined.

Class IsIncludedOnce (Σ Σ' : StoreTy) : Type := is_included_once : Σ ⊑ Σ'.
Hint Mode IsIncludedOnce + + : typeclass_instances.

Instance IsIncludedOnce_ext {T} Σ : IsIncludedOnce Σ (T :: Σ).
Proof. apply store_ext_incl. Defined.

Class IsIncluded (Σ Σ' : StoreTy) : Type := is_included : Σ ⊑ Σ'.
Hint Mode IsIncluded + + : typeclass_instances.

Instance IsIncluded_refl Σ : IsIncluded Σ Σ := refl_incl.
Instance IsIncluded_trans Σ Σ' Σ'' : IsIncludedOnce Σ Σ' -> IsIncluded Σ' Σ'' -> IsIncluded Σ Σ''.
Proof. intros H H'. exact (trans_incl H H'). Defined.

Equations wk {Σ Σ' P} {W : Weakenable P} (p : P Σ) {incl : IsIncluded Σ Σ'} : P Σ' :=
  wk p := weaken incl p.

Equations bind_ext {Σ Γ} {P Q : StoreTy -> Type} (f : M Γ P Σ) (g : ∀ {Σ'} `{IsIncluded Σ Σ'}, P Σ' -> M Γ Q Σ') : M Γ Q Σ :=
  bind_ext f g E μ with f E μ :=
    { | None := None;
      | Some (Σ', μ', x, ext) with g _ ext x (weaken_env ext E) μ' :=
          { | None := None;
            | Some (_, μ'', y, ext') := Some (_, μ'', y, ext ⊚ ext') } }.

Infix ">>='" := bind_ext (at level 20, left associativity).

Equations eval_ext (n : nat) {Γ Σ t} (e : Expr Γ t) : M Γ (Val t) Σ :=
  eval_ext 0 _                := timeout;
  eval_ext (S k) tt           := ret val_unit;
  eval_ext (S k) true         := ret val_true;
  eval_ext (S k) false        := ret val_false;
  eval_ext (S k) (ite b t f)  := eval_ext k b >>=' λ{ | _ | ext | val_true => eval_ext k t;
                                                      | _ | ext | val_false => eval_ext k f };

  eval_ext (S k) (var x)      := getEnv >>=' fun {Σ ext} E => ret (lookup E x);
  eval_ext (S k) (abs x)      := getEnv >>=' fun {Σ ext} E => ret (val_closure x E);
  eval_ext (S k) (app (Γ:=Γ) e1 e2) :=
      eval_ext k e1 >>=' λ{ | _ | ext | val_closure e' E => 
      eval_ext k e2 >>=' fun {Σ' ext'} v => usingEnv (all_cons v (wk E)) (eval_ext k e')};
  eval_ext (S k) (new e)      := eval_ext k e >>=' fun {Σ ext} v => storeM v;
  eval_ext (S k) (deref l)    := eval_ext k l >>=' λ{ | _ | ext | val_loc l => derefM l };
  eval_ext (S k) (assign l e) := eval_ext k l >>=' λ{ | _ | ext | val_loc l =>
                                 eval_ext k e >>=' λ{ | _ | ext | v => updateM (wk l) (wk v) }}.

Equations strength {Σ Γ} {P Q : StoreTy -> Type} {w : Weakenable Q} (m : M Γ P Σ) (q : Q Σ) : M Γ (P ⊛ Q) Σ :=
  strength m q E μ with m E μ => {
    | None => None;
    | Some (Σ, μ', p, ext) => Some (Σ, μ', storepred_pair p (weaken ext q), ext) }.

Infix "^" := strength.
(* TODO improve pattern matching lambda to have implicit arguments implicit.
   Hard because Coq does not keep the implicit status of bind's [g] argument. *)

Equations eval (n : nat) {Γ Σ t} (e : Expr Γ t) : M Γ (Val t) Σ :=
  eval 0 _                := timeout;
  eval (S k) tt           := ret val_unit;
  eval (S k) true         := ret val_true;
  eval (S k) false        := ret val_false;
  eval (S k) (ite b t f)  := eval k b >>= λ{ | _ | val_true => eval k t;
                                             | _ | val_false => eval k f };

  eval (S k) (var x)      := getEnv >>= fun Σ E => ret (lookup E x);
  eval (S k) (abs x)      := getEnv >>= fun Σ E => ret (val_closure x E);
  eval (S k) (app (Γ:=Γ) e1 e2) :=
      eval k e1 >>= λ{ | _ | val_closure e' E =>
                             (eval k e2 ^ E) >>= fun Σ' '(storepred_pair v E) => usingEnv (all_cons v E) (eval k e')};
  eval (S k) (new e)      := eval k e >>= fun Σ v => storeM v;
  eval (S k) (deref l)    := eval k l >>= λ{ | _ | val_loc l => derefM l };
  eval (S k) (assign l e) := eval k l >>= λ{ | _ | val_loc l =>
                             (eval k e ^ l) >>= λ{ | _ | storepred_pair v l => updateM l v }}.

Definition idu : Expr [] (unit ⇒ unit) :=
  abs (var here).

Definition idapp : Expr [] unit := app idu tt.

Eval vm_compute in eval 100 idapp all_nil all_nil.

Definition neg : Expr [] (bool ⇒ bool) :=
  abs (ite (var here) false true).

Definition letref {t u} (v : Expr [] t) (b : Expr [ref t] u) : Expr [] u :=
  app (abs b) (new v).

Equations weaken_expr {Γ Γ' t u} (e1 : Expr (Γ ++ Γ') t) : Expr (Γ ++ u :: Γ') t :=
  weaken_expr tt             := tt;
  weaken_expr true           := true;
  weaken_expr false          := false;
  weaken_expr (ite b t f)    := ite (weaken_expr b) (weaken_expr t) (weaken_expr f);
  weaken_expr (var x)        := var _;
  weaken_expr (abs (t:=t) x) := abs (weaken_expr (Γ := t :: Γ) x);
  weaken_expr (app e1 e2)    := app (weaken_expr e1) (weaken_expr e2);
  weaken_expr (new e)        := new (weaken_expr e);
  weaken_expr (deref l)      := deref (weaken_expr l);
  weaken_expr (assign l e)   := assign (weaken_expr l) (weaken_expr e).
Next Obligation.
  clear weaken_expr.
  induction Γ in Γ', u, x |- *. now apply there. simpl.
  depelim x. constructor. apply there. apply IHΓ. apply x.
Defined.

Definition seq {Γ u} (e1 : Expr Γ unit) (e2 : Expr Γ u) : Expr Γ u :=
  app (abs (weaken_expr (Γ := []) e2)) e1.

(* let x = ref true in
   x := false; !x *)

Definition letupdate : Expr [] bool :=
  letref true (seq (assign (var here) false) (deref (var here))).

Eval vm_compute in eval 100 letupdate all_nil all_nil.








(*
Inductive eval_sem {Γ : Ctx} {Σ} {env : Env Γ Σ} : forall {t : Ty}, Expr Γ t -> Val t Σ -> Prop :=
| eval_tt (e : Expr Γ unit) : eval_sem e val_unit
| eval_var t (i : t ∈ Γ) : eval_sem (var i) (lookup env i)
| eval_abs {t u} (b : Expr (t :: Γ) u) : eval_sem (abs b) (val_closure b env)
| eval_app {t u} (f : Expr Γ (t ⇒ u)) b' (a : Expr Γ t) v :
    eval_sem f (val_closure b' env) ->
    eval_sem a v ->
    forall u, @eval_sem (t :: Γ) _ (all_cons v env) _ b' u ->
    eval_sem (app f a) u.



Lemma eval_correct {n} Γ Σ (μ : Store Σ) t (e : Expr Γ t) env v : eval n e env μ = Some v ->
                                                                  @eval_sem _ _ (weaken v.2.2.2 env) _ e (v.2.2.1).
Proof.
  induction n. intros; discriminate.
  destruct e; simp eval; try intros [= <-]; simpl; try constructor.
  admit. admit.



  pose proof (fun_elim (f:=eval)).
  specialize (H (fun n Γ Σ t e m => forall env v μ, m env μ = Some v -> @eval_sem _ _ (weaken v.2.2.2 env) _ e v.2.2.1)
                (fun n Γ Σ t u f a v m => forall env v',
                     @eval_sem _ _ env _ f v.2.2.1 -> m env = Some v' -> @eval_sem _ env _ (app f a) v')).
  rapply H; clear; intros.
  discriminate.
  noconf H. constructor.
  noconf H. constructor.

  noconf H. constructor.

  unfold bind in H1.
  destruct (eval n e0 env) eqn:Heq.
  specialize (H _ _ Heq).
  specialize (H0 v0 _ _ H H1). apply H0.
  discriminate.

  unfold bind in H2.
  destruct (eval k arg env) eqn:Heq.
  specialize (H _ _ Heq).
  unfold usingEnv in H2. specialize (H0 v (all_cons v a) v').
  econstructor; eauto.
Admitted.*)