open import Agda.Primitive using (lsuc)

module MonotonicState {s} {S : Set s}
                      (_≼_ : S → S → Set s)
                      (≼-trans : ∀ {s₁ s₂ s₃ : S} → s₁ ≼ s₂ → s₂ ≼ s₃ → s₁ ≼ s₃) where

open import Data.Product using (Σ; _×_; _,_)

open import Utils using (_⊗_; _,_)

-- Sometimes, we need a state monad whose values depend on the state. However,
-- one trouble with such monads is that as the state evolves, old values
-- in scope are over the 'old' state, and don't get updated accordingly.
-- Apparently, a related version of this problem is called 'demonic bind'.
--
-- One solution to the problem is to also witness some kind of relationtion
-- between the input and output states. Using this relationship makes it possible
-- to 'bring old values up to speed'.
--
-- Motivated primarily by constructing a Control Flow Graph, the 'relationship'
-- I've chosen is a 'less-than' relation. Thus, 'MonotonicState' is just
-- a (dependent) state "monad" that also witnesses that the state keeps growing.

MonotonicState : (S → Set s) → Set s
MonotonicState T = (s₁ : S) → Σ S (λ s₂ → T s₂ × s₁ ≼ s₂)

-- It's not a given that the (arbitrary) _≼_ relationship can be used for
-- updating old values. The Relaxable typeclass represents type constructor
-- that support the operation.

record Relaxable (T : S → Set s) : Set (lsuc s) where
    field relax : ∀ {s₁ s₂ : S} → s₁ ≼ s₂ → T s₁ → T s₂

instance
    ProdRelaxable : ∀ {P : S → Set s} {Q : S → Set s} →
                    {{ PRelaxable : Relaxable P }} → {{ QRelaxable : Relaxable Q }} →
                    Relaxable (P ⊗ Q)
    ProdRelaxable {{pr}} {{qr}} = record
        { relax = (λ { g₁≼g₂ (p , q) →
            ( Relaxable.relax pr g₁≼g₂ p
            , Relaxable.relax qr g₁≼g₂ q) }
            )
        }

-- In general, the "MonotonicState monad" is not even a monad; it's not
-- even applicative. The trouble is that functions in general cannot be
-- 'relaxed', and to apply an 'old' function to a 'new' value, you'd thus
-- need to un-relax the value (which also isn't possible in general).
--
-- However, we _can_ combine pairs from two functions into a tuple, which
-- would equivalent to the applicative operation if functions were relaxable.
--
-- TODO: Now that I think about it, the swapped version of the applicative
-- operation is possible, since it doesn't require lifting functions.

infixr 4 _⟨⊗⟩_
_⟨⊗⟩_ : ∀ {T₁ T₂ : S → Set s} {{ _ : Relaxable T₁ }} →
        MonotonicState T₁ → MonotonicState T₂ → MonotonicState (T₁ ⊗ T₂)
_⟨⊗⟩_ {{r}} f₁ f₂ s
    with (s' , (t₁ , s≼s')) ← f₁ s
    with (s'' , (t₂ , s'≼s'')) ← f₂ s' =
    (s'' , ((Relaxable.relax r s'≼s'' t₁ , t₂) , ≼-trans s≼s' s'≼s''))

infixl 4 _update_
_update_ : ∀ {T : S → Set s} {{ _ : Relaxable T }} →
         MonotonicState T → (∀ (s : S) → T s → Σ S (λ s' → s ≼ s')) →
         MonotonicState T
_update_ {{r}} f mod s
    with (s' , (t , s≼s')) ← f s
    with (s'' , s'≼s'') ← mod s' t =
    (s'' , ((Relaxable.relax r s'≼s'' t , ≼-trans s≼s' s'≼s'')))

infixl 4 _map_
_map_ : ∀ {T₁ T₂ : S → Set s} →
        MonotonicState T₁ → (∀ (s : S) → T₁ s → T₂ s) → MonotonicState T₂
_map_ f fn s = let (s' , (t₁ , s≼s')) = f s in (s' , (fn s' t₁ , s≼s'))


-- To reason about MonotonicState instances, we need predicates over their
-- values. But such values are dependent, so our predicates need to accept
-- the state as argument, too.

DependentPredicate : (S → Set s) → Set (lsuc s)
DependentPredicate T = ∀ (s₁ : S) → T s₁ → Set s

Both : {T₁ T₂ : S → Set s} → DependentPredicate T₁ → DependentPredicate T₂ →
       DependentPredicate (T₁ ⊗ T₂)
Both P Q = (λ { s (t₁ , t₂) → (P s t₁ × Q s t₂) })

-- Since monotnic functions keep adding on to the state, proofs of
-- predicates over their outputs go stale fast (they describe old values of
-- the state). To keep them relevant, we need them to still hold on 'bigger
-- states'. We call such predicates monotonic as well, since they respect the
-- ordering relation.

MonotonicPredicate : ∀ {T : S → Set s} {{ _ : Relaxable T }} →
                     DependentPredicate T → Set s
MonotonicPredicate {T} {{r}} P = ∀ (s₁ s₂ : S) (t₁ : T s₁) (s₁≼s₂ : s₁ ≼ s₂) →
                           P s₁ t₁ → P s₂ (Relaxable.relax r s₁≼s₂ t₁)

-- A MonotonicState "monad" m has a certain property if its ouputs satisfy that
-- property for all inputs.

always : ∀ {T : S → Set s} → DependentPredicate T → MonotonicState T → Set s
always P m = ∀ s₁ → let (s₂ , t , _) = m s₁ in P s₂ t

⟨⊗⟩-reason : ∀ {T₁ T₂ : S → Set s} {{ _ : Relaxable T₁ }}
             {P : DependentPredicate T₁} {Q : DependentPredicate T₂}
             {P-Mono : MonotonicPredicate P}
             {m₁ : MonotonicState T₁} {m₂ : MonotonicState T₂} →
             always P m₁ → always Q m₂ → always (Both P Q) (m₁ ⟨⊗⟩ m₂)
⟨⊗⟩-reason {P-Mono = P-Mono} {m₁ = m₁} {m₂ = m₂} aP aQ s
    with p ← aP s
    with (s' , (t₁ , s≼s')) ← m₁ s
    with q ← aQ s'
    with (s'' , (t₂ , s'≼s'')) ← m₂ s' = (P-Mono _ _ _ s'≼s'' p , q)