Choice.lean

inductive Fin2 : (n : Nat) → Type
  | zero : Fin2 (.succ n)
  | succ (i : Fin2 n) : Fin2 n.succ

def Fin2.elim0 {C : Fin2 .zero → Sort u} : ∀ i, C i := nofun

def Fin2.cases {C : Fin2 (.succ n) → Sort u} (zero : C .zero) (succ : ∀ i, C (.succ i)) : ∀ i, C i
  | .zero => zero
  | .succ i => succ i

def congrArg₂' {α : Sort u} {β : α → Prop} {γ : Sort v} (f : ∀ x, β x → γ) {x₁ x₂ : α} (h : x₁ = x₂) {y₁ : β x₁} {y₂ : β x₂} : f x₁ y₁ = f x₂ y₂ :=
  match x₂, h with
  | _, rfl => rfl


def finite_choice
  {α : Fin2 n → Sort u}
  {r : ∀ i, α i → α i → Prop}
  (f : ∀ i, Quot (r i))
  {β : Sort v}
  (g : (∀ i, α i) → β)
  (h : ∀ f₁ f₂, (∀ i, f₁ i = f₂ i ∨ r i (f₁ i) (f₂ i)) → g f₁ = g f₂)
: β :=
  match n with
  | .zero => g Fin2.elim0
  | .succ _ => f .zero |>.lift
    (fun fz => finite_choice (fun i => f i.succ) (fun fs => g <| Fin2.cases fz fs) fun _ _ hs => h _ _ <| Fin2.cases (.inl rfl) hs)
    fun _ _ hz => congrArg₂' (finite_choice _) <| funext fun _ => by exact h _ _ <| Fin2.cases (.inr hz) fun _ => .inl rfl

theorem finite_choice_mk f {β g h} : @finite_choice n α r (fun i => .mk _ (f i)) β g h = g f :=
  match n with
  | .zero => congrArg g <| funext Fin2.elim0
  | .succ _ => (finite_choice_mk fun i => f i.succ).trans <| congrArg g <| funext <| Fin2.cases rfl fun _ => rfl


opaque ULower (α : Sort u) : Type
unsafe def ULower.down : α → ULower α := unsafeCast
unsafe def ULower.up : ULower α → α := unsafeCast


private unsafe def choice.impl
  {α : Sort u}
  [DecidableEq α]
  {β : α → Sort v}
  {s : ∀ x, Setoid (β x)}
  (f : ∀ x, Quotient (s x))
  {γ : Sort w}
  (g : (∀ x, β x) → γ)
  (_ : ∀ f₁ f₂, (∀ x, (s x).r (f₁ x) (f₂ x)) → g f₁ = g f₂)
: γ :=
  ULower.up <| unsafeBaseIO do
  let entriesRef ← IO.mkRef #[]
  return .down <| g fun x =>
    ULower.up <| unsafeBaseIO do
    let entries ← entriesRef.get
    if let some entry := entries.find? (fun entry : ULower (PSigma β) => entry.up.fst = x) then
      return .down <| unsafeCast entry.up.snd
    let y := unsafeCast (f x)
    entriesRef.set <| entries.push <| .down ⟨x, y⟩
    return .down y

@[implemented_by choice.impl]
def choice
  {α : Sort u}
  [DecidableEq α]
  {β : α → Sort v}
  {s : ∀ x, Setoid (β x)}
  (f : ∀ x, Quotient (s x))
  {γ : Sort w}
  (g : (∀ x, β x) → γ)
  (h : ∀ f₁ f₂, (∀ x, (s x).r (f₁ x) (f₂ x)) → g f₁ = g f₂)
: γ :=
  g fun x => Classical.choose (f x).exists_rep

theorem choice_mk [DecidableEq α] {β s} f {γ g h} : @choice α _ β s (fun x => .mk _ (f x)) γ g h = g f :=
  h _ _ fun x => Quotient.exact <| Classical.choose_spec (Quotient.mk (s x) (f x)).exists_rep


private unsafe def countable_choice.impl
  {α : Nat → Sort u}
  {s : ∀ n, Setoid (α n)}
  (f : ∀ n, Quotient (s n))
  {β : Sort v}
  (g : (∀ n, α n) → β)
  (_ : ∀ f₁ f₂, (∀ n, (s n).r (f₁ n) (f₂ n)) → g f₁ = g f₂)
: β :=
  g (unsafeCast f)

@[implemented_by countable_choice.impl]
def countable_choice :
  {α : Nat → Sort u} →
  {s : ∀ n, Setoid (α n)} →
  (f : ∀ n, Quotient (s n)) →
  {β : Sort v} →
  (g : (∀ n, α n) → β) →
  (h : ∀ f₁ f₂, (∀ n, (s n).r (f₁ n) (f₂ n)) → g f₁ = g f₂) →
β :=
  @choice _ _

theorem countable_choice_mk : ∀ {α s} f {β g h}, @countable_choice α s (fun i => .mk _ (f i)) β g h = g f :=
  @choice_mk _ _