{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE QualifiedDo #-}

module Control.Monad.Indexed.Cont
  ( -- * Abstract delimited control
    Shifty (..),
    Stacked (..),
    stack,
    pop,
    pop_,
    push,
    (@),

    -- * Comonad-to-indexed-monad transformer
    ContW (..),
    shift0,
    yield,
    yield_,

    -- * Experimental combinators (subject to radical change)
    abort,
    capture,
    handle,
    pullback,
  )
where

import Control.Comonad
import Control.Comonad.Store
import Control.Monad.Indexed ((*>))
import Control.Monad.Indexed qualified as Indexed
import Prelude hiding (Applicative (..), Monad (..), MonadFail (..))
import Prelude qualified

newtype ContW w r r' a = ContW {forall (w :: * -> *) r r' a. ContW w r r' a -> w (a -> r') -> r
runContW :: w (a -> r') -> r}
  deriving stock ((forall a b. (a -> b) -> ContW w r r' a -> ContW w r r' b)
-> (forall a b. a -> ContW w r r' b -> ContW w r r' a)
-> Functor (ContW w r r')
forall a b. a -> ContW w r r' b -> ContW w r r' a
forall a b. (a -> b) -> ContW w r r' a -> ContW w r r' b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
forall (w :: * -> *) r r' a b.
Functor w =>
a -> ContW w r r' b -> ContW w r r' a
forall (w :: * -> *) r r' a b.
Functor w =>
(a -> b) -> ContW w r r' a -> ContW w r r' b
$cfmap :: forall (w :: * -> *) r r' a b.
Functor w =>
(a -> b) -> ContW w r r' a -> ContW w r r' b
fmap :: forall a b. (a -> b) -> ContW w r r' a -> ContW w r r' b
$c<$ :: forall (w :: * -> *) r r' a b.
Functor w =>
a -> ContW w r r' b -> ContW w r r' a
<$ :: forall a b. a -> ContW w r r' b -> ContW w r r' a
Functor)

deriving via (Indexed.FromIndexed (ContW w) r r) instance (Comonad w) => Prelude.Applicative (ContW w r r)

deriving via (Indexed.FromIndexed (ContW w) r r) instance (Comonad w) => Prelude.Monad (ContW w r r)

instance (Comonad w) => Indexed.Applicative (ContW w) where
  pure :: forall a i. a -> ContW w i i a
pure a
x = (w (a -> i) -> i) -> ContW w i i a
forall (w :: * -> *) r r' a. (w (a -> r') -> r) -> ContW w r r' a
ContW ((w (a -> i) -> i) -> ContW w i i a)
-> (w (a -> i) -> i) -> ContW w i i a
forall a b. (a -> b) -> a -> b
$ \w (a -> i)
k -> w (a -> i) -> a -> i
forall a. w a -> a
forall (w :: * -> *) a. Comonad w => w a -> a
extract w (a -> i)
k a
x

instance (Comonad w) => Indexed.Monad (ContW w) where
  ContW w (a -> j) -> i
a >>= :: forall i j a k1 b.
ContW w i j a -> (a -> ContW w j k1 b) -> ContW w i k1 b
>>= a -> ContW w j k1 b
f = (w (b -> k1) -> i) -> ContW w i k1 b
forall (w :: * -> *) r r' a. (w (a -> r') -> r) -> ContW w r r' a
ContW ((w (b -> k1) -> i) -> ContW w i k1 b)
-> (w (b -> k1) -> i) -> ContW w i k1 b
forall a b. (a -> b) -> a -> b
$ \w (b -> k1)
wk -> w (a -> j) -> i
a (w (b -> k1) -> a -> j
k' (w (b -> k1) -> a -> j) -> w (b -> k1) -> w (a -> j)
forall a b. (w a -> b) -> w a -> w b
forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
`extend` w (b -> k1)
wk)
    where
      k' :: w (b -> k1) -> a -> j
