import for_mathlib.snake_lemma3
import for_mathlib.les_homology
import for_mathlib.snake_lemma_naturality

noncomputable theory

open category_theory category_theory.limits

namespace category_theory

section

local attribute [-instance] category_theory.prod

@[elab_as_eliminator]
lemma preorder_prod_induction {C D : Type*} [preorder C] [preorder D]
  {motive : Π ⦃i j : C × D⦄ (f : i ⟶ j), Prop}
  (comp : ∀ {i j k : C × D} (f : i ⟶ j) (g : j ⟶ k), motive f → motive g → motive (f ≫ g))
  (H1 : ∀ (i : C) {j k : D} (f : j ≤ k), @motive (i,j) (i,k) (hom_of_le $ ⟨le_rfl, f⟩))
  (H2 : ∀ {i j : C} (k : D) (f : i ≤ j), @motive (i,k) (j,k) (hom_of_le $ ⟨f, le_rfl⟩))
  ⦃i j : C × D⦄ (f : i ⟶ j) : motive f :=
begin
  cases i with i1 i2, cases j with j1 j2,
  convert comp _ _ (H1 i1 f.le.2) (H2 j2 f.le.1),
end

end

variables {C D : Type*} [category C] [category D]

@[elab_as_eliminator]
lemma prod_induction
  {motive : Π ⦃i j : C × D⦄ (f : i ⟶ j), Prop}
  (comp : ∀ {i j k : C × D} (f : i ⟶ j) (g : j ⟶ k), motive f → motive g → motive (f ≫ g))
  (H1 : ∀ (i : C) {j k : D} (f : j ⟶ k), @motive (i,j) (i,k) (𝟙 i, f))
  (H2 : ∀ {i j : C} (k : D) (f : i ⟶ j), @motive (i,k) (j,k) (f, 𝟙 k))
  ⦃i j : C × D⦄ (f : i ⟶ j) : motive f :=
begin
  let f1 : (i.1, i.2) ⟶ (i.1, j.2) := (𝟙 i.1, f.2),
  let f2 : (i.1, j.2) ⟶ (j.1, j.2) := (f.1, 𝟙 j.2),
  have hf : f = f1 ≫ f2,
  { ext; simp only [prod_comp_fst, prod_comp_snd, category.id_comp, category.comp_id], },
  rw hf, cases i, cases j,
  apply comp; apply_assumption,
end

@[elab_as_eliminator]
lemma fin_induction (n : ℕ)
  {motive : Π ⦃i j : fin (n+1)⦄ (f : i ≤ j), Prop}
  (id : ∀ i, motive (le_refl i))
  (comp : ∀ {i j k : fin (n+1)} (f : i ≤ j) (g : j ≤ k), motive f → motive g → motive (f.trans g : i ≤ k))
  (Hsucc : ∀ (i : fin n), @motive i.cast_succ i.succ (le_of_lt $ by { rw fin.cast_succ_lt_iff_succ_le }))
  ⦃i j : fin (n+1)⦄ (f : i ≤ j) : motive f :=
begin
  revert f,
  refine fin.induction_on j _ _; clear j,
  { intro f, have hi : i = 0, { erw eq_bot_iff, exact f }, subst i, convert id _, },
  { intros j IH f,
    obtain (hij|rfl|hij) := lt_trichotomy i j.succ,
    { rw ← fin.le_cast_succ_iff at hij,
      convert comp _ _ (IH hij) (Hsucc j), },
    { convert id _, },
    { exact (f.not_lt hij).elim } }
end

end category_theory

variables {C 𝓐 : Type*} [category C] [category 𝓐] [abelian 𝓐]

namespace homological_complex

variables {ι : Type*} {c : complex_shape ι}

local notation x `⟶[`D`]` y := D.map (snake_diagram.hom x y)

def cast_horizontal (i : fin 4) (j : fin 2) : snake_diagram := (i,j.cast_succ)
def cast_vertical (i : fin 3) (j : fin 3) : snake_diagram := (i.cast_succ,j)
def succ_horizontal (i : fin 4) (j : fin 2) : snake_diagram := (i, j.succ)
def succ_vertical (i : fin 3) (j : fin 3) : snake_diagram := (i.succ,j)
def to_succ_horizontal (i : fin 4) (j : fin 2) :
  cast_horizontal i j ⟶ succ_horizontal i j := snake_diagram.hom _ _
def to_succ_vertical ( i : fin 3) (j : fin 3) :
  cast_vertical i j ⟶ succ_vertical i j := snake_diagram.hom _ _

lemma snake_diagram_induction
  {motive : Π ⦃i j : snake_diagram⦄ (f : i ⟶ j), Prop}
  (id : ∀ i : snake_diagram, motive (𝟙 i))
  (comp : ∀ (i j k : snake_diagram) (f : i ⟶ j) (g : j ⟶ k),
    motive f → motive g → motive (f ≫ g))
  (succ_horizontal : ∀ (i : fin 4) (j : fin 2),
    motive (to_succ_horizontal i j))
  (succ_vertical : ∀ (i : fin 3) (j : fin 3),
    motive (to_succ_vertical i j)) ⦃i j : snake_diagram⦄ (f : i ⟶ j) : motive f :=
begin
  apply category_theory.preorder_prod_induction comp; clear f i j,
  { intros i,
    refine @category_theory.fin_induction 2
      (λ j k f, motive (hom_of_le $ (⟨le_refl i, f⟩ : (i,j) ≤ (i,k)))) _ _ _,
    { intros j, convert id _, },
    { intros i' j k f g hf hg, convert comp _ _ _ _ _ hf hg, },
    { intros j, convert succ_horizontal i j } },
  { intros i j k, revert i j,
    refine @category_theory.fin_induction 3
      (λ i j f, motive (hom_of_le $ (⟨f, le_refl k⟩ : (i,k) ≤ (j,k)))) _ _ _,
    { intros j, convert id _, },
    { intros i' j k f g hf hg, convert comp _ _ _ _ _ hf hg, },
    { intros i, convert succ_vertical i k } },
end

variables
  {X Y Z : C ⥤ homological_complex 𝓐 c} (f : X ⟶ Y) (g : Y ⟶ Z)
  (H : ∀ c i, short_exact ((f.app c).f i) ((g.app c).f i))
  {c₁ c₂ : C} (φ : c₁ ⟶ c₂) (i j : ι) (hij : c.rel i j)

def mk_snake_diagram_nat_trans_app : Π (e : snake_diagram),
  (snake (f.app c₁) (g.app c₁) (H _) i j hij).snake_diagram.obj e ⟶
  (snake (f.app c₂) (g.app c₂) (H _) i j hij).snake_diagram.obj e
| ⟨⟨0,_⟩,⟨0,_⟩⟩ := (homology_functor _ _ i).map (X.map φ)
| ⟨⟨0,_⟩,⟨1,_⟩⟩ := (homology_functor _ _ i).map (Y.map φ)
| ⟨⟨0,_⟩,⟨2,_⟩⟩ := (homology_functor _ _ i).map (Z.map φ)
| ⟨⟨1,_⟩,⟨0,_⟩⟩ := (mod_boundaries_functor _).map (X.map φ)
| ⟨⟨1,_⟩,⟨1,_⟩⟩ := (mod_boundaries_functor _).map (Y.map φ)
| ⟨⟨1,_⟩,⟨2,_⟩⟩ := (mod_boundaries_functor _).map (Z.map φ)
| ⟨⟨2,_⟩,⟨0,_⟩⟩ := (cycles_functor _ _ _).map (X.map φ)
| ⟨⟨2,_⟩,⟨1,_⟩⟩ := (cycles_functor _ _ _).map (Y.map φ)
| ⟨⟨2,_⟩,⟨2,_⟩⟩ := (cycles_functor _ _ _).map (Z.map φ)
| ⟨⟨3,_⟩,⟨0,_⟩⟩ := (homology_functor _ _ j).map (X.map φ)
| ⟨⟨3,_⟩,⟨1,_⟩⟩ := (homology_functor _ _ j).map (Y.map φ)
| ⟨⟨3,_⟩,⟨2,_⟩⟩ := (homology_functor _ _ j).map (Z.map φ)
| _ := 0 -- impossible case
.

def mk_snake_diagram_nat_trans_hor :
  ∀ (a : fin 4) (b : fin 2),
  (snake (f.app c₁) (g.app c₁) (H _) i j hij).snake_diagram.map (to_succ_horizontal a b) ≫
    mk_snake_diagram_nat_trans_app f g H φ i j hij (succ_horizontal a b) =
    mk_snake_diagram_nat_trans_app f g H φ i j hij (cast_horizontal a b) ≫
    (snake (f.app c₂) (g.app c₂) (H _) i j hij).snake_diagram.map (to_succ_horizontal a b)
| ⟨0,_⟩ ⟨0,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_f0, ← category_theory.functor.map_comp] }, rw nat_trans.naturality, }
| ⟨0,_⟩ ⟨1,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_g0, ← category_theory.functor.map_comp] }, rw nat_trans.naturality, }
| ⟨0,_⟩ ⟨n+2,h⟩ := by { exfalso, rw [nat.succ_lt_succ_iff, nat.succ_lt_succ_iff] at h, exact nat.not_lt_zero n h }
| ⟨1,_⟩ ⟨0,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_f1, ← category_theory.functor.map_comp] }, rw nat_trans.naturality, }
| ⟨1,_⟩ ⟨1,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_g1, ← category_theory.functor.map_comp] }, rw nat_trans.naturality, }
| ⟨1,_⟩ ⟨n+2,h⟩ := by { exfalso, rw [nat.succ_lt_succ_iff, nat.succ_lt_succ_iff] at h, exact nat.not_lt_zero n h }
| ⟨2,_⟩ ⟨0,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_f2, ← category_theory.functor.map_comp] }, rw nat_trans.naturality, }
| ⟨2,_⟩ ⟨1,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_g2, ← category_theory.functor.map_comp] }, rw nat_trans.naturality, }
| ⟨2,_⟩ ⟨n+2,h⟩ := by { exfalso, rw [nat.succ_lt_succ_iff, nat.succ_lt_succ_iff] at h, exact nat.not_lt_zero n h }
| ⟨3,_⟩ ⟨0,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_f3, ← category_theory.functor.map_comp] }, rw nat_trans.naturality, }
| ⟨3,_⟩ ⟨1,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_g3, ← category_theory.functor.map_comp] }, rw nat_trans.naturality, }
| ⟨3,_⟩ ⟨n+2,h⟩ := by { exfalso, rw [nat.succ_lt_succ_iff, nat.succ_lt_succ_iff] at h, exact nat.not_lt_zero n h }
| ⟨n+4,h⟩ _   := by { exfalso, repeat { rw [nat.succ_lt_succ_iff] at h }, exact nat.not_lt_zero n h }
.

def mk_snake_diagram_nat_trans_ver :
  ∀ (a b : fin 3),
  (snake (f.app c₁) (g.app c₁) (H _) i j hij).snake_diagram.map (to_succ_vertical a b) ≫
    mk_snake_diagram_nat_trans_app f g H φ i j hij (succ_vertical a b) =
    mk_snake_diagram_nat_trans_app f g H φ i j hij (cast_vertical a b) ≫
    (snake (f.app c₂) (g.app c₂) (H _) i j hij).snake_diagram.map (to_succ_vertical a b)
