compute.lean

import data.set.basic
import data.nat.lattice
import tactic

universes u v

@[ext]
structure sval (α : Type u) :=
(s : set α)
(singleton : ∃ x, s = {x})

prefix `!`:100 := sval

noncomputable def sval.val {α : Type u} (v : !α) : α := v.singleton.some

@[simp] lemma sval.val_spec {α : Type u} (v : !α) : {v.val} = v.s := v.singleton.some_spec.symm

@[simp] lemma sval.val_spec' {α : Type u} (v : !α) : v.val ∈ v.s :=
by { rw ←sval.val_spec, apply set.mem_singleton, }

def sval.incl {α : Type u} (x : α) : !α := ⟨{x}, x, rfl⟩

@[simp] lemma sval.incl_s {α : Type u} (x : α) : (sval.incl x).s = {x} := rfl

lemma sval.incl_prop {α : Type u} (v : !α) : ∃ x, v = sval.incl x :=
begin
  cases v with s hs,
  obtain ⟨x, hx⟩ := hs,
  subst s,
  use x,
  refl,
end

@[simp]
lemma sval.incl_injective {α : Type u} (x y : α) : sval.incl x = sval.incl y ↔ x = y :=
by simp [sval.incl]

@[simp] lemma sval.mem_s {α : Type u} (v : !α) (x : α) : sval.incl x = v ↔ x ∈ v.s :=
begin
  obtain ⟨y, hy⟩ := v.incl_prop,
  simp [hy],
end

@[simp] lemma sval.mem_s' {α : Type u} (v : !α) (x : α) : v = sval.incl x ↔ x ∈ v.s :=
begin
  rw eq_comm,
  simp,
end

@[simp] lemma sval.eq_val {α : Type u} (v : !α) (x : α) : x = v.val ↔ x ∈ v.s :=
begin
  split,
  { rintro rfl, convert set.mem_singleton _, simp, },
  { intro h, rw ←set.mem_singleton_iff, simp [h], },
end

@[simp] lemma sval.val_eq {α : Type u} (v : !α) (x : α) : v.val = x ↔ x ∈ v.s :=
begin
  rw eq_comm,
  simp,
end

lemma sval.exists {α : Type u} (v : !α) : ∃ x, x ∈ v.s :=
begin
  obtain ⟨x, h⟩ := v.singleton,
  use x,
  simp [h],
end

@[simp] lemma sval.val_incl {α : Type u} (x : α) : (sval.incl x).val = x :=
begin
  simp,
end

@[simp] lemma sval.incl_val {α : Type u} (v : !α) : sval.incl v.val = v :=
begin
  obtain ⟨x, rfl⟩ := sval.incl_prop v,
  simp,
end

@[simp] lemma sval.incl_some' {α : Type u} (v : !α) : sval.incl v.exists.some = v :=
begin
  ext x,
  cases v with s hs,
  obtain ⟨x', rfl⟩ := hs,
  simp,
  generalize_proofs h,
  rw h.some_spec,
end

lemma sval.eq_exists_some {α : Type u} (v : !α) : v.exists.some = v.val :=
by rw [←sval.incl_injective, sval.incl_val, sval.incl_some']

@[elab_as_eliminator]
lemma sval.elim {α : Type u} {P : !α → Prop} (h : ∀ x, P (sval.incl x)) (v : !α) : P v :=
begin
  cases v with q hq,
  cases hq with x hx,
  subst hx,
  exact h x,
end

@[elab_as_eliminator]
def sval.rec' {α : Type u} {β : Π (v : !α), Type*} (h : Π x, β (sval.incl x)) (v : !α) : !β v :=
{ s := {y | begin let z := h v.val, rw sval.incl_val at z, exact y = z, end},
  singleton := by simp }

@[simp] lemma sval.rec'_incl {α : Type u} {β : Π (v : !α), Type*} (h : Π x, β (sval.incl x)) (x : α) :
  sval.rec' h (sval.incl x) = sval.incl (h x) :=
begin
  ext y,
  simp [sval.rec'],
  split; rintro rfl,
  { have : ∀ x' (h' : x' = x) (h'' : β (sval.incl x') = β (sval.incl x)), cast h'' (h x') = h x,
    { rintros _ rfl, simp, },
    apply this,
    simp, },
  { have : ∀ x' (h' : x' = x) (h'' : β (sval.incl x') = β (sval.incl x)), h x = cast h'' (h x'),
    { rintros _ rfl, simp, },
    apply this,
    simp, },
