--  An abstract formalization of "isomorphism is equality up to relabeling"
-- -------------------------------------------------------------------------
--
-- See `README.md` for more info.
--
-- Helpers for the construction of functors into `universeStructure`.

import Structure.Basic
import Structure.Forgetfulness

open Morphisms
open HasStructure
open Structure
open GeneralizedFunctor
open StructureFunctor
open Forgetfulness
open SetoidStructureFunctor
open SetoidStructureEquiv

set_option autoBoundImplicitLocal false

-- Aliases for functors into `universeStructure`.

@[reducible] def UniverseFunctor (S : Structure) := StructureFunctor S universeStructure

-- A special limitation of `UniverseFunctor S` for a generic `S` is that since equivalences of
-- equivalences in `S` are propositions, we can only map instances of `S` to setoid structures. The
-- definition of `UniverseFunctor` itself does not encode this knowledge, but sometimes we want to
-- restrict ourselves to this very frequent case. We can do this by replacing `UniverseFunctor` with
-- `SetoidUniverseFunctor` and then coercing it into a `UniverseFunctor`.

structure SetoidUniverseFunctor (S : Structure) :=
(map     : S → universeStructure)
(functor : Functor (T := universeStructure) (setoidStructure ∘ map))

namespace SetoidUniverseFunctor

def universeFunctor {S : Structure} (F : SetoidUniverseFunctor S) : UniverseFunctor S :=
{ map     := setoidStructure ∘ F.map,
  functor := F.functor }

instance {S : Structure} : Coe (SetoidUniverseFunctor S) (UniverseFunctor S) := ⟨universeFunctor⟩