| ⟨0,_⟩ ⟨0,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_a0] }, erw nat_trans.naturality, refl }
| ⟨0,_⟩ ⟨1,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_b0] }, erw nat_trans.naturality, refl }
| ⟨0,_⟩ ⟨2,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_c0] }, erw nat_trans.naturality, refl }
| ⟨0,_⟩ ⟨n+3,h⟩ := by { exfalso, repeat { rw [nat.succ_lt_succ_iff] at h }, exact nat.not_lt_zero n h }
| ⟨1,_⟩ ⟨0,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_a1] }, erw nat_trans.naturality, refl }
| ⟨1,_⟩ ⟨1,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_b1] }, erw nat_trans.naturality, refl }
| ⟨1,_⟩ ⟨2,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_c1] }, erw nat_trans.naturality, refl }
| ⟨1,_⟩ ⟨n+3,h⟩ := by { exfalso, repeat { rw [nat.succ_lt_succ_iff] at h }, exact nat.not_lt_zero n h }
| ⟨2,_⟩ ⟨0,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_a2] }, erw nat_trans.naturality, refl }
| ⟨2,_⟩ ⟨1,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_b2] }, erw nat_trans.naturality, refl }
| ⟨2,_⟩ ⟨2,_⟩ := by { repeat { erw [snake_diagram.mk_functor_map_c2] }, erw nat_trans.naturality, refl }
| ⟨2,_⟩ ⟨n+3,h⟩ := by { exfalso, repeat { rw [nat.succ_lt_succ_iff] at h }, exact nat.not_lt_zero n h }
| ⟨n+3,h⟩ _   := by { exfalso, repeat { rw [nat.succ_lt_succ_iff] at h }, exact nat.not_lt_zero n h }
.

-- TODO: Make a general construction, similar to `snake_diagram.mk_functor`
def mk_snake_diagram_nat_trans :
  (snake (f.app c₁) (g.app c₁) (H _) i j hij).snake_diagram ⟶
  (snake (f.app c₂) (g.app c₂) (H _) i j hij).snake_diagram :=
{ app := λ e, mk_snake_diagram_nat_trans_app f g H φ i j hij e,
  naturality' := begin
    apply snake_diagram_induction,
    { intro, simp only [category_theory.functor.map_id, category.id_comp, category.comp_id] },
    { intros i j k f g h1 h2, simp only [functor.map_comp, category.assoc, h2, reassoc_of h1] },
    { exact mk_snake_diagram_nat_trans_hor f g H φ i j hij },
    { exact mk_snake_diagram_nat_trans_ver f g H φ i j hij },
  end }

lemma δ_natural :
  δ (f.app c₁) (g.app c₁) (H _) i j hij ≫ (homology_functor _ _ j).map (X.map φ) =
    (homology_functor _ _ i).map (Z.map φ) ≫ δ (f.app c₂) (g.app c₂) (H _) i j hij :=
begin
  let η := mk_snake_diagram_nat_trans f g H φ i j hij,
  apply (snake_lemma.δ_natural η _ _).symm,
end

end homological_complex


import mcl.defs
import mcl.rhl
import mcl.compute_list

open parlang
open parlang.state
open parlang.thread_state
open mcl
open mcl.rhl

inductive op (sig : signature)
| store {t} {dim} (var : string) (idx : vector (expression sig type.int) dim) (h₁ : type_of (sig.val var) = t) (h₂ : ((sig.val var).type).dim = dim) : op
| compute_list (computes : list (memory (parlang_mcl_tlocal sig) → memory (parlang_mcl_tlocal sig))) : op

def ts_updates {sig : signature} : list (op sig) → thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig) → thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig)
| [] ts := ts
| (op.store var idx h₁ h₂ :: ops) ts := ts_updates ops $ thread_state.tlocal_to_shared var idx h₁ h₂ ts
| (op.compute_list computes :: ops) ts := ts_updates ops $ compute_list computes ts

lemma ts_update_compute_list {sig : signature} (ups : list (op sig)) (computes) : ts_updates (op.compute_list computes :: ups) = ts_updates ups ∘ compute_list computes := by refl

lemma ts_update_split {sig : signature} (up) (ups : list (op sig)) : ts_updates (list.reverse (up :: ups)) = ts_updates [up] ∘ ts_updates (list.reverse ups) := begin
    funext ts,
    rw list.reverse_cons,
    induction (list.reverse ups) generalizing ts,
    {
        refl,
    }, {
        simp,
        cases hd;
        rw ts_updates;
        apply ih,
    }
end

lemma ts_updates_tlocal {sig : signature}  {ts : thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig)} {updates} (m loads stores) :
(ts_updates updates ts).tlocal = (ts_updates updates { tlocal := ts.tlocal, loads := loads, stores := stores, shared := m }).tlocal := begin
    sorry,
end

lemma ts_updates_nil {sig : signature} (f : thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig) → thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig)) :
ts_updates [] ∘ f = f := begin
    refl,
end

@[simp]
lemma ts_updates_store {sig : signature} {dim} {idx : vector (expression sig type.int) dim} {var t} {h₁ : type_of (sig.val var) = t} {h₂} {updates} (f : thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig) → thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig)) :
ts_updates updates ∘ thread_state.tlocal_to_shared var idx h₁ h₂ ∘ f = ts_updates (op.store var idx h₁ h₂ :: updates) ∘ f := begin
    refl,
end

@[simp]
lemma ts_updates_compute {sig : signature} {g} {updates} (f : thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig) → thread_state (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig)) :
ts_updates updates ∘ compute g ∘ f = ts_updates (op.compute_list [g] :: updates) ∘ f := begin
    refl,
end

@[simp]
lemma ts_updates_merge_computes_list {sig : signature} {updates} {com com' : list (memory (parlang_mcl_tlocal sig) → memory (parlang_mcl_tlocal sig))} :
ts_updates (op.compute_list com :: op.compute_list com' :: updates) = ts_updates (op.compute_list (com ++ com') :: updates) := begin
    sorry,
end

@[simp]
lemma compute_list_stores' {sig : signature} {n} {tid : fin n} {ac : vector bool n} {computes}
{s : state n (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig)} :
(vector.nth ((map_active_threads ac (ts_updates [op.compute_list computes]) s).threads) tid).stores = (vector.nth s.threads tid).stores := begin
    by_cases h : ac.nth tid = tt,
    {
        simp [ts_updates, map_active_threads_nth_ac h],
    }, {
        rw ←map_active_threads_nth_inac h,
    }
end

@[simp]
lemma compute_list_loads' {sig : signature} {n} {tid : fin n} {ac : vector bool n} {computes}
{s : state n (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig)} :
(vector.nth ((map_active_threads ac (ts_updates [op.compute_list computes]) s).threads) tid).loads = (vector.nth s.threads tid).loads := begin
    by_cases h : ac.nth tid = tt,
    {
        simp [ts_updates, map_active_threads_nth_ac h],
    }, {
        rw ←map_active_threads_nth_inac h,
    }
end

@[simp]
lemma compute_list_shared' {sig : signature} {n} {tid : fin n} {ac : vector bool n} {computes}
{s : state n (memory $ parlang_mcl_tlocal sig) (parlang_mcl_shared sig)} :
(vector.nth ((map_active_threads ac (ts_updates [op.compute_list computes]) s).threads) tid).shared = (vector.nth s.threads tid).shared := begin
    by_cases h : ac.nth tid = tt,
    {
        simp [ts_updates, map_active_threads_nth_ac h],
    }, {
        rw ←map_active_threads_nth_inac h,
    }
end


-- TODO: Adapt to `HasIdentity`:
-- Add type classes to "upgrade" a meta-relation to a relation,
-- and especially to upgrade instance equivalences to an equality-like recursor
-- (see `IsIdentity` below).
#exit

import UniverseAbstractions.Axioms.Universes
import UniverseAbstractions.Axioms.Universe.Functors
import UniverseAbstractions.Axioms.Universe.Products
import UniverseAbstractions.Axioms.Universe.Equivalences
import UniverseAbstractions.Axioms.Universe.DependentTypes.Properties
import UniverseAbstractions.Axioms.Universe.DependentTypes.DependentFunctors
import UniverseAbstractions.Axioms.Universe.DependentTypes.DependentProducts
import UniverseAbstractions.Lemmas.DerivedFunctors
import UniverseAbstractions.Lemmas.DerivedProductFunctors
import UniverseAbstractions.Notation

set_option autoBoundImplicitLocal false
--set_option pp.universes true

universe u v w w' w''

class HasRelations (U : Universe.{u}) [HasFunOp.{u, w} U] [HasInternalProducts.{u, w} U]
                   (V : Universe.{v}) extends
  HasDependentFunctors.{u, v, v, w', w''} U V V : Type (max 1 u v w w' w'')

