------------------------------------------------------------------------
-- Fixpoint combinators
------------------------------------------------------------------------

{-# OPTIONS --without-K #-}

module Partiality-monad.Inductive.Fixpoints where

open import Equality.Propositional
open import Interval using (ext; ⟨ext⟩; bad-ext)
open import Logical-equivalence using (_⇔_)
open import Prelude hiding (; head; tail)

open import Bijection equality-with-J using (_↔_)
import Equivalence equality-with-J as Eq
open import Function-universe equality-with-J hiding (id; _∘_)
open import H-level equality-with-J
open import H-level.Closure equality-with-J
open import Monad equality-with-J
open import Univalence-axiom equality-with-J

open import Partiality-algebra using (Partiality-algebra)
import Partiality-algebra.Fixpoints as PAF
open import Partiality-algebra.Monotone
open import Partiality-algebra.Omega-continuous
open import Partiality-algebra.Pi as Pi hiding (at)
open import Partiality-monad.Inductive
open import Partiality-monad.Inductive.Monad
open import Partiality-monad.Inductive.Monotone using ([_⊥→_⊥]⊑)
open import Partiality-monad.Inductive.Omega-continuous

------------------------------------------------------------------------
-- The fixpoint combinator machinery instantiated with the partiality
-- monad

module _ {a} {A : Set a} where

  open PAF.Fix {P = partiality-algebra A} public
  open [_⊥→_⊥]

  -- A variant of fix.

  fix′ : (A  A )  A 
  fix′ f = fix (monotone-function (f ))

  -- This variant also produces a kind of least fixpoint.

  fix′-is-fixpoint-combinator :
    (f : A  A ) 
    fix′ f  fix′ f >>= f
  fix′-is-fixpoint-combinator f =
    fix′ f                               ≡⟨⟩
    fix (monotone-function (f ))        ≡⟨ fix-is-fixpoint-combinator (f ) ⟩∎
    fix (monotone-function (f )) >>= f  

  fix′-is-least :
    (f : A  A ) 
     x  x >>= f  x  fix′ f  x
  fix′-is-least f = fix-is-least (monotone-function (f ))

  -- Taking steps that are "twice as large" does not affect the end
  -- result.

  fix′-∘ : (f : A  A )  fix′  x  f x >>= f)  fix′ f
  fix′-∘ f =
    fix′  x  f x >>= f)                                    ≡⟨⟩

    fix (monotone-function ((λ x  f x >>= f) ))             ≡⟨ cong fix (_↔_.to equality-characterisation-monotone λ x 

        (x >>= λ y  f y >>= f)                                    ≡⟨ associativity x f f ⟩∎
        x >>= f >>= f                                              ) 

    fix (monotone-function (f ) ∘⊑ monotone-function (f ))  ≡⟨ fix-∘ (monotone-function (f )) ⟩∎

    fix′ f                                                    

  -- Unary Scott induction for fix′.

  fix′-induction₁ :
     {p}
    (P : A   Set p)
    (P⊥ : P never)
    (P⨆ :  s  (∀ n  P (s [ n ]))  P ( s))
    (f : A  A ) 
    (∀ x  P x  P (x >>= f)) 
    P (fix′ f)
  fix′-induction₁ P P⊥ P⨆ f =
    fix-induction₁ P P⊥ P⨆ ([_⊥→_⊥].monotone-function (f ))

-- N-ary Scott induction.

module _ {a p} n
  (As : Fin n  Set a)
  (P : (∀ i  As i )  Set p)
  (P⊥ : P  _  never))
  (P⨆ :  ss  (∀ n  P  i  ss i [ n ]))  P (  ss))
  where

  open PAF.N-ary n As  i  partiality-algebra (As i)) P P⊥ P⨆ public

  fix′-induction :
    (fs :  i  As i  As i ) 
    (∀ xs  P xs  P  i  xs i >>= fs i)) 
    P (fix′  fs)
  fix′-induction fs =
    fix-induction  i  [_⊥→_⊥].monotone-function (fs i ))