def toSetoidUniverseFunctor {S : Structure} (F : UniverseFunctor S) : SetoidUniverseFunctor S :=
{ map     := F.map,
  functor := comp.genFun' F.map F.functor toSetoidStructureEquiv.genFun }

def compFun {S T : Structure} (F : StructureFunctor S T) (G : SetoidUniverseFunctor T) : SetoidUniverseFunctor S :=
{ map     := G.map ∘ F.map,
  functor := comp.genFun' F.map F.functor G.functor }

def constFun {S : Structure} (T : Structure) : SetoidUniverseFunctor S :=
{ map     := Function.const (IsType.type S) T,
  functor := const.genFun (setoidStructure T) }

end SetoidUniverseFunctor

-- The function `setoidStructure`, which truncates the equivalences of a structure to setoids, is a
-- `UniverseFunctor`.

def structureToSetoidStructureFunctor' : SetoidUniverseFunctor universeStructure :=
{ map     := id,
  functor := toSetoidStructureEquiv.genFun }

def setoidUniverseFunctor {S : Structure} (F : UniverseFunctor S) : SetoidUniverseFunctor S :=
SetoidUniverseFunctor.compFun F structureToSetoidStructureFunctor'

def structureToSetoidStructureFunctor : UniverseFunctor universeStructure :=
SetoidUniverseFunctor.universeFunctor structureToSetoidStructureFunctor'

-- A helper structure for constructing a `SetoidUniverseFunctor`, and thus indirectly also a
-- `UniverseFunctor`.
--
-- Instead of directly specifying the resulting structure equivalence, we only give its `toFun` because
-- the `invFun` can be obtained by inverting the input equivalence. By stating the functor axioms in terms
-- of `toFun`, we can ensure that the result is indeed a functor.
--
-- The limitation that we can only define functors to setoid structures is reflected in
-- `SetoidUniverseFunctorDesc` in two ways:
-- * `toFun` yields a `SetoidStructureFunctor` because for a given equivalence `e`, it would be
--   unrealistic to obtain a regular `StructureFunctor`.
-- * The `respects...` functions return proofs. Especially, `respectsSetoid` can only return a proof since
--   it takes a proof as an input.

structure SetoidUniverseFunctorDesc (S : Structure) where
(map                                                : S → Structure)
(toFun          {a b   : S}                         : a ≃ b → SetoidStructureFunctor (map a) (map b))
(respectsSetoid {a b   : S} {e₁ e₂ : a ≃ b}         : e₁ ≈ e₂ → ∀ T, toFun e₁ T ≈ toFun e₂ T)
(respectsComp   {a b c : S} (e : a ≃ b) (f : b ≃ c) : ∀ T, toFun (f • e) T ≈ toFun f (toFun e T))
(respectsId     (a     : S)                         : ∀ T, toFun (id_ a) T ≈ T)

namespace SetoidUniverseFunctorDesc

variable {S : Structure} (D : SetoidUniverseFunctorDesc S)

def targetLeftInv {a b : S} (e : a ≃ b) : D.toFun e⁻¹ ⊙ D.toFun e ≃ @idFun (setoidStructure (D.map a)) :=
makeToSetoidStructureFunctorEquiv (λ T => let h₁ := D.respectsComp e e⁻¹ T;
                                          let h₂ := D.respectsSetoid (leftInv e) T;
                                          let h₃ := Setoid.trans (Setoid.symm h₁) h₂;
                                          let h₄ := D.respectsId a T;
                                          Setoid.trans h₃ h₄)

def targetRightInv {a b : S} (e : a ≃ b) : D.toFun e ⊙ D.toFun e⁻¹ ≃ @idFun (setoidStructure (D.map b)) :=
makeToSetoidStructureFunctorEquiv (λ T => let h₁ := D.respectsComp e⁻¹ e T;
                                          let h₂ := D.respectsSetoid (rightInv e) T;
                                          let h₃ := Setoid.trans (Setoid.symm h₁) h₂
                                          let h₄ := D.respectsId b T;
                                          Setoid.trans h₃ h₄)

def targetEquiv {a b : S} (e : a ≃ b) : SetoidStructureEquiv (D.map a) (D.map b) :=
{ toFun  := D.toFun e,
  invFun := D.toFun e⁻¹,
  isInv  := makeSetoidStructureFunctorInverse (targetLeftInv D e) (targetRightInv D e) }

theorem targetRespectsEquiv {a b : S} {e₁ e₂ : a ≃ b} :
  e₁ ≈ e₂ → targetEquiv D e₁ ≈ targetEquiv D e₂ :=
λ h => ⟨makeSetoidStructureEquivEquiv (D.respectsSetoid h)
                                      (D.respectsSetoid (inv_congrArg h))⟩

theorem targetRespectsComp {a b c : S} (e : a ≃ b) (f : b ≃ c) :
  targetEquiv D (f • e) ≈ StructureEquiv.trans (targetEquiv D e) (targetEquiv D f) :=
⟨makeSetoidStructureEquivEquiv (D.respectsComp e f)
                               (λ T => let h₁ := D.respectsComp f⁻¹ e⁻¹ T;
                                       let h₂ := D.respectsSetoid (compInv e f) T;
                                       Setoid.trans h₂ h₁)⟩

theorem targetRespectsId (a : S) :
  targetEquiv D (id_ a) ≈ StructureEquiv.refl (setoidStructure (D.map a)) :=
⟨makeSetoidStructureEquivEquiv (D.respectsId a)
                               (λ T => let h₁ := D.respectsId a T;
                                       let h₂ := D.respectsSetoid (idInv a) T;
                                       Setoid.trans h₂ h₁)⟩

theorem targetRespectsInv {a b : S} (e : a ≃ b) :
  targetEquiv D e⁻¹ ≈ StructureEquiv.symm (targetEquiv D e) :=
⟨makeSetoidStructureEquivEquiv (λ T => Setoid.refl (D.toFun e⁻¹ T))
                               (D.respectsSetoid (invInv e))⟩

def setoidUniverseFunctor : SetoidUniverseFunctor S :=
{ map     := D.map,
  functor := { mapEquiv  := targetEquiv D,
               isFunctor := { respectsEquiv := targetRespectsEquiv D,
                              respectsComp  := targetRespectsComp   D,
                              respectsId    := targetRespectsId     D,
                              respectsInv   := targetRespectsInv    D } } }

def universeFunctor : UniverseFunctor S := SetoidUniverseFunctor.universeFunctor (setoidUniverseFunctor D)

end SetoidUniverseFunctorDesc

instance {S : Structure} : Coe (SetoidUniverseFunctorDesc S) (SetoidUniverseFunctor S) := ⟨SetoidUniverseFunctorDesc.setoidUniverseFunctor⟩
instance {S : Structure} : Coe (SetoidUniverseFunctorDesc S) (UniverseFunctor S)       := ⟨SetoidUniverseFunctorDesc.universeFunctor⟩

-- A 2-functor specifically between universe structures. I.e. we have a functor between structures and a
-- functor between equivalences.

structure UniverseStructureFunctor where
(map                                                       : Structure → Structure)
(mapEquiv      {S T   : Structure}                         : S ≃ T → map S ≃ map T)
(respectsEquiv {S T   : Structure}                         : GeneralizedFunctor.Functor (S := StructureEquiv.equivStructure S T) (T := StructureEquiv.equivStructure (map S) (map T)) mapEquiv)
(respectsComp  {S T U : Structure} (e : S ≃ T) (f : T ≃ U) : mapEquiv (f • e) ≃ mapEquiv f • mapEquiv e)
-- TODO: Why don't `≃` and `id_` work as they should here?
(respectsId    (S     : Structure)                         : StructureEquiv.EquivEquiv (mapEquiv (id_ S)) (StructureEquiv.refl (map S)))
(respectsInv   {S T   : Structure} (e : S ≃ T)             : mapEquiv e⁻¹ ≃ (mapEquiv e)⁻¹)

namespace UniverseStructureFunctor

instance universeStructureFunctorCoeFun : CoeFun UniverseStructureFunctor (λ _ => Structure → Structure) := ⟨UniverseStructureFunctor.map⟩

variable (F : UniverseStructureFunctor)

def congrArg {S T : Structure} (e : S ≃ T) : F S ≃ F T := F.mapEquiv e

def universeFunctor : UniverseFunctor universeStructure :=
{ map     := F.map,
  functor := { mapEquiv  := F.mapEquiv,
               isFunctor := { respectsEquiv := λ ⟨η⟩ => ⟨F.respectsEquiv η⟩,
                              respectsComp  := λ e f => ⟨F.respectsComp  e f⟩,
                              respectsId    := λ S   => ⟨F.respectsId    S⟩,
                              respectsInv   := λ e   => ⟨F.respectsInv   e⟩ } } }

end UniverseStructureFunctor

instance : Coe UniverseStructureFunctor (UniverseFunctor universeStructure) := ⟨UniverseStructureFunctor.universeFunctor⟩

-- Similar to `SetoidUniverseFunctorDesc`, but constructs a `UniverseStructureFunctor` producing regular
-- structures without any truncation to setoids. This is possible because equivalences of structure
-- equivalences also come in a non-setoid form.

structure UniverseStructureFunctorDesc where
(map                                                       : Structure → Structure)
(toFun         {S T   : Structure}                         : S ≃ T → StructureFunctor (map S) (map T))
(respectsEquiv {S T   : Structure}                         : GeneralizedFunctor.Functor (S := StructureEquiv.equivStructure S T) (T := functorStructure (map S) (map T)) toFun)
(respectsComp  {S T U : Structure} (e : S ≃ T) (f : T ≃ U) : toFun (f • e) ≃ toFun f ⊙ toFun e)
(respectsCompNat {S T U : Structure} {e₁ e₂ : S ≃ T} {f₁ f₂ : T ≃ U} (η : e₁ ≃ e₂) (θ : f₁ ≃ f₂) :
   compFun.congrArg (respectsEquiv η) (respectsEquiv θ) • respectsComp e₁ f₁ ≈ respectsComp e₂ f₂ • respectsEquiv (StructureEquiv.comp_congrArg η θ))
(respectsId    (S     : Structure)                         : toFun (id_ S) ≃ @idFun (map S))

namespace UniverseStructureFunctorDesc

variable (D : UniverseStructureFunctorDesc)

-- TODO: We can probably make use of `toFunFunctor` somehow in order to make the missing proofs easier.

def toFunFunctor (S T : Structure) : StructureFunctor (StructureEquiv.equivStructure S T) (functorStructure (D.map S) (D.map T)) :=
{ map     := D.toFun,
  functor := D.respectsEquiv }

def targetLeftInv {S T : Structure} (e : S ≃ T) : D.toFun e⁻¹ ⊙ D.toFun e ≃ @idFun (D.map S) :=
let η₁ := FunctorEquiv.symm (D.respectsComp e e⁻¹);
let η₂ := D.respectsEquiv (StructureEquiv.leftInv' e);
FunctorEquiv.trans (FunctorEquiv.trans η₁ η₂) (D.respectsId S)

def targetRightInv {S T : Structure} (e : S ≃ T) : D.toFun e ⊙ D.toFun e⁻¹ ≃ @idFun (D.map T) :=
let η₁ := FunctorEquiv.symm (D.respectsComp e⁻¹ e);
let η₂ := D.respectsEquiv (StructureEquiv.rightInv' e);
FunctorEquiv.trans (FunctorEquiv.trans η₁ η₂) (D.respectsId T)

-- This might simplify some proofs.
theorem targetInv {S T : Structure} (e : S ≃ T) :
  targetRightInv D e ≈ StructureEquiv.symm_symm e ▸ targetLeftInv D e⁻¹ :=
sorry

def targetEquiv {S T : Structure} (e : S ≃ T) : StructureEquiv (D.map S) (D.map T) :=
{ toFun  := D.toFun e,
  invFun := D.toFun e⁻¹,
  isInv  := { leftInv  := targetLeftInv  D e,
              rightInv := targetRightInv D e,
              lrCompat := sorry,
              rlCompat := sorry } }

def targetRespectsEquiv {S T : Structure} {e₁ e₂ : S ≃ T} (η : e₁ ≃ e₂) :
  targetEquiv D e₁ ≃ targetEquiv D e₂ :=
{ toFunEquiv    := D.respectsEquiv η,
  invFunEquiv   := D.respectsEquiv (StructureEquiv.inv_congrArg η),
  leftInvEquiv  := sorry,
  rightInvEquiv := sorry }

def targetEquiv.functor {S T : Structure} :
  GeneralizedFunctor.Functor (S := StructureEquiv.equivStructure S T) (T := StructureEquiv.equivStructure (D.map S) (D.map T)) (targetEquiv D) :=
{ mapEquiv  := targetRespectsEquiv D,
  isFunctor := sorry }

def targetRespectsComp {S T U : Structure} (e : S ≃ T) (f : T ≃ U) :
  targetEquiv D (StructureEquiv.trans e f) ≃ StructureEquiv.trans (targetEquiv D e) (targetEquiv D f) :=
{ toFunEquiv    := D.respectsComp e   f,
  invFunEquiv   := D.respectsComp f⁻¹ e⁻¹,
  leftInvEquiv  := sorry,
  rightInvEquiv := sorry }

def targetRespectsId (S : Structure) :
  targetEquiv D (StructureEquiv.refl S) ≃ StructureEquiv.refl (D.map S) :=
{ toFunEquiv    := D.respectsId S,
  invFunEquiv   := D.respectsId S,
  leftInvEquiv  := sorry,
  rightInvEquiv := sorry }

def targetRespectsInv {S T : Structure} (e : S ≃ T) :
  targetEquiv D (StructureEquiv.symm e) ≃ StructureEquiv.symm (targetEquiv D e) :=
{ toFunEquiv    := FunctorEquiv.refl (D.toFun e⁻¹),
  invFunEquiv   := D.respectsEquiv (StructureEquiv.invInv e),
  leftInvEquiv  := sorry,
  rightInvEquiv := sorry }

def universeStructureFunctor : UniverseStructureFunctor :=
{ map            := D.map,
  mapEquiv       := targetEquiv         D,
  respectsEquiv  := targetEquiv.functor D,
  respectsComp   := targetRespectsComp  D,
  respectsId     := targetRespectsId    D,
  respectsInv    := targetRespectsInv   D }

def universeFunctor : UniverseFunctor universeStructure := UniverseStructureFunctor.universeFunctor (universeStructureFunctor D)

end UniverseStructureFunctorDesc

instance : Coe UniverseStructureFunctorDesc UniverseStructureFunctor := ⟨UniverseStructureFunctorDesc.universeStructureFunctor⟩
instance : Coe UniverseStructureFunctorDesc (UniverseFunctor universeStructure) := ⟨UniverseStructureFunctorDesc.universeFunctor⟩


-- Copyright (c) 2018 Michael Jendrusch. All rights reserved.
import data.equiv.basic
import category_theory.category
import category_theory.functor
import category_theory.products
import category_theory.natural_isomorphism
import ..monoidal_category
import ..monoid_object
open category_theory
open tactic

universes u v

namespace category_theory.monoidal

section

open monoidal_category

inductive nat_hom (m n : ℕ) : Type
| mk : nat_hom

@[simp] lemma equality {m n : ℕ} (f : nat_hom m n) : f = nat_hom.mk m n :=
by cases f; refl

@[simp] def nat_id (X : ℕ) : nat_hom X X :=
nat_hom.mk X X
@[simp] def nat_comp (X Y Z : ℕ) (f : nat_hom X Y) (g : nat_hom Y Z) : nat_hom X Z :=
nat_hom.mk X Z
@[simp] def nat_tensor_obj (X Y : ℕ) : ℕ := X + Y
@[simp] def nat_tensor_hom (A B C D : ℕ) (f : nat_hom A B) (g : nat_hom C D) :
  nat_hom (A + C) (B + D) := nat_hom.mk (A + C) (B + D)

instance naturals : monoidal_category (nat) :=
{ hom  := λ X Y, nat_hom X Y,
  id   := nat_id,
  comp := nat_comp,
  id_comp' := by tidy; rw equality f,
  comp_id' := by tidy; rw equality f,
  tensor_obj := nat_tensor_obj,
  tensor_hom := nat_tensor_hom,
  tensor_unit := nat.zero,
  left_unitor := λ X,
    { hom := nat_hom.mk (nat_tensor_obj 0 X) X,
      inv := nat_hom.mk X (nat_tensor_obj 0 X) },
  right_unitor := λ X,
    { hom := nat_hom.mk (nat_tensor_obj X 0) X,
      inv := nat_hom.mk X (nat_tensor_obj X 0) },
  associator := λ X Y Z,
    { hom := nat_hom.mk (nat_tensor_obj (nat_tensor_obj X Y) Z)
                        (nat_tensor_obj X (nat_tensor_obj Y Z)),
      inv := nat_hom.mk (nat_tensor_obj X (nat_tensor_obj Y Z))
                        (nat_tensor_obj (nat_tensor_obj X Y) Z)} }

end

instance nat_monoid_object (n : nat) : monoid_object n :=
{ unit    := nat_hom.mk 0 n,
  product := nat_hom.mk (n + n) n }

instance nat_comonoid_object (n : nat) : comonoid_object n :=
{ counit    := nat_hom.mk n 0,
  coproduct := nat_hom.mk n (n + n) }

instance nat_frobenius_object (n : nat) : frobenius_object n := {}

end category_theory.monoidal


import mll

def sequent := list Form_occ

instance : has_append sequent := ⟨list.append⟩
instance : has_mem Form_occ sequent := ⟨list.mem⟩

inductive proof : sequent → Type
| ax {A pi ni}                   : proof [(~A,pi), (A,ni)]
| cut (A) {i j} {Γ Γ' Δ Δ'}      : proof (Γ ++ [(A,i)] ++ Γ') → proof (Δ ++ [(~A,j)] ++ Δ') → proof (Γ++Γ'++Δ++Δ')
| tensor {A B ai bi ci} {Γ Γ' Δ Δ'} : proof (Γ ++ [(A,ai)] ++ Γ') → proof (Δ ++ [(B,bi)] ++ Δ') → proof (Γ++Γ'++ [(A ⊗ B,ci)] ++Δ++Δ')
| par {A B ai bi ci} {Γ Γ'}         : proof (Γ ++ [(A,ai),(B,bi)] ++ Γ') → proof (Γ ++ [(A ⅋ B,ci)] ++ Γ')
| ex {Ai Bi} {Γ Γ'}          : proof (Γ ++ [Ai,Bi] ++ Γ') → proof (Γ ++ [Bi,Ai] ++ Γ')

def mem_Form_occ_proof {Γ : sequent} : Form_occ → proof Γ → Prop :=
begin
  intros Ai π,
  induction π,
  case proof.ax : A pi ni { exact Ai = (A,pi)  }

end

instance {Γ : sequent} : has_mem Form_occ (proof Γ) := ⟨∃ Γ⟩

inductive proof_net : sequent → Type
| mk {Γ : sequent} (ps : proof_structure) : (Π A ∈ Γ, { i : ℕ // (A,i) ∈ ps ∧ ∀ Δ ∈ ps.links, ¬premise (A,i) Δ }) → proof_net Γ

instance {Γ : sequent} : has_coe (proof_net Γ) proof_structure := ⟨by rintro ⟨Γ,ps,_⟩; exact ps⟩

def relabel_Link (f : ℕ → ℕ) : Link → Link
| (Link.ax pi ni A) := Link.ax (f pi) (f ni) A
| (Link.cut pi ni A) := Link.cut (f pi) (f ni) A
| (Link.tensor ai bi ci A B) := Link.tensor (f ai) (f bi) (f ci) A B
| (Link.par ai bi ci A B) := Link.par (f ai) (f bi) (f ci) A B

lemma relabel_valid {l f} (hf : function.injective f): valid_link l → valid_link (relabel_Link f l) :=
begin
  cases l,
  case Link.ax : pi ni A { rintro ⟨_⟩, constructor, },
  case Link.cut : pi ni A { rintro ⟨_⟩, constructor, },
  case Link.tensor : ai bi ci A B {
    rintro ⟨_⟩, constructor, rintro e, injection e with e₁ e₂, apply ᾰ_ᾰ, congr, assumption, exact hf e₂, },
  case Link.par : ai bi ci A B {
    rintro ⟨_⟩, constructor, rintro e, injection e with e₁ e₂, apply ᾰ_ᾰ, congr, assumption, exact hf e₂, },
end

lemma relabel_injective {f} (hf : function.injective f) : function.injective (relabel_Link f) :=
by rintros ⟨l₁⟩ ⟨l₂⟩; intros h; injection h; congr; repeat {refl <|> assumption <|> apply hf}

lemma relabel_premise {l f D i } (hf : function.injective f) : premise (D,i) (relabel_Link f l) → ∃ j, f j = i ∧ premise (D,j) l :=
begin
  cases l,
  case Link.ax : pi ni A { rintro ⟨_⟩ },
  case Link.cut : pi ni A { rintro ⟨_⟩, exact ⟨pi,rfl,premise.cut_pos⟩, exact ⟨ni,rfl,premise.cut_neg⟩, },
  case Link.tensor : ai bi ci A B {
    rintro ⟨_⟩, exact ⟨ai,rfl,premise.tensor_left⟩, exact ⟨bi,rfl,premise.tensor_right⟩, },
  case Link.par : ai bi ci A B {
    rintro ⟨_⟩, exact ⟨ai,rfl,premise.par_left⟩, exact ⟨bi,rfl,premise.par_right⟩, },
end

lemma relabel_conclusion {l f D i } (hf : function.injective f) : conclusion (D,i) (relabel_Link f l) → ∃j, f j = i ∧ conclusion (D,j) l :=
begin
  cases l,
  case Link.ax : pi ni A {
    rintro ⟨_⟩, exact ⟨pi,rfl,conclusion.ax_pos⟩, exact ⟨ni,rfl,conclusion.ax_neg⟩, },
  case Link.cut : pi ni A { rintro ⟨_⟩ },
  case Link.tensor : ai bi ci A B { rintro ⟨_⟩, exact ⟨ci,rfl,conclusion.tensor⟩ },
  case Link.par : ai bi ci A B { rintro ⟨_⟩, exact ⟨ci,rfl,conclusion.par⟩ },
end

lemma relabel_mem {Δ f A i} (hf : function.injective f) : (A,i) ∈ (relabel_Link f Δ) → ∃ j, f j = i ∧ (A,j) ∈ Δ :=
begin
  intro h, cases h with h h,
  rcases (relabel_premise hf h) with ⟨j, ⟨fji,pΔ⟩⟩, refine ⟨j,fji,mem_Link.prem pΔ⟩,
  rcases (relabel_conclusion hf h) with ⟨j, ⟨fji,pΔ⟩⟩, refine ⟨j,fji,mem_Link.con pΔ⟩,
end

def proof_structure.relabel (ps : proof_structure) (f : ℕ → ℕ) (hf : function.injective f) : proof_structure :=
⟨set.image (relabel_Link f) ps.links,
  by rintros l ⟨l', ⟨hl',⟨_⟩⟩⟩; exact relabel_valid hf (ps.valid l' hl'),
begin
  rintros ⟨A,i⟩ _ _ ⟨k₁, ⟨hk₁,⟨_⟩⟩⟩ ⟨k₂, ⟨hk₂,⟨_⟩⟩⟩,
  intros pk₁ pk₂,
  congr,
  rcases relabel_premise hf pk₁ with ⟨j,hfj,u₁⟩,
  rcases relabel_premise hf pk₂ with ⟨j',hfj',u₂⟩,
  have : j' = j, rw ←hfj at hfj', exact hf hfj', rw this at u₂,
  exact ps.prem_unique (A,j) _ _ hk₁ hk₂ u₁ u₂
end
,
begin
  rintros ⟨A,i⟩ _ _ ⟨k₁, ⟨hk₁,⟨_⟩⟩⟩ ⟨k₂, ⟨hk₂,⟨_⟩⟩⟩,
  intros pk₁ pk₂,
  congr,
  rcases relabel_conclusion hf pk₁ with ⟨j,hfj,u₁⟩,
  rcases relabel_conclusion hf pk₂ with ⟨j',hfj',u₂⟩,
  have : j' = j, rw ←hfj at hfj', exact hf hfj', rw this at u₂,
  exact ps.con_unique (A,j) _ _ hk₁ hk₂ u₁ u₂
end⟩

def separators {α β} (f g : α → β) : Prop := ∀ x y, f x ≠ g y

lemma sep_even_odd : separators (λ x, 2 * x) (λ x, 2 * x + 1) :=
  λ x y, nat.two_mul_ne_two_mul_add_one

def disjoint_of_separators {ps₁ ps₂ : proof_structure} {f g} (hf hg) : separators f g → disjoint { Ai | Ai ∈ (ps₁.relabel f hf) } { Ai | Ai ∈ (ps₂.relabel g hg) } :=
begin
  rintros s ⟨A,i⟩ ⟨⟨Δ₁,⟨Δ₁', hΔ₁', ⟨_⟩⟩,h₁⟩,⟨Δ₂,⟨Δ₂', hΔ₂', ⟨_⟩⟩,h₂⟩⟩,
  rcases (relabel_mem hf h₁) with ⟨j₁,hfg,h₁⟩,
  rcases (relabel_mem hg h₂) with ⟨j₂,⟨_⟩,h₂⟩,
  exact s j₁ j₂ hfg,
end

def proof_net.disjoint {Γ Δ} : proof_net Γ → proof_net Δ → Prop :=
  by rintro ⟨_,ps₁,_⟩ ⟨_,ps₂,_⟩; exact disjoint {Ai | Ai ∈ ps₁} {Ai | Ai ∈ ps₂}

def net_links_ax (pi ni A) : set Link :=
  {Link.ax pi ni A}

def net_links_tensor (ai bi ci A B) (sA sB : set Link) : set Link :=
  {Link.tensor ai bi ci A B} ∪ sA ∪ sB

def net_links_par (ai bi ci A B) (s : set Link) : set Link :=
  {Link.par ai bi ci A B} ∪ s

def net_links_cut (pi ni A) (sA snA : set Link) : set Link :=
  {Link.cut pi ni A} ∪ sA ∪ snA

def proof_net_ax (A) : proof_net [~A,A] :=
⟨
  ⟨{Link.ax 0 0 A},
  by rintro l ⟨h⟩; exact valid_link.ax,
  by rintro Ai Δ₁ Δ₂ ⟨_⟩ ⟨_⟩; finish,
  by rintro Ai Δ₁ Δ₂ ⟨_⟩ ⟨_⟩; finish ⟩
,
  begin
    rintro B Bmem,
    refine ⟨0,_,_⟩, rcases Bmem with ⟨_,_⟩, use Link.ax 0 0 A, simp, exact mem_Link.con conclusion.ax_neg,

      -- exact ⟨0,Link.ax 0 0 A,by simp,conclusion.ax_neg,_⟩, rintro Δ' ⟨_⟩ ⟨_⟩,
    rcases H with ⟨⟨_⟩⟩,
      exact ⟨0,Link.ax 0 0 A,by simp,conclusion.ax_pos,_⟩, rintro Δ' ⟨_⟩ ⟨_⟩,
    cases H,
  end
⟩

def proof_net_tensor {Γ Γ' A B Δ Δ'} (pnA : proof_net (Γ ++ [A] ++ Γ')) (pnB : proof_net (Δ ++ [B] ++ Δ')) : pnA.disjoint pnB → proof_net (Γ ++ Γ' ++ [A ⊗ B] ++ Δ ++ Δ') :=
begin
  rcases pnA with ⟨_,psA, hA⟩,
  rcases pnB with ⟨_,psB, hB⟩,
  intro dAB,
  specialize hA A (by refine list.mem_append_left _ (list.mem_append_right _ (list.mem_cons_self A list.nil))),
  specialize hB B,
  cases hA with ai ΔA hA,

end


theorem ex (v v' : Nat) (ty : Nat → String) (h : some v = some v') : some (ty v) = some (ty v') := by
  simp_all


import algebra.homology.short_complex.short_exact
import algebra.add_torsor

noncomputable theory

namespace category_theory

open limits category

universes v u

variables (C : Type u) [category.{v} C]

@[derive category]
def short_exact_sequence [has_zero_morphisms C] :=
full_subcategory (λ (S : short_complex C), S.short_exact)

namespace short_exact_sequence

variable {C}

section

variables [has_zero_morphisms C] (S : short_exact_sequence C)

abbreviation short_complex : short_complex C := S.1

lemma short_exact : S.short_complex.short_exact := S.2

lemma exact : S.short_complex.exact := S.short_exact.exact

instance : mono S.short_complex.f := S.short_exact.mono_f
instance : epi S.short_complex.g := S.short_exact.epi_g

end

instance five_lemma [preadditive C] [balanced C]
  {S₁ S₂ : short_exact_sequence C} (φ : S₁ ⟶ S₂)
  [is_iso φ.τ₁] [is_iso φ.τ₃] : is_iso φ.τ₂ :=
begin
  rw is_iso_iff_mono_and_epi,
  refine ⟨_, _⟩,
  { rw preadditive.mono_iff_cancel_zero,
    intros A f hf,
    let f' := S₁.short_exact.lift f
      (by simp only [assoc, ← cancel_mono φ.τ₃, ← φ.comm₂₃, reassoc_of hf, zero_comp]),
    have hf' : f' ≫ _ = _ := S₁.short_exact.lift_f _ _,
    have hf'' : f' = 0,
    { simp only [← cancel_mono φ.τ₁, ← cancel_mono S₂.short_complex.f, assoc, φ.comm₁₂,
        reassoc_of hf', hf, zero_comp], },
    rw [← hf', hf'', zero_comp], },
  { rw preadditive.epi_iff_cancel_zero,
    intros A f hf,
    let f' := S₂.short_exact.desc f
      (by simp only [← cancel_epi φ.τ₁, φ.comm₁₂_assoc, hf, comp_zero]),
    have hf' : _ ≫ f' = _ := S₂.short_exact.g_desc _ _,
    have hf'' : f' = 0,
    { simp only [← cancel_epi φ.τ₃, ← cancel_epi S₁.short_complex.g, ← φ.comm₂₃_assoc,
        hf', hf, comp_zero], },
    rw [← hf', hf'', comp_zero], },
end

end short_exact_sequence

namespace abelian

variables {C} (A B : C)

structure extension [has_zero_morphisms C] :=
(X : C)
(i : B ⟶ X)
(p : X ⟶ A)
(w : i ≫ p = 0)
(ex : (short_complex.mk _ _ w).short_exact)

namespace extension

section

variables {A B} [has_zero_morphisms C]

instance (E : extension A B) : mono E.i := E.ex.mono_f
instance (E : extension A B) : epi E.p := E.ex.epi_g

@[ext]
structure hom (E₁ E₂ : extension A B) :=
(τ : E₁.X ⟶ E₂.X)
(commi' : E₁.i ≫ τ = E₂.i . obviously)
(commp' : τ ≫ E₂.p = E₁.p . obviously)

restate_axiom hom.commi'
restate_axiom hom.commp'
attribute [simp, reassoc] w hom.commi hom.commp

@[simps]
def hom.id (E : extension A B) : hom E E :=
{ τ := 𝟙 _, }

@[simps]
def hom.comp {E₁ E₂ E₃ : extension A B} (φ : hom E₁ E₂) (φ' : hom E₂ E₃) : hom E₁ E₃ :=
{ τ := φ.τ ≫ φ'.τ, }

instance : category (extension A B) :=
{ hom := hom,
  id := hom.id,
  comp := λ E₁ E₂ E₃, hom.comp, }

@[simps]
def hom.mk' {E₁ E₂ : extension A B} (τ : E₁.X ⟶ E₂.X)
  (commi : E₁.i ≫ τ = E₂.i) (commp : τ ≫ E₂.p = E₁.p) : E₁ ⟶ E₂ :=
{ τ := τ,
  commi' := commi,
  commp' := commp, }

@[simp]
lemma comp_τ {E₁ E₂ E₃ : extension A B} (φ : E₁ ⟶ E₂) (φ' : E₂ ⟶ E₃) :
  (φ ≫ φ').τ = φ.τ ≫ φ'.τ := rfl

@[simp]
lemma id_τ (E : extension A B) :
  hom.τ (𝟙 E) = 𝟙 E.X := rfl

variables (A B)

@[simps]
def to_short_exact_sequence_functor : extension A B ⥤ short_exact_sequence C :=
{ obj := λ E, ⟨short_complex.mk E.i E.p E.w, E.ex⟩,
  map := λ E₁ E₂ φ,
  { τ₁ := 𝟙 _,
    τ₂ := φ.τ,
    τ₃ := 𝟙 _, },
  map_comp' := λ E₁ E₂ E₃ φ φ', begin
    ext,
    { dsimp, erw short_complex.comp_τ₁, dsimp, simp only [comp_id], },
    { refl, },
    { dsimp, erw short_complex.comp_τ₃, dsimp, simp only [comp_id], },
  end, }

instance : faithful (to_short_exact_sequence_functor A B) :=
⟨λ E₁ E₂ f₁ f₂ eq, begin
  ext,
  simpa only using congr_arg short_complex.hom.τ₂ eq,
end⟩

instance (E₁ E₂ : extension A B) (f : E₁ ⟶ E₂) :
  is_iso ((to_short_exact_sequence_functor A B).map f).τ₁ :=
by { dsimp, apply_instance, }

instance (E₁ E₂ : extension A B) (f : E₁ ⟶ E₂) :
  is_iso ((to_short_exact_sequence_functor A B).map f).τ₃ :=
by { dsimp, apply_instance, }

end

section preadditive

variables [has_zero_object C] [preadditive C]
  (A B) [has_binary_biproduct B A]

@[simps]
def trivial : extension A B :=
{ X := biprod B A,
  i := biprod.inl,
  p := biprod.snd,
  w := biprod.inl_snd,
  ex := short_complex.splitting.short_exact
    { r := biprod.fst,
      s := biprod.inr,
      f_r := by tidy,
      s_g := by tidy,
      id := by tidy, }, }

end preadditive

variable [abelian C]

variables {A B} {E₁ E₂ : extension A B}

instance (f : E₁ ⟶ E₂) : is_iso f.τ :=
(infer_instance : is_iso ((to_short_exact_sequence_functor A B).map f).τ₂)

instance (f : E₁ ⟶ E₂) : is_iso f :=
⟨begin
  refine ⟨⟨inv f.τ, _, _⟩, _, _⟩,
  tidy,
end⟩

@[simp, reassoc]
lemma iso_hom_inv_τ (e : E₁ ≅ E₂) : e.hom.τ ≫ e.inv.τ = 𝟙 _ :=
by rw [← comp_τ, e.hom_inv_id, id_τ]

@[simp, reassoc]
lemma iso_inv_hom_τ (e : E₁ ≅ E₂) : e.inv.τ ≫ e.hom.τ = 𝟙 _ :=
by rw [← comp_τ, e.inv_hom_id, id_τ]

@[simps]
instance has_vadd : has_vadd (A ⟶ B) (E₁ ⟶ E₂) :=
{ vadd := λ g f,
  { τ := E₁.p ≫ g ≫ E₂.i + f.τ, }, }

instance : add_action (A ⟶ B) (E₁ ⟶ E₂) :=
{ zero_vadd := by tidy,
  add_vadd := λ g₁ g₂ f, begin
    ext,
    simp only [has_vadd_vadd_τ, preadditive.add_comp, preadditive.comp_add, add_assoc],
  end, }

def hom.vsub (f₁ f₂ : E₁ ⟶ E₂) : A ⟶ B :=
begin
  let g₀ := E₂.ex.lift (f₁.τ - f₂.τ) (by simp),
  have hg₀ : g₀ ≫ E₂.i = _ := E₂.ex.lift_f _ _,
  exact E₁.ex.desc g₀ begin
    dsimp,
    simp only [← cancel_mono E₂.i, assoc, hg₀, preadditive.comp_sub,
      hom.commi, sub_self, zero_comp],
  end,
end

lemma hom.p_vsub_i (f₁ f₂ : E₁ ⟶ E₂) : E₁.p ≫ hom.vsub f₁ f₂ ≫ E₂.i = f₁.τ - f₂.τ :=
begin
  dsimp [hom.vsub],
  rw [← assoc, E₁.ex.g_desc, E₂.ex.lift_f],
end

instance has_vsub : has_vsub (A ⟶ B) (E₁ ⟶ E₂) :=
{ vsub := hom.vsub, }

@[simp, reassoc]
lemma p_has_vsub_vsub_i (f₁ f₂ : E₁ ⟶ E₂) :
  E₁.p ≫ (f₁ -ᵥ f₂) ≫ E₂.i = f₁.τ - f₂.τ :=
hom.p_vsub_i f₁ f₂

@[simp]
lemma vsub_vadd (f₁ f₂ : E₁ ⟶ E₂) :
  (f₁ -ᵥ f₂ : A ⟶ B) +ᵥ f₂ = f₁ :=
begin
  ext,
  simp only [has_vadd_vadd_τ, p_has_vsub_vsub_i, sub_add_cancel],
end

@[simp]
lemma vadd_vsub (g : A ⟶ B) (f : E₁ ⟶ E₂) :
  g +ᵥ f -ᵥ f = g :=
by rw [← cancel_mono E₂.i, ← cancel_epi E₁.p, p_has_vsub_vsub_i, has_vadd_vadd_τ, add_sub_cancel]

@[simps]
def iso_trivial_equiv (e : extension A B) :
  (e ≅ trivial A B) ≃ (short_complex.mk _ _ e.w).splitting :=
{ to_fun := λ φ,
  { r := φ.hom.τ ≫ biprod.fst,
    s := biprod.inr ≫ φ.inv.τ,
    f_r := by simp only [hom.commi_assoc, trivial_i, biprod.inl_fst],
    s_g := by simp only [assoc, hom.commp, trivial_p, biprod.inr_snd],
    id := begin
      dsimp,
      rw [← cancel_epi φ.inv.τ, ← cancel_mono φ.hom.τ],
      simp only [assoc, preadditive.comp_add, iso_inv_hom_τ_assoc, hom.commp_assoc,
        trivial_p, preadditive.add_comp, hom.commi, trivial_i, iso_inv_hom_τ, comp_id],
      erw [comp_id],
      dsimp,
      rw biprod.total,
    end },
  inv_fun := λ s, as_iso
  { τ := biprod.lift s.r e.p,
    commi' := begin
      ext,
      { simp only [assoc, biprod.lift_fst, trivial_i, biprod.inl_fst, s.f_r], },
      { simp only [w, assoc, biprod.lift_snd, trivial_i, biprod.inl_snd], },
    end, },
  left_inv := λ φ, begin
    ext,
    { tidy, },
    { dsimp,
      simpa only [biprod.lift_snd] using φ.hom.commp'.symm, },
  end,
  right_inv := λ s, short_complex.splitting.ext_r _ _ (by simp), }

@[simps]
def pull {A' : C} (E : extension A B) (π : A' ⟶ A) : extension A' B :=
{ X := pullback E.p π,
  i := pullback.lift E.i 0 (by simp),
  p := pullback.snd,
  w := pullback.lift_snd _ _ _,
  ex := short_complex.short_exact.of_f_is_kernel begin
    refine limits.kernel_fork.is_limit.of_ι _ _
      (λ Z x hx, E.ex.lift (x ≫ pullback.fst)
        (by { dsimp at hx ⊢, rw [assoc, pullback.condition, reassoc_of hx, zero_comp], })) _ _,
    { intros Z x hx,
      ext,
      { simp only [assoc, pullback.lift_fst, short_complex.short_exact.lift_f], },
      { simp only [assoc, pullback.lift_snd, comp_zero, hx], }, },
    { intros Z x hx m hm,
      simpa only [← cancel_mono E.i, assoc, short_complex.short_exact.lift_f,
        pullback.lift_fst] using hm =≫ pullback.fst, },
  end, }

@[simps]
def pull_short_complex {A' : C} (E : extension A B) (π : A' ⟶ A) :
  short_complex.mk _ _ (E.pull π).w ⟶ short_complex.mk _ _ E.w :=
{ τ₁ := 𝟙 _,
  τ₂ := pullback.fst,
  τ₃ := π,
  comm₂₃' := pullback.condition, }

@[simps]
def pull_functor {A A' : C} (π : A' ⟶ A) (B : C) : extension A B ⥤ extension A' B :=
{ obj := λ E, E.pull π,
  map := λ E₁ E₂ f,
  { τ := pullback.map _ _ _ _ f.τ (𝟙 A') (𝟙 A) (by simp) (by simp), }, }

def pull_functor_id (A B : C) : pull_functor (𝟙 A) B ≅ 𝟭 _ :=
nat_iso.of_components
  (λ E, as_iso
    { τ := pullback.fst,
      commp' := by { dsimp, rw [pullback.condition, comp_id], }, })
  (by tidy)

def pull_functor_comp {A A' A'' : C} (π : A' ⟶ A) (π' : A'' ⟶ A') (B : C) :
  pull_functor π B ⋙ pull_functor π' B ≅ pull_functor (π' ≫ π) B :=
nat_iso.of_components
  (λ E, as_iso
    { τ := pullback.lift (pullback.fst ≫ pullback.fst) pullback.snd
        (by erw [assoc, pullback.condition, pullback.condition_assoc]), })
  (by tidy)

@[simps]
def push {B' : C} (E : extension A B) (ι : B ⟶ B') : extension A B' :=
{ X := pushout E.i ι,
  i := pushout.inr,
  p := pushout.desc E.p 0 (by simp),
  w := pushout.inr_desc _ _ _,
  ex := short_complex.short_exact.of_g_is_cokernel begin
    refine limits.cokernel_cofork.is_colimit.of_π _ _
      (λ Z x hx, E.ex.desc (pushout.inl ≫ x)
      (by { dsimp at hx ⊢, rw [pushout.condition_assoc, hx, comp_zero], })) _ _,
    { intros A x hx,
      ext,
      { simp only [pushout.inl_desc_assoc, E.ex.g_desc (pushout.inl ≫ x)], },
      { simp only [pushout.inr_desc_assoc, zero_comp, hx], }, },
    { intros Z x hx m hm,
      rw [← cancel_epi E.p, E.ex.g_desc (pushout.inl ≫ x), ← hm, pushout.inl_desc_assoc], },
  end, }

@[simps]
def push_short_complex {B' : C} (E : extension A B) (ι : B ⟶ B') :
  short_complex.mk _ _ E.w ⟶ short_complex.mk _ _ (E.push ι).w :=
{ τ₁ := ι,
  τ₂ := pushout.inl,
  τ₃ := 𝟙 _,
  comm₁₂' := pushout.condition.symm, }

@[simps]
def push_functor (A : C) {B B' : C} (ι : B ⟶ B') : extension A B ⥤ extension A B' :=
{ obj := λ E, E.push ι,
  map := λ E₁ E₂ f,
  { τ := pushout.map _ _ _ _ f.τ (𝟙 B') (𝟙 B) (by simp) (by simp), }, }

def push_functor_id (A B : C) : push_functor A (𝟙 B) ≅ 𝟭 _ :=
iso.symm (nat_iso.of_components
  (λ E, as_iso
    { τ := pushout.inl,
      commi' := by { dsimp, rw [pushout.condition, id_comp], }, })
  (by tidy))

def push_functor_comp (A : C) {B B' B'' : C} (ι : B ⟶ B') (ι' : B' ⟶ B'') :
  push_functor A ι ⋙ push_functor A ι' ≅ push_functor A (ι ≫ ι') :=
iso.symm (nat_iso.of_components
  (λ E, as_iso
    { τ := pushout.desc (pushout.inl ≫ pushout.inl) pushout.inr
        (by rw [pushout.condition_assoc, pushout.condition, assoc]), })
  (by tidy))

def pull_functor_comm_push_functor {A A' B B' : C} (π : A' ⟶ A) (ι : B ⟶ B') :
  pull_functor π B ⋙ push_functor A' ι ≅
    push_functor A ι ⋙ pull_functor π B' :=
nat_iso.of_components
  (λ E, as_iso
    { τ := pushout.desc
        (pullback.map _ _ _ _ pushout.inl (𝟙 A') (𝟙 A) (by tidy) (by simp))
        (pullback.lift pushout.inr 0
          (by { dsimp, simp only [pushout.inr_desc, zero_comp], }))
        begin
          ext,
          { dsimp, simp [pushout.condition], },
          { dsimp, simp, },
        end, })
  (by tidy)

end extension

variable [abelian C]

def extensions := quotient (is_isomorphic_setoid (extension A B))

def extensions_map_src {A A' : C} (π : A' ⟶ A) (B : C) : extensions A B → extensions A' B :=
quot.map (extension.pull_functor π B).obj begin
  rintro E₁ E₂ ⟨e⟩,
  exact ⟨(extension.pull_functor π B).map_iso e⟩,
end

def extensions_map_tgt (A : C) {B B' : C} (ι : B ⟶ B') : extensions A B → extensions A B' :=
quot.map (extension.push_functor A ι).obj begin
  rintro E₁ E₂ ⟨e⟩,
  exact ⟨(extension.push_functor A ι).map_iso e⟩,
end

lemma extensions_map_src_id (A B : C) :
  extensions_map_src (𝟙 A) B = id :=
begin
  ext E,
  obtain ⟨E, rfl⟩ := quotient.surjective_quotient_mk' E,
  exact quot.sound ⟨(extension.pull_functor_id A B).app E⟩,
end

lemma extensions_map_src_comp {A A' A'' : C} (π' : A'' ⟶ A') (π : A' ⟶ A) (B : C) :
  extensions_map_src π' B ∘ extensions_map_src π B = extensions_map_src (π' ≫ π) B :=
begin
  ext E,
  obtain ⟨E, rfl⟩ := quotient.surjective_quotient_mk' E,
  exact quot.sound ⟨(extension.pull_functor_comp π π' B).app E⟩,
end

lemma extensions_map_tgt_id (A B : C) :
  extensions_map_tgt A (𝟙 B) = id :=
begin
  ext E,
  obtain ⟨E, rfl⟩ := quotient.surjective_quotient_mk' E,
  exact quot.sound ⟨(extension.push_functor_id A B).app E⟩,
end

lemma extensions_map_tgt_comp (A : C) {B B' B'' : C} (ι : B ⟶ B') (ι' : B' ⟶ B'') :
  extensions_map_tgt A ι' ∘ extensions_map_tgt A ι = extensions_map_tgt A (ι ≫ ι') :=
begin
  ext E,
  obtain ⟨E, rfl⟩ := quotient.surjective_quotient_mk' E,
  exact quot.sound ⟨(extension.push_functor_comp A ι ι').app E⟩,
end

lemma extensions_map_tgt_comp_map_src {A A' B B' : C} (π : A' ⟶ A) (ι : B ⟶ B') :
  extensions_map_tgt A' ι ∘ extensions_map_src π B =
    extensions_map_src π B' ∘ extensions_map_tgt A ι :=
begin
  ext E,
  obtain ⟨E, rfl⟩ := quotient.surjective_quotient_mk' E,
  exact quot.sound ⟨(extension.pull_functor_comm_push_functor π ι).app E⟩,
end

variable (C)

@[simps]
def extensions_functor : C ⥤ Cᵒᵖ ⥤ Type (max u v) :=
{ obj := λ B,
  { obj := λ A, extensions A.unop B,
    map := λ A A' π, extensions_map_src π.unop B,
    map_id' := λ A, extensions_map_src_id A.unop B,
    map_comp' := λ A A' A'' π π', (extensions_map_src_comp π'.unop π.unop B).symm, },
  map := λ B B' ι,
  { app := λ A, extensions_map_tgt A.unop ι,
    naturality' := λ A A' π, extensions_map_tgt_comp_map_src π.unop ι, },
  map_id' := λ B, begin
    ext A : 2,
    exact extensions_map_tgt_id A.unop B,
  end,
  map_comp' := λ B B' B'' ι ι', begin
    ext A : 2,
    exact (extensions_map_tgt_comp A.unop ι ι').symm,
  end }

end abelian

end category_theory


import MLIR.Dialects.ToyModel
import MLIR.Dialects.BuiltinModel
import MLIR.Semantics.Fitree
import MLIR.Semantics.Verifier
import MLIR.Semantics.SSAEnv
import MLIR.Semantics.UB
import MLIR.Semantics.TensorElem
import MLIR.Util.Metagen
import MLIR.Util.Reduce

import MLIR.AST
import MLIR.EDSL

import Lean

open MLIR.AST

/- To be automatically generated -/

inductive ToyOp: Type → Type :=
  | Constant:
      (cst_D: DimList) → (cst_τ: MLIRTy) → (cst: TensorLiteral cst_D cst_τ) →
      ToyOp (RankedTensor cst_D cst_τ)
  | Transpose:
      (τ: MLIRTy) → (n m: Nat) →
      RankedTensor [Dimension.Known n, Dimension.Known m] τ →
      ToyOp (RankedTensor [Dimension.Known m, Dimension.Known n] τ)
  | Reshape:
      (τ: MLIRTy) → (D D': DimList) → (H: D.known) → (H': D'.known) →
      (Hprod: D'.prod = D.prod) →
      RankedTensor D τ →
      ToyOp (RankedTensor D' τ)

/- To be automatically generated (hopefully; basically this is the
   verification stuff) -/

def toy_semantics_op (op: Op builtin):
    Fitree (UBE +' SSAEnvE builtin +' ToyOp) Unit :=
  match op with
  | Op.mk "toy.constant" [(res, builtin.tensor D₁ τ₁)] [] [] attrs =>
      match AttrDict.find attrs "value" with
      | some (builtin.dense_tensor_attr elem D₂ τ₂) =>
          match TensorLiteral.ofTensorElem elem D₁ τ₁ with
          | none =>
              raiseUB s!"{op}"
          | some t_lit => do
              let t ← Fitree.trigger <| ToyOp.Constant D₁ τ₁ t_lit
              SSAEnv.set? (builtin.tensor D₁ τ₁) res t
      | _ =>
          raiseUB s!"{op}"

  | Op.mk "toy.transpose" [(res, τ₂)] [(t_name, builtin.tensor D τ)] [] _ =>
      match D with
      | [Dimension.Known n, Dimension.Known m] => do
          let t ← Fitree.trigger (SSAEnvE.Get (builtin.tensor
                  [Dimension.Known n, Dimension.Known m] τ) t_name);
          let t' ← Fitree.trigger (ToyOp.Transpose τ n m t);
          SSAEnv.set? (builtin.tensor [Dimension.Known m, Dimension.Known n] τ)
            (some res) t'
      | _ =>
          raiseUB s!"{op}"

  | Op.mk "toy.reshape" [(res, builtin.tensor D' τ₂)]
        [(t_name, builtin.tensor D τ₁)] [] _ =>
      if H: τ₁ = τ₂
        ∧ DimList.known D
        ∧ DimList.known D'
        ∧ DimList.prod D' = DimList.prod D then do
        let t ← Fitree.trigger (SSAEnvE.Get (builtin.tensor D τ₁) t_name);
        let t' ← Fitree.trigger (ToyOp.Reshape τ₁ D D'
                H.2.1 H.2.2.1 H.2.2.2 t);
        let t': RankedTensor D' τ₂ := cast (by rw [H.1]) t';
        SSAEnv.set? (builtin.tensor D' τ₂) (some res) t'
      else
        raiseUB s!"{op}"

  | _ => raiseUB s!"{op}"

-- TODO: toy_semantics_bb: handle basic block arguments
@[simp]
def toy_semantics_region: Region builtin →
      Fitree (UBE +' (SSAEnvE builtin) +' ToyOp) Unit
  | Region.mk name args [] =>
      Fitree.ret ()
  | Region.mk name args (op1::ops) =>
      List.foldr (fun t acc => Fitree.bind acc (fun _ => t))
                 (toy_semantics_op op1)
                 (ops.map toy_semantics_op)

/- Manually specified: ToyOp event handler -/

def ToyOp.handle {E}: ToyOp ~> Fitree E :=
  fun _ e => match e with
  | ToyOp.Constant D τ t_lit =>
      return RankedTensor.ofTensorLiteral t_lit
  | ToyOp.Transpose α n m t =>
      return transpose t
  | ToyOp.Reshape α D D' H H' Hprod t =>
      return reshape D' H H' Hprod t

-- Interpretation in context

def interp_toy {E} (t: Fitree (ToyOp +' E) R): Fitree E R :=
  t.interp (Fitree.case ToyOp.handle (fun T => @Fitree.trigger E E T _))

@[simp]
def run_toy (t: Fitree (UBE +' SSAEnvE builtin +' ToyOp) Unit)
    (env: SSAEnv builtin): Fitree Void1 (Unit × SSAEnv builtin) :=
  Fitree.interp ToyOp.handle (interpSSA' (interpUB'! t) env)

/-
### Examples and testing
-/

-- TODO: Can we infer the builtin in there?
def transpose_stmt: Op builtin := [mlir_op|
  %t2 = "toy.transpose"(%t1): (tensor<2×4×i32>) -> (tensor<4×2×i32>)
]

def constant_stmt: Op builtin := [mlir_op|
  %t = "toy.constant"() {value=dense<[[1,2],[3,4]]>: tensor<2×2×i32>}:
    () -> (tensor<2×2×i32>)
]

def double_transpose: Region builtin := [mlir_region| {
  ^dbl:
    %t2 = "toy.transpose"(%t1): (tensor<2×4×i32>) -> (tensor<4×2×i32>)
    %t3 = "toy.transpose"(%t2): (tensor<4×2×i32>) -> (tensor<2×4×i32>)
}]

#eval Fitree.run <| run_toy (toy_semantics_op transpose_stmt) SSAEnv.empty

#eval Fitree.run <| run_toy (toy_semantics_op constant_stmt) SSAEnv.empty

theorem double_transpose_correct:
  ∀ (t1: RankedTensor [.Known 2, .Known 4] .i32),
    run_toy (toy_semantics_region double_transpose)
      (SSAEnv.One [("t1", ⟨builtin.tensor [.Known 2, .Known 4] .i32, t1⟩)])
    =
    Fitree.ret ((), SSAEnv.One [
      (SSAVal.SSAVal "t1", ⟨builtin.tensor [.Known 2, .Known 4] .i32, t1⟩),
      (SSAVal.SSAVal "t2", ⟨builtin.tensor [.Known 4, .Known 2] .i32,
                           transpose t1⟩),
      (SSAVal.SSAVal "t3", ⟨builtin.tensor [.Known 2, .Known 4] .i32, t1⟩)
    ]) := by
  intros t1
  simp [double_transpose, toy_semantics_region, toy_semantics_op]; simp_itree
  simp [interpUB'!]; simp_itree
  simp [interpSSA', Fitree.interpState, SSAEnvE.handle]; simp_itree
  simp [SSAEnv.get, SSAEnv.getT, SSAEnv.set]; simp_itree
  simp [SSAEnv.get, SSAEnv.getT, SSAEnv.set]; simp_itree
  rw [transpose_involutive]


/-
Copyright (c) 2020 Scott Morrison. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Scott Morrison, Bhavik Mehta
-/
import category_theory.limits.preserves.basic
import category_theory.limits.shapes.equalizers
import category_theory.limits.shapes.strong_epi
import category_theory.limits.shapes.pullbacks

/-!
# Definitions and basic properties of regular monomorphisms and epimorphisms.

A regular monomorphism is a morphism that is the equalizer of some parallel pair.

We give the constructions
* `split_mono → regular_mono` and
* `regular_mono → mono`
as well as the dual constructions for regular epimorphisms. Additionally, we give the
construction
* `regular_epi ⟶ strong_epi`.

-/

noncomputable theory

namespace category_theory
open category_theory.limits

universes v₁ u₁ u₂

variables {C : Type u₁} [category.{v₁} C]

variables {X Y : C}

/-- A regular monomorphism is a morphism which is the equalizer of some parallel pair. -/
class regular_mono (f : X ⟶ Y) :=
(Z : C)
(left right : Y ⟶ Z)
(w : f ≫ left = f ≫ right)
(is_limit : is_limit (fork.of_ι f w))

attribute [reassoc] regular_mono.w

/-- Every regular monomorphism is a monomorphism. -/
@[priority 100]
instance regular_mono.mono (f : X ⟶ Y) [regular_mono f] : mono f :=
mono_of_is_limit_parallel_pair regular_mono.is_limit

instance equalizer_regular (g h : X ⟶ Y) [has_limit (parallel_pair g h)] :
  regular_mono (equalizer.ι g h) :=
{ Z := Y,
  left := g,
  right := h,
  w := equalizer.condition g h,
  is_limit := fork.is_limit.mk _ (λ s, limit.lift _ s) (by simp) (λ s m w, by { ext1, simp [←w] }) }

/-- Every split monomorphism is a regular monomorphism. -/
@[priority 100]
instance regular_mono.of_split_mono (f : X ⟶ Y) [split_mono f] : regular_mono f :=
{ Z     := Y,
  left  := 𝟙 Y,
  right := retraction f ≫ f,
  w     := by tidy,
  is_limit := split_mono_equalizes f }

/-- If `f` is a regular mono, then any map `k : W ⟶ Y` equalizing `regular_mono.left` and
    `regular_mono.right` induces a morphism `l : W ⟶ X` such that `l ≫ f = k`. -/
def regular_mono.lift' {W : C} (f : X ⟶ Y) [regular_mono f] (k : W ⟶ Y)
  (h : k ≫ (regular_mono.left : Y ⟶ @regular_mono.Z _ _ _ _ f _) = k ≫ regular_mono.right) :
  {l : W ⟶ X // l ≫ f = k} :=
fork.is_limit.lift' regular_mono.is_limit _ h

/--
The second leg of a pullback cone is a regular monomorphism if the right component is too.

See also `pullback.snd_of_mono` for the basic monomorphism version, and
`regular_of_is_pullback_fst_of_regular` for the flipped version.
-/
def regular_of_is_pullback_snd_of_regular {P Q R S : C} {f : P ⟶ Q} {g : P ⟶ R} {h : Q ⟶ S}
  {k : R ⟶ S} [hr : regular_mono h] (comm : f ≫ h = g ≫ k)
  (t : is_limit (pullback_cone.mk _ _ comm)) :
regular_mono g :=
{ Z := hr.Z,
  left := k ≫ hr.left,
  right := k ≫ hr.right,
  w := by rw [← reassoc_of comm, ← reassoc_of comm, hr.w],
  is_limit :=
  begin
    apply fork.is_limit.mk' _ _,
    intro s,
    have l₁ : (fork.ι s ≫ k) ≫ regular_mono.left = (fork.ι s ≫ k) ≫ regular_mono.right,
      rw [category.assoc, s.condition, category.assoc],
    obtain ⟨l, hl⟩ := fork.is_limit.lift' hr.is_limit _ l₁,
    obtain ⟨p, hp₁, hp₂⟩ := pullback_cone.is_limit.lift' t _ _ hl,
    refine ⟨p, hp₂, _⟩,
    intros m w,
    have z : m ≫ g = p ≫ g := w.trans hp₂.symm,
    apply t.hom_ext,
    apply (pullback_cone.mk f g comm).equalizer_ext,
    { erw [← cancel_mono h, category.assoc, category.assoc, comm, reassoc_of z] },
    { exact z },
  end }

/--
The first leg of a pullback cone is a regular monomorphism if the left component is too.

See also `pullback.fst_of_mono` for the basic monomorphism version, and
`regular_of_is_pullback_snd_of_regular` for the flipped version.
-/
def regular_of_is_pullback_fst_of_regular {P Q R S : C} {f : P ⟶ Q} {g : P ⟶ R} {h : Q ⟶ S}
  {k : R ⟶ S} [hr : regular_mono k] (comm : f ≫ h = g ≫ k)
  (t : is_limit (pullback_cone.mk _ _ comm)) :
regular_mono f :=
regular_of_is_pullback_snd_of_regular comm.symm (pullback_cone.flip_is_limit t)

/-- A regular monomorphism is an isomorphism if it is an epimorphism. -/
lemma is_iso_of_regular_mono_of_epi (f : X ⟶ Y) [regular_mono f] [e : epi f] : is_iso f :=
@is_iso_limit_cone_parallel_pair_of_epi _ _ _ _ _ _ _ regular_mono.is_limit e

/-- A regular epimorphism is a morphism which is the coequalizer of some parallel pair. -/
class regular_epi (f : X ⟶ Y) :=
(W : C)
(left right : W ⟶ X)
(w : left ≫ f = right ≫ f)
(is_colimit : is_colimit (cofork.of_π f w))

attribute [reassoc] regular_epi.w

/-- Every regular epimorphism is an epimorphism. -/
@[priority 100]
instance regular_epi.epi (f : X ⟶ Y) [regular_epi f] : epi f :=
epi_of_is_colimit_parallel_pair regular_epi.is_colimit

instance coequalizer_regular (g h : X ⟶ Y) [has_colimit (parallel_pair g h)] :
  regular_epi (coequalizer.π g h) :=
{ W := X,
  left := g,
  right := h,
  w := coequalizer.condition g h,
  is_colimit := cofork.is_colimit.mk _ (λ s, colimit.desc _ s) (by simp)
    (λ s m w, by { ext1, simp [←w] }) }

/-- Every split epimorphism is a regular epimorphism. -/
@[priority 100]
instance regular_epi.of_split_epi (f : X ⟶ Y) [split_epi f] : regular_epi f :=
{ W     := X,
  left  := 𝟙 X,
  right := f ≫ section_ f,
  w     := by tidy,
  is_colimit := split_epi_coequalizes f }

/-- If `f` is a regular epi, then every morphism `k : X ⟶ W` coequalizing `regular_epi.left` and
    `regular_epi.right` induces `l : Y ⟶ W` such that `f ≫ l = k`. -/
def regular_epi.desc' {W : C} (f : X ⟶ Y) [regular_epi f] (k : X ⟶ W)
  (h : (regular_epi.left : regular_epi.W f ⟶ X) ≫ k = regular_epi.right ≫ k) :
  {l : Y ⟶ W // f ≫ l = k} :=
cofork.is_colimit.desc' (regular_epi.is_colimit) _ h

/--
The second leg of a pushout cocone is a regular epimorphism if the right component is too.

See also `pushout.snd_of_epi` for the basic epimorphism version, and
`regular_of_is_pushout_fst_of_regular` for the flipped version.
-/
def regular_of_is_pushout_snd_of_regular
  {P Q R S : C} {f : P ⟶ Q} {g : P ⟶ R} {h : Q ⟶ S} {k : R ⟶ S}
  [gr : regular_epi g] (comm : f ≫ h = g ≫ k) (t : is_colimit (pushout_cocone.mk _ _ comm)) :
regular_epi h :=
{ W := gr.W,
  left := gr.left ≫ f,
  right := gr.right ≫ f,
  w := by rw [category.assoc, category.assoc, comm, reassoc_of gr.w],
  is_colimit :=
  begin
    apply cofork.is_colimit.mk' _ _,
    intro s,
    have l₁ : gr.left ≫ f ≫ s.π = gr.right ≫ f ≫ s.π,
      rw [← category.assoc, ← category.assoc, s.condition],
    obtain ⟨l, hl⟩ := cofork.is_colimit.desc' gr.is_colimit (f ≫ cofork.π s) l₁,
    obtain ⟨p, hp₁, hp₂⟩ := pushout_cocone.is_colimit.desc' t _ _ hl.symm,
    refine ⟨p, hp₁, _⟩,
    intros m w,
    have z := w.trans hp₁.symm,
    apply t.hom_ext,
    apply (pushout_cocone.mk _ _ comm).coequalizer_ext,
    { exact z },
    { erw [← cancel_epi g, ← reassoc_of comm, ← reassoc_of comm, z], refl },
  end }

/--
The first leg of a pushout cocone is a regular epimorphism if the left component is too.

See also `pushout.fst_of_epi` for the basic epimorphism version, and
`regular_of_is_pushout_snd_of_regular` for the flipped version.
-/
def regular_of_is_pushout_fst_of_regular
  {P Q R S : C} {f : P ⟶ Q} {g : P ⟶ R} {h : Q ⟶ S} {k : R ⟶ S}
  [fr : regular_epi f] (comm : f ≫ h = g ≫ k) (t : is_colimit (pushout_cocone.mk _ _ comm)) :
regular_epi k :=
regular_of_is_pushout_snd_of_regular comm.symm (pushout_cocone.flip_is_colimit t)

/-- A regular epimorphism is an isomorphism if it is a monomorphism. -/
lemma is_iso_of_regular_epi_of_mono (f : X ⟶ Y) [regular_epi f] [m : mono f] : is_iso f :=
@is_iso_limit_cocone_parallel_pair_of_epi _ _ _ _ _ _ _ regular_epi.is_colimit m

@[priority 100]
instance strong_epi_of_regular_epi (f : X ⟶ Y) [regular_epi f] : strong_epi f :=
{ epi := by apply_instance,
  has_lift :=
  begin
    introsI,
    have : (regular_epi.left : regular_epi.W f ⟶ X) ≫ u = regular_epi.right ≫ u,
    { apply (cancel_mono z).1,
      simp only [category.assoc, h, regular_epi.w_assoc] },
    obtain ⟨t, ht⟩ := regular_epi.desc' f u this,
    exact arrow.has_lift.mk ⟨t, ht, (cancel_epi f).1
      (by simp only [←category.assoc, ht, ←h, arrow.mk_hom, arrow.hom_mk'_right])⟩,
  end }

end category_theory


-- A minimal implementation of lenses.
universe variables u v

-- This tactic proves (∀v s, get (set v s) = v by case splitting on s
-- and using reflexivity.
meta def lens.get_set_tactic : tactic unit := do
  tactic.intro `v,
  e ← tactic.intro `s,
  tactic.cases e [],
  tactic.reflexivity

-- This tactic proves (∀v s, set (get s) = s by case splitting on s
-- and using reflexivity.
meta def lens.set_get_tactic : tactic unit := do
  e ← tactic.intro `s,
  tactic.cases e [],
  tactic.reflexivity

-- This tactic proves (∀u v s, set u (set v s) = set u s by case splitting on s
-- and using reflexivity.
meta def lens.set_set_tactic : tactic unit := do
  tactic.intro `u,
  tactic.intro `v,
  e ← tactic.intro `s,
  tactic.cases e [],
  tactic.reflexivity

-- A lens is a structure with a getter, setter, and axioms about their effect.
structure lens (S : Type u) (α : Type v) :=
  (get : S → α)
  (set : α → S → S)
  (get_set :  ∀ (v : α) (s : S), get (set v s) = v . lens.get_set_tactic)
  (set_get :  ∀ (s : S), set (get s) s = s         . lens.set_get_tactic)
  (set_set :  ∀ (u v : α) (s : S), set u (set v s) = set u s . lens.set_set_tactic)

@[simp]
lemma get_set_cancel {S : Type u} {α : Type v} (l : lens S α)
: ∀ (v : α) (s : S), l.get (l.set v s) = v := l.get_set

@[simp]
lemma set_get_cancel {S : Type u} {α : Type v} (l : lens S α)
: ∀ (s : S), l.set (l.get s) s = s := l.set_get

@[simp]
lemma set_set_cancel {S : Type u} {α : Type v} (l : lens S α)
: ∀ (u v : α) (s : S), l.set u (l.set v s) = l.set u s := l.set_set

namespace lens

variables { s t : Type u }
variables { a b : Type v }

-- This applies a function to the value of a lens on an object and updates the state.
def over (l : lens s a) (f : a → a) (x : s) : s := lens.set l (f (lens.get l x)) x

end lens

infixr ` .~ `:4 := lens.set

infixr ` %~ `:4 := lens.over

@[reducible]
def call {a b} : a → (a → b) → b := λx f, f x

infixl ` & `:2 := call

@[simp]
theorem call_elim {α : Type u} {β : Type v} (f : α → β) (x : α) : (x & f) = f x :=  by simp [call]


-- /-
-- Copyright (c) 2020 Wojciech Nawrocki. All rights reserved.
-- Released under Apache 2.0 license as described in the file LICENSE.
-- Authors: Wojciech Nawrocki
-- -/

-- import category_theory.limits.shapes.finite_products
-- import category_theory.limits.shapes.binary_products
-- import category_theory.limits.shapes.terminal
-- import tactic.rcases
-- import pfin

-- /-! # Stuff that should be in the catthy library. -/
-- namespace category_theory

-- universes w w₁
-- -- def discrete.equiv_of_iso {J : Type w} {K : Type w₁} (h : J ≃ K) : (discrete J ≌ discrete K) :=
-- -- equivalence.mk
-- --   (functor.of_function h.to_fun) -- C ⥤ D
-- --   (functor.of_function h.inv_fun) -- D ⥤ C
-- --   { hom := {
-- --       app := λ X, begin
-- --         apply eq_to_hom,
-- --         simp only [functor.id_obj, functor.of_function_obj, functor.comp_obj],
-- --         exact (h.left_inv X).symm,
-- --       end,
-- --       naturality' := λ X Y f, dec_trivial },
-- --     inv := {
-- --       app := λ X, begin
-- --         apply eq_to_hom,
-- --         simp only [functor.id_obj, functor.of_function_obj, functor.comp_obj],
-- --         exact h.left_inv X,
-- --       end,
-- --       naturality' := λ X Y f, dec_trivial },
-- --     hom_inv_id' := by ext1; exact dec_trivial,
-- --     inv_hom_id' := by ext1; exact dec_trivial }
-- --   { hom := {
-- --       app := λ X, begin
-- --         apply eq_to_hom,
-- --         simp only [functor.id_obj, functor.of_function_obj, functor.comp_obj],
-- --         exact h.right_inv X
-- --       end,
-- --       naturality' := λ X Y f, dec_trivial },
-- --     inv := {
-- --       app := λ X, begin
-- --         apply eq_to_hom,
-- --         simp only [functor.id_obj, functor.of_function_obj, functor.comp_obj],
-- --         exact (h.right_inv X).symm,
-- --       end,
-- --       naturality' := λ X Y f, dec_trivial },
-- --     hom_inv_id' := by ext1; exact dec_trivial,
-- --     inv_hom_id' := by ext1; exact dec_trivial }

-- namespace limits

-- universes v u
-- variables {C : Type u} [𝒞 : category.{v} C]
-- include 𝒞

-- lemma prod.lift_uniq {X Y Z : C} [has_limit (pair X Y)] (f : Z ⟶ X) (g : Z ⟶ Y) (m : Z ⟶ X ⨯ Y)
--   (hLeft : m ≫ prod.fst = f) (hRight : m ≫ prod.snd = g)
--   : m = prod.lift f g :=
-- begin
--   apply limit.hom_ext,
--   intro j,
--   cases hLeft, cases hRight, cases j,
--     simp only [limit.lift_π, binary_fan.mk_π_app_left],
--   simp only [limit.lift_π, binary_fan.mk_π_app_right],
-- end

-- end limits
-- end category_theory

-- /-!
-- # Constructing finite products from binary products and a terminal object

-- If a category has all binary products, and a terminal object, then it has all finite products.
-- -/

-- namespace category_theory.limits
-- open category_theory

-- universes v u
-- variables {C : Type u} [𝒞 : category.{v} C]
-- include 𝒞

-- -- We hide the "implementation details" inside a namespace
-- namespace has_finite_products_of_binary_products_and_terminal_object

-- @[reducible]
-- def match_statement_lol [has_binary_products.{v} C]
--   {n : ℕ} (F: discrete (pfin (nat.succ n)) ⥤ C) (limF' : has_limit (discrete.lift pfin.succ ⋙ F))
--   : Π (j: pfin (nat.succ n)), F.obj ⟨0, nat.succ_pos n⟩ ⨯ limF'.cone.X ⟶ F.obj j
-- | ⟨0, _⟩ := prod.fst
-- | w@⟨nat.succ j, _⟩ := prod.snd ≫
--   limF'.cone.π.app (w.pred (λ h, nat.succ_ne_zero j (pfin.veq_of_eq h)))

-- set_option eqn_compiler.zeta true
-- def has_limit_for_pfin_diagram [has_binary_products.{v} C] [has_terminal.{v} C]
-- : Π {n: ℕ} (F: (discrete (pfin n)) ⥤ C)
-- , has_limit F
-- | 0 F :=
--   -- In the base case, the category of cones over a diagram of shape ∅ is simply 𝒞, so
--   -- the limit cone is 𝒞's terminal object.
--   let absurdJ (x : pfin 0) : false := x.elim0 in
--   let myCone : cone F :=
--     { X := terminal C,
--       π := nat_trans.of_homs (λ j, (absurdJ j).elim) } in
--   { cone := myCone,
--     is_limit :=
--       { lift := λ s, terminal.from s.X
--       , fac' := λ s j, (absurdJ j).elim
--       , uniq' := λ s m h, dec_trivial } }

-- | (nat.succ n) F :=
--   -- In the inductive case, we construct a limit cone with apex (F 0) ⨯ (apex of smaller limit cone)
--   -- where the smaller cone is obtained from the below functor.
--   let F' : discrete (pfin n) ⥤ C := discrete.lift pfin.succ ⋙ F in
--   let limF' : has_limit F' := has_limit_for_pfin_diagram F' in
--   let myCone : cone F :=
--     { X := (F.obj ⟨0, nat.succ_pos n⟩) ⨯ limF'.cone.X
--     , π := nat_trans.of_homs (match_statement_lol F limF') } in -- TODO(WN): using an actual match statement here
--                                                                 -- is hard to unfold later, but would obv be nicer.
--   { cone := myCone,
--     is_limit :=
--       { lift := λ s,
--           -- Show that s.X is also the apex of a cone over F' ..
--           let s' : cone F' :=
--             { X := s.X
--             , π := nat_trans.of_homs (λ j, s.π.app j.succ) } in
--           -- .. in order to get from s.X to limF'.cone.X in the right morphism
--           -- using the fact that limF' is a limit cone over F'.
--           prod.lift
--             (s.π.app $ ⟨0, nat.succ_pos n⟩)
--             (eq_to_hom rfl ≫ limF'.is_limit.lift s')
--       -- Show that lift is in fact a morphism of cones from s into myCone.
--       , fac' := λ s j, begin
--         rcases j with ⟨j, hj⟩, cases j;
--         simp only [category.id_comp, nat_trans.of_homs_app, eq_to_hom_refl, match_statement_lol,
--           prod.lift_fst, limit.lift_π_assoc, is_limit.fac, nat_trans.of_homs_app,
--           binary_fan.mk_π_app_right], congr
--       end
--       -- Show that lift is the unique morphism into myCone.
--       , uniq' := λ s m h, begin
--         have h0 := h ⟨0, nat.succ_pos n⟩,
--         simp [match_statement_lol] at h0,
--         let s' : cone F' :=
--           { X := s.X
--           , π := nat_trans.of_homs (λ j, s.π.app j.succ) },
--         have hS : m ≫ prod.snd = eq_to_hom rfl ≫ limF'.is_limit.lift s',
--         { -- m ≫ prod.snd is a morphism of cones over F' into limF'.X ..
--           have hN : ∀ (j: discrete (pfin n)), (m ≫ prod.snd) ≫ limF'.cone.π.app j = s'.π.app j,
--           { intro j,
--             unfold_projs, simp [(h j.succ).symm],
--             rcases j with ⟨j, hj⟩, refl },
--           -- .. and therefore unique.
--           have hUniq' : m ≫ prod.snd = limF'.is_limit.lift s',
--           from limF'.is_limit.uniq' s' (m ≫ prod.snd) hN,
--           simp only [hUniq', category.id_comp, eq_to_hom_refl] },
--         exact prod.lift_uniq _ _ _ h0 hS
--       end } }
-- set_option eqn_compiler.zeta false

-- end has_finite_products_of_binary_products_and_terminal_object

-- open has_finite_products_of_binary_products_and_terminal_object

-- -- TODO(WN): instance or def? Is there another way one might want to construct limits of shape pfin?
-- instance has_limits_of_shape_pfin [has_binary_products.{v} C] [has_terminal.{v} C] (n : ℕ)
--   : @has_limits_of_shape (discrete $ pfin n) _ C 𝒞 :=
-- ⟨λ F, has_limit_for_pfin_diagram F⟩

-- -- TODO(WN): trunc? #22
-- def has_trunc_finite_products [has_binary_products.{v} C] [has_terminal.{v} C]
--   {J : Type v} [fintype J] [decidable_eq J]
--   : trunc (has_limits_of_shape (discrete J) C) :=
-- trunc.lift_on (fintype.equiv_pfin J)
--   (λ h,
--     let hIso : discrete (pfin $ fintype.card J) ≌ discrete J :=
--       discrete.equiv_of_iso h.symm in
--     let limsPfin : @has_limits_of_shape (discrete (pfin $ fintype.card J)) _ C 𝒞 :=
--       by apply_instance in
--     trunc.mk $ has_limits_of_shape_of_equivalence hIso)
--   (λ a b, trunc.eq _ _)

-- end category_theory.limits


import for_mathlib.category_theory.functor.shift
import category_theory.localization.predicate

noncomputable theory

open category_theory category_theory.category

namespace category_theory

@[simps]
instance localization.lifting.L_comp {C H D : Type*} [category C] [category H] [category D]
  (L : C ⥤ H) (W : morphism_property C) [L.is_localization W] (F : H ⥤ D) :
  localization.lifting L W (L ⋙ F) F :=
⟨iso.refl _⟩

namespace functor

namespace has_comm_shift

section

variables {C H D : Type*} [category C] [category H] [category D]
  {L : C ⥤ H} {F : C ⥤ D} {G : H ⥤ D} (e : L ⋙ G ≅ F)
  (W : morphism_property C) [L.is_localization W] {A : Type*} [add_monoid A]
  [has_shift C A] [has_shift D A] [has_shift H A]
  [F.has_comm_shift A] [L.has_comm_shift A]

include e W

def of_localization.iso (a : A) : shift_functor H a ⋙ G ≅ G ⋙ shift_functor D a :=
localization.lift_nat_iso L W (L ⋙ (shift_functor H a ⋙ G)) (L ⋙ G ⋙ shift_functor D a) _ _
  ((functor.associator _ _ _).symm ≪≫ iso_whisker_right (L.comm_shift_iso a).symm _ ≪≫
    functor.associator _ _ _ ≪≫ iso_whisker_left _ e ≪≫ F.comm_shift_iso a ≪≫
    iso_whisker_right e.symm  _ ≪≫ functor.associator _ _ _)

@[simp]
lemma of_localization.iso_hom_app_obj (a : A) (X : C) :
  (of_localization.iso e W a).hom.app (L.obj X) =
    G.map ((L.comm_shift_iso a).inv.app X) ≫ e.hom.app ((shift_functor C a).obj X) ≫
      (F.comm_shift_iso a).hom.app X ≫ (shift_functor D a).map (e.inv.app X) :=
begin
  dsimp [of_localization.iso],
  simp only [localization.lift_nat_trans_app, localization.lifting.L_comp_iso, iso.refl_hom,
    nat_trans.id_app, nat_trans.comp_app, associator_inv_app, whisker_right_app,
    associator_hom_app, whisker_left_app, comp_id, id_comp, iso.refl_inv, assoc],
  erw id_comp,
end

@[simp]
lemma of_localization.iso_inv_app_obj (a : A) (X : C) :
  (of_localization.iso e W a).inv.app (L.obj X) =
    (shift_functor D a).map (e.hom.app X) ≫  (F.comm_shift_iso a).inv.app X ≫
    e.inv.app ((shift_functor C a).obj X) ≫ G.map ((L.comm_shift_iso a).hom.app X) :=
begin
  dsimp [of_localization.iso],
  simp only [assoc, localization.lift_nat_trans_app, localization.lifting.L_comp_iso,
    iso.refl_hom, nat_trans.id_app, nat_trans.comp_app, associator_inv_app, whisker_right_app,
    whisker_left_app, associator_hom_app, comp_id, id_comp, iso.refl_inv],
  erw id_comp,
end

variable (A)

@[simps]
def of_localization : G.has_comm_shift A :=
{ iso := of_localization.iso e W,
  iso_zero := begin
    ext1,
    apply localization.nat_trans_ext L W,
    intro X,
    simp only [of_localization.iso_hom_app_obj, comm_shift.unit_hom_app, iso.symm_hom,
      iso.symm_inv, monoidal_functor.ε_iso_hom, F.comm_shift_iso_zero A,
      L.comm_shift_iso_zero A, comm_shift.unit_inv_app, assoc, functor.map_comp,
      ← nat_trans.naturality, ← nat_trans.naturality_assoc, id_map, comp_map],
    congr' 1,
    dsimp,
    simp only [e.hom_inv_id_app_assoc, ← functor.map_comp_assoc, ← functor.map_comp,
      ε_hom_inv_app],
    erw [functor.map_id, functor.map_id, id_comp],
  end,
  iso_add := λ a b, begin
    ext1,
    apply localization.nat_trans_ext L W,
    intro X,
    simp only [of_localization.iso_hom_app_obj, comm_shift.add_hom_app, iso.symm_hom,
      map_comp, iso.symm_inv, monoidal_functor.μ_iso_hom, assoc, μ_naturality,
      functor.comm_shift_iso_add, comm_shift.add_inv_app],
    erw [← nat_trans.naturality_assoc, ← nat_trans.naturality_assoc],
    dsimp,
    simp only [of_localization.iso_hom_app_obj, assoc],
    nth_rewrite 3 ← G.map_comp_assoc,
    erw [← (shift_functor D b).map_comp_assoc, e.inv_hom_id_app, functor.map_id, id_comp,
      ← L.map_comp, μ_hom_inv_app, L.map_id, G.map_id, id_comp],
    refl,
  end, }

include A

lemma of_localization.respects_comm_shift :
  by { letI := of_localization e W A,
    exact e.hom.respects_comm_shift A } :=
begin
  letI := of_localization e W A,
  refine ⟨λ a, _⟩,
  ext X,
  simp only [nat_trans.comp_app, comp_hom_app, whisker_right_app, assoc, whisker_left_app],
  erw of_localization.iso_hom_app_obj e W a X,
  simp only [assoc, ← G.map_comp_assoc, iso.hom_inv_id_app, ← functor.map_comp,
    iso.inv_hom_id_app],
  dsimp,
  simp only [map_id, id_comp, comp_id],
end

end

section

variables {C H D A : Type*}
  [category C] [category H] [category D] [add_monoid A]
  [has_shift C A] [has_shift D A] [has_shift H A] (W : morphism_property C)
  {L : C ⥤ H} {F : C ⥤ D} [L.has_comm_shift A] [F.has_comm_shift A]
  [L.is_localization W] (hF : W.is_inverted_by F)

instance localization_lift_has_comm_shift :
  (localization.lift F hF L).has_comm_shift A :=
of_localization (localization.fac F hF L) W A

instance localization_fac_hom_respects_comm_shift :
  (localization.fac F hF L).hom.respects_comm_shift A :=
of_localization.respects_comm_shift _ _ _

end
end has_comm_shift

end functor

end category_theory


import Smt

theorem falsum : !false := by
  smt
  simp_all


/-
Copyright (c) 2021 Microsoft Corporation. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Leonardo de Moura
-/
import Lean.Elab.PreDefinition.Structural.Basic

namespace Lean.Elab.Structural
open Meta

private def getIndexMinPos (xs : Array Expr) (indices : Array Expr) : Nat := Id.run do
  let mut minPos := xs.size
  for index in indices do
    match xs.indexOf? index with
    | some pos => if pos.val < minPos then minPos := pos.val
    | _        => pure ()
  return minPos

-- Indices can only depend on other indices
private def hasBadIndexDep? (ys : Array Expr) (indices : Array Expr) : MetaM (Option (Expr × Expr)) := do
  for index in indices do
    let indexType ← inferType index
    for y in ys do
      if !indices.contains y && (← dependsOn indexType y.fvarId!) then
        return some (index, y)
  return none

-- Inductive datatype parameters cannot depend on ys
private def hasBadParamDep? (ys : Array Expr) (indParams : Array Expr) : MetaM (Option (Expr × Expr)) := do
  for p in indParams do
    let pType ← inferType p
    for y in ys do
      if ← dependsOn pType y.fvarId! then
        return some (p, y)
  return none

private def throwStructuralFailed : MetaM α :=
  throwError "structural recursion cannot be used"

private def orelse' (x y : M α) : M α := do
  let saveState ← get
  orelseMergeErrors x (do set saveState; y)

/--
  Try to find an argument that is structurally smaller in every recursive application.
  We use this argument to justify termination using the auxiliary `brecOn` construction.

  We give preference for arguments that are *not* indices of inductive types of other arguments.
  See issue #837 for an example where we can show termination using the index of an inductive family, but
  we don't get the desired definitional equalities.

  We perform two passes. In the first-pass, we only consider arguments that are not indices.
  In the second pass, we consider them.

  TODO: explore whether there are better solutions, and whether there are other ways to break the heuristic used
  for creating the smart unfolding auxiliary definition.
-/
partial def findRecArg (numFixed : Nat) (xs : Array Expr) (k : RecArgInfo → M α) : M α := do
  /- Collect arguments that are indices. See comment above. -/
  let indicesRef : IO.Ref FVarIdSet ← IO.mkRef {}
  for x in xs do
    let xType ← inferType x
    /- Traverse all sub-expressions in the type of `x` -/
    forEachExpr xType fun e =>
      /- If `e` is an inductive family, we store in `indicesRef` all variables in `xs` that occur in "index positions". -/
      matchConstInduct e.getAppFn (fun _ => pure ()) fun info _ => do
        if info.numIndices > 0 && info.numParams + info.numIndices == e.getAppNumArgs then
          for arg in e.getAppArgs[info.numParams:] do
            forEachExpr arg fun e => do
              if e.isFVar && xs.any (· == e) then
                indicesRef.modify fun indices => indices.insert e.fvarId!
  let indices ← indicesRef.get
  /- We perform two passes. See comment above. -/
  let rec go (i : Nat) (firstPass : Bool) : M α := do
    if h : i < xs.size then
      let x := xs.get ⟨i, h⟩
      trace[Elab.definition.structural] "findRecArg x: {x}, firstPass: {firstPass}"
      let localDecl ← getFVarLocalDecl x
      if localDecl.isLet then
        throwStructuralFailed
      else if firstPass == indices.contains localDecl.fvarId then
        go (i+1) firstPass
      else
        let xType ← whnfD localDecl.type
        matchConstInduct xType.getAppFn (fun _ => go (i+1) firstPass) fun indInfo us => do
        if !(← hasConst (mkBRecOnName indInfo.name)) then
          go (i+1) firstPass
        else if indInfo.isReflexive && !(← hasConst (mkBInductionOnName indInfo.name)) && !(← isInductivePredicate indInfo.name) then
          go (i+1) firstPass
        else
          let indArgs    := xType.getAppArgs
          let indParams  := indArgs.extract 0 indInfo.numParams
          let indIndices := indArgs.extract indInfo.numParams indArgs.size
          if !indIndices.all Expr.isFVar then
            orelse'
              (throwError "argument #{i+1} was not used because its type is an inductive family and indices are not variables{indentExpr xType}")
              (go (i+1) firstPass)
          else if !indIndices.allDiff then
            orelse'
              (throwError "argument #{i+1} was not used because its type is an inductive family and indices are not pairwise distinct{indentExpr xType}")
              (go (i+1) firstPass)
          else
            let indexMinPos := getIndexMinPos xs indIndices
            let numFixed    := if indexMinPos < numFixed then indexMinPos else numFixed
            let fixedParams := xs.extract 0 numFixed
            let ys          := xs.extract numFixed xs.size
            match (← hasBadIndexDep? ys indIndices) with
            | some (index, y) =>
              orelse'
                (throwError "argument #{i+1} was not used because its type is an inductive family{indentExpr xType}\nand index{indentExpr index}\ndepends on the non index{indentExpr y}")
                (go (i+1) firstPass)
            | none =>
              match (← hasBadParamDep? ys indParams) with
              | some (indParam, y) =>
                orelse'
                  (throwError "argument #{i+1} was not used because its type is an inductive datatype{indentExpr xType}\nand parameter{indentExpr indParam}\ndepends on{indentExpr y}")
                  (go (i+1) firstPass)
              | none =>
                let indicesPos := indIndices.map fun index => match ys.indexOf? index with | some i => i.val | none => unreachable!
                orelse'
                  (mapError
                    (k { fixedParams := fixedParams
                         ys          := ys
                         pos         := i - fixedParams.size
                         indicesPos  := indicesPos
                         indName     := indInfo.name
                         indLevels   := us
                         indParams   := indParams
                         indIndices  := indIndices
                         reflexive := indInfo.isReflexive
                         indPred := ←isInductivePredicate indInfo.name })
                    (fun msg => m!"argument #{i+1} was not used for structural recursion{indentD msg}"))
                  (go (i+1) firstPass)
    else if firstPass then
      go (i := numFixed) (firstPass := false)
    else
      throwStructuralFailed

  go (i := numFixed) (firstPass := true)

end Lean.Elab.Structural


import Yatima.Lean.Utils
import Yatima.ContAddr.ContAddrM
import YatimaStdLib.RBMap
import Lurk.LightData

namespace Yatima.ContAddr

scoped instance : HMul Ordering Ordering Ordering where hMul
  | .gt, _ => .gt
  | .lt, _ => .lt
  | .eq, x => x

def concatOrds : List Ordering → Ordering :=
  List.foldl (· * ·) .eq

open IR
open Std (RBMap)

/-- Defines an ordering for Lean universes -/
def cmpLevel (x : Lean.Level) (y : Lean.Level) : ContAddrM Ordering :=
  match x, y with
  | .mvar .., _ => throw $ .unfilledLevelMetavariable x
  | _, .mvar .. => throw $ .unfilledLevelMetavariable y
  | .zero, .zero => return .eq
  | .zero, _ => return .lt
  | _, .zero => return .gt
  | .succ x, .succ y => cmpLevel x y
  | .succ .., _ => return .lt
  | _, .succ .. => return .gt
  | .max lx ly, .max rx ry => (· * ·) <$> cmpLevel lx rx <*> cmpLevel ly ry
  | .max .., _ => return .lt
  | _, .max .. => return .gt
  | .imax lx ly, .imax rx ry => (· * ·) <$> cmpLevel lx rx <*> cmpLevel ly ry
  | .imax .., _ => return .lt
  | _, .imax .. => return .gt
  | .param x, .param y => do
    let lvls := (← read).univCtx
    match (lvls.indexOf? x), (lvls.indexOf? y) with
    | some xi, some yi => return (compare xi yi)
    | none,    _       => throw $ .levelNotFound x lvls
    | _,       none    => throw $ .levelNotFound y lvls

/-- Content-addresses a Lean universe level and adds it to the store -/
def contAddrUniv : Lean.Level → ContAddrM Univ
  | .zero => pure .zero
  | .succ u => return .succ (← contAddrUniv u)
  | .max a b  => return .max  (← contAddrUniv a) (← contAddrUniv b)
  | .imax a b => return .imax (← contAddrUniv a) (← contAddrUniv b)
  | .param name => do
    let lvls := (← read).univCtx
    match lvls.indexOf? name with
    | some n => pure $ .var n
    | none   => throw $ .levelNotFound name lvls
  | l@(.mvar ..) => throw $ .unfilledLevelMetavariable l

/-- Retrieves a Lean constant from the environment by its name -/
def getLeanConstant (name : Lean.Name) : ContAddrM Lean.ConstantInfo := do
  match (← read).constMap.find? name with
  | some const => pure const
  | none => throw $ .unknownConstant name

def isInternalRec (expr : Lean.Expr) (name : Lean.Name) : Bool :=
  match expr with
  | .forallE _ t e _  => match e with
    | .forallE ..  => isInternalRec e name
    | _ => isInternalRec t name -- t is the major premise
  | .app e .. => isInternalRec e name
  | .const n .. => n == name
  | _ => false

mutual

partial def contAddrConst (const : Lean.ConstantInfo) : ContAddrM Lurk.F := do
  match (← get).env.consts.find? const.name with
  | some hash => pure hash
  | none => match const with
    | .defnInfo val => withLevelsAndReset val.levelParams $ contAddrDefinition val
    | .inductInfo val => withLevelsAndReset val.levelParams $ contAddrInductive val
    | .ctorInfo val => do
      match ← getLeanConstant val.induct with
      | .inductInfo ind => discard $ contAddrConst (.inductInfo ind)
      | const => throw $ .invalidConstantKind const.name "inductive" const.ctorName
      contAddrConst const
    | .recInfo val => do
      match ← getLeanConstant val.getInduct with
      | .inductInfo ind => discard $ contAddrConst (.inductInfo ind)
      | const => throw $ .invalidConstantKind const.name "inductive" const.ctorName
      contAddrConst const
    -- The rest adds the constants to the cache one by one
    | const => withLevelsAndReset const.levelParams do
      let obj ← match const with
        | .defnInfo _ | .inductInfo _ | .ctorInfo _ | .recInfo _ => unreachable!
        | .axiomInfo val =>
          pure $ .axiom ⟨val.levelParams.length, ← contAddrExpr val.type⟩
        | .thmInfo val =>
          -- Theorems are never truly recursive
          pure $ .theorem ⟨val.levelParams.length, ← contAddrExpr val.type,
            ← contAddrExpr val.value⟩
        | .opaqueInfo val =>
          let recrs := .single val.name 0
          pure $ .opaque ⟨val.levelParams.length, ← contAddrExpr val.type,
            ← withRecrs recrs $ contAddrExpr val.value⟩
        | .quotInfo val =>
          pure $ .quotient ⟨val.levelParams.length, ← contAddrExpr val.type, val.kind⟩
      let hash ← commit obj
      addConstToEnv const.name hash
      return hash

partial def contAddrDefinition (struct : Lean.DefinitionVal) : ContAddrM Lurk.F := do
  -- If the mutual size is one, simply content address the single definition
  if struct.all matches [_] then
    let hash ← commit $ .definition
      (← withRecrs (.single struct.name 0) $ definitionToIR struct)
    addConstToEnv struct.name hash
    return hash

  -- Collecting and sorting all definitions in the mutual block
  let mutualDefs ← struct.all.mapM fun name => do
    match ← getLeanConstant name with
    | .defnInfo defn => pure defn
    | const => throw $ .invalidConstantKind const.name "definition" const.ctorName
  let mutualDefs ← sortDefs [mutualDefs]

  -- Building the `recrCtx`
  let mut recrCtx := default
  for (i, ds) in mutualDefs.enum do
    for d in ds do
      recrCtx := recrCtx.insert d.name i

  let definitions ← withRecrs recrCtx $ mutualDefs.mapM (·.mapM definitionToIR)

  -- Building and storing the block
  let definitionsIr := (definitions.map (match ·.head? with
    | some d => [d] | none => [])).join
  let blockHash ← commit $ .mutDefBlock definitionsIr
  addBlockToEnv blockHash

  -- While iterating on the definitions from the mutual block, we need to track
  -- the correct objects to return
  let mut ret? : Option Lurk.F := none

  for name in struct.all do
    -- Storing and caching the definition projection
    -- Also adds the constant to the array of constants
    let some idx := recrCtx.find? name | throw $ .cantFindMutDefIndex name
    let hash ← commit $ .definitionProj ⟨blockHash, idx⟩
    addConstToEnv name hash
    if struct.name == name then ret? := some hash

  match ret? with
  | some ret => return ret
  | none => throw $ .constantNotContentAddressed struct.name

partial def definitionToIR (defn : Lean.DefinitionVal) : ContAddrM Definition :=
  return ⟨defn.levelParams.length, ← contAddrExpr defn.type,
    ← contAddrExpr defn.value, defn.safety == .partial⟩

/--
Content-addresses an inductive and all inductives in the mutual block as a
mutual block, even if the inductive itself is not in a mutual block.

Content-addressing an inductive involves content-addressing its associated
constructors and recursors, hence the lenght of this function.
-/
partial def contAddrInductive (initInd : Lean.InductiveVal) : ContAddrM Lurk.F := do
  -- `mutualConsts` is the list of the names of all constants associated with an inductive block
  -- it has the form: ind₁ ++ ctors₁ ++ recrs₁ ++ ... ++ indₙ ++ ctorsₙ ++ recrsₙ
  let mut inds := []
  let mut indCtors := []
  let mut indRecs := []
  let mut nameData : RBMap Name (List Name × List Name) compare := .empty
  for indName in initInd.all do
    match ← getLeanConstant indName with
    | .inductInfo ind =>
      let indRecrs := ((← read).constMap.childrenOfWith ind.name
        fun c => match c with | .recInfo _ => true | _ => false).map (·.name)
      inds := inds ++ [indName]
      indCtors := indCtors ++ ind.ctors
      indRecs := indRecs ++ indRecrs
      nameData := nameData.insert indName (ind.ctors, indRecrs)
    | const => throw $ .invalidConstantKind const.name "inductive" const.ctorName

  -- `mutualConsts` is the list of the names of all constants associated with an
  -- inductive block: the inductives themselves, the constructors and the recursors
  let mutualConsts := inds ++ indCtors ++ indRecs

  let recrCtx := mutualConsts.enum.foldl (init := default)
    fun acc (i, n) => acc.insert n i

  -- This part will build the inductive block and add all inductives,
  -- constructors and recursors to `consts`
  let irInds ← initInd.all.mapM fun name => do match ← getLeanConstant name with
    | .inductInfo ind => withRecrs recrCtx do pure $ (← inductiveToIR ind)
    | const => throw $ .invalidConstantKind const.name "inductive" const.ctorName
  let blockHash ← commit $ .mutIndBlock irInds
  addBlockToEnv blockHash

  -- While iterating on the inductives from the mutual block, we need to track
  -- the correct objects to return
  let mut ret? : Option Lurk.F := none
  for (indIdx, indName) in initInd.all.enum do
    -- Store and cache inductive projections
    let name := indName
    let hash ← commit $ .inductiveProj ⟨blockHash, indIdx⟩
    addConstToEnv name hash
    if name == initInd.name then ret? := some hash

    let some (ctors, recrs) := nameData.find? indName
      | throw $ .cantFindMutDefIndex indName

    for (ctorIdx, ctorName) in ctors.enum do
      -- Store and cache constructor projections
      let hashes ← commit $ .constructorProj ⟨blockHash, indIdx, ctorIdx⟩
      addConstToEnv ctorName hashes

    for (recrIdx, recrName) in recrs.enum do
      -- Store and cache recursor projections
      let hashes ← commit $ .recursorProj ⟨blockHash, indIdx, recrIdx⟩
      addConstToEnv recrName hashes

  match ret? with
  | some ret => return ret
  | none => throw $ .constantNotContentAddressed initInd.name

partial def inductiveToIR (ind : Lean.InductiveVal) : ContAddrM Inductive := do
  let leanRecs := (← read).constMap.childrenOfWith ind.name
    fun c => match c with | .recInfo _ => true | _ => false
  let (recs, ctors) ← leanRecs.foldrM (init := ([], []))
    fun r (recs, ctors) => match r with
      | .recInfo rv =>
        if isInternalRec rv.type ind.name then do
          let (thisRec, thisCtors) := ← internalRecToIR ind.ctors r
          pure (thisRec :: recs, thisCtors)
        else do
          let thisRec ← externalRecToIR r
          pure (thisRec :: recs, ctors)
      | _ => throw $ .nonRecursorExtractedFromChildren r.name
  let (struct, unit) ← if ind.isRec || ind.numIndices != 0 then pure (false, false) else
    match ctors with
    -- Structures can only have one constructor
    | [ctor] => pure (true, ctor.fields == 0)
    | _ => pure (false, false)
  return ⟨ind.levelParams.length, ← contAddrExpr ind.type, ind.numParams, ind.numIndices,
    -- NOTE: for the purpose of extraction, the order of `ctors` and `recs` MUST
    -- match the order used in `recrCtx`
    ctors, recs, ind.isRec, ind.isReflexive, struct, unit⟩

partial def internalRecToIR (ctors : List Lean.Name) :
    Lean.ConstantInfo → ContAddrM (Recursor × List Constructor)
  | .recInfo rec => withLevels rec.levelParams do
    let typ ← contAddrExpr rec.type
    let (retCtors, retRules) ← rec.rules.foldrM (init := ([], []))
      fun r (retCtors, retRules) => do
        if ctors.contains r.ctor then
          let (ctor, rule) ← recRuleToIR r
          pure $ (ctor :: retCtors, rule :: retRules)
        else pure (retCtors, retRules) -- this is an external recursor rule
    let recr := ⟨rec.levelParams.length, typ, rec.numParams, rec.numIndices,
      rec.numMotives, rec.numMinors, retRules, rec.k, true⟩
    return (recr, retCtors)
  | const => throw $ .invalidConstantKind const.name "recursor" const.ctorName

partial def recRuleToIR (rule : Lean.RecursorRule) : ContAddrM $ Constructor × RecursorRule := do
  let rhs ← contAddrExpr rule.rhs
  match ← getLeanConstant rule.ctor with
  | .ctorInfo ctor => withLevels ctor.levelParams do
    let typ ← contAddrExpr ctor.type
    let ctor := ⟨ctor.levelParams.length, typ, ctor.cidx, ctor.numParams, ctor.numFields⟩
    pure (ctor, ⟨rule.nfields, rhs⟩)
  | const => throw $ .invalidConstantKind const.name "constructor" const.ctorName

partial def externalRecToIR : Lean.ConstantInfo → ContAddrM Recursor
  | .recInfo rec => withLevels rec.levelParams do
    let typ ← contAddrExpr rec.type
    let rules ← rec.rules.mapM externalRecRuleToIR
    return ⟨rec.levelParams.length, typ, rec.numParams, rec.numIndices,
      rec.numMotives, rec.numMinors, rules, rec.k, false⟩
  | const => throw $ .invalidConstantKind const.name "recursor" const.ctorName

partial def externalRecRuleToIR (rule : Lean.RecursorRule) : ContAddrM RecursorRule :=
  return ⟨rule.nfields, ← contAddrExpr rule.rhs⟩

/--
Content-addresses a Lean expression and adds it to the store.

Constants are the tricky case, for which there are two possibilities:
* The constant belongs to `recrCtx`, representing a recursive call. Those are
encoded as variables with indexes that go beyond the bind indexes
* The constant doesn't belong to `recrCtx`, meaning that it's not a recursion
and thus we can contAddr the actual constant right away
-/
partial def contAddrExpr : Lean.Expr → ContAddrM Expr
  | .mdata _ e => contAddrExpr e
  | expr => match expr with
    | .bvar idx => do match (← read).bindCtx.get? idx with
      -- Bound variables must be in the bind context
      | some _ => return .var idx []
      | none => throw $ .invalidBVarIndex idx
    | .sort lvl => return .sort $ ← contAddrUniv lvl
    | .const name lvls => do
      let univs ← lvls.mapM contAddrUniv
      match (← read).recrCtx.find? name with
      | some i => -- recursing!
        let idx := (← read).bindCtx.length + i
        return .var idx univs
      | none => return .const (← contAddrConst $ ← getLeanConstant name) univs
    | .app fnc arg => return .app (← contAddrExpr fnc) (← contAddrExpr arg)
    | .lam name typ bod _ =>
      return .lam (← contAddrExpr typ) (← withBinder name $ contAddrExpr bod)
    | .forallE name dom img _ =>
      return .pi (← contAddrExpr dom) (← withBinder name $ contAddrExpr img)
    | .letE name typ exp bod _ =>
      return .letE (← contAddrExpr typ) (← contAddrExpr exp)
        (← withBinder name $ contAddrExpr bod)
    | .lit lit => return .lit lit
    | .proj _ idx exp => return .proj idx (← contAddrExpr exp)
    | .fvar ..  => throw $ .freeVariableExpr expr
    | .mvar ..  => throw $ .metaVariableExpr expr
    | .mdata .. => throw $ .metaDataExpr expr

/--
A name-irrelevant ordering of Lean expressions.
`weakOrd` contains the best known current mutual ordering
-/
partial def cmpExpr (weakOrd : Std.RBMap Name Nat compare) :
    Lean.Expr → Lean.Expr → ContAddrM Ordering
  | e@(.mvar ..), _ => throw $ .unfilledExprMetavariable e
  | _, e@(.mvar ..) => throw $ .unfilledExprMetavariable e
  | e@(.fvar ..), _ => throw $ .freeVariableExpr e
  | _, e@(.fvar ..) => throw $ .freeVariableExpr e
  | .mdata _ x, .mdata _ y  => cmpExpr weakOrd x y
  | .mdata _ x, y  => cmpExpr weakOrd x y
  | x, .mdata _ y  => cmpExpr weakOrd x y
  | .bvar x, .bvar y => return (compare x y)
  | .bvar .., _ => return .lt
  | _, .bvar .. => return .gt
  | .sort x, .sort y => cmpLevel x y
  | .sort .., _ => return .lt
  | _, .sort .. => return .gt
  | .const x xls, .const y yls => do
    let univs ← concatOrds <$> (xls.zip yls).mapM fun (x,y) => cmpLevel x y
    if univs != .eq then return univs
    match weakOrd.find? x, weakOrd.find? y with
    | some nx, some ny => return compare nx ny
    | none, some _ => return .gt
    | some _, none => return .lt
    | none, none =>
      return compare (← contAddrConst $ ← getLeanConstant x)
        (← contAddrConst $ ← getLeanConstant y)
  | .const .., _ => return .lt
  | _, .const .. => return .gt
  | .app xf xa, .app yf ya =>
    (· * ·) <$> cmpExpr weakOrd xf yf <*> cmpExpr weakOrd xa ya
  | .app .., _ => return .lt
  | _, .app .. => return .gt
  | .lam _ xt xb _, .lam _ yt yb _ =>
    (· * ·) <$> cmpExpr weakOrd xt yt <*> cmpExpr weakOrd xb yb
  | .lam .., _ => return .lt
  | _, .lam .. => return .gt
  | .forallE _ xt xb _, .forallE _ yt yb _ =>
    (· * ·) <$> cmpExpr weakOrd xt yt <*> cmpExpr weakOrd xb yb
  | .forallE .., _ => return .lt
  | _, .forallE .. => return .gt
  | .letE _ xt xv xb _, .letE _ yt yv yb _ =>
    (· * · * ·) <$> cmpExpr weakOrd xt yt <*> cmpExpr weakOrd xv yv <*> cmpExpr weakOrd xb yb
  | .letE .., _ => return .lt
  | _, .letE .. => return .gt
  | .lit x, .lit y =>
    return if x < y then .lt else if x == y then .eq else .gt
  | .lit .., _ => return .lt
  | _, .lit .. => return .gt
  | .proj _ nx tx, .proj _ ny ty => do
    let ts ← cmpExpr weakOrd tx ty
    return concatOrds [compare nx ny, ts]

/-- AST comparison of two Lean definitions.
  `weakOrd` contains the best known current mutual ordering -/
partial def cmpDef (weakOrd : Std.RBMap Name Nat compare)
  (x : Lean.DefinitionVal) (y : Lean.DefinitionVal) :
    ContAddrM Ordering := do
  let ls := compare x.levelParams.length y.levelParams.length
  let ts ← cmpExpr weakOrd x.type y.type
  let vs ← cmpExpr weakOrd x.value y.value
  return concatOrds [ls, ts, vs]

/-- AST equality between two Lean definitions.
  `weakOrd` contains the best known current mutual ordering -/
@[inline] partial def eqDef (weakOrd : Std.RBMap Name Nat compare)
    (x y : Lean.DefinitionVal) : ContAddrM Bool :=
  return (← cmpDef weakOrd x y) == .eq

/--
`sortDefs` recursively sorts a list of mutual definitions into weakly equal blocks.
At each stage, we take as input the current best approximation of known weakly equal
blocks as a List of blocks, hence the `List (List DefinitionVal)` as the argument type.
We recursively take the input blocks and resort to improve the approximate known
weakly equal blocks, obtaining a sequence of list of blocks:
```
dss₀ := [startDefs]
dss₁ := sortDefs dss₀
dss₂ := sortDefs dss₁
dss₍ᵢ₊₁₎ := sortDefs dssᵢ ...
```
Initially, `startDefs` is simply the list of definitions we receive from `DefinitionVal.all`;
since there is no order yet, we treat it as one block all weakly equal. On the other hand,
at the end, there is some point where `dss₍ᵢ₊₁₎ := dssᵢ`, then we have hit a fixed point
and we may end the sorting process. (We claim that such a fixed point exists, although
technically we don't really have a proof.)

On each iteration, we hope to improve our knowledge of weakly equal blocks and use that
knowledge in the next iteration. e.g. We start with just one block with everything in it,
but the first sort may differentiate the one block into 3 blocks. Then in the second
iteration, we have more information than than first, since the relationship of the 3 blocks
gives us more information; this information may then be used to sort again, turning 3 blocks
into 4 blocks, and again 4 blocks into 6 blocks, etc, until we have hit a fixed point.
This step is done in the computation of `newDss` and then comparing it to the original `dss`.

Two optimizations:

1. `names := enum dss` records the ordering information in a map for faster access.
    Directly using `List.findIdx?` on dss is slow and introduces `Option` everywhere.
    `names` is used as a custom comparison in `ds.sortByM (cmpDef names)`.
2. `normDss/normNewDss`. We want to compare if two lists of blocks are equal.
    Technically blocks are sets and their order doesn't matter, but we have encoded
    them as lists. To fix this, we sort the list by name before comparing. Note we
    could maybe also use `List (RBTree ..)` everywhere, but it seemed like a hassle.
-/
partial def sortDefs (dss : List (List Lean.DefinitionVal)) :
    ContAddrM (List (List Lean.DefinitionVal)) := do
  let enum (ll : List (List Lean.DefinitionVal)) :=
    Std.RBMap.ofList (ll.enum.map fun (n, xs) => xs.map (·.name, n)).join
  let weakOrd := enum dss _
  let newDss ← (← dss.mapM fun ds =>
    match ds with
    | []  => unreachable!
    | [d] => pure [[d]]
    | ds  => do pure $ (← List.groupByM (eqDef weakOrd) $
      ← ds.sortByM (cmpDef weakOrd))).joinM

  -- must normalize, see comments
  let normDss    := dss.map    fun ds => ds.map (·.name) |>.sort
  let normNewDss := newDss.map fun ds => ds.map (·.name) |>.sort
  if normDss == normNewDss then return newDss
  else sortDefs newDss

end

/-- Iterates over a list of `Lean.ConstantInfo`, triggering their content-addressing -/
def contAddrM (delta : List Lean.ConstantInfo) : ContAddrM Unit := do
  delta.forM fun c => if !c.isUnsafe then discard $ contAddrConst c else pure ()
  if (← read).persist then dumpData (← get).ldonHashState LDONHASHCACHE

/--
Content-addresses the "delta" of an environment, that is, the content that is
added on top of the imports.

Important: constants with open references in their expressions are filtered out.
Open references are variables that point to names which aren't present in the
`Lean.ConstMap`.
-/
def contAddr (constMap : Lean.ConstMap) (delta : List Lean.ConstantInfo)
    (quick persist : Bool) : IO $ Except ContAddrError ContAddrState := do
  let persist := if quick then false else persist
  let ldonHashState ←
    if quick then pure default
    else pure $ (← loadData LDONHASHCACHE).getD default
  if persist then IO.FS.createDirAll STOREDIR
  match ← StateT.run (ReaderT.run (contAddrM delta)
    (.init constMap quick persist)) (.init ldonHashState) with
  | (.ok _, stt) => return .ok stt
  | (.error e, _) => return .error e

end Yatima.ContAddr