namespace HasRelations

  open HasProducts HasInternalProducts HasCompFunProp'

  variable {U : Universe.{u}} [HasFunOp.{u, w} U] [HasInternalProducts.{u, w} U]

  def relMap {V : Universe.{v}} [HasRelations.{u, v, w, w', w''} U V] {A : U} (r : A ⊓ A → V) :
    A → A → V :=
  λ a b => r (intro a b)

  def propMap {V : Universe.{v}} [HasRelations.{u, v, w, w', w''} U V] {A : U} (r : A → A → V) :
    A ⊓ A → V :=
  λ P => r (fst P) (snd P)

  def DefRel (A : U) (V : Universe.{v}) [HasRelations.{u, v, w, w', w''} U V] (r : A → A → V) :=
  A ⊓ A ⟿[propMap r] V
  notation:20 A:21 " ⤐[" r:0 "] " V:21 => HasRelations.DefRel A V r

  def Relation (A : U) (V : Universe.{v}) [HasRelations.{u, v, w, w', w''} U V] := A ⊓ A ⟶ ⌊V⌋
  infixr:20 " ⤐ " => HasRelations.Relation

  variable {V : Universe} [HasRelations U V] {A : U}

  instance coeRel : CoeFun (A ⤐ V) (λ _ => A → A → V) := ⟨λ θ => relMap θ.p⟩

  def defExtractABFun : (A ⊓ A) ⊓ A ⟶{λ P => intro (fst (fst P)) (snd (fst P))} A ⊓ A :=
  fstFun (A ⊓ A) A
  ◄ λ _ => by simp

  def defExtractBCFun : (A ⊓ A) ⊓ A ⟶{λ P => intro (snd (fst P)) (snd P)} A ⊓ A :=
  elim₃LFun (HasSubLinearFunOp.constFun A (introFunFun A A))
  ◄ λ _ => by simp [elim₃LFun]

  def defExtractACFun : (A ⊓ A) ⊓ A ⟶{λ P => intro (fst (fst P)) (snd P)} A ⊓ A :=
  elim₃LFun (HasLinearFunOp.swapFunFun (HasSubLinearFunOp.constFun A (introFunFun A A)))
  ◄ λ _ => by simp [elim₃LFun]

  @[reducible] def extractABFun' : (A ⊓ A) ⊓ A ⟶' A ⊓ A := HasFunctoriality.fromDefFun defExtractABFun
  @[reducible] def extractBCFun' : (A ⊓ A) ⊓ A ⟶' A ⊓ A := HasFunctoriality.fromDefFun defExtractBCFun
  @[reducible] def extractACFun' : (A ⊓ A) ⊓ A ⟶' A ⊓ A := HasFunctoriality.fromDefFun defExtractACFun

  variable [HasCompFunProp' U U V] (θ : A ⤐ V)

  class HasRefl where
  (reflPi : Π compProp (dupIntroFun' A) θ)

  def HasRefl.refl [HasRefl θ] (a : A) : θ a a := reflPi a

  variable [HasInternalFunctors V] [HasFunProp U V V V]

  class HasTrans where
  (transPi : Π {compProp extractABFun' θ ⟶ {compProp extractBCFun' θ ⟶ compProp extractACFun' θ}})

  @[simp] theorem simp_extractAB (a b c : A) :
    let P := intro₃L a b c;
    θ (fst (fst P)) (snd (fst P)) = θ a b :=
  by simp

  @[simp] theorem simp_extractBC (a b c : A) :
    let P := intro₃L a b c;
    θ (snd (fst P)) (snd P) = θ b c :=
  by simp

  @[simp] theorem simp_extractAC (a b c : A) :
    let P := intro₃L a b c;
    θ (fst (fst P)) (snd P) = θ a c :=
  by simp

  def HasTrans.trans [HasTrans θ] (a b c : A) : θ a b ⟶ θ b c ⟶ θ a c :=
  simp_extractAB θ a b c ▸ simp_extractBC θ a b c ▸ simp_extractAC θ a b c ▸ transPi (intro₃L a b c)

  class IsPreorder extends HasRefl θ, HasTrans θ

  variable [HasInternalProducts V] [HasInternalEquivalences V] [HasEquivProp U V V]

  class HasSymm where
  (symmPi : Π {θ ⟶ compProp (commFun' A A) θ})

  @[simp] theorem simp_swap (a b : A) :
    let P := intro a b;
    θ (snd P) (fst P) = θ b a :=
  by simp

  def HasSymm.symm [HasSymm θ] (a b : A) : θ a b ⟶ θ b a :=
  simp_swap θ a b ▸ symmPi (intro a b)

  class HasTransEquiv [HasTrans θ] [HasSymm θ] where
  (defTransEquiv    {a b : A} (f : θ a b) (c : A) :
      θ b c ⟷{HasTrans.trans θ a b c f, HasTrans.trans θ b a c (HasSymm.symm θ a b f)} θ a c)
  (defTransEquivFun (a b c : A)                   :
      θ a b ⟶{λ f => HasEquivalences.fromDefEquiv (defTransEquiv f c)} (θ b c ⟷ θ a c))

  class HasSymmEquiv [HasSymm θ] where
  (defSymmEquiv (a b : A) : θ a b ⟷{HasSymm.symm θ a b, HasSymm.symm θ b a} θ b a)

  class IsEquivalence extends IsPreorder θ, HasSymm θ, HasTransEquiv θ

  def substRel (φ : A ⟶ ⌊V⌋) : A ⤐ V :=
  {compProp (fstFun' A A) φ ⟷ compProp (sndFun' A A) φ}

  @[simp] theorem simp_subst_refl (φ : A ⟶ ⌊V⌋) :
    compProp (dupIntroFun' A) (substRel φ) = {φ ⟷ φ} :=
  sorry

  class HasIdEquivPi (φ : A ⟶ ⌊V⌋) where
  [hasIdFun   : HasIdFun V]
  [hasIdEquiv : HasIdEquiv V V]
  (F          : Π{λ a => HasIdEquiv.idEquiv (φ a)} {φ ⟷ φ})

  instance substRel.hasRefl (φ : A ⟶ ⌊V⌋) [h : HasIdEquivPi φ] :
    HasRefl (substRel φ) :=
  ⟨simp_subst_refl φ ▸ h.F⟩

  class IsSubstitution extends HasRefl θ where
  (substPi (φ : A ⟶ ⌊V⌋) [HasIdEquivPi φ] : Π {θ ⟶ substRel φ})

  @[simp] theorem simp_apply_fst (φ : A ⟶ ⌊V⌋) (a b : A) :
    φ (fst (intro a b)) = φ a :=
  by simp

  @[simp] theorem simp_apply_snd (φ : A ⟶ ⌊V⌋) (a b : A) :
    φ (snd (intro a b)) = φ b :=
  by simp

  def IsSubstitution.subst [IsSubstitution θ] (φ : A ⟶ ⌊V⌋) [HasIdEquivPi φ] (a b : A) :
    θ a b ⟶ (φ a ⟷ φ b) :=
  simp_apply_fst φ a b ▸ simp_apply_snd φ a b ▸ substPi φ (intro a b)

  class IsIdentity extends HasRefl θ where
  (elimPi (ξ : A ⤐ V) [HasRefl ξ] : Π {θ ⟶ ξ})

  namespace IsIdentity

    variable [IsIdentity θ]

    def elim (ξ : A ⤐ V) [HasRefl ξ] (a b : A) : θ a b ⟶ ξ a b :=
    elimPi ξ (HasProducts.intro a b)

    instance isSubstitution : IsSubstitution θ :=
    { substPi := λ φ => elimPi (substRel φ) }

  end IsIdentity

end HasRelations


/-
In this file, we explore the use of the tagless final style [1]
to encode SSA semantics.

[1] https://okmij.org/ftp/tagless-final/
-/
import Lean
open Lean

namespace WriterT
def WriterT (m: Type _ -> Type _) (a: Type _) := m (a × String)

def WriterT.run (wm: WriterT m a ): m (a × String) := wm

def WriterT.mk (x: m (a × String)): WriterT m a := x

instance [Functor m]: Functor (WriterT m) where
  map f w := Functor.map (f := m) (fun (a, log) => (f a, log)) w

instance [Pure m]: Pure (WriterT m) where
  pure x := pure (f := m) (x, "")

instance [Monad m]: Seq (WriterT m) where
   seq mx my := WriterT.mk do
    let wx <- mx
    let wy <- (my ())
    let wb := wx.fst wy.fst
    return (wb, wx.snd ++ wy.snd)

instance [Monad m] : SeqLeft (WriterT m) where
   seqLeft mx my := WriterT.mk do
    let wx <- mx
    let wy <- (my ())
    return (wx.fst, wx.snd ++ wy.snd)

instance [Monad m] : SeqRight (WriterT m) where
   seqRight mx my := WriterT.mk do
    let wx <- mx
    let wy <- (my ())
    return (wy.fst, wx.snd  ++ wy.snd )

def WriterT.bindCont [Bind m] [Pure m] (k: α → WriterT m β) (x: α × String):
    WriterT m β := WriterT.mk do
  let y ← k x.fst
  return (y.fst, x.snd ++ y.snd)

def WriterT.bind [Bind m] [Pure m] (wma: WriterT m α) (a2wmb: α → WriterT m β):
    WriterT m β :=
  WriterT.mk do
    let x <- wma
    WriterT.bindCont a2wmb x

instance [Bind m] [Pure m]: Bind (WriterT m) where
  bind wma a2wmb := WriterT.bind wma a2wmb

def WriterT.lift [Monad m] {α : Type u} (ma: m α): WriterT m α :=
  Bind.bind (m := m) ma (fun a => return (a, ""))

instance [Monad m]: MonadLift m (WriterT m) where
  monadLift := WriterT.lift

instance : MonadFunctor m (WriterT m) where
  monadMap f := f

instance [Monad m] : Applicative (WriterT m) where
  pure := Pure.pure
  seqLeft := SeqLeft.seqLeft
  seqRight := SeqRight.seqRight

instance [Monad m]: Monad (WriterT m) where
  pure := Pure.pure
  bind := Bind.bind
  map  := Functor.map

def logWriterT [Monad m] (s: String): WriterT.{u} m PUnit.{u+1} :=
  pure (f := m) (.unit, s)

end WriterT

namespace Fitree
open WriterT
/- Extendable effect families -/

abbrev to1 (E: Type → Type u) (F: Type → Type v) :=
  ∀ T, E T → F T
abbrev sum1 (E F: Type → Type) :=
  fun T => E T ⊕ F T
inductive Void1: Type → Type :=

infixr:40 " ~> " => to1
infixr:60 " +' " => sum1

class Member (E: Type → Type) (F: Type → Type) where
  inject : E ~> F

instance MemberId {E}: Member E E where
  inject := (fun _ => id)

instance MemberSumL {E F G} [Member E F]: Member E (F +' G) where
  inject T := Sum.inl ∘ Member.inject T

instance MemberSumR {E F G} [Member E G]: Member E (F +' G) where
  inject T := Sum.inr ∘ Member.inject T

def Sum.cases {α β γ} (fα: α → γ) (fβ: β → γ): (α ⊕ β) → γ
  | .inl a => fα a
  | .inr b => fβ b

instance MemberSum {E F G H} [Member E G] [Member F H]:
    Member (E +' F) (G +' H) where
  inject T := Sum.cases (Member.inject T) (Member.inject T)

instance MemberVoid1 {E}:
    Member Void1 E where
  inject _ e := nomatch e

-- Effects can now be put in context automatically by typeclass resolution
example E:      Member E E := inferInstance
example E F:    Member E (E +' F) := inferInstance
example E F:    Member E (F +' (F +' E)) := inferInstance
example E F G:  Member (E +' F) (E +' F +' G) := inferInstance

/- The monadic domain; essentially finite Interaction Trees -/

inductive Fitree (E: Type → Type) (R: Type) where
  | Ret (r: R): Fitree E R
  | Vis {T: Type} (e: E T) (k: T → Fitree E R): Fitree E R

def Fitree.ret {E R}: R → Fitree E R :=
  Fitree.Ret

def Fitree.trigger {E: Type → Type} {F: Type → Type} {T} [Member E F]
    (e: E T): Fitree F T :=
  Fitree.Vis (Member.inject _ e) Fitree.ret

def Fitree.bind {E R T} (t: Fitree E T) (k: T → Fitree E R) :=
  match t with
  | Ret r => k r
  | Vis e k' => Vis e (fun r => bind (k' r) k)

instance {E}: Monad (Fitree E) where
  pure := Fitree.ret
  bind := Fitree.bind

-- Since we only use finite ITrees, we can actually run them when they're
-- fully interpreted (which leaves only the Ret constructor)
def Fitree.run {R}: Fitree Void1 R → R
  | Ret r => r
  | Vis e _ => nomatch e

@[simp] theorem Fitree.run_ret:
  Fitree.run (Fitree.ret r) = r := rfl

def Fitree.translate {E F R} (f: E ~> F): Fitree E R → Fitree F R
  | Ret r => Ret r
  | Vis e k => Vis (f _ e) (fun r => translate f (k r))

@[simp] theorem Fitree.translate_ret:
  Fitree.translate f (Fitree.ret r) = Fitree.ret r := rfl
@[simp] theorem Fitree.translate_vis:
    Fitree.translate f (Vis e k) = Vis (f _ e) (fun r => translate f (k r)) :=
  rfl

def Fitree.case (h₁: E ~> G) (h₂: F ~> G): E +' F ~> G :=
  fun R ef => match ef with
  | Sum.inl e => h₁ R e
  | Sum.inr f => h₂ R f

@[simp] theorem Fitree.case_left:
  Fitree.case h₁ h₂ _ (Sum.inl e) = h₁ _ e := rfl
@[simp] theorem Fitree.case_right:
  Fitree.case h₁ h₂ _ (Sum.inr e) = h₂ _ e := rfl

/-
### Monadic interpretation
-/

def Fitree.interp {M} [Monad M] {E} (h: E ~> M) {R}: Fitree E R → M R
  | .Ret r => pure r
  | .Vis e k => Bind.bind (h _ e) (fun t => interp h (k t))

def Fitree.interp' {E F} (h: E ~> Fitree Void1) {R} (t: Fitree (E +' F) R):
    Fitree F R :=
  interp (Fitree.case
    (fun _ e => (h _ e).translate $ fun _ e => nomatch e)
    (fun _ e => Fitree.trigger e)) t

-- Interp `F` by lifting into a monad transformer (this is used when
-- interpreting `E +' F` into the monad)
def Fitree.liftHandler {F M} [MonadLiftT (Fitree F) M]: F ~> M := fun R e =>
  monadLift (Fitree.trigger e: Fitree F R)

-- Interpretation into various predefined monads. These are predefined so that
-- rewriting theorems that expose the monad structure can be provided.

def Fitree.interpState {M S} [Monad M] {E} (h: E ~> StateT S M):
    forall {R}, Fitree E R → StateT S M R :=
  interp h

def Fitree.interpWriter {M} [Monad M] {E} (h: E ~> WriterT M):
    forall {R}, Fitree E R → WriterT M R :=
  interp h

def Fitree.interpOption {M} [Monad M] {E} (h: E ~> OptionT M):
    forall {R}, Fitree E R → OptionT M R :=
  interp h

def Fitree.interpExcept {M ε} [Monad M] {E} (h: E ~> ExceptT ε M) {R}:
    Fitree E R → ExceptT ε M R :=
  interp h

/-
### Combinator identities

The following theorems act as the main interface for computation on ITrees. We
don't unfold definitions because Lean 4 doesn't yet have the match-unfolding
behavior of Coq's `simpl` tactic, and runs into performance issues as unfolded
terms grow larger. Instead, we aggressively rewrite the following simplifying
equalities.
-/

@[simp] theorem Fitree.bind_ret:
  Fitree.bind (Fitree.ret r) k = k r := rfl

@[simp] theorem Fitree.bind_Ret:
  Fitree.bind (Fitree.Ret r) k = k r := rfl

@[simp] theorem Fitree.bind_ret':
    Fitree.bind t (fun r => Fitree.ret r) = t := by
  induction t with
  | Ret _ => rfl
  | Vis _ _ ih => simp [bind, ih]

@[simp] theorem Fitree.bind_Ret':
    Fitree.bind t (fun r => Fitree.Ret r) = t := by
  induction t with
  | Ret _ => rfl
  | Vis _ _ ih => simp [bind, ih]

@[simp] theorem Fitree.bind_bind:
    Fitree.bind (Fitree.bind t k) k' =
    Fitree.bind t (fun x => Fitree.bind (k x) k') := by
  induction t with
  | Ret _ => rfl
  | Vis _ _ ih => simp [bind, ih]

@[simp] theorem Fitree.pure_is_ret:
  @Pure.pure (Fitree E) _ _ r = Fitree.ret r := rfl

@[simp] theorem Fitree.bind_is_bind:
  @Bind.bind (Fitree E) _ _ _  t k = Fitree.bind t k := rfl

@[simp] theorem Fitree.StateT_bind_is_bind (k: T → S → Fitree E (R × S)):
  StateT.bind (m := Fitree E) t k =
    fun s => Fitree.bind (t s) (fun (x,s) => k x s) := rfl

@[simp] theorem Fitree.WriterT_bind_is_bind (k: T → Fitree E (R × String)):
  WriterT.bind (m := Fitree E) t k =
    Fitree.bind t (WriterT.bindCont k) := rfl

@[simp] theorem Fitree.OptionT_bind_is_bind (k: T → Fitree E (Option R)):
  OptionT.bind (m := Fitree E) t k =
    Fitree.bind t (fun
      | some x => k x
      | none => Fitree.ret none) := rfl

@[simp] theorem Fitree.ExceptT_bind_is_bind (k: T → Fitree E (Except ε R)):
  ExceptT.bind (m := Fitree E) t k = Fitree.bind t (ExceptT.bindCont k) := rfl

@[simp] theorem Fitree.liftHandler_StateT_is_StateT_lift:
  @Fitree.liftHandler F (StateT S (Fitree F)) _ _ e =
  fun s => Fitree.bind (Fitree.trigger e) (fun x => Fitree.ret (x, s)) := rfl

@[simp] theorem Fitree.liftHandler_WriterT_is_WriterT_lift:
  @Fitree.liftHandler F (WriterT (Fitree F)) _ _ e =
  Fitree.bind (Fitree.trigger e) (fun x => Fitree.ret (x, "")) := rfl

@[simp] theorem Fitree.liftHandler_OptionT_is_OptionT_lift:
  @Fitree.liftHandler F (OptionT (Fitree F)) _ _ e =
  Fitree.bind (Fitree.trigger e) (fun x => Fitree.ret (some x)) := rfl

@[simp] theorem Fitree.liftHandler_ExceptT_is_ExceptT_lift:
  @Fitree.liftHandler F (ExceptT ε (Fitree F)) _ _ e =
  Fitree.bind (Fitree.trigger e) (fun x => Fitree.ret (Except.ok x)) := rfl

@[simp] theorem Member.injectId:
  @Member.inject E E MemberId _ e = e := rfl

@[simp] theorem Member.injectSumL [Member E F]:
  @Member.inject E (F +' G) MemberSumL _ e = Sum.inl (Member.inject _ e) := rfl

@[simp] theorem Member.injectSumR [Member E G]:
  @Member.inject E (F +' G) MemberSumR _ e = Sum.inr (Member.inject _ e) := rfl

@[simp] theorem Member.injectSum_inl [Member E G] [Member F H]:
  @Member.inject (E +' F) (G +' H) MemberSum _ (Sum.inl e) =
    Sum.inl (Member.inject _ e) := rfl

@[simp] theorem Member.injectSum_inr [Member E G] [Member F H]:
  @Member.inject (E +' F) (G +' H) MemberSum _ (Sum.inr e) =
    Sum.inr (Member.inject _ e) := rfl

-- Interpretatin identities

@[simp] theorem Fitree.interp_ret:
  Fitree.interp h (Fitree.ret r) = Fitree.ret r := rfl

@[simp] theorem Fitree.interp_Ret:
  Fitree.interp h (Fitree.Ret r) = Fitree.ret r := rfl

@[simp] theorem Fitree.interp_Vis:
  Fitree.interp h (Fitree.Vis e k) =
  Fitree.bind (h _ e) (fun x => Fitree.interp h (k x)) := rfl

@[simp] theorem Fitree.interp'_ret:
  @Fitree.interp' E F h _ (Fitree.ret r) = Fitree.ret r := rfl

@[simp] theorem Fitree.interp'_Vis_left:
  @Fitree.interp' E F h _ (Fitree.Vis (Sum.inl e) k) =
  Fitree.bind (Fitree.translate (fun _ e => nomatch e) (h _ e))
              (fun x => Fitree.interp' h (k x)) := rfl

@[simp] theorem Fitree.interp'_Vis_right:
  @Fitree.interp' E F h _ (Fitree.Vis (Sum.inr e) k) =
  Fitree.bind (Fitree.trigger e)
              (fun x => Fitree.interp' h (k x)) := rfl

@[simp] theorem Fitree.interpState_ret:
  Fitree.interpState h (Fitree.ret r) = (fun s => Fitree.ret (r, s)) := rfl

@[simp] theorem Fitree.interpState_Vis {M S} [Monad M] (h: E ~> StateT S M):
  Fitree.interpState h (Fitree.Vis e k) =
  StateT.bind (h _ e) (fun x => Fitree.interpState h (k x)) := rfl

@[simp] theorem Fitree.interpWriter_ret:
  Fitree.interpWriter h (Fitree.ret r) = Fitree.ret (r, "") := rfl

@[simp] theorem Fitree.interpWriter_Vis {M} [Monad M] (h: E ~> WriterT M):
  Fitree.interpWriter h (Fitree.Vis e k) =
  WriterT.bind (h _ e) (fun x => Fitree.interpWriter h (k x)) := rfl

@[simp] theorem Fitree.interpOption_ret:
  Fitree.interpOption h (Fitree.ret r) = Fitree.ret (some r) := rfl

@[simp] theorem Fitree.interpOption_Vis {M} [Monad M] (h: E ~> OptionT M):
  Fitree.interpOption h (Fitree.Vis e k) =
  OptionT.bind (h _ e) (fun x => Fitree.interpOption h (k x)) := rfl

@[simp] theorem Fitree.interpExcept_ret:
  Fitree.interpExcept h (Fitree.ret r) = Fitree.ret (.ok r) := rfl

@[simp] theorem Fitree.interpExcept_Vis {M ε} [Monad M] (h: E ~> ExceptT ε M):
  Fitree.interpExcept h (Fitree.Vis e k) =
  ExceptT.bind (h _ e) (fun x => Fitree.interpExcept h (k x)) := rfl

-- We don't assume [LawfulMonad M] so we can't simplify the continuation. But
-- when it's an ITree the other simp lemmas will do it anyway.
@[simp] theorem Fitree.interp_trigger [Member E F] [Monad M] (e: E T):
  Fitree.interp (M := M) (E := F) h (Fitree.trigger e) =
  Bind.bind (h _ (Member.inject _ e)) (fun x => pure x) := rfl

@[simp] theorem Fitree.interp'_trigger_left (e: E R):
  @Fitree.interp' E F h _ (@Fitree.trigger (E +' F) _ _ MemberId (Sum.inl e)) =
  Fitree.bind
    (Fitree.translate (fun _ e => nomatch e) (h _ (Member.inject _ e)))
    (fun x => pure x) := rfl

@[simp] theorem Fitree.interp'_trigger_right [Member G F]:
  @Fitree.interp' E F h _ (@Fitree.trigger (E +' G) (E +' F) _ _ (Sum.inr e)) =
  Fitree.trigger e := rfl

-- The following theorems are only applied manually

theorem Fitree.run_bind {T R} (t: Fitree Void1 T) (k: T → Fitree Void1 R):
    run (bind t k) = run (k (run t)) :=
  match t with
  | Ret _ => rfl
  | Vis e _ => nomatch e

theorem Fitree.interp_bind:
    Fitree.interp h (Fitree.bind t k) =
    Fitree.bind (Fitree.interp h t) (fun x => Fitree.interp h (k x)) := by
  induction t with
  | Ret _ => rfl
  | Vis _ _ ih => simp [bind, ih]

theorem Fitree.interp'_bind:
    Fitree.interp' h (Fitree.bind t k) =
    Fitree.bind (Fitree.interp' h t) (fun x => Fitree.interp' h (k x)) := by
  simp [interp', interp_bind]

-- Specialized interp_bind lemmas that unfold the monadic structure and expose
-- the Fitree.bind directly rather than the monadic Bind.bind

theorem Fitree.interpState_bind (h: E ~> StateT S (Fitree F)) (t: Fitree E R):
    Fitree.interpState h (Fitree.bind t k) s =
    Fitree.bind (Fitree.interpState h t s)
      (fun (x,s') => Fitree.interpState h (k x) s') := by
  revert s
  induction t with
  | Ret _ => intros s; rfl
  | Vis _ _ ih =>
    simp [interpState] at *
    simp [interp, Bind.bind, StateT.bind]
    simp [ih]

example {F R}: WriterT (Fitree F) R = Fitree F (R × String) := by
  simp [WriterT]

theorem Fitree.interpWriter_bind (h: E ~> WriterT (Fitree F))
  (t: Fitree E T) (k: T → Fitree E R):
    Fitree.interpWriter h (Fitree.bind t k) =
    Fitree.bind (Fitree.interpWriter h t) fun (x,s₁) =>
      Fitree.bind (Fitree.interpWriter h (k x)) fun (y,s₂) =>
        Fitree.ret (y,s₁++s₂) := by
  induction t with
  | Ret _ =>
      simp [bind, interpWriter]
      have h₁: forall x, "" ++ x = x := by
        simp [HAppend.hAppend, Append.append, String.append]
        simp [List.nil_append]
      simp [h₁]
      have h₂: forall (α β: Type) (x: α × β), (x.fst, x.snd) = x := by simp
      simp [h₂]
  | Vis _ _ ih =>
      simp [interpWriter] at *
      simp [interp, Bind.bind, WriterT.bindCont, WriterT.mk]
      have h: forall (x y z: String), x ++ (y ++ z) = x ++ y ++ z := by
        simp [HAppend.hAppend, Append.append, String.append]
        simp [List.append_assoc]
      simp [ih, h]

theorem Fitree.interpOption_bind (h: E ~> OptionT (Fitree F))
  (t: Fitree E T) (k: T → Fitree E R):
    Fitree.interpOption h (Fitree.bind t k) =
    Fitree.bind (Fitree.interpOption h t) fun x? =>
      match x? with
      | some x => Fitree.interpOption h (k x)
      | none => Fitree.ret none := by
  induction t with
  | Ret _ => rfl
  | Vis _ _ ih =>
      simp [interpOption] at *
      simp [interp, bind, Bind.bind, OptionT.bind, OptionT.mk]
      -- I can't get a bind (match) → match (bind) theorem to rewrite, so...
      have fequal2 α β (f g: α → β) x y: f = g → x = y → f x = g y :=
        fun h₁ h₂ => by simp [h₁, h₂]
      apply fequal2; rfl; funext x
      cases x <;> simp [ih]

theorem Fitree.interpExcept_bind (h: E ~> ExceptT ε (Fitree F))
  (t: Fitree E T) (k: T → Fitree E R):
    Fitree.interpExcept h (Fitree.bind t k) =
    Fitree.bind (Fitree.interpExcept h t) fun x? =>
      match x? with
      | .error ε => Fitree.ret (.error ε)
      | .ok x => Fitree.interpExcept h (k x) := by
  induction t with
  | Ret _ => rfl
  | Vis _ _ ih =>
      simp [interpExcept] at *
      simp [interp, bind, Bind.bind]
      simp [ExceptT.bind, ExceptT.mk, ExceptT.bindCont]
      -- See above
      have fequal2 α β (f g: α → β) x y: f = g → x = y → f x = g y :=
        fun h₁ h₂ => by simp [h₁, h₂]
      apply fequal2; rfl; funext x
      cases x <;> simp [ih]

-- This theorem has the drawback of hiding the continuation of `bind` into the
-- `Vis` node, which blocks other theorems like `Fitree.bind_bind`.
theorem Fitree.bind_trigger [Member E F] (e: E T) (k: T → Fitree F R):
  Fitree.bind (Fitree.trigger e) k = Fitree.Vis (Member.inject _ e) k := rfl

/-
### Other properties
-/

inductive Fitree.noEventL {E F R}: Fitree (E +' F) R → Prop :=
  | Ret r: noEventL (Ret r)
  | Vis f k: (∀ t, noEventL (k t)) → noEventL (Vis (Sum.inr f) k)

end Fitree

namespace Exp

-- https://okmij.org/ftp/tagless-final/course/lecture.pdf
inductive Exp where
| Lit: Int -> Exp
| Neg: Exp -> Exp
| Add: Exp -> Exp -> Exp

def Exp.eval: Exp -> Int
| .Lit i => i
| .Neg e => -1 * e.eval
| .Add e e' => e.eval + e'.eval

class ExpSYM (repr: Type) where
  lit: Int -> repr
  neg: repr -> repr
  add: repr -> repr -> repr
  -- neg_involutive: (a: repr) -> neg (neg a) = a

instance : ExpSYM Int where
  lit i := i
  neg i := (-i)
  add i j := i + j

instance : ExpSYM String where
  lit i := toString i
  neg i := s!"(neg {i})"
  add i i' := s!"(add {i} {i'})"
end Exp

namespace Tree
inductive Tree where
| Leaf: String -> Tree
| Node: String -> List Tree -> Tree
deriving BEq

open Exp

-- Serialize Exp into Tree

instance : ExpSYM Tree where
  lit n := .Node "Lit" [.Leaf (toString n)]
  neg e := .Node "Neg" [e]
  add e e' := .Node "Add" [e, e']

def fromTree {repr: Type} [ExpSYM repr] : Tree -> Except String repr
| .Node "Lit" [.Leaf n] => do
   Except.ok (ExpSYM.lit 42) -- TODO: convert from string to nat.
| .Node "Neg" [e] => do
       return (ExpSYM.neg (<- fromTree e))
| .Node "Add" [e, e'] => do
   return ExpSYM.add (<- fromTree e) (<-
fromTree e')
| _t => Except.error "incorrect tree"

end Tree

namespace PushNeg
open Exp

def Exp.pushNeg: Exp -> Exp
| .Lit v => .Lit v
| .Neg (.Lit v) => .Neg (.Lit v)
| .Neg (.Neg e) => Exp.pushNeg e
| .Neg (.Add e e') => .Add (Exp.pushNeg e) (Exp.pushNeg e')
| .Add e e' => .Add (Exp.pushNeg e) (Exp.pushNeg e')

inductive Ctx where
| Pos: Ctx
| Neg: Ctx

instance {repr: Type} [ExpSYM repr] : ExpSYM (Ctx -> repr) where
  lit n := fun ctx => match ctx with
    | .Pos => ExpSYM.lit n
    | .Neg => ExpSYM.neg (ExpSYM.lit n)
  neg e := fun ctx => match ctx with
    | .Pos => e .Neg
    | .Neg => e .Pos
  add e1 e2 := fun ctx =>  ExpSYM.add (e1 ctx) (e2 ctx)
end PushNeg

namespace HO -- higher order tagless final

class Symantics (repr: Type -> Type) where
  int : Int -> repr Int
  add: repr Int -> repr Int -> repr Int
  lam: (repr a -> repr b) -> repr (a -> b)
  app: repr (a -> b) -> repr a -> repr b

structure R (a: Type) where
  val : a

instance : Symantics R where
  int i := { val := i }
  add i j := { val := i.val + j.val }
  lam f := { val := fun a =>  (f (R.mk a)).val }
  app f a := R.mk $ f.val a.val

class BoolSYM (repr: Type -> Type) where
  bool: Bool -> repr Bool
  leq : repr Int -> repr Int -> repr Bool
  if_: repr Bool -> repr a -> repr a -> repr a

instance : BoolSYM R where
 bool b := R.mk b
 leq a a' := R.mk (a.val <= a'.val)
 if_ cond t e := R.mk $ if cond.val then t.val else e.val

class FixSYM (repr: Type -> Type) where
  fix: (repr a -> repr a) -> repr a

-- lol
partial instance : FixSYM R where
  fix := sorry

-- h is heaps

inductive IR (h: Type _ -> Type _): Type _ -> Type _ where
| int: Int -> IR h Int
| add: IR h t -> IR h t -> IR h t
| var: h t -> IR h t
-- | lam: (IR h t1 -> IR h t2) -> IR h (t1 -> t2) -- non-positive occurence, cannot be encoded in initial style!

end HO

namespace SSA
/-
inductive BB (repr: Type _ -> Type _ ): Type _ -> Type _ where
| entry: String -> BB repr a -> BB repr a -- begin a bb
| seq: BB repr a -> BB repr b -> BB repr b
| op: repr a -> BB repr a -- operation
| ret: repr a -> BB repr a -- only place where problem occurs.
| condbr: repr Bool -> String -> String -> BB repr Unit
| br: String -> BB repr Unit

class BBSemantics (repr: Type _ -> Type _) where
  bb: BB repr a -> repr a

structure R (a: Type) where
  val : a

instance : BBSemantics R where
  bb repr := match repr with
             | .entry name rest =>

-/

inductive Op: Type _ -> Type _ where
| add: Int -> Int -> Op Int
| lt: Int -> Int -> Op Bool
| const: Int -> Op Int

class OpSYM (repr: Type -> Type) where
  add: Int -> Int -> repr Int
  lt: Int -> Int -> repr Bool
  const: Int -> repr Int

instance : OpSYM Op where
  add := .add
  lt := .lt
  const := .const

structure BBName where
  name: String

structure BBRef (a: Type _) where
  name: String

-- class BBRefSYM (repr: Type -> Type) := String -> repr a

-- Terminator has single type for interprocedural control flow.
-- Inside and Outside
-- k for things that are unknown, in the grand CPS style
-- BB intra inter.
inductive Terminator: Type _ -> Type _ where
| br: BBRef i -> i -> Terminator Unit
| ret: o -> Terminator o
| condbr: Bool -> (BBRef i × i) -> (BBRef i' × i') -> Terminator Unit

class TerminatorSYM (repr: Type _ -> Type _) where
  br: BBRef i -> i -> repr Unit
  ret: o -> repr o
  condbr: Bool -> (BBRef i × i) -> (BBRef i' × i') -> repr Unit

instance : TerminatorSYM Terminator where
  br := .br
  ret := .ret
  condbr := .condbr

-- BB has three two type: one for interprocedural control flow
-- one for intraprocedural control flow
-- Inside and Outside
-- BB intra inter.
-- BB <input-type> <interprocedural-out-type>
-- O: type of ops
-- T: type of terminators.
inductive BB (O: Type _ -> Type _) (T: Type _ -> Type _): Type _ -> Type _ -> Type _ where
| begin: (i -> BB O T Unit o) -> BB O T i o
| seq: O a -> (a -> BB O T Unit o) -> BB O T Unit o
| terminator: T o -> BB O T Unit o

class  BBSYM (bbRepr: Type _ -> Type _ -> Type _)
  (opRepr: Type _ -> Type _)
  (terminatorRepr: Type _ -> Type _)
  extends OpSYM opRepr, TerminatorSYM terminatorRepr where
  begin: (i -> bbRepr Unit o) -> bbRepr i o
  seq: (opRepr a) -> (a -> bbRepr Unit o) -> bbRepr Unit o
  terminator: (terminatorRepr o) -> bbRepr Unit o

-- instance of Symantics for BB.
instance [OpSYM O] [TerminatorSYM T]: BBSYM (BB O T) O T where
  begin := BB.begin
  seq := BB.seq
  terminator := BB.terminator

-- build a BB which takes 'Int' input, produces 'Int' output.
def prog0 : BB Op Terminator Int Int :=
  .begin (fun input =>
    .seq (.const 4) (fun j =>
    .seq (.add input j) (fun k =>
    .terminator (.ret k)
  )))

namespace RegionBuilder
-- Build a region
-- The list of types is the labels that have been defined.
inductive RegionBuilder
  (O: Type _ -> Type _)
  (T: Type _ -> Type _): List (Σ (i: Type), BBRef i) -> Type _ -> Type _ -> Type _ where
| lbl: ((ref: BBRef i) ->
   RegionBuilder O T (⟨ i, ref ⟩::ris) ri ro) -- if you want a label,
   ->  RegionBuilder O T ris ri ro -- I can then forget about the `i` and remember that the `is` have been defined
                                   -- you have an obligation to define it in the output
| define: (ref: BBRef i) -> BB O T i o -> RegionBuilder O T ris ri ro
    -> RegionBuilder O T (⟨i,ref⟩::ris) ri ro -- define defines an `i`.
| empty: RegionBuilder O T [] ri ro -- empty region defines no BBS.

def prog1: RegionBuilder Op Terminator [] Int Int :=
  .lbl (i := Int) (fun entry =>
     .define entry (.begin fun i =>
      .terminator (.ret i)
      ) .empty)

-- takes an int as input, produces an int as output
-- entry(input):
--   br loop (input, 0)
-- loop(i, k):
--   knew := k + 1
--   inew := i + 1
--   exit := knew == 10
--   condbr exit(inew), loop(inew, knew)
-- exit(inew):
--   ret inew
def prog2: RegionBuilder Op Terminator [] Int Int :=
  .lbl (i := Int) (fun entrybb =>
  .lbl (i := Int × Int) (fun loopbb =>
  .lbl (i := Int) (fun exitbb =>
     .define exitbb (.begin fun inew => .terminator (.ret inew)) $
     .define loopbb (.begin fun args =>
       .seq (.add 1 args.fst) (fun knew =>
       .seq (.add 1 args.snd) (fun inew =>
       .seq (.lt knew 10) (fun isExit =>
       -- .terminator (.ret knew)))) -- (.condbr isExit ⟨exitbb, inew⟩, ⟨loopbb, (inew, knew)⟩)))))
       .terminator (.condbr isExit (exitbb, inew) (loopbb, (inew, knew))))))
     ) $
     .define entrybb (.begin fun input =>
      .terminator (.br loopbb (input, 0))) $
     .empty)))
#reduce prog2
end RegionBuilder

namespace Region

end Region

end SSA

namespace StructuredSSA
/-
We flatten Op, BasicBlock, Region into a single Def'.
We need three notions:
(1) Running some semantic value (R), labelled by a label (L)
(3) creating a new scope
(2) invoking control flow (C) to a label (L)
(2) sequentially composing two defs

-/

inductive Producer: Type -> Type where
inductive Consumer: Type -> Type where
inductive ProducerConsumer: Type -> Type -> Type where

-- op: dataflow
-- bb: ? (Ill defined concept)
-- br, condbr: control flow.
-- CFG: control flow

-- backwards dataflow graph.
inductive Dataflow (D: Type -> Type -> Type) (C: Type -> Type -> Type): Type _ -> Type _ where
| val: O -> Dataflow D C O
| df: D I O -> (I -> Dataflow D C O) -> Dataflow D C O

-- forwards control flow graph.
inductive Controlflow (D: Type -> Type -> Type) (C: Type -> Type -> Type): Type _ -> Type _ where
| controldep: (I -> Dataflow D C O) -> Controlflow D C I -- Create phi nodes / control flow dependent values.
| cf: C I BLANK  /- instruction condbr in conbr b bb1, bb2 (I = Bool) -/
   -- -> (I -> Dataflow D C O' × Controlflow D C I') /- function that maps true ->bb1(x), false -> bb2(x), and shows how to produce (x, y) when mapping. -/
   -> (I ->  Controlflow D C I') /- function that maps true ->bb1, false -> bb2 -/

inductive Void where

abbrev Unit2 (a: Type) := a
abbrev Void2 (_a: Type) := Void

inductive OpD : Type -> Type where -- tagged by output type
| add: Int -> Int -> OpD Int
| neg: Int -> OpD Int

abbrev Op O := Dataflow OpD Void2 O

inductive TerminatorC : Type -> Type where -- tagged by input type
|  br: TerminatorC Unit
|  condbr: TerminatorC Bool

-- A basic block is obtained by taking the data flow of an Op and the control flow of a Terminator
abbrev BasicBlock I := Controlflow OpD TerminatorC I

inductive Adapt (D: Type -> Type _) (C: Type -> Type _): Type _ -> Type _
| adapt: D O ->  C I -> (O -> I) -> Adapt D C O

-- A region adapts
abbrev Region O := Adapt BasicBlock  BasicBlock O

end StructuredSSA

namespace PartialFunctionReasoning

-- @[mlirdBy "factorial"]
opaque factorial: Int -> Int
-- axiom factorial_succ: ∀ (n: Nat), factorial (Int.ofNat (Nat.succ n))
axiom factorial_rec: ∀ (i: Nat), factorial (Int.ofNat (Nat.succ i)) = (Nat.succ i) * factorial (Int.ofNat i)
axiom factorial_zero: factorial (Int.ofNat 0) = 1

def terminating_factorial (n: Nat): Nat :=
  match n with
  | 0 => 1
  | n' + 1 => n * terminating_factorial n'

#check Nat
theorem agree: forall (n: Nat), terminating_factorial n = factorial (.ofNat n) := by {
  intros n;
  induction n;
  case zero =>  {
  simp [factorial_zero, terminating_factorial];
  }
  case succ n H => {
   simp[terminating_factorial, H];
   simp[factorial_rec];
   rewrite [<- H];
   sorry
  }

}

end PartialFunctionReasoning

namespace PartialEvaluator
-- Section 4.6
end PartialEvaluator


/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Sebastian Ullrich, Leonardo de Moura
-/
prelude
import Init.SimpLemmas
import Init.Control.Except
import Init.Control.StateRef

open Function

@[simp] theorem monadLift_self [Monad m] (x : m α) : monadLift x = x :=
  rfl

class LawfulFunctor (f : Type u → Type v) [Functor f] : Prop where
  map_const          : (Functor.mapConst : α → f β → f α) = Functor.map ∘ const β
  id_map   (x : f α) : id <$> x = x
  comp_map (g : α → β) (h : β → γ) (x : f α) : (h ∘ g) <$> x = h <$> g <$> x

export LawfulFunctor (map_const id_map comp_map)

attribute [simp] id_map

@[simp] theorem id_map' [Functor m] [LawfulFunctor m] (x : m α) : (fun a => a) <$> x = x :=
  id_map x

class LawfulApplicative (f : Type u → Type v) [Applicative f] extends LawfulFunctor f : Prop where
  seqLeft_eq  (x : f α) (y : f β)     : x <* y = const β <$> x <*> y
  seqRight_eq (x : f α) (y : f β)     : x *> y = const α id <$> x <*> y
  pure_seq    (g : α → β) (x : f α)   : pure g <*> x = g <$> x
  map_pure    (g : α → β) (x : α)     : g <$> (pure x : f α) = pure (g x)
  seq_pure    {α β : Type u} (g : f (α → β)) (x : α) : g <*> pure x = (fun h => h x) <$> g
  seq_assoc   {α β γ : Type u} (x : f α) (g : f (α → β)) (h : f (β → γ)) : h <*> (g <*> x) = ((@comp α β γ) <$> h) <*> g <*> x
  comp_map g h x := by
    repeat rw [← pure_seq]
    simp [seq_assoc, map_pure, seq_pure]

export LawfulApplicative (seqLeft_eq seqRight_eq pure_seq map_pure seq_pure seq_assoc)

attribute [simp] map_pure seq_pure

@[simp] theorem pure_id_seq [Applicative f] [LawfulApplicative f] (x : f α) : pure id <*> x = x := by
  simp [pure_seq]

class LawfulMonad (m : Type u → Type v) [Monad m] extends LawfulApplicative m : Prop where
  bind_pure_comp (f : α → β) (x : m α) : x >>= pure ∘ f = f <$> x
  bind_map       {α β : Type u} (f : m (α → β)) (x : m α) : f >>= (. <$> x) = f <*> x
  pure_bind      (x : α) (f : α → m β) : pure x >>= f = f x
  bind_assoc     (x : m α) (f : α → m β) (g : β → m γ) : x >>= f >>= g = x >>= fun x => f x >>= g
  map_pure g x    := by rw [← bind_pure_comp, pure_bind]
  seq_pure g x    := by rw [← bind_map]; simp [map_pure, bind_pure_comp]
  seq_assoc x g h := by
    -- TODO: support for applying `symm` at `simp` arguments
    let bind_pure_comp_symm {α β : Type u} (f : α → β) (x : m α) : f <$> x = x >>= pure ∘ f := by
      rw [bind_pure_comp]
    let bind_map_symm {α β : Type u} (f : m (α → (β : Type u))) (x : m α) : f <*> x = f >>= (. <$> x) := by
      rw [bind_map]
    simp[bind_pure_comp_symm, bind_map_symm, bind_assoc, pure_bind]

export LawfulMonad (bind_pure_comp bind_map pure_bind bind_assoc)
attribute [simp] pure_bind bind_assoc

@[simp] theorem bind_pure [Monad m] [LawfulMonad m] (x : m α) : x >>= pure = x := by
  show x >>= pure ∘ id = x
  rw [bind_pure_comp, id_map]

theorem map_eq_pure_bind [Monad m] [LawfulMonad m] (f : α → β) (x : m α) : f <$> x = x >>= fun a => pure (f a) := by
  rw [← bind_pure_comp]

theorem seq_eq_bind_map {α β : Type u} [Monad m] [LawfulMonad m] (f : m (α → β)) (x : m α) : f <*> x = f >>= (. <$> x) := by
  rw [← bind_map]

theorem bind_congr [Bind m] {x : m α} {f g : α → m β} (h : ∀ a, f a = g a) : x >>= f = x >>= g := by
  simp [funext h]

@[simp] theorem bind_pure_unit [Monad m] [LawfulMonad m] {x : m PUnit} : (x >>= fun _ => pure ⟨⟩) = x := by
  have : (x >>= fun _ => pure ⟨⟩) = (x >>= pure) := by
    apply bind_congr; intro u
    cases u; simp
  rw [bind_pure] at this
  assumption

theorem map_congr [Functor m] {x : m α} {f g : α → β} (h : ∀ a, f a = g a) : (f <$> x : m β) = g <$> x := by
  simp [funext h]

theorem seq_eq_bind {α β : Type u} [Monad m] [LawfulMonad m] (mf : m (α → β)) (x : m α) : mf <*> x = mf >>= fun f => f <$> x := by
  rw [bind_map]

theorem seqRight_eq_bind [Monad m] [LawfulMonad m] (x : m α) (y : m β) : x *> y = x >>= fun _ => y := by
  rw [seqRight_eq]; simp [map_eq_pure_bind, seq_eq_bind_map]

theorem seqLeft_eq_bind [Monad m] [LawfulMonad m] (x : m α) (y : m β) : x <* y = x >>= fun a => y >>= fun _ => pure a := by
  rw [seqLeft_eq]; simp [map_eq_pure_bind, seq_eq_bind_map]

/- Id -/

namespace Id

@[simp] theorem map_eq (x : Id α) (f : α → β) : f <$> x = f x := rfl
@[simp] theorem bind_eq (x : Id α) (f : α → id β) : x >>= f = f x := rfl
@[simp] theorem pure_eq (a : α) : (pure a : Id α) = a := rfl

instance : LawfulMonad Id := by
  refine' { .. } <;> intros <;> rfl

end Id

/- ExceptT -/

namespace ExceptT

theorem ext [Monad m] {x y : ExceptT ε m α} (h : x.run = y.run) : x = y := by
  simp [run] at h
  assumption

@[simp] theorem run_pure [Monad m] : run (pure x : ExceptT ε m α) = pure (Except.ok x) := rfl

@[simp] theorem run_lift [Monad m] (x : m α) : run (ExceptT.lift x : ExceptT ε m α) = (Except.ok <$> x : m (Except ε α)) := rfl

@[simp] theorem run_throw [Monad m] : run (throw e : ExceptT ε m β) = pure (Except.error e) := rfl

@[simp] theorem run_bind_lift [Monad m] [LawfulMonad m] (x : m α) (f : α → ExceptT ε m β) : run (ExceptT.lift x >>= f : ExceptT ε m β) = x >>= fun a => run (f a) := by
  simp[ExceptT.run, ExceptT.lift, bind, ExceptT.bind, ExceptT.mk, ExceptT.bindCont, map_eq_pure_bind]

@[simp] theorem bind_throw [Monad m] [LawfulMonad m] (f : α → ExceptT ε m β) : (throw e >>= f) = throw e := by
  simp [throw, throwThe, MonadExceptOf.throw, bind, ExceptT.bind, ExceptT.bindCont, ExceptT.mk]

theorem run_bind [Monad m] (x : ExceptT ε m α)
        : run (x >>= f : ExceptT ε m β)
          =
          run x >>= fun
                     | Except.ok x => run (f x)
                     | Except.error e => pure (Except.error e) :=
  rfl

@[simp] theorem lift_pure [Monad m] [LawfulMonad m] (a : α) : ExceptT.lift (pure a) = (pure a : ExceptT ε m α) := by
  simp [ExceptT.lift, pure, ExceptT.pure]

@[simp] theorem run_map [Monad m] [LawfulMonad m] (f : α → β) (x : ExceptT ε m α)
    : (f <$> x).run = Except.map f <$> x.run := by
  simp [Functor.map, ExceptT.map, map_eq_pure_bind]
  apply bind_congr
  intro a; cases a <;> simp [Except.map]

protected theorem seq_eq {α β ε : Type u} [Monad m] (mf : ExceptT ε m (α → β)) (x : ExceptT ε m α) : mf <*> x = mf >>= fun f => f <$> x :=
  rfl

protected theorem bind_pure_comp [Monad m] [LawfulMonad m] (f : α → β) (x : ExceptT ε m α) : x >>= pure ∘ f = f <$> x := by
  intros; rfl

protected theorem seqLeft_eq {α β ε : Type u} {m : Type u → Type v} [Monad m] [LawfulMonad m] (x : ExceptT ε m α) (y : ExceptT ε m β) : x <* y = const β <$> x <*> y := by
  show (x >>= fun a => y >>= fun _ => pure a) = (const (α := α) β <$> x) >>= fun f => f <$> y
  rw [← ExceptT.bind_pure_comp]
  apply ext
  simp [run_bind]
  apply bind_congr
  intro
  | Except.error _ => simp
  | Except.ok _ =>
    simp [map_eq_pure_bind]; apply bind_congr; intro b;
    cases b <;> simp [comp, Except.map, const]

protected theorem seqRight_eq [Monad m] [LawfulMonad m] (x : ExceptT ε m α) (y : ExceptT ε m β) : x *> y = const α id <$> x <*> y := by
  show (x >>= fun _ => y) = (const α id <$> x) >>= fun f => f <$> y
  rw [← ExceptT.bind_pure_comp]
  apply ext
  simp [run_bind]
  apply bind_congr
  intro a; cases a <;> simp

instance [Monad m] [LawfulMonad m] : LawfulMonad (ExceptT ε m) where
  id_map         := by intros; apply ext; simp
  map_const      := by intros; rfl
  seqLeft_eq     := ExceptT.seqLeft_eq
  seqRight_eq    := ExceptT.seqRight_eq
  pure_seq       := by intros; apply ext; simp [ExceptT.seq_eq, run_bind]
  bind_pure_comp := ExceptT.bind_pure_comp
  bind_map       := by intros; rfl
  pure_bind      := by intros; apply ext; simp [run_bind]
  bind_assoc     := by intros; apply ext; simp [run_bind]; apply bind_congr; intro a; cases a <;> simp

end ExceptT

/- ReaderT -/

namespace ReaderT

theorem ext [Monad m] {x y : ReaderT ρ m α} (h : ∀ ctx, x.run ctx = y.run ctx) : x = y := by
  simp [run] at h
  exact funext h

@[simp] theorem run_pure [Monad m] (a : α) (ctx : ρ) : (pure a : ReaderT ρ m α).run ctx = pure a := rfl

@[simp] theorem run_bind [Monad m] (x : ReaderT ρ m α) (f : α → ReaderT ρ m β) (ctx : ρ)
    : (x >>= f).run ctx = x.run ctx >>= λ a => (f a).run ctx := rfl

@[simp] theorem run_map [Monad m] (f : α → β) (x : ReaderT ρ m α) (ctx : ρ)
    : (f <$> x).run ctx = f <$> x.run ctx := rfl

@[simp] theorem run_monadLift [MonadLiftT n m] (x : n α) (ctx : ρ)
    : (monadLift x : ReaderT ρ m α).run ctx = (monadLift x : m α) := rfl

@[simp] theorem run_monadMap [Monad m] [MonadFunctor n m] (f : {β : Type u} → n β → n β) (x : ReaderT ρ m α) (ctx : ρ)
    : (monadMap @f x : ReaderT ρ m α).run ctx = monadMap @f (x.run ctx) := rfl

@[simp] theorem run_read [Monad m] (ctx : ρ) : (ReaderT.read : ReaderT ρ m ρ).run ctx = pure ctx := rfl

@[simp] theorem run_seq {α β : Type u} [Monad m] [LawfulMonad m] (f : ReaderT ρ m (α → β)) (x : ReaderT ρ m α) (ctx : ρ) : (f <*> x).run ctx = (f.run ctx <*> x.run ctx) := by
  rw [seq_eq_bind (m := m)]; rfl

@[simp] theorem run_seqRight [Monad m] [LawfulMonad m] (x : ReaderT ρ m α) (y : ReaderT ρ m β) (ctx : ρ) : (x *> y).run ctx = (x.run ctx *> y.run ctx) := by
  rw [seqRight_eq_bind (m := m)]; rfl

@[simp] theorem run_seqLeft [Monad m] [LawfulMonad m] (x : ReaderT ρ m α) (y : ReaderT ρ m β) (ctx : ρ) : (x <* y).run ctx = (x.run ctx <* y.run ctx) := by
  rw [seqLeft_eq_bind (m := m)]; rfl

instance [Monad m] [LawfulMonad m] : LawfulMonad (ReaderT ρ m) where
  id_map         := by intros; apply ext; intros; simp
  map_const      := by intros; rfl
  seqLeft_eq     := by intros; apply ext; intros; simp; apply LawfulApplicative.seqLeft_eq
  seqRight_eq    := by intros; apply ext; intros; simp; apply LawfulApplicative.seqRight_eq
  pure_seq       := by intros; apply ext; intros; simp; apply LawfulApplicative.pure_seq
  bind_pure_comp := by intros; apply ext; intros; simp; apply LawfulMonad.bind_pure_comp
  bind_map       := by intros; rfl
  pure_bind      := by intros; apply ext; intros; simp
  bind_assoc     := by intros; apply ext; intros; simp

end ReaderT

/- StateRefT -/

instance [Monad m] [LawfulMonad m] : LawfulMonad (StateRefT' ω σ m) :=
  inferInstanceAs (LawfulMonad (ReaderT (ST.Ref ω σ) m))

/- StateT -/

namespace StateT

theorem ext {x y : StateT σ m α} (h : ∀ s, x.run s = y.run s) : x = y :=
  funext h

@[simp] theorem run'_eq [Monad m] (x : StateT σ m α) (s : σ) : run' x s = (·.1) <$> run x s :=
  rfl

@[simp] theorem run_pure [Monad m] (a : α) (s : σ) : (pure a : StateT σ m α).run s = pure (a, s) := rfl

@[simp] theorem run_bind [Monad m] (x : StateT σ m α) (f : α → StateT σ m β) (s : σ)
    : (x >>= f).run s = x.run s >>= λ p => (f p.1).run p.2 := by
  simp [bind, StateT.bind, run]
  apply bind_congr
  intro p; cases p; rfl

@[simp] theorem run_map {α β σ : Type u} [Monad m] [LawfulMonad m] (f : α → β) (x : StateT σ m α) (s : σ) : (f <$> x).run s = (fun (p : α × σ) => (f p.1, p.2)) <$> x.run s := by
  simp [Functor.map, StateT.map, run, map_eq_pure_bind]
  apply bind_congr
  intro p; cases p; rfl

@[simp] theorem run_get [Monad m] (s : σ)    : (get : StateT σ m σ).run s = pure (s, s) := rfl

@[simp] theorem run_set [Monad m] (s s' : σ) : (set s' : StateT σ m PUnit).run s = pure (⟨⟩, s') := rfl

@[simp] theorem run_modify [Monad m] (f : σ → σ) (s : σ) : (modify f : StateT σ m PUnit).run s = pure (⟨⟩, f s) := rfl

@[simp] theorem run_modifyGet [Monad m] (f : σ → α × σ) (s : σ) : (modifyGet f : StateT σ m α).run s = pure ((f s).1, (f s).2) := by
  simp [modifyGet, MonadStateOf.modifyGet, StateT.modifyGet, run]; cases f s <;> rfl

@[simp] theorem run_lift {α σ : Type u} [Monad m] (x : m α) (s : σ) : (StateT.lift x : StateT σ m α).run s = x >>= fun a => pure (a, s) := rfl

@[simp] theorem run_bind_lift {α σ : Type u} [Monad m] [LawfulMonad m] (x : m α) (f : α → StateT σ m β) (s : σ) : (StateT.lift x >>= f).run s = x >>= fun a => (f a).run s := by
  simp [StateT.lift, StateT.run, bind, StateT.bind]

@[simp] theorem run_monadLift {α σ : Type u} [Monad m] [MonadLiftT n m] (x : n α) (s : σ) : (monadLift x : StateT σ m α).run s = (monadLift x : m α) >>= fun a => pure (a, s) := rfl

@[simp] theorem run_monadMap [Monad m] [MonadFunctor n m] (f : {β : Type u} → n β → n β) (x : StateT σ m α) (s : σ)
    : (monadMap @f x : StateT σ m α).run s = monadMap @f (x.run s) := rfl

@[simp] theorem run_seq {α β σ : Type u} [Monad m] [LawfulMonad m] (f : StateT σ m (α → β)) (x : StateT σ m α) (s : σ) : (f <*> x).run s = (f.run s >>= fun fs => (fun (p : α × σ) => (fs.1 p.1, p.2)) <$> x.run fs.2) := by
  show (f >>= fun g => g <$> x).run s = _
  simp

@[simp] theorem run_seqRight [Monad m] [LawfulMonad m] (x : StateT σ m α) (y : StateT σ m β) (s : σ) : (x *> y).run s = (x.run s >>= fun p => y.run p.2) := by
  show (x >>= fun _ => y).run s = _
  simp

@[simp] theorem run_seqLeft {α β σ : Type u} [Monad m] [LawfulMonad m] (x : StateT σ m α) (y : StateT σ m β) (s : σ) : (x <* y).run s = (x.run s >>= fun p => y.run p.2 >>= fun p' => pure (p.1, p'.2)) := by
  show (x >>= fun a => y >>= fun _ => pure a).run s = _
  simp

theorem seqRight_eq [Monad m] [LawfulMonad m] (x : StateT σ m α) (y : StateT σ m β) : x *> y = const α id <$> x <*> y := by
  apply ext; intro s
  simp [map_eq_pure_bind]
  apply bind_congr; intro p; cases p
  simp [Prod.ext]

theorem seqLeft_eq [Monad m] [LawfulMonad m] (x : StateT σ m α) (y : StateT σ m β) : x <* y = const β <$> x <*> y := by
  apply ext; intro s
  simp [map_eq_pure_bind]

instance [Monad m] [LawfulMonad m] : LawfulMonad (StateT σ m) where
  id_map         := by intros; apply ext; intros; simp[Prod.ext]
  map_const      := by intros; rfl
  seqLeft_eq     := seqLeft_eq
  seqRight_eq    := seqRight_eq
  pure_seq       := by intros; apply ext; intros; simp
  bind_pure_comp := by intros; apply ext; intros; simp; apply LawfulMonad.bind_pure_comp
  bind_map       := by intros; rfl
  pure_bind      := by intros; apply ext; intros; simp
  bind_assoc     := by intros; apply ext; intros; simp

end StateT


/-
Copyright (c) 2020 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Meta.Transform
import Lean.Meta.Tactic.Replace
import Lean.Meta.Tactic.Util
import Lean.Meta.Tactic.Clear
import Lean.Meta.Tactic.Simp.Types
import Lean.Meta.Tactic.Simp.Rewrite

namespace Lean.Meta
namespace Simp

builtin_initialize congrHypothesisExceptionId : InternalExceptionId ←
  registerInternalExceptionId `congrHypothesisFailed

def throwCongrHypothesisFailed : MetaM α :=
  throw <| Exception.internal congrHypothesisExceptionId

def Result.getProof (r : Result) : MetaM Expr := do
  match r.proof? with
  | some p => return p
  | none   => mkEqRefl r.expr

private def mkEqTrans (r₁ r₂ : Result) : MetaM Result := do
  match r₁.proof? with
  | none => return r₂
  | some p₁ => match r₂.proof? with
    | none    => return { r₂ with proof? := r₁.proof? }
    | some p₂ => return { r₂ with proof? := (← Meta.mkEqTrans p₁ p₂) }

private def mkCongrFun (r : Result) (a : Expr) : MetaM Result :=
  match r.proof? with
  | none   => return { expr := mkApp r.expr a, proof? := none }
  | some h => return { expr := mkApp r.expr a, proof? := (← Meta.mkCongrFun h a) }

private def mkCongr (r₁ r₂ : Result) : MetaM Result :=
  let e := mkApp r₁.expr r₂.expr
  match r₁.proof?, r₂.proof? with
  | none,     none   => return { expr := e, proof? := none }
  | some h,  none    => return { expr := e, proof? := (← Meta.mkCongrFun h r₂.expr) }
  | none,    some h  => return { expr := e, proof? := (← Meta.mkCongrArg r₁.expr h) }
  | some h₁, some h₂ => return { expr := e, proof? := (← Meta.mkCongr h₁ h₂) }

private def mkImpCongr (r₁ r₂ : Result) : MetaM Result := do
  let e ← mkArrow r₁.expr r₂.expr
  match r₁.proof?, r₂.proof? with
  | none,     none   => return { expr := e, proof? := none }
  | _,        _      => return { expr := e, proof? := (← Meta.mkImpCongr (← r₁.getProof) (← r₂.getProof)) } -- TODO specialize if bootleneck

private def reduceProj (e : Expr) : MetaM Expr := do
  match (← reduceProj? e) with
  | some e => return e
  | _      => return e

private def reduceProjFn? (e : Expr) : SimpM (Option Expr) := do
  matchConst e.getAppFn (fun _ => pure none) fun cinfo _ => do
    match (← getProjectionFnInfo? cinfo.name) with
    | none => return none
    | some projInfo =>
      if projInfo.fromClass then
        if (← read).simpLemmas.isDeclToUnfold cinfo.name then
          -- We only unfold class projections when the user explicitly requested them to be unfolded.
          -- Recall that `unfoldDefinition?` has support for unfolding this kind of projection.
          withReducibleAndInstances <| unfoldDefinition? e
        else
          return none
      else
        -- `structure` projection
        match (← unfoldDefinition? e) with
        | none   => pure none
        | some e =>
          match (← reduceProj? e.getAppFn) with
          | some f => return some (mkAppN f e.getAppArgs)
          | none   => return none

private def reduceFVar (cfg : Config) (e : Expr) : MetaM Expr := do
  if cfg.zeta then
    match (← getFVarLocalDecl e).value? with
    | some v => return v
    | none   => return e
  else
    return e

private def unfold? (e : Expr) : SimpM (Option Expr) := do
  let f := e.getAppFn
  if !f.isConst then
    return none
  let fName := f.constName!
  if (← isProjectionFn fName) then
    return none -- should be reduced by `reduceProjFn?`
  if (← read).simpLemmas.isDeclToUnfold e.getAppFn.constName! then
    withDefault <| unfoldDefinition? e
  else
    return none

private partial def reduce (e : Expr) : SimpM Expr := withIncRecDepth do
  let cfg := (← read).config
  if cfg.beta then
    let e' := e.headBeta
    if e' != e then
      return (← reduce e')
  -- TODO: eta reduction
  if cfg.proj then
    match (← reduceProjFn? e) with
    | some e => return (← reduce e)
    | none   => pure ()
  if cfg.iota then
    match (← reduceRecMatcher? e) with
    | some e => return (← reduce e)
    | none   => pure ()
  match (← unfold? e) with
  | some e => reduce e
  | none => return e

private partial def dsimp (e : Expr) : M Expr := do
  transform e (post := fun e => return TransformStep.done (← reduce e))

partial def simp (e : Expr) : M Result := withIncRecDepth do
  let cfg ← getConfig
  if (← isProof e) then
    return { expr := e }
  if cfg.memoize then
    if let some result := (← get).cache.find? e then
      return result
  simpLoop { expr := e }

where
  simpLoop (r : Result) : M Result := do
    let cfg ← getConfig
    if (← get).numSteps > cfg.maxSteps then
      throwError "simp failed, maximum number of steps exceeded"
    else
      let init := r.expr
      modify fun s => { s with numSteps := s.numSteps + 1 }
      match (← pre r.expr) with
      | Step.done r   => cacheResult cfg r
      | Step.visit r' =>
        let r ← mkEqTrans r r'
        let r ← mkEqTrans r (← simpStep r.expr)
        match (← post r.expr) with
        | Step.done r'  => cacheResult cfg (← mkEqTrans r r')
        | Step.visit r' =>
          let r ← mkEqTrans r r'
          if cfg.singlePass || init == r.expr then
            cacheResult cfg r
          else
            simpLoop r

  simpStep (e : Expr) : M Result := do
    match e with
    | Expr.mdata _ e _ => simp e
    | Expr.proj ..     => pure { expr := (← reduceProj e) }
    | Expr.app ..      => simpApp e
    | Expr.lam ..      => simpLambda e
    | Expr.forallE ..  => simpForall e
    | Expr.letE ..     => simpLet e
    | Expr.const ..    => simpConst e
    | Expr.bvar ..     => unreachable!
    | Expr.sort ..     => pure { expr := e }
    | Expr.lit ..      => pure { expr := e }
    | Expr.mvar ..     => pure { expr := (← instantiateMVars e) }
    | Expr.fvar ..     => pure { expr := (← reduceFVar (← getConfig) e) }

  congrDefault (e : Expr) : M Result :=
    withParent e <| e.withApp fun f args => do
      let infos := (← getFunInfoNArgs f args.size).paramInfo
      let mut r ← simp f
      let mut i := 0
      for arg in args do
        trace[Debug.Meta.Tactic.simp] "app [{i}] {infos.size} {arg} hasFwdDeps: {infos[i].hasFwdDeps}"
        if i < infos.size && !infos[i].hasFwdDeps then
          r ← mkCongr r (← simp arg)
        else if (← whnfD (← inferType r.expr)).isArrow then
          r ← mkCongr r (← simp arg)
        else
          r ← mkCongrFun r (← dsimp arg)
        i := i + 1
      return r

  /- Return true iff processing the given congruence lemma hypothesis produced a non-refl proof. -/
  processCongrHypothesis (h : Expr) : M Bool := do
    forallTelescopeReducing (← inferType h) fun xs hType => withNewLemmas xs do
      let lhs ← instantiateMVars hType.appFn!.appArg!
      let r ← simp lhs
      let rhs := hType.appArg!
      rhs.withApp fun m zs => do
        let val ← mkLambdaFVars zs r.expr
        unless (← isDefEq m val) do
          throwCongrHypothesisFailed
        unless (← isDefEq h (← mkLambdaFVars xs (← r.getProof))) do
          throwCongrHypothesisFailed
        return r.proof?.isSome

  /- Try to rewrite `e` children using the given congruence lemma -/
  tryCongrLemma? (c : CongrLemma) (e : Expr) : M (Option Result) := withNewMCtxDepth do
    trace[Debug.Meta.Tactic.simp.congr] "{c.theoremName}, {e}"
    let lemma ← mkConstWithFreshMVarLevels c.theoremName
    let (xs, bis, type) ← forallMetaTelescopeReducing (← inferType lemma)
    if c.hypothesesPos.any (· ≥ xs.size) then
      return none
    let lhs := type.appFn!.appArg!
    let rhs := type.appArg!
    if (← isDefEq lhs e) then
      let mut modified := false
      for i in c.hypothesesPos do
        let x := xs[i]
        try
          if (← processCongrHypothesis x) then
            modified := true
        catch _ =>
          trace[Meta.Tactic.simp.congr] "processCongrHypothesis {c.theoremName} failed {← inferType x}"
          return none
      unless modified do
        trace[Meta.Tactic.simp.congr] "{c.theoremName} not modified"
        return none
      unless (← synthesizeArgs c.theoremName xs bis (← read).discharge?) do
        trace[Meta.Tactic.simp.congr] "{c.theoremName} synthesizeArgs failed"
        return none
      let eNew ← instantiateMVars rhs
      let proof ← instantiateMVars (mkAppN lemma xs)
      return some { expr := eNew, proof? := proof }
    else
      return none

  congr (e : Expr) : M Result := do
    let f := e.getAppFn
    if f.isConst then
      let congrLemmas ← getCongrLemmas
      let cs := congrLemmas.get f.constName!
      for c in cs do
        match (← tryCongrLemma? c e) with
        | none   => pure ()
        | some r => return r
      congrDefault e
    else
      congrDefault e

  simpApp (e : Expr) : M Result := do
    let e ← reduce e
    if !e.isApp then
      simp e
    else
      congr e

  simpConst (e : Expr) : M Result :=
    return { expr := (← reduce e) }

  withNewLemmas {α} (xs : Array Expr) (f : M α) : M α := do
    if (← getConfig).contextual then
      let mut s ← getSimpLemmas
      let mut updated := false
      for x in xs do
        if (← isProof x) then
          s ← s.add #[] x
          updated := true
      if updated then
        withSimpLemmas s f
      else
        f
    else
      f

  simpLambda (e : Expr) : M Result :=
    withParent e <| lambdaTelescope e fun xs e => withNewLemmas xs do
      let r ← simp e
      let eNew ← mkLambdaFVars xs r.expr
      match r.proof? with
      | none   => return { expr := eNew }
      | some h =>
        let p ← xs.foldrM (init := h) fun x h => do
          mkFunExt (← mkLambdaFVars #[x] h)
        return { expr := eNew, proof? := p }

  simpArrow (e : Expr) : M Result := do
    trace[Debug.Meta.Tactic.simp] "arrow {e}"
    let p := e.bindingDomain!
    let q := e.bindingBody!
    let rp ← simp p
    trace[Debug.Meta.Tactic.simp] "arrow [{(← getConfig).contextual}] {p} [{← isProp p}] -> {q} [{← isProp q}]"
    if (← (← getConfig).contextual <&&> isProp p <&&> isProp q) then
      trace[Debug.Meta.Tactic.simp] "ctx arrow {rp.expr} -> {q}"
      withLocalDeclD e.bindingName! rp.expr fun h => do
        let s ← getSimpLemmas
        let s ← s.add #[] h
        withSimpLemmas s do
          let rq ← simp q
          match rq.proof? with
          | none    => mkImpCongr rp rq
          | some hq =>
            let hq ← mkLambdaFVars #[h] hq
            return { expr := (← mkArrow rp.expr rq.expr), proof? := (← mkImpCongrCtx (← rp.getProof) hq) }
    else
      mkImpCongr rp (← simp q)

  simpForall (e : Expr) : M Result := withParent e do
    trace[Debug.Meta.Tactic.simp] "forall {e}"
    if e.isArrow then
      simpArrow e
    else if (← isProp e) then
      withLocalDecl e.bindingName! e.bindingInfo! e.bindingDomain! fun x => withNewLemmas #[x] do
        let b := e.bindingBody!.instantiate1 x
        let rb ← simp b
        let eNew ← mkForallFVars #[x] rb.expr
        match rb.proof? with
        | none   => return { expr := eNew }
        | some h => return { expr := eNew, proof? := (← mkForallCongr (← mkLambdaFVars #[x] h)) }
    else
      return { expr := (← dsimp e) }

  simpLet (e : Expr) : M Result := do
    if (← getConfig).zeta then
      match e with
      | Expr.letE _ _ v b _ => return { expr := b.instantiate1 v }
      | _ => unreachable!
    else
      -- TODO: simplify nondependent let-decls
      return { expr := (← dsimp e) }

  cacheResult (cfg : Config) (r : Result) : M Result := do
    if cfg.memoize then
      modify fun s => { s with cache := s.cache.insert e r }
    return r

def main (e : Expr) (ctx : Context) (methods : Methods := {}) : MetaM Result := do
  withReducible do
    simp e methods ctx |>.run' {}

namespace DefaultMethods
mutual
  partial def discharge? (e : Expr) : SimpM (Option Expr) := do
    let ctx ← read
    if ctx.dischargeDepth >= ctx.config.maxDischargeDepth then
      trace[Meta.Tactic.simp.discharge] "maximum discharge depth has been reached"
      return none
    else
      withReader (fun ctx => { ctx with dischargeDepth := ctx.dischargeDepth + 1 }) do
        let r ← simp e methods
        if r.expr.isConstOf ``True then
          try
            return some (← mkOfEqTrue (← r.getProof))
          catch _ =>
            return none
        else
          return none

  partial def pre (e : Expr) : SimpM Step :=
    preDefault e discharge?

  partial def post (e : Expr) : SimpM Step :=
    postDefault e discharge?

  partial def methods : Methods :=
    { pre := pre, post := post, discharge? := discharge? }
end
end DefaultMethods

end Simp

def simp (e : Expr) (ctx : Simp.Context) : MetaM Simp.Result := do profileitM Exception "simp" (← getOptions) do
  Simp.main e ctx (methods := Simp.DefaultMethods.methods)

/-- See `simpTarget`. This method assumes `mvarId` is not assigned, and we are already using `mvarId`s local context. -/
def simpTargetCore (mvarId : MVarId) (ctx : Simp.Context) : MetaM (Option MVarId) := do
  let target ← instantiateMVars (← getMVarType mvarId)
  let r ← simp target ctx
  if r.expr.isConstOf ``True then
    match r.proof? with
    | some proof => assignExprMVar mvarId  (← mkOfEqTrue proof)
    | none => assignExprMVar mvarId (mkConst ``True.intro)
    return none
  else
    match r.proof? with
    | some proof => replaceTargetEq mvarId r.expr proof
    | none =>
      if target != r.expr then
        replaceTargetDefEq mvarId r.expr
      else
        return mvarId

/--
  Simplify the given goal target (aka type). Return `none` if the goal was closed. Return `some mvarId'` otherwise,
  where `mvarId'` is the simplified new goal. -/
def simpTarget (mvarId : MVarId) (ctx : Simp.Context) : MetaM (Option MVarId) :=
  withMVarContext mvarId do
    checkNotAssigned mvarId `simp
    simpTargetCore mvarId ctx

/--
  Simplify `prop` (which is inhabited by `proof`). Return `none` if the goal was closed. Return `some (proof', prop')`
  otherwise, where `proof' : prop'` and `prop'` is the simplified `prop`.

  This method assumes `mvarId` is not assigned, and we are already using `mvarId`s local context. -/
def simpStep (mvarId : MVarId) (proof : Expr) (prop : Expr) (ctx : Simp.Context) : MetaM (Option (Expr × Expr)) := do
  let r ← simp prop ctx
  if r.expr.isConstOf ``False then
    match r.proof? with
    | some eqProof => assignExprMVar mvarId (← mkFalseElim (← getMVarType mvarId) (← mkEqMP eqProof proof))
    | none => assignExprMVar mvarId (← mkFalseElim (← getMVarType mvarId) proof)
    return none
  else
    match r.proof? with
    | some eqProof => return some ((← mkEqMP eqProof proof), r.expr)
    | none =>
      if r.expr != prop then
        return some ((← mkExpectedTypeHint proof r.expr), r.expr)
      else
        return some (proof, r.expr)

end Lean.Meta