path: root/Nominal/Instances/Abstraction.lean

import Mathlib
import Nominal.Support
import Nominal.Equivariance
import Nominal.Category
import Nominal.Permutations
import Nominal.Instances.Basic

/-!
# Name Abstractions

Given a nominal set X, we define the set `[𝔸]X` of abstractions as pairs in 𝔸 × X modulo alpha-equivalence.
-/

variable (𝔸 : Type) [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸]

/-!
## Raw Structure

We first define the raw type of name abstractions as a pair (𝔸, X).

This structure, along with the obvious permutation action, forms a nominal set where supp (a, x) = supp x ∪ {a}.
-/
section RawStructure

/-- The structure forming a name abstraction. -/
structure NameAbstraction' (𝔸 : Type) [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] (X : Type) [NominalSet 𝔸 X] where
  /-- The atom that is abstracted away. -/
  atom : 𝔸
  /-- The body from which the atom is abstracted. -/
  body : X

variable (X : Type) [NominalSet 𝔸 X]

attribute [ext] NameAbstraction'

namespace NameAbstraction'

instance : SMul (FPerm 𝔸) (NameAbstraction' 𝔸 X) where
  smul π x := .mk (π • x.atom) (π • x.body)

@[simp] lemma smul_atom (π : FPerm 𝔸) (x : NameAbstraction' 𝔸 X) : (π • x).atom = π • x.atom := rfl

@[simp] lemma smul_body (π : FPerm 𝔸) (x : NameAbstraction' 𝔸 X) : (π • x).body = π • x.body := rfl

instance : PermSet 𝔸 (NameAbstraction' 𝔸 X) where
  one_smul := by intro _; ext <;> simp [one_smul]
  mul_smul := by intro _ _ _; ext <;> simp [mul_smul]

private def equiv : EquivariantEquiv 𝔸 (NameAbstraction' 𝔸 X) (𝔸 × X) :=
  { toFun x := ⟨x.atom, x.body⟩
    invFun x := ⟨x.fst, x.snd⟩
    toFun_equivariant := by intro x π; ext <;> simp
  }

instance : NominalSet 𝔸 (NameAbstraction' 𝔸 X) :=
  NominalSet.fromEquivariantEquiv 𝔸 (equiv 𝔸 X)

@[simp] lemma supp {x : NameAbstraction' 𝔸 X} : NominalSet.supp 𝔸 x = NominalSet.supp 𝔸 x.body ∪ {x.atom} := by
  have h : NominalSet.supp 𝔸 x = NominalSet.supp 𝔸 (x.atom, x.body) :=
    EquivariantEquiv.supp_eq 𝔸 (equiv 𝔸 X) x
  simp only [Prod.supp, NominalSet.supp_atom] at h
  rwa [Set.union_comm] at h

@[simp] lemma fresh {a : 𝔸} {x : NameAbstraction' 𝔸 X} : a # x ↔ a ≠ x.atom ∧ a # x.body := by
  simp [NominalSet.IsFresh]

end NameAbstraction'

end RawStructure

/-!
## Alpha-Equivalence

We continue to define alpha-equivalence on this structure.
-/
section AlphaEquivalence

namespace NameAbstraction

variable {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] {X : Type} [NominalSet 𝔸 X]

/-- (a, x) ≈ (b, y) if (c a) x = (c b) y for some fresh c. -/
def AlphaEq (x y : NameAbstraction' 𝔸 X) : Prop :=
  И c, ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body

/-- Alternatively define alpha equivalence as an equivariant relation.

This is used to show some properties using the Equivariance.for_fresh_* lemmas about the freshness quantifier. -/
private def AlphaEqPred (c : 𝔸) (xy : NameAbstraction' 𝔸 X × NameAbstraction' 𝔸 X) : Prop :=
  ⸨c, xy.fst.atom⸩ • xy.fst.body = ⸨c, xy.snd.atom⸩ • xy.snd.body

private lemma alpha_eq_iff_pred {x y : NameAbstraction' 𝔸 X} : AlphaEq x y = (И c, AlphaEqPred c (x, y)) := by
  apply NominalSet.for_fresh_ext
  simp [AlphaEqPred]

private lemma alpha_eq_pred_equivariant : Rel.Equivariant 𝔸 (AlphaEqPred : 𝔸 → NameAbstraction' 𝔸 X × NameAbstraction' 𝔸 X → Prop) := by
  intro a x π
  rw [AlphaEqPred, AlphaEqPred]
  intro h
  have h' : π • ⸨a, x.fst.atom⸩ • x.fst.body = π • ⸨a, x.snd.atom⸩ • x.snd.body := by congr
  simp only [smul_smul, FPerm.mul_swap_eq_swap_mul] at h'
  simpa [smul_smul]

private lemma alpha_eq_iff_exists' {x y : NameAbstraction' 𝔸 X} : AlphaEq x y ↔ ∃ c, c # x ∧ c # y ∧ ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body := by
  suffices (∃ c, c # x ∧ c # y ∧ ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body) = (∃ c, c # (x, y) ∧ AlphaEqPred c (x, y)) by
    rw [alpha_eq_iff_pred, this]
    apply Rel.Equivariant.for_fresh_iff_exists
    exact alpha_eq_pred_equivariant
  simp [Prod.fresh, AlphaEqPred, @and_assoc]

private lemma alpha_eq_iff_all' {x y : NameAbstraction' 𝔸 X} : AlphaEq x y ↔ ∀ c, c # x ∧ c # y → ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body := by
  suffices (∀ c, c # x ∧ c # y → ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body) = (∀ c, c # (x, y) → AlphaEqPred c (x, y)) by
    rw [alpha_eq_iff_pred, this]
    apply Rel.Equivariant.for_fresh_iff_all
    exact alpha_eq_pred_equivariant
  simp [Prod.fresh, AlphaEqPred]

private lemma alpha_eq_iff_cases' {x y : NameAbstraction' 𝔸 X} : AlphaEq x y ↔ (x.atom = y.atom ∧ x.body = y.body) ∨ (x.atom ≠ y.atom ∧ x.atom # y.body ∧ x.body = ⸨x.atom, y.atom⸩ • y.body) := by
  constructor
  · rw [alpha_eq_iff_exists']
    intro h
    obtain ⟨c, ⟨hcx, ⟨hcy, hc⟩⟩⟩ := h
    by_cases h_eq : x.atom = y.atom
    · left; constructor
      · exact h_eq
      · calc
          x.body = ⸨c, x.atom⸩ • ⸨c, x.atom⸩ • x.body := by rw [← mul_smul, FPerm.swap_mul_self, one_smul]
          _      = ⸨c, x.atom⸩ • ⸨c, y.atom⸩ • y.body := by rw [hc]
          _      = ⸨c, y.atom⸩ • ⸨c, y.atom⸩ • y.body := by rw [h_eq]
          _      = y.body                             := by rw [← mul_smul, FPerm.swap_mul_self, one_smul]
    · right; constructor
      · exact h_eq
      · simp at hcx
        simp at hcy
        have h_atom : (⸨c, y.atom⸩ * ⸨c, x.atom⸩) • c # (⸨c, y.atom⸩ * ⸨c, x.atom⸩) • x.body :=
          NominalSet.fresh_equivariant 𝔸 _ _ _ hcx.right
        have h_swap : ⸨c, x.atom⸩ * ⸨c, y.atom⸩ = ⸨x.atom, y.atom⸩ * ⸨c, x.atom⸩ := by
          ext a
          by_cases a = c <;> by_cases a = x.atom <;> by_cases a = y.atom
          all_goals simp_all [FPerm.swap_apply_of_ne_of_ne, hcy.left, hcx.left, Ne.symm hcy.left, Ne.symm hcx.left, Ne.symm h_eq]
        rw [mul_smul, PermSet.atom_swap, PermSet.atom_swap, Equiv.swap_apply_left, Equiv.swap_apply_of_ne_of_ne (Ne.symm hcx.left) h_eq] at h_atom
        rw [mul_smul, hc, ← mul_smul, FPerm.swap_mul_self, one_smul] at h_atom
        constructor
        · exact h_atom
        · calc
            x.body = ⸨c, x.atom⸩ • ⸨c, x.atom⸩ • x.body      := by rw [← mul_smul, FPerm.swap_mul_self, one_smul]
            _      = (⸨c, x.atom⸩ * ⸨c, y.atom⸩) • y.body    := by rw [mul_smul, hc]
            _      = ⸨x.atom, y.atom⸩ • ⸨c, x.atom⸩ • y.body := by rw [← mul_smul, h_swap]
            _      = ⸨x.atom, y.atom⸩ • y.body               := by rw [NominalSet.fresh_swap_fix 𝔸 hcy.right h_atom]
  · rw [alpha_eq_iff_all']
    intro h c ⟨hcx, hcy⟩
    cases h with
    | inl h =>
      rw [h.left, h.right]
    | inr h =>
      simp at hcx
      simp at hcy
      have h_swap : ⸨c, x.atom⸩ * ⸨x.atom, y.atom⸩ = ⸨c, y.atom⸩ * ⸨c, x.atom⸩ := by
        ext a
        by_cases a = c <;> by_cases a = x.atom <;> by_cases a = y.atom
        all_goals simp_all [FPerm.swap_apply_of_ne_of_ne, hcy.left, hcx.left, Ne.symm hcy.left, Ne.symm hcx.left, Ne.symm h.left]
      rw [h.right.right, ← mul_smul, h_swap, mul_smul, NominalSet.fresh_swap_fix 𝔸 hcy.right h.right.left]

private lemma alpha_eq_refl (x : NameAbstraction' 𝔸 X) : AlphaEq x x := by
  rw [alpha_eq_iff_cases']; left
  trivial

private lemma alpha_eq_symm {x y : NameAbstraction' 𝔸 X} : AlphaEq x y → AlphaEq y x := by
  rw [alpha_eq_iff_cases', alpha_eq_iff_cases']
  intro h
  cases h with
  | inl h =>
    left; exact ⟨h.left.symm, h.right.symm⟩
  | inr h =>
    right; constructor
    · exact h.left.symm
    · constructor
      · have h : ⸨x.atom, y.atom⸩ • x.atom # x.body := by
          rw [h.right.right]
          exact NominalSet.fresh_equivariant 𝔸 _ _ _ h.right.left
        simp at h; exact h
      · rw [h.right.right, ← mul_smul, FPerm.swap_comm, FPerm.swap_mul_self, one_smul]

private lemma alpha_eq_trans {x y z : NameAbstraction' 𝔸 X} : AlphaEq x y → AlphaEq y z → AlphaEq x z := by
  intro hxy hyz
  have h : AlphaEq x y ∧ AlphaEq y z := ⟨hxy, hyz⟩
  simp only [AlphaEq, AlphaEq, NominalSet.for_fresh_and] at h
  apply NominalSet.for_fresh_imp
  · exact h
  · intro c ⟨hxy, hyz⟩
    rw [hxy, hyz]

private lemma alpha_eq_equivariant : Rel.Equivariant 𝔸 (AlphaEq : NameAbstraction' 𝔸 X → NameAbstraction' 𝔸 X → Prop) := by
  intro x y π h
  rw [alpha_eq_iff_exists'] at *
  obtain ⟨c, ⟨hcx, ⟨hcy, hc⟩⟩⟩ := h
  use π • c
  constructor
  · apply NominalSet.fresh_equivariant; exact hcx
  · constructor
    · apply NominalSet.fresh_equivariant; exact hcy
    · -- XXX: This is largely similar to alpha_eq_pred...
      have h' : π • ⸨c, x.atom⸩ • x.body = π • ⸨c, y.atom⸩ • y.body := by congr
      simp only [smul_smul, FPerm.mul_swap_eq_swap_mul] at h'
      simpa [smul_smul]

end NameAbstraction

end AlphaEquivalence

/-!
## Quotient Type

We then define `NameAbstraction` as the quotient type of this structure modulo alpha-equivalence.
-/
section Quotient

variable (X : Type) [NominalSet 𝔸 X]

/-- Associate the name abstraction structure with alpha-equivalence. -/
instance NameAbstraction.setoid : EquivariantSetoid 𝔸 (NameAbstraction' 𝔸 X) where
  r := AlphaEq
  iseqv := ⟨alpha_eq_refl, alpha_eq_symm, alpha_eq_trans⟩
  equivariant := alpha_eq_equivariant

/-- Define name abstractions as the quotient set (𝔸 × X)/≈. -/
def NameAbstraction := Quotient (NameAbstraction.setoid 𝔸 X).toSetoid

namespace NameAbstraction

/-- The equivalence class of a name abstraction (a, x), often written `⟪a⟫ x` -/
def mk {𝔸 : Type} [Infinite 𝔸] [Countable 𝔸] [DecidableEq 𝔸] {X : Type} [NominalSet 𝔸 X] (a : 𝔸) (x : X) : NameAbstraction 𝔸 X :=
  Quotient.mk (setoid 𝔸 X).toSetoid (.mk a x)

/-- We write `⟪a⟫ x` instead of the usual `⟨a⟩ x` because `⟨⟩` is already heavily used in Lean for constructors. -/
notation "⟪" a "⟫ " x:60 => mk a x

instance : NominalSet 𝔸 (NameAbstraction 𝔸 X) := Quotient.instNominalSet 𝔸

@[simp]
lemma smul {π : FPerm 𝔸} {a : 𝔸} {x : X } : π • ⟪a⟫ x = ⟪π • a⟫ π • x := by congr

private lemma alpha_eq {x y : NameAbstraction' 𝔸 X} : x ≈ y ↔ AlphaEq x y := by rfl

/-- Two name abstractions are equal iff they are alpha-equivalent. -/
lemma eq_iff_alpha_eq {a b : 𝔸} {x y : X} : ⟪a⟫ x = ⟪b⟫ y ↔ (⟨a, x⟩ : NameAbstraction' 𝔸 X) ≈ (⟨b, y⟩ : NameAbstraction' 𝔸 X) := by
  rw [mk, mk, Quotient.eq_iff_equiv]

/-- The freshness quantifier in the AlphaEq definition is equivalent to "there exists some c # (x, y)". -/
lemma alpha_eq_iff_exists {x y : NameAbstraction' 𝔸 X} : x ≈ y ↔ ∃ c, c # x ∧ c # y ∧ ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body := by
  rw [alpha_eq, alpha_eq_iff_exists']

@[inherit_doc alpha_eq_iff_exists]
lemma eq_iff_exists {a b : 𝔸} {x y : X} : ⟪a⟫ x = ⟪b⟫ y ↔ ∃ c, c ≠ a ∧ c ≠ b ∧ c # x ∧ c # y ∧ ⸨c, a⸩ • x = ⸨c, b⸩ • y := by
  simp only [eq_iff_alpha_eq, alpha_eq_iff_exists, NameAbstraction'.fresh, ne_eq]
  ac_nf

/-- The freshness quantifier in the AlphaEq definition is equivalent to "for all c # (x, y)". -/
lemma alpha_eq_iff_all {x y : NameAbstraction' 𝔸 X} : x ≈ y ↔ ∀ c, c # x ∧ c # y → ⸨c, x.atom⸩ • x.body = ⸨c, y.atom⸩ • y.body := by
  rw [alpha_eq, alpha_eq_iff_all']

@[inherit_doc alpha_eq_iff_all]
lemma eq_iff_all {a b : 𝔸} {x y : X} : ⟪a⟫ x = ⟪b⟫ y ↔ ∀ c, c ≠ a ∧ c ≠ b ∧ c # x ∧ c # y → ⸨c, a⸩ • x = ⸨c, b⸩ • y := by
  simp only [eq_iff_alpha_eq, alpha_eq_iff_all, NameAbstraction'.fresh, ne_eq]
  ac_nf

/-- An alternative definition of α-equivalence:

For ⟪a⟫ x = ⟪b⟫ y, we either have (a, x) = (b, y) or x is just (a b) y, requiring that a is not in the support of y. -/
lemma alpha_eq_iff_cases {x y : NameAbstraction' 𝔸 X} : x ≈ y ↔ (x.atom = y.atom ∧ x.body = y.body) ∨ (x.atom ≠ y.atom ∧ x.atom # y.body ∧ x.body = ⸨x.atom, y.atom⸩ • y.body) := by
  rw [alpha_eq, alpha_eq_iff_cases']

@[inherit_doc alpha_eq_iff_cases]
lemma eq_iff_cases {a b : 𝔸} {x y : X} : ⟪a⟫ x = ⟪b⟫ y ↔ (a = b ∧ x = y) ∨ (a ≠ b ∧ a # y ∧ x = ⸨a, b⸩ • y) := by
  simp [eq_iff_alpha_eq, alpha_eq_iff_cases]

private lemma raw_supp_swap {a b : 𝔸} {x y : X} : ⟪a⟫ x = ⟪b⟫ y → NominalSet.supp 𝔸 (⟨b, y⟩ : NameAbstraction' 𝔸 X) = {b} ∪ (NominalSet.supp 𝔸 x \ {a}) := by
  intro haeq
  rw [eq_iff_cases] at haeq
  cases haeq with
  | inl h =>
    rw [h.left, h.right]
    simp only [NameAbstraction'.supp, Set.union_singleton, Set.union_diff_self, Set.singleton_union]
  | inr h =>
    simp only [NameAbstraction'.supp, Set.union_comm]
    rw [h.right.right, NominalSet.supp_equivariant 𝔸 y ⸨a, b⸩]
    ext c
    by_cases h₁ : c = b
    · simp [h₁]
    · by_cases h₂ : c = a
      · simp [h₂, show a ∉ NominalSet.supp 𝔸 y by exact h.right.left]
      · simp [h₁, h₂, ← PermSet.atom_set_swap_ne_ne]

@[simp]
lemma supp {a : 𝔸} {x : X} : NominalSet.supp 𝔸 (⟪a⟫ x) = NominalSet.supp 𝔸 x \ {a} := by
  rw [mk, Quotient.supp]
  ext b
  constructor
  · by_cases h_eq : a = b
    · -- The interesting case: Show that we don't need a in the support.
      subst h_eq; contrapose; intro _
      simp only [NameAbstraction'.supp, Set.union_singleton, Set.mem_insert_iff, Set.mem_setOf_eq, Classical.not_imp, not_forall, not_or]
      have hfin : (NominalSet.supp 𝔸 x).Finite := NominalSet.supp_finite 𝔸
      obtain ⟨c, hc⟩ := Infinite.exists_notMem_finset (hfin.toFinset ∪ {a})
      simp only [Finset.union_singleton, Finset.mem_insert, Set.Finite.mem_toFinset, not_or] at hc
      use ⟨c, ⸨a, c⸩ • x⟩
      have haeq : ⟪c⟫ ⸨a, c⸩ • x = ⟪a⟫ x := by
        rw [eq_iff_cases]
        right
        exact ⟨hc.left, hc.right, by rw [FPerm.swap_comm]⟩
      have hraw : a ∉ NominalSet.supp 𝔸 (⟨c, ⸨a, c⸩ • x⟩ : NameAbstraction' 𝔸 X) := by
        rw [raw_supp_swap 𝔸 X haeq.symm]
        simp [Ne.symm hc.left]
      simp only [NameAbstraction'.supp, Set.union_singleton, Set.mem_insert_iff, not_or] at hraw
      exact ⟨Quotient.exact haeq, Ne.symm hc.left, hraw.right⟩
    · -- For all other cases, it suffices that supp (⟪a⟫ x) ⊆ supp (a, x) ∪ {a}.
      intro h
      specialize h ⟨a, x⟩ (by rfl)
      simp only [NameAbstraction'.supp, Set.union_singleton, Set.mem_insert_iff, Ne.symm h_eq, false_or] at h
      rw [Set.mem_diff_singleton]
      exact ⟨h, Ne.symm h_eq⟩
  · -- Show that this is actually the least support.
    intro h x' hx'
    suffices NominalSet.supp 𝔸 x' = {x'.atom} ∪ NominalSet.supp 𝔸 x \ {a} by
      rw [this]
      exact Set.mem_union_right {x'.atom} h
    apply raw_supp_swap
    apply Quotient.sound
    exact (Setoid.symm hx')

@[simp]
lemma fresh {a b : 𝔸} {x : X} : b # ⟪a⟫ x ↔ b # x ∨ b = a := by
  simp only [NominalSet.IsFresh, supp, Set.mem_diff, Set.mem_singleton_iff, not_and_or, Decidable.not_not]

end NameAbstraction

end Quotient

/-!
## Abstraction Functor

Name abstraction can also be considered as a functor `Nom ⥤ Nom`.
-/
section Functor

open CategoryTheory

namespace NameAbstraction

variable {X Y Z : Type} [NominalSet 𝔸 X] [NominalSet 𝔸 Y] [NominalSet 𝔸 Z]

private def map' (f : X →ₑ[𝔸] Y) : NameAbstraction 𝔸 X → NameAbstraction 𝔸 Y :=
  Quotient.lift (fun x => ⟪x.atom⟫ f x.body) (by
    intro x y hxy
    rw [alpha_eq_iff_exists] at hxy
    obtain ⟨c, ⟨hcx, hcy, hc⟩⟩ := hxy
    rw [NameAbstraction'.fresh] at hcx
    rw [NameAbstraction'.fresh] at hcy
    rw [eq_iff_exists]
    use c
    have hc' : ⸨c,x.atom⸩ • f x.body = ⸨c,y.atom⸩ • f y.body := by
      rw [show f x.body = f.toFun x.body from rfl, ← f.equivariant]
      rw [show f y.body = f.toFun y.body from rfl, ← f.equivariant]
      congr
    exact
      ⟨hcx.left,
       hcy.left,
       Function.Equivariant.preserve_fresh 𝔸 f.equivariant hcx.right,
       Function.Equivariant.preserve_fresh 𝔸 f.equivariant hcy.right,
       hc'⟩)

private lemma map_equivariant (f : X →ₑ[𝔸] Y) : Function.Equivariant 𝔸 (map' 𝔸 f) := by
  apply Quotient.ind
  intro x π
  rw [← Quotient.mk_equivariant]
  simp only [map', Quotient.lift_mk, NameAbstraction'.smul_atom, NameAbstraction'.smul_body, smul]
  rw [show f (π • x.body) = f.toFun (π • x.body) from rfl, f.equivariant]
  rfl

/-- Lifts an equivariant function to name abstractions. -/
def map (f : X →ₑ[𝔸] Y) : NameAbstraction 𝔸 X →ₑ[𝔸] NameAbstraction 𝔸 Y :=
  ⟨map' 𝔸 f, map_equivariant 𝔸 f⟩

/-- Application of equivariant functions lifted to name abstractions. -/
@[simp]
lemma map_eq {f : X →ₑ[𝔸] Y} {a : 𝔸} {x : X} : map 𝔸 f (⟪a⟫ x) = ⟪a⟫ f x := rfl

private lemma map_id : map 𝔸 (EquivariantMap.id 𝔸 : X →ₑ[𝔸] X) = EquivariantMap.id 𝔸 := by
  suffices ∀ (x : NameAbstraction 𝔸 X), map' 𝔸 (EquivariantMap.id 𝔸) x = x by ext x; exact this x
  apply Quotient.ind
  intro x
  simp [map', NameAbstraction.mk]

private lemma map_comp (f : Y →ₑ[𝔸] Z) (g : X →ₑ[𝔸] Y) : map 𝔸 (EquivariantMap.comp 𝔸 f g) = EquivariantMap.comp 𝔸 (map 𝔸 f) (map 𝔸 g) := by
  suffices ∀ (x : NameAbstraction 𝔸 X), map' 𝔸 (EquivariantMap.comp 𝔸 f g) x = (map' 𝔸 f) ((map' 𝔸 g) x) by ext x; exact this x
  apply Quotient.ind
  intro x
  simp [map', NameAbstraction.mk]

/-- Name abstractions form a functor `Nom ⥤ Nom`. -/
def Functor : Nom 𝔸 ⥤ Nom 𝔸 :=
  { obj X := ⟨NameAbstraction 𝔸 X⟩
    map f := ⟨NameAbstraction.map 𝔸 f, NameAbstraction.map_equivariant 𝔸 f⟩
    map_id := fun _ => map_id 𝔸
    map_comp f g := map_comp 𝔸 g f }

end NameAbstraction

end Functor