Formalization of Schröder–Bernstein theorem in Lean, based on an exercise from Mathematics in Lean

schroeder_bernstein.lean

--
-- Inspired by https://leanprover-community.github.io/mathematics_in_lean/C04_Sets_and_Functions.html#the-schroder-bernstein-theorem
--

import Mathlib

open Set
open Function

noncomputable section
open Classical
variable {α β : Type*} [Nonempty β]

section

variable (f : α → β) (g : β → α)

def directIter: ℕ → Set α
  | 0 => (Set.range g)ᶜ
  | n + 1 => (g ∘ f) '' (directIter n)

def directPoints :=
  ⋃ i, directIter f g i

def h (x : α) : β :=
  if x ∈ directPoints f g then f x else (invFun g) x

theorem sb_surjective (injg : Injective g) : Surjective (h f g) := by
  intro b
  simp only [h]
  by_cases gb_direct : g b ∈ directPoints f g
  · /- If g(b) is a part of a direct chain, it must have been added
    to the direct set via some a -(f)-> b -(g)-> g(b) iteration.
    By picking that a, we get h(a) = f(a) = b, proving b is reachable. -/
    simp only [directPoints, mem_iUnion] at *
    rcases gb_direct with ⟨n, hn⟩
    rcases n with _ | n
    · tauto
    · rcases hn with ⟨a, a_direct, ha⟩
      use a
      rw [if_pos (by use n)]
      exact injg ha
  /- If g(b) is not a part of some direct chain [1], pick a := g(b).
  For that a, h(a) =(by [1])= inv_g(a) =(by def.)= inv_g(g(b)).
  Thus, proving inv_g(g(b)) = b would be enough to show b is reachable.
  To cancel them out, we'll wrap both sides with g (thanks, injectivity). -/
  use g b
  rw [if_neg gb_direct]
  apply injg
  apply invFun_eq
  use b

theorem g_right_inv {a : α} (ha : a ∉ directPoints f g) : g (invFun g a) = a := by
  /- If a is not a part of the direct chain, it must have been excluded
  when the set was originally being constructed. The elements that were
  excluded were the elements reachable via g. So a must be reachable via g. -/
  have : ∃ b, g b = a := by
    simp only [directPoints, mem_iUnion] at ha
    contrapose! ha
    use 0
    simp only [directIter, mem_compl_iff, mem_range]
    tauto
  apply invFun_eq this

theorem heq_imp_same_directness {a1 a2 : α} (ha1 : a1 ∈ directPoints f g) (heq : h f g a1 = h f g a2) : a2 ∈ directPoints f g := by
  /- If h(a1) = h(a2), either both must be direct or neither is. To see this, we'll unfold
  the definition of h. If one was g^-1(a1) but the other was f(a2), a1 and a2 would be only
  one f->g hop away from each other, forcing both a1 and a2 to be direct (a contradiction). -/
  simp only [h] at *
  rw [if_pos ha1] at heq
  by_cases ha2 : a2 ∈ directPoints f g
  · tauto
  · rw [if_neg ha2] at heq
    simp only [directPoints, mem_iUnion] at ha1
    rcases ha1 with ⟨i, hi⟩
    simp only [directPoints, mem_iUnion]
    use i + 1, a1, hi
    rw [comp_apply, heq]
    apply g_right_inv f g ha2

theorem sb_injective (injf : Injective f) : Injective (h f g) := by
  intro a1 a2 heq
  by_cases a1_direct : a1 ∈ directPoints f g
  · by_cases a2_direct : a2 ∈ directPoints f g
    · rw [h, h, if_pos a1_direct, if_pos a2_direct] at heq
      exact injf heq
    · have := heq_imp_same_directness f g a1_direct heq
      contradiction
  · by_cases a2_direct : a2 ∈ directPoints f g
    · have := heq_imp_same_directness f g a2_direct heq.symm
      contradiction
    · rw [h, h, if_neg a1_direct, if_neg a2_direct] at heq
      have ginv_a1 := g_right_inv f g a1_direct
      have ginv_a2 := g_right_inv f g a2_direct
      rw [← ginv_a1, ← ginv_a2, heq]

theorem schroeder_bernstein {f : α → β} {g : β → α} (hf : Injective f) (hg : Injective g) :
    ∃ h : α → β, Bijective h :=
  ⟨h f g, sb_injective f g hf, sb_surjective f g hg⟩

end