-- | This module exposes a simple notion of a Staged computation (or a freer
-- monad) to be used when modifying mockchain runs with Ltl formulas.
module Cooked.Staged
  ( Staged (..),
    interpStaged,
  )
where

import Control.Monad
import Data.Kind

-- | The freer monad on @op@. We think of this as the AST of a computation with
-- operations of types @op a@.
data Staged (op :: Type -> Type) :: Type -> Type where
  Return :: a -> Staged op a
  Instr :: op a -> (a -> Staged op b) -> Staged op b

instance Functor (Staged op) where
  fmap :: forall a b. (a -> b) -> Staged op a -> Staged op b
fmap a -> b
f (Return a
x) = b -> Staged op b
forall a (op :: * -> *). a -> Staged op a
Return (b -> Staged op b) -> b -> Staged op b
forall a b. (a -> b) -> a -> b
$ a -> b
f a
x
  fmap a -> b
f (Instr op a
op a -> Staged op a
cont) = op a -> (a -> Staged op b) -> Staged op b
forall (op :: * -> *) a b.
op a -> (a -> Staged op b) -> Staged op b
Instr op a
op ((a -> b) -> Staged op a -> Staged op b
forall a b. (a -> b) -> Staged op a -> Staged op b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f (Staged op a -> Staged op b)
-> (a -> Staged op a) -> a -> Staged op b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Staged op a
cont)

instance Applicative (Staged op) where
  pure :: forall a. a -> Staged op a
pure = a -> Staged op a
forall a (op :: * -> *). a -> Staged op a
Return
  <*> :: forall a b. Staged op (a -> b) -> Staged op a -> Staged op b
(<*>) = Staged op (a -> b) -> Staged op a -> Staged op b
forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad (Staged op) where
  (Return a
x) >>= :: forall a b. Staged op a -> (a -> Staged op b) -> Staged op b
>>= a -> Staged op b
f = a -> Staged op b
f a
x
  (Instr op a
i a -> Staged op a
m) >>= a -> Staged op b
f = op a -> (a -> Staged op b) -> Staged op b
forall (op :: * -> *) a b.
op a -> (a -> Staged op b) -> Staged op b
Instr op a
i (a -> Staged op a
m (a -> Staged op a) -> (a -> Staged op b) -> a -> Staged op b
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> a -> Staged op b
f)

-- | Interprets a staged computation given a interpreter of the builtins
interpStaged :: (Monad m) => (forall a. op a -> m a) -> forall a. Staged op a -> m a
interpStaged :: forall (m :: * -> *) (op :: * -> *).
Monad m =>
(forall a. op a -> m a) -> forall a. Staged op a -> m a
interpStaged forall a. op a -> m a
_ (Return a
a) = a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a
interpStaged forall a. op a -> m a
interpBuiltin (Instr op a
op a -> Staged op a
cont) = op a -> m a
forall a. op a -> m a
interpBuiltin op a
op m a -> (a -> m a) -> m a
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (forall a. op a -> m a) -> forall a. Staged op a -> m a
forall (m :: * -> *) (op :: * -> *).
Monad m =>
(forall a. op a -> m a) -> forall a. Staged op a -> m a
interpStaged op a -> m a
forall a. op a -> m a
interpBuiltin (Staged op a -> m a) -> (a -> Staged op a) -> a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Staged op a
cont