k' w (b -> k1)
wk a
x = ContW w j k1 b -> w (b -> k1) -> j
forall (w :: * -> *) r r' a. ContW w r r' a -> w (a -> r') -> r
runContW (a -> ContW w j k1 b
f a
x) w (b -> k1)
wk

shift0 :: (Comonad w) => (w (a -> r') -> ContW w r k k) -> ContW w r r' a
shift0 :: forall (w :: * -> *) a r' r k.
Comonad w =>
(w (a -> r') -> ContW w r k k) -> ContW w r r' a
shift0 w (a -> r') -> ContW w r k k
f = (w (a -> r') -> r) -> ContW w r r' a
forall (w :: * -> *) r r' a. (w (a -> r') -> r) -> ContW w r r' a
ContW ((w (a -> r') -> r) -> ContW w r r' a)
-> (w (a -> r') -> r) -> ContW w r r' a
forall a b. (a -> b) -> a -> b
$ \w (a -> r')
wk -> ContW w r k k -> w (k -> k) -> r
forall (w :: * -> *) r r' a. ContW w r r' a -> w (a -> r') -> r
runContW (w (a -> r') -> ContW w r k k
f w (a -> r')
wk) ((k -> k) -> (a -> r') -> k -> k
forall a b. a -> b -> a
const k -> k
forall a. a -> a
id ((a -> r') -> k -> k) -> w (a -> r') -> w (k -> k)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> w (a -> r')
wk)

yield :: (Comonad w) => (w (a -> r) -> r) -> ContW w r r a
yield :: forall (w :: * -> *) a r.
Comonad w =>
(w (a -> r) -> r) -> ContW w r r a
yield w (a -> r) -> r
act = (w (a -> r) -> ContW w r r r) -> ContW w r r a
forall (w :: * -> *) a r' r k.
Comonad w =>
(w (a -> r') -> ContW w r k k) -> ContW w r r' a
shift0 ((w (a -> r) -> ContW w r r r) -> ContW w r r a)
-> (w (a -> r) -> ContW w r r r) -> ContW w r r a
forall a b. (a -> b) -> a -> b
$ r -> ContW w r r r
forall a i. a -> ContW w i i a
forall {k} (f :: k -> k -> * -> *) a (i :: k).
Applicative f =>
a -> f i i a
Indexed.pure (r -> ContW w r r r)
-> (w (a -> r) -> r) -> w (a -> r) -> ContW w r r r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. w (a -> r) -> r
act

yield_ :: (Comonad w) => (w r -> r) -> ContW w r r ()
yield_ :: forall (w :: * -> *) r. Comonad w => (w r -> r) -> ContW w r r ()
yield_ w r -> r
act = (w (() -> r) -> ContW w r r r) -> ContW w r r ()
forall (w :: * -> *) a r' r k.
Comonad w =>
(w (a -> r') -> ContW w r k k) -> ContW w r r' a
shift0 ((w (() -> r) -> ContW w r r r) -> ContW w r r ())
-> (w (() -> r) -> ContW w r r r) -> ContW w r r ()
forall a b. (a -> b) -> a -> b
$ \w (() -> r)
wk -> r -> ContW w r r r
forall a i. a -> ContW w i i a
forall {k} (f :: k -> k -> * -> *) a (i :: k).
Applicative f =>
a -> f i i a
Indexed.pure (r -> ContW w r r r) -> r -> ContW w r r r
forall a b. (a -> b) -> a -> b
$ w r -> r
act (((() -> r) -> () -> r
forall a b. (a -> b) -> a -> b
$ ()) ((() -> r) -> r) -> w (() -> r) -> w r
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> w (() -> r)
wk)

class (Stacked m) => Shifty m where
  shift :: ((a -> r') -> m r k k) -> m r r' a

instance (Comonad w) => Shifty (ContW w) where
  shift :: forall a r' r k. ((a -> r') -> ContW w r k k) -> ContW w r r' a
shift (a -> r') -> ContW w r k k
f = (w (a -> r') -> ContW w r k k) -> ContW w r r' a
forall (w :: * -> *) a r' r k.
Comonad w =>
(w (a -> r') -> ContW w r k k) -> ContW w r r' a
shift0 ((w (a -> r') -> ContW w r k k) -> ContW w r r' a)
-> (w (a -> r') -> ContW w r k k) -> ContW w r r' a
forall a b. (a -> b) -> a -> b
$ \w (a -> r')
wk -> (a -> r') -> ContW w r k k
f (w (a -> r') -> a -> r'
forall a. w a -> a
forall (w :: * -> *) a. Comonad w => w a -> a
extract w (a -> r')
wk)

class Stacked m where
  shift_ :: (r' -> m r r'' r'') -> m r r' ()

instance (Comonad w) => Stacked (ContW w) where
  shift_ :: forall r' r r''. (r' -> ContW w r r'' r'') -> ContW w r r' ()
shift_ r' -> ContW w r r'' r''
f = ((() -> r') -> ContW w r r'' r'') -> ContW w r r' ()
forall a r' r k. ((a -> r') -> ContW w r k k) -> ContW w r r' a
forall {k} (m :: k -> * -> * -> *) a r' (r :: k) k.
Shifty m =>
((a -> r') -> m r k k) -> m r r' a
shift (\() -> r'
k -> r' -> ContW w r r'' r''
f (() -> r'
k ()))

stack :: (Indexed.Applicative m, Stacked m) => (j -> i) -> m i j ()
stack :: forall (m :: * -> * -> * -> *) j i.
(Applicative m, Stacked m) =>
(j -> i) -> m i j ()
stack j -> i
f = (j -> m i i i) -> m i j ()
forall r' r r''. (r' -> m r r'' r'') -> m r r' ()
forall {k} (m :: k -> * -> * -> *) r' (r :: k) r''.
Stacked m =>
(r' -> m r r'' r'') -> m r r' ()
shift_ ((j -> m i i i) -> m i j ()) -> (j -> m i i i) -> m i j ()
forall a b. (a -> b) -> a -> b
$ \j
k -> i -> m i i i
forall a i. a -> m i i a
forall {k} (f :: k -> k -> * -> *) a (i :: k).
Applicative f =>
a -> f i i a
Indexed.pure (i -> m i i i) -> i -> m i i i
forall a b. (a -> b) -> a -> b
$ j -> i
f j
k

(@) :: (Indexed.Applicative m, Stacked m) => m (a -> i) j b -> a -> m i j b
m (a -> i) j b
act @ :: forall (m :: * -> * -> * -> *) a i j b.
(Applicative m, Stacked m) =>
m (a -> i) j b -> a -> m i j b
@ a
a = a -> m i (a -> i) ()
forall (m :: * -> * -> * -> *) a i.
(Applicative m, Stacked m) =>
a -> m i (a -> i) ()
push a
a m i (a -> i) () -> m (a -> i) j b -> m i j b
forall i j a k1 b. m i j a -> m j k1 b -> m i k1 b
forall {k} (f :: k -> k -> * -> *) (i :: k) (j :: k) a (k1 :: k) b.
Applicative f =>
f i j a -> f j k1 b -> f i k1 b
*> m (a -> i) j b
act

infixl 9 @

pop :: (Indexed.Applicative m, Shifty m) => m (a -> i) i a
pop :: forall (m :: * -> * -> * -> *) a i.
(Applicative m, Shifty m) =>
m (a -> i) i a
pop = ((a -> i) -> m (a -> i) (a -> i) (a -> i)) -> m (a -> i) i a
forall a r' r k. ((a -> r') -> m r k k) -> m r r' a
forall {k} (m :: k -> * -> * -> *) a r' (r :: k) k.
Shifty m =>
((a -> r') -> m r k k) -> m r r' a
shift (((a -> i) -> m (a -> i) (a -> i) (a -> i)) -> m (a -> i) i a)
-> ((a -> i) -> m (a -> i) (a -> i) (a -> i)) -> m (a -> i) i a
forall a b. (a -> b) -> a -> b
$ \a -> i
k -> (a -> i) -> m (a -> i) (a -> i) (a -> i)
forall a i. a -> m i i a
forall {k} (f :: k -> k -> * -> *) a (i :: k).
Applicative f =>
a -> f i i a
Indexed.pure (\a
a -> a -> i
k a
a)

push :: (Indexed.Applicative m, Stacked m) => a -> m i (a -> i) ()
push :: forall (m :: * -> * -> * -> *) a i.
(Applicative m, Stacked m) =>
a -> m i (a -> i) ()
push a
a = ((a -> i) -> m i i i) -> m i (a -> i) ()
forall r' r r''. (r' -> m r r'' r'') -> m r r' ()
forall {k} (m :: k -> * -> * -> *) r' (r :: k) r''.
Stacked m =>
(r' -> m r r'' r'') -> m r r' ()
shift_ (((a -> i) -> m i i i) -> m i (a -> i) ())
-> ((a -> i) -> m i i i) -> m i (a -> i) ()
forall a b. (a -> b) -> a -> b
$ \a -> i
k -> i -> m i i i
forall a i. a -> m i i a
forall {k} (f :: k -> k -> * -> *) a (i :: k).
Applicative f =>
a -> f i i a
Indexed.pure (a -> i
k a
a)

pop_ :: (Indexed.Applicative m, Stacked m) => m (a -> i) i ()
pop_ :: forall (m :: * -> * -> * -> *) a i.
(Applicative m, Stacked m) =>
m (a -> i) i ()
pop_ = (i -> m (a -> i) (a -> i) (a -> i)) -> m (a -> i) i ()
forall r' r r''. (r' -> m r r'' r'') -> m r r' ()
forall {k} (m :: k -> * -> * -> *) r' (r :: k) r''.
Stacked m =>
(r' -> m r r'' r'') -> m r r' ()
shift_ ((i -> m (a -> i) (a -> i) (a -> i)) -> m (a -> i) i ())
-> (i -> m (a -> i) (a -> i) (a -> i)) -> m (a -> i) i ()
forall a b. (a -> b) -> a -> b
$ \i
k -> (a -> i) -> m (a -> i) (a -> i) (a -> i)
forall a i. a -> m i i a
forall {k} (f :: k -> k -> * -> *) a (i :: k).
Applicative f =>
a -> f i i a
Indexed.pure (\a
_a -> i
k)

----------------------------------------------------------------------------------------
--
-- Experimental
--
----------------------------------------------------------------------------------------

abort :: (Indexed.Applicative m, Shifty m) => r -> m r r' a
abort :: forall (m :: * -> * -> * -> *) r r' a.
(Applicative m, Shifty m) =>
r -> m r r' a
abort r
a = ((a -> r') -> m r r r) -> m r r' a
forall a r' r k. ((a -> r') -> m r k k) -> m r r' a
forall {k} (m :: k -> * -> * -> *) a r' (r :: k) k.
Shifty m =>
((a -> r') -> m r k k) -> m r r' a
shift (((a -> r') -> m r r r) -> m r r' a)
-> ((a -> r') -> m r r r) -> m r r' a
forall a b. (a -> b) -> a -> b
$ \a -> r'
_ -> r -> m r r r
forall a i. a -> m i i a
forall {k} (f :: k -> k -> * -> *) a (i :: k).
Applicative f =>
a -> f i i a
Indexed.pure r
a

capture :: (Comonad w) => ContW w b r' r' -> ContW w r r (w b)
capture :: forall (w :: * -> *) b r' r.
Comonad w =>
ContW w b r' r' -> ContW w r r (w b)
capture (ContW w (r' -> r') -> b
a) = (w (w b -> r) -> r) -> ContW w r r (w b)
forall (w :: * -> *) r r' a. (w (a -> r') -> r) -> ContW w r r' a
ContW ((w (w b -> r) -> r) -> ContW w r r (w b))
-> (w (w b -> r) -> r) -> ContW w r r (w b)
forall a b. (a -> b) -> a -> b
$ \w (w b -> r)
wk -> w (w b -> r) -> w b -> r
forall a. w a -> a
forall (w :: * -> *) a. Comonad w => w a -> a
extract w (w b -> r)
wk (w (r' -> r') -> b
a (w (r' -> r') -> b) -> w (r' -> r') -> w b
forall a b. (w a -> b) -> w a -> w b
forall (w :: * -> *) a b. Comonad w => (w a -> b) -> w a -> w b
`extend` ((r' -> r') -> (w b -> r) -> r' -> r'
forall a b. a -> b -> a
const r' -> r'
forall a. a -> a
id ((w b -> r) -> r' -> r') -> w (w b -> r) -> w (r' -> r')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> w (w b -> r)
wk))

handle :: (Comonad w) => ContW (StoreT k w) r r' a -> ContW w k r' a -> ContW w r r' a
handle :: forall (w :: * -> *) k r r' a.
Comonad w =>
ContW (StoreT k w) r r' a -> ContW w k r' a -> ContW w r r' a
handle (ContW StoreT k w (a -> r') -> r
inner) (ContW w (a -> r') -> k
handler) =
  (w (a -> r') -> r) -> ContW w r r' a
forall (w :: * -> *) r r' a. (w (a -> r') -> r) -> ContW w r r' a
ContW ((w (a -> r') -> r) -> ContW w r r' a)
-> (w (a -> r') -> r) -> ContW w r r' a
forall a b. (a -> b) -> a -> b
$ \w (a -> r')
wk -> StoreT k w (a -> r') -> r
inner (w (k -> a -> r') -> k -> StoreT k w (a -> r')
forall s (w :: * -> *) a. w (s -> a) -> s -> StoreT s w a
StoreT ((a -> r') -> k -> a -> r'
forall a b. a -> b -> a
const ((a -> r') -> k -> a -> r') -> w (a -> r') -> w (k -> a -> r')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> w (a -> r')
wk) (w (a -> r') -> k
handler w (a -> r')
wk))

pullback :: (Comonad w) => (forall x. w x -> v x) -> ContW v r r' a -> ContW w r r' a
pullback :: forall (w :: * -> *) (v :: * -> *) r r' a.
Comonad w =>
(forall x. w x -> v x) -> ContW v r r' a -> ContW w r r' a
pullback forall x. w x -> v x
f (ContW v (a -> r') -> r
a) = (w (a -> r') -> r) -> ContW w r r' a
forall (w :: * -> *) r r' a. (w (a -> r') -> r) -> ContW w r r' a
ContW ((w (a -> r') -> r) -> ContW w r r' a)
-> (w (a -> r') -> r) -> ContW w r r' a
forall a b. (a -> b) -> a -> b
$ \w (a -> r')
wk -> v (a -> r') -> r
a (w (a -> r') -> v (a -> r')
forall x. w x -> v x
f w (a -> r')
wk)