------------------------------------------------------------------------
-- The fixpoint combinator machinery instantiated with dependent
-- functions to the partiality monad

module _ {a b} (A : Set a) (B : A  Set b) where

  -- Partial function transformers.

  Trans : Set (a  b)
  Trans = ((x : A)  B x )  ((x : A)  B x )

  -- A partiality algebra.

  Pi : Partiality-algebra (a  b) (a  b) ((x : A)  B x)
  Pi = Π A (partiality-algebra  B)

  -- Monotone partial function transformers.

  Trans-⊑ : Set (a  b)
  Trans-⊑ = [ Pi  Pi ]⊑

  private

    sanity-check : Trans-⊑  Trans
    sanity-check = [_⟶_]⊑.function

  -- Partial function transformers that are ω-continuous.

  Trans-ω : Set (a  b)
  Trans-ω = [ Pi  Pi ]

module _ {a b} {A : Set a} {B : A  Set b} where

  module Trans-⊑ = [_⟶_]⊑ {P₁ = Pi A B} {P₂ = Pi A B}
  module Trans-ω = [_⟶_]  {P₁ = Pi A B} {P₂ = Pi A B}

  private
    module F = PAF.Fix {P = Pi A B}

  -- Applies an increasing sequence of functions to a value.

  at :
    ( λ (f :   (x : A)  B x )   n x  f n x  f (suc n) x) 
    (x : A)   λ (f :   B x )   n  f n  f (suc n)
  at = Pi.at (partiality-algebra  B)

  -- Repeated application of a monotone partial function transformer
  -- to the least element.

  app→ : Trans-⊑ A B    ((x : A)  B x )
  app→ = F.app

  -- The increasing sequence fix→-sequence f consists of the elements
  -- const never, Trans-⊑.function f (const never), and so on. This
  -- sequence is used in the implementation of fix→.

  fix→-sequence : Trans-⊑ A B 
                  Partiality-algebra.Increasing-sequence (Pi A B)
  fix→-sequence = F.fix-sequence

  -- A fixpoint combinator.

  fix→ : Trans-⊑ A B  ((x : A)  B x )
  fix→ = F.fix

  -- The fixpoint combinator produces fixpoints for ω-continuous
  -- arguments.

  fix→-is-fixpoint-combinator :
    (f : Trans-ω A B) 
    fix→ (Trans-ω.monotone-function f) 
    Trans-ω.function f (fix→ (Trans-ω.monotone-function f))
  fix→-is-fixpoint-combinator = F.fix-is-fixpoint-combinator

  -- The result of the fixpoint combinator is pointwise smaller than
  -- or equal to every "pointwise post-fixpoint".

  fix→-is-least :
    (f : Trans-⊑ A B) 
     g  (∀ x  Trans-⊑.function f g x  g x) 
       x  fix→ f x  g x
  fix→-is-least = F.fix-is-least

  -- Taking steps that are "twice as large" does not affect the end
  -- result.

  fix→-∘ : (f : Trans-⊑ A B)  fix→ (f ∘⊑ f)  fix→ f
  fix→-∘ = F.fix-∘

-- N-ary Scott induction.

module _
  {a b p} n
  (As : Fin n  Set a)
  (Bs : (i : Fin n)  As i  Set b)
  (P : (∀ i  (x : As i)  Bs i x )  Set p)
  (P⊥ : P  _ _  never))
  (P⨆ : (ss :  i (x : As i)  Increasing-sequence (Bs i x)) 
        (∀ n  P  i xs  ss i xs [ n ])) 
        P  i xs   (ss i xs)))
  (fs :  i  Trans-⊑ (As i) (Bs i)) where

  private
    module N = PAF.N-ary n
                  i  (x : As i)  Bs i x)
                  i  Π (As i) (partiality-algebra  Bs i))
                 P
                 P⊥
                  _  P⨆ _)
                 fs

  -- Generalised.

  fix→-induction′ :
    (∀ n  P  i  app→ (fs i) n) 
           P  i  app→ (fs i) (suc n))) 
    P  i  fix→ (fs i))
  fix→-induction′ = N.fix-induction′

  -- Basic.

  fix→-induction :
    ((gs :  i (x : As i)  Bs i x ) 
     P gs  P  i  Trans-⊑.function (fs i) (gs i))) 
    P  i  fix→ (fs i))
  fix→-induction = N.fix-induction