end

@[elab_as_eliminator]
def sval.rec'' {α : Type u} {β : Π (v : !α), Type*} (h : Π x, !β (sval.incl x)) (v : !α) : !β v :=
{ s := {y | begin let z := h v.val, rw sval.incl_val at z, exact sval.incl y = z, end},
  singleton := begin
    refine sval.elim (λ x, _) v,
    simp,
    refine sval.elim (λ y h', _) (h x) (_ : ∀ x' (h' : x = x'), cast _ (h x) = (h x')),
    subst h', swap, rintro _ rfl, simp,
    specialize h' x rfl,
    use y,
    ext z, simp,
    suffices : ∀ x' (h' : x' = x) (h'' : !β (sval.incl x') = !β (sval.incl x)), z ∈ (cast h'' (h x')).s ↔ z = y,
    { apply this,
      simp, },
    rintro _ rfl,
    simp [←h'],
  end }

@[simp] lemma sval.rec''_incl {α : Type u} {β : Π (v : !α), Type*} (h : Π x, !β (sval.incl x)) (x : α) :
  sval.rec'' h (sval.incl x) = h x :=
begin
  ext y,
  simp [sval.rec''],
  generalize_proofs,
  suffices : ∀ x' (h' : x' = x) (h'' : !β (sval.incl x') = !β (sval.incl x)), y ∈ (cast h'' (h x')).s ↔ y ∈ (h x).s,
  { apply this,
    simp, },
  rintro _ rfl,
  simp,
end

/-- Create a proposition that refers to the value contained in `v` -/
def sval.with_val {α : Type u} (v : !α) (h : Π (x : α), Prop) : Prop :=
∀ x, x ∈ v.s → h x

@[simp] lemma sval.incl_with_val {α : Type u} (x : α) (h : Π (x : α), Prop) :
  (sval.incl x).with_val h = h x :=
begin
  simp [sval.incl, sval.with_val],
end

def sval.set_unwrap {α : Type u} : set (!α) ≃ set α :=
{ to_fun := λ s, {x | sval.incl x ∈ s},
  inv_fun := λ s, {v | v.with_val (λ x, x ∈ s)},
  left_inv := begin
    intro s,
    ext v,
    refine sval.elim (λ x, _) v,
    simp,
  end,
  right_inv := begin
    intro s,
    ext x,
    simp,
  end }

@[simp] lemma sval.set_unwrap_mem {α : Type u} (s : set (!α)) (x : α) : x ∈ sval.set_unwrap s ↔ sval.incl x ∈ s := iff.rfl

@[simp] lemma sval.set_wrap_mem {α : Type u} (s : set α) (v : !α) : v ∈ sval.set_unwrap.symm s ↔ v.val ∈ s :=
begin
  refine sval.elim (λ x, _) v,
  simp [sval.set_unwrap, sval.with_val],
end

def sval.set_remove {α : Type u} : !set α ≃ set α :=
{ to_fun := λ s, {x | ∃ s', s' ∈ s.s ∧ x ∈ s'},
  inv_fun := sval.incl,
  left_inv := begin
    intro s,
    ext s',
    simp,
    split,
    { rintro rfl,
      obtain ⟨s', hs'⟩ := s.incl_prop,
      simp [hs'], },
    { intro h,
      ext x,
      simp,
      split,
      { intro h',
        exact ⟨s', h, h'⟩, },
      { rintro ⟨s'', hs'', hx⟩,
        obtain ⟨a,b⟩ := s.incl_prop,
        simp [b] at h hs'',
        subst s', subst s'',
        assumption, }, },
  end,
  right_inv := begin
    intro s,
    ext x,
    simp,
  end }

@[simps]
instance sval.monad : monad (λ α, !α) :=
{ pure := @sval.incl,
  bind := λ α β v f, by { refine sval.rec'' (λ x, _) v, exact f x } }

def sval.dpush {α : Type*} {β : α → Type*}: !(Π (x : α), β x) ≃ (Π (x : α), !β x) :=
{ to_fun := λ f x, (do f' ← f,
                       pure (f' x)),
  inv_fun := λ f, ⟨{f' | ∀ x, f' x ∈ (f x).s}, begin
    use λ x, (f x).val,
    ext f',
    simp,
    split,
    { intro h,
      ext x,
      simp [h], },
    { rintro rfl x,
      simp, },
  end⟩,
  left_inv := λ f, begin
    refine sval.elim (λ f', _) f,
    ext x,
    simp,
    split,
    { intro h,
      ext x',
      exact h x', },
    { rintro rfl,
      intro,
      refl, },
  end,
  right_inv := λ f, begin
    ext1 x,
    simp [sval.rec''],
    have : ∀ x, ∃ y, f x = sval.incl y := λ x, (f x).incl_prop,
    choose g h using this,
    simp [h],
    convert_to (sval.incl g).val x = g x,
    swap, simp,
    congr' 1,
    ext g',
    simp,
    rw function.funext_iff,
  end }

def sval.push {α β : Type*} : !(α → β) ≃ (α → !β) := sval.dpush

class computable {α : Type*} (v : !α) :=
(compute : α)
(compute_spec : compute ∈ v.s)

def sval.compute {α : Type*} (v : !α) [computable v] : α := computable.compute v

lemma sval.compute_spec {α : Type*} (v : !α) [computable v] : v.compute ∈ v.s := computable.compute_spec

instance {α : Type*} (v : !α) : subsingleton (computable v) :=
⟨begin
  rintros ⟨a, ha⟩ ⟨b, hb⟩,
  congr,
  obtain ⟨x, hx⟩ := v.singleton,
  rw hx at ha hb,
  simp at ha hb,
  subst a, subst b,
end⟩

noncomputable
def computable_eqv {α : Type*} : !α ≃ α :=
{ to_fun := λ v, v.val,
  inv_fun := sval.incl,
  left_inv := begin
    intro v,
    refine sval.elim (λ x, _) v,
    simp,
  end,
  right_inv := begin
    intro x,
    simp,
  end }

instance computable_incl {α : Type*} (x : α) : computable (sval.incl x) :=
{ compute := x,
  compute_spec := by simp }

instance mem_incl_decidable {α : Type*} [decidable_eq α] (x : α) : decidable_pred (∈ (sval.incl x).s) :=
begin
  intro n,
  simp [sval.incl],
  apply_instance,
end

lemma compute_incl' {α : Type*} (x : α) : (sval.incl x).compute = x := rfl

@[simp]
lemma compute_incl {α : Type*} (x : α) (h : computable (sval.incl x)) :
  @sval.compute _ (sval.incl x) h = x :=
begin
  convert compute_incl' x,
end

@[simp]
lemma compute_eq {α : Type*} (v : !α) [computable v] : v.compute = v.val :=
by simp [v.compute_spec]

@[simp]
lemma incl_compute {α : Type*} (v : !α) [h : computable v] : sval.incl v.compute = v :=
begin
  simp,
end

def computable_eqv' {α : Type*} [∀ (v : !α), computable v] : !α ≃ α :=
{ to_fun := λ v, v.compute,
  inv_fun := sval.incl,
  left_inv := λ v, by simp,
  right_inv := λ x, by simp }

/-- ite, but without decidability -/
def ite' {α : Type*} (c : Prop) (x y : !α) : !α :=
⟨{v | (c → v ∈ x.s) ∧ (¬c → v ∈ y.s)},
begin
  by_cases c,
  { use x.val, simp [h], },
  { use y.val, simp [h], },
end⟩

/-- When the proposition is decidable, compute the ite' with ite. -/
instance ite'.compute {α : Type*} (c : Prop) [decidable c] (x y : !α) [computable x] [computable y] : computable (ite' c x y) :=
⟨ite c x.compute y.compute,
begin
  unfold ite',
  dsimp,
  simp {contextual:=tt},
end⟩

instance nat.computable (v : !ℕ) [decidable_pred (∈ v.s)] : computable v :=
{ compute := nat.find v.exists,
  compute_spec := nat.find_spec v.exists }

def Inf' (s : set ℕ) : !ℕ := ⟨{Inf s}, _, rfl⟩

instance nat.Inf'_computable (s : set ℕ) [decidable_pred (∈ s)] [h : decidable s.nonempty] : computable (Inf' s) :=
begin
  tactic.unfreeze_local_instances,
  simp [Inf'],
  cases h,
  use 0, simp, symmetry, rw nat.Inf_eq_zero, right, rwa set.not_nonempty_iff_eq_empty at h,
  use nat.find h,
  simp [nat.Inf_def h],
  congr,
end

instance nat.Inf'_interval (a : ℕ) : computable (Inf' (set.Ici a)) :=
⟨a, by simp [Inf']⟩

#eval (Inf' (set.Ici 100)).compute

noncomputable
def foo : ℕ := classical.choice infer_instance

def foo2 : !ℕ := ⟨{foo}, _, rfl⟩

noncomputable def nth' (p : ℕ → Prop) : ℕ → ℕ
| n := Inf { i : ℕ | p i ∧ ∀ k < n, nth' k < i }

def nth (p : ℕ → Prop) (n : ℕ) : !ℕ := ⟨{nth' p n}, _, rfl⟩

abbreviation finset' (α : Type*) := !finset α

namespace finset'
variables {α : Type*}

instance is_set_like : set_like (finset' α) α :=
⟨λ f, sval.set_remove (functor.map coe f),
begin
  intro s,
  refine sval.elim (λ s, _) s,
  intro t,
  refine sval.elim (λ t, _) t,
  simp [functor.map],
end⟩

@[simp] lemma mem_to_finset {s : finset' α} {x : α} : x ∈ s.val ↔ x ∈ s :=
begin
  change _ ↔ x ∈ sval.set_remove _,
  simp [sval.set_remove, sval.set_unwrap],
  generalize_proofs h,
  split,
  { intro h',
    use s.val,
    simp [h', sval.rec''], },
  { rintro ⟨t, ha, hx⟩,
    simp [sval.rec''] at ha,
    simpa [←ha] using hx, },
end

@[simp] lemma mem_incl {s : finset α} (x : α) : x ∈ sval.incl s ↔ x ∈ s :=
begin
  rw ←mem_to_finset,
  rw sval.val_incl,
end

@[ext] theorem ext {s t : finset' α} (h : ∀ x, x ∈ s ↔ x ∈ t) : s = t := set_like.ext h

@[simp] lemma coe_to_finset (s : finset' α) : (s.val : set α) = s :=
by { ext, simp }

def union (s t : finset' α) : finset' α :=
⟨{u | by { classical, exact u = s.val ∪ t.val }}, _, rfl⟩

instance : has_union (finset' α) := ⟨union⟩

lemma mem_union {s t : finset' α} (x : α) : x ∈ s ∪ t ↔ x ∈ s ∨ x ∈ t :=
begin
  change x ∈ sval.incl _ ↔ _,
  simp,
end

instance (s t : finset' α) [computable s] [computable t] [decidable_eq α] : computable (s ∪ t) :=
{ compute := s.compute ∪ t.compute,
  compute_spec := begin
    ext x,
    simp,
  end }

abbreviation a : finset' ℕ := sval.incl {1,2}
abbreviation b : finset' ℕ := sval.incl {2,3,4}

#eval (a ∪ b).compute

end finset'