partial.lean

universes u v

@[reducible] def partial.fuel := nat
private def partial.huge := 1000

/- Remark: I could not use (reader_t partial.fuel (except unit α)) because of universe
   constraints. `partial.fuel` is at (Type 0), but α is in (Type u) -/
structure partial (α : Type u) :=
(run_prv : partial.fuel → except unit α) -- TODO: mark as private
/-
By marking partial.run_prv, partial.mk and partial.rec as private,
we make sure this module is the only one that have access to the internal implementation.
During code generation, we ignore the fuel and treat (partial α) as (id α).
-/

/- Compiler implements `partial` using the identity monad -/
private def partial.pure (α : Type) (a : α) : partial α :=
{run_prv := λ _, except.ok a}

/- Compiler implements `partial` using the identity monad -/
private def partial.bind (α β : Type) (a : partial α) (b : α → partial β) : partial β :=
{run_prv := λ f, match a.run_prv f with
                 | except.ok v := (b v).run_prv f
                 | _           := except.error ()
                 end}

instance : monad partial :=
{pure := @partial.pure, bind := @partial.bind}

variables {α : Type u} {β : α → Type v}

private def partial.fix_core
    (F : Π (d₁ : partial.fuel), (Π (d₂ : partial.fuel), d₂ < d₁ → Π (a : α), except unit (β a)) → Π (a : α), except unit (β a))
    (d : partial.fuel) (a : α) : except unit (β a) :=
well_founded.fix nat.lt_wf F d a

private def partial.fix_aux
    (F : Π (d₁ : partial.fuel), (Π (a : α), except unit (β a)) → Π (a : α), except unit (β a))
    (d : partial.fuel) (a : α) : except unit (β a) :=
partial.fix_core (λ d₁ g a₁, if h : d₁ = 0 then except.error () else F d₁ (λ a₂, g (nat.pred d₁) (nat.pred_lt h) a₂) a₁) d a

private def partial.pack (f : partial.fuel → Π (a : α), except unit (β a)) : Π (x : α), partial (β x) :=
λ x : α, {run_prv := λ d, f d x}

private def partial.unpack (F : Π (x : α), (Π y : α, partial (β y)) → partial (β x)) :
                           Π (d₁ : partial.fuel), (Π (a : α), except unit (β a)) → Π (a : α), except unit (β a) :=
λ d₁ g a, (F a (λ b, {run_prv := λ _, g b})).run_prv d₁

/- Compiler implements this function as
   fix F x = F x (fix F). -/
def partial.fix (F : Π (x : α), (Π y : α, partial (β y)) → partial (β x)) (x : α) : partial (β x) :=
partial.pack (partial.fix_aux (partial.unpack F)) x

def partial.run_with (a : partial α) (d : partial.fuel) : except unit α :=
a.run_prv d

/- Compiler implements this function as `some a` -/
def partial.run (a : partial α) : option α :=
(a.run_with partial.huge).to_option

/-
Example:
foo a := if a = 10 then return a else foo (a+1)
-/
def foo (a : nat) : partial nat :=
@partial.fix nat (λ _, nat) (λ a F, if a = 10 then return a else F (a+1)) a

#eval (foo 1).run  -- some 10
#eval (foo 20).run -- none

def find_core (f : nat → bool) (n : nat) : partial {x : nat // f x = tt} :=
@partial.fix nat (λ _, {x : nat // f x = tt})
  (λ x F, if h : f x = tt then return ⟨x, h⟩ else F (x+1))
  n

def find (f : nat → bool) : partial {x : nat // f x = tt} :=
find_core f 0

#eval (find (λ x, to_bool (x > 10 ∧ x % 6 = 1))).run -- (some 13)

def foo2 (a : nat) : partial {x : nat // x ≥ a ∧ x ≤ 10} :=
@partial.fix nat (λ a, {x : nat // x ≥ a ∧ x ≤ 10})
   (λ a' F,
      if h : a' ≤ 10 then return ⟨a', and.intro (nat.le_refl a') h⟩
      else do ⟨n, h₂⟩ ← F (a'+1),
              have aux₁ : n ≥ a' + 1, from h₂.1,
              have aux₂ : a' > 10, from nat.gt_of_not_le h,
              have aux₃ : a' + 1 > 10 + 1, from nat.add_lt_add_right aux₂ 1,
              have aux₄ : n > 10 + 1, from nat.lt_of_lt_of_le aux₃ aux₁,
              have aux₅ : n ≤ 10, from h₂.2,
              have aux₆ : n > 10, from nat.lt_of_succ_lt aux₄,
              have aux₇ : ¬ (n ≤ 10), from nat.not_le_of_gt aux₆,
              return ⟨n, absurd aux₅ aux₇⟩) a