-- Unary Scott induction.

fix→-induction₁ :
   {a b p} {A : Set a} {B : A  Set b}
  (P : ((x : A)  B x )  Set p) 
  P  _  never) 
  ((s : (x : A)  Increasing-sequence (B x)) 
   (∀ n  P  x  s x [ n ])) 
   P  x   (s x))) 
  (f : Trans-⊑ A B) 
  (∀ g  P g  P (Trans-⊑.function f g)) 
  P (fix→ f)
fix→-induction₁ P P⊥ P⨆ = PAF.fix-induction₁ P P⊥  _  P⨆ _)

------------------------------------------------------------------------
-- Some combinators that can be used to construct ω-continuous partial
-- function transformers, for use with fix→

-- A type used by these combinators.

record Partial
         {a b c}
         (A : Set a) (B : A  Set b) (C : Set c) :
         Set (a  b  lsuc c) where
  field
    -- A function that is allowed to make "recursive calls" of type
    -- (x : A) → B x ⊥.

    function : ((x : A)  B x )  C 

    -- The function must be monotone.

    monotone :  {rec₁ rec₂} 
               (∀ x  rec₁ x  rec₂ x) 
               function rec₁  function rec₂

  -- The function can be turned into an increasing sequence.

  sequence : ((x : A)  Increasing-sequence (B x)) 
             Increasing-sequence C
  sequence recs =
       n  function  x  recs x [ n ]))
    ,  n  monotone  x  increasing (recs x) n))

  field
    -- The function must be ω-continuous in the following sense.
    --
    -- The proof can make use of propositional extensionality. This
    -- assumption is included so that the monad instance can be
    -- defined without a corresponding assumption.

    ω-continuous : Propositional-extensionality c 
                   (recs : (x : A)  Increasing-sequence (B x)) 
                   function (  recs)   (sequence recs)

-- The first two arguments of Partial specify the domain and codomain
-- of "recursive calls".

rec :  {a b} {A : Set a} {B : A  Set b} 
      (x : A)  Partial A B (B x)
rec x = record
  { function     = _$ x
  ; monotone     = _$ x
  ; ω-continuous = λ _ _  refl
  }

-- Turns certain Partial-valued functions into monotone partial
-- function transformers.

transformer :  {a b} {A : Set a} {B : A  Set b} 
              ((x : A)  Partial A B (B x))  Trans-⊑ A B
transformer f = record
  { function = λ g     x  function (f x) g
  ; monotone = λ g₁⊑g₂ x  monotone (f x) g₁⊑g₂
  }
  where
  open Partial

-- Turns certain Partial-valued functions into ω-continuous partial
-- function transformers (assuming propositional extensionality).

transformer-ω :  {a b} {A : Set a} {B : A  Set b} 
                Propositional-extensionality b 
                ((x : A)  Partial A B (B x))  Trans-ω A B
transformer-ω prop-ext f = record
  { monotone-function = transformer f
  ; ω-continuous      = λ _  ⟨ext⟩ λ x  ω-continuous (f x) prop-ext _
  }
  where
  open Partial

-- A fixpoint combinator.

fixP :  {a b} {A : Set a} {B : A  Set b} 
       ((x : A)  Partial A B (B x))  ((x : A)  B x )
fixP {A = A} {B} =
  ((x : A)  Partial A B (B x))  ↝⟨ transformer 
  Trans-⊑ A B                    ↝⟨ fix→ ⟩□
  ((x : A)  B x )              

-- The fixpoint combinator produces fixpoints (assuming propositional
-- extensionality).

fixP-is-fixpoint-combinator :
   {a b} {A : Set a} {B : A  Set b} 
  Propositional-extensionality b 
  (f : (x : A)  Partial A B (B x)) 
  fixP f  flip (Partial.function  f) (fixP f)
fixP-is-fixpoint-combinator prop-ext =
  fix→-is-fixpoint-combinator  transformer-ω prop-ext

-- The result of the fixpoint combinator is pointwise smaller than
-- or equal to every "pointwise post-fixpoint".

fixP-is-least :
   {a b} {A : Set a} {B : A  Set b}
  (f : (x : A)  Partial A B (B x)) 
   g  (∀ x  Partial.function (f x) g  g x) 
     x  fixP f x  g x
fixP-is-least = fix→-is-least  transformer

-- Composition of certain Partial-valued functions.

infixr 40 _∘P_

_∘P_ :
   {a b} {A : Set a} {B : A  Set b} 
  ((x : A)  Partial A B (B x)) 
  ((x : A)  Partial A B (B x)) 
  ((x : A)  Partial A B (B x))
Partial.function     ((f ∘P g) x) = λ h 
                                      Partial.function (f x) λ y 
                                      Partial.function (g y) h
Partial.monotone     ((f ∘P g) x) = λ hyp 
                                      Partial.monotone (f x) λ y 
                                      Partial.monotone (g y) hyp
Partial.ω-continuous ((f ∘P g) x) = λ prop-ext recs 
  function (f x)  y  function (g y) (  recs))  ≡⟨ cong (function (f x)) (⟨ext⟩ λ y  ω-continuous (g y) prop-ext recs) 
  function (f x)  y   (sequence (g y) recs))    ≡⟨ ω-continuous (f x) prop-ext  y  sequence (g y) recs) ⟩∎
   (sequence (f x) λ y  sequence (g y) recs)      
  where
  open Partial

-- Taking steps that are "twice as large" does not affect the end
-- result.

fixP-∘ :  {a b} {A : Set a} {B : A  Set b}
         (f : (x : A)  Partial A B (B x)) 
         fixP (f ∘P f)  fixP f
fixP-∘ f =
  fix→ (transformer (f ∘P f))            ≡⟨⟩
  fix→ (transformer f ∘⊑ transformer f)  ≡⟨ fix→-∘ (transformer f) ⟩∎
  fix→ (transformer f)                   

-- N-ary Scott induction.

module _
  {a b p} n
  (As : Fin n  Set a)
  (Bs : (i : Fin n)  As i  Set b)
  (P : (∀ i (x : As i)  Bs i x )  Set p)
  (P⊥ : P  _ _  never))
  (P⨆ : (ss :  i (x : As i)  Increasing-sequence (Bs i x)) 
        (∀ n  P  i xs  ss i xs [ n ])) 
        P  i xs   (ss i xs)))
  (fs :  i (x : As i)  Partial (As i) (Bs i) (Bs i x)) where

  -- Generalised.

  fixP-induction′ :
    (∀ n  P  i  app→ (transformer (fs i)) n) 
           P  i  app→ (transformer (fs i)) (suc n))) 
    P  i  fixP (fs i))
  fixP-induction′ =
    fix→-induction′ n As Bs P P⊥ P⨆ (transformer  fs)

  -- Basic.

  fixP-induction :
    ((gs :  i (x : As i)  Bs i x ) 
     P gs  P  i xs  Partial.function (fs i xs) (gs i))) 
    P  i  fixP (fs i))
  fixP-induction =
    fix→-induction n As Bs P P⊥ P⨆ (transformer  fs)

-- Unary Scott induction.

fixP-induction₁ :
   {a b p} {A : Set a} {B : A  Set b}
  (P : ((x : A)  B x )  Set p) 
  P  _  never) 
  ((s : (x : A)  Increasing-sequence (B x)) 
   (∀ n  P  x  s x [ n ])) 
   P  x   (s x))) 
  (f : (x : A)  Partial A B (B x)) 
  (∀ g  P g  P  x  Partial.function (f x) g)) 
  P (fixP f)
fixP-induction₁ P P⊥ P⨆ f =
  fix→-induction₁ P P⊥ P⨆ (transformer f)

-- Equality characterisation lemma for Partial.

equality-characterisation-Partial :
   {a b c} {A : Set a} {B : A  Set b} {C : Set c}
    {f g : Partial A B C} 
  (∀ rec  Partial.function f rec  Partial.function g rec) 
  f  g
equality-characterisation-Partial {f = f} {g} =
  (∀ rec  function f rec  function g rec)  ↔⟨ Eq.extensionality-isomorphism bad-ext 
  function f  function g                    ↝⟨ ignore-propositional-component
                                                  (Σ-closure 1
                                                     (implicit-Π-closure ext 1 λ _ 
                                                      implicit-Π-closure ext 1 λ _ 
                                                      Π-closure          ext 1 λ _ 
                                                      ⊑-propositional) λ _ 
                                                     Π-closure ext 1 λ _ 
                                                     Π-closure ext 1 λ _ 
                                                     ⊥-is-set _ _) 
  (function f , _)  (function g , _)        ↔⟨ Eq.≃-≡ (Eq.↔⇒≃ lemma) ⟩□
  f  g                                      
  where
  open Partial

  lemma : Partial _ _ _
            
           λ _   λ (_ :  {rec₁ rec₂}  _  _)  _
  lemma = record
    { surjection = record
      { logical-equivalence = record
        { to   = λ x  function x , monotone x , ω-continuous x
        ; from = λ { (f , m , ω)  record
                      { function     = f
                      ; monotone     = λ {rec₁ rec₂}  m {rec₁ = rec₁} {rec₂ = rec₂}
                      ; ω-continuous = ω
                      } }
        }
      ; right-inverse-of = λ _  refl
      }
    ; left-inverse-of = λ _  refl
    }

-- "Non-recursive" computations can be converted into possibly
-- recursive ones.

non-recursive :
   {a b c} {A : Set a} {B : A  Set b} {C : Set c} 
  C   Partial A B C
non-recursive x = record
  { function     = const x
  ; monotone     = const (x )
  ; ω-continuous = λ _ _ 
      x             ≡⟨ sym ⨆-const ⟩∎
       (constˢ x)  
  }

instance

  -- Partial A B is a monad.

  partial-raw-monad :  {a b c} {A : Set a} {B : A  Set b} 
                      Raw-monad (Partial {c = c} A B)
  Raw-monad.return partial-raw-monad x   = non-recursive (return x)
  Raw-monad._>>=_  partial-raw-monad x f = record
    { function     = λ rec 
                       function x rec >>=′ λ y 
                       function (f y) rec
    ; monotone     = λ rec⊑rec 
                       monotone x rec⊑rec >>=-mono λ y 
                       monotone (f y) rec⊑rec
    ; ω-continuous = λ prop-ext recs 
        (function x (  recs) >>=′ λ y  function (f y) (  recs))     ≡⟨ cong₂ _>>=′_ (ω-continuous x prop-ext recs)
                                                                                         (⟨ext⟩ λ y  ω-continuous (f y) prop-ext recs) 
        ( (sequence x recs) >>=′ λ y   (sequence (f y) recs))         ≡⟨ ⨆->>= 

         (  n 
               function x  z  recs z [ n ]) >>=′ λ y 
                (sequence (f y) recs))
          , _
          )                                                              ≡⟨ ⨆>>=⨆≡⨆>>= prop-ext prop-ext
                                                                                       (sequence x recs)
                                                                                        y  sequence (f y) recs) ⟩∎
         (  n  function x  z  recs z [ n ]) >>=′ λ y 
                   function (f y)  z  recs z [ n ]))
          , _
          )                                                              
    }
    where
    open Partial

  partial-monad :  {a b c} {A : Set a} {B : A  Set b} 
                  Monad (Partial {c = c} A B)
  Monad.raw-monad partial-monad = partial-raw-monad

  Monad.left-identity partial-monad _ f =
    _↔_.to equality-characterisation-Partial
       h  left-identity _  y  Partial.function (f y) h))

  Monad.right-identity partial-monad _ =
    _↔_.to equality-characterisation-Partial
       _  right-identity _)

  Monad.associativity partial-monad x _ _ =
    _↔_.to equality-characterisation-Partial
       h  associativity (Partial.function x h) _ _)