The State Comonad

StateComonad.hs

-- http://comonad.com/reader/2018/the-state-comonad/
-- https://www.reddit.com/r/haskell/comments/7oav51/i_made_a_monad_that_i_havent_seen_before_and_i/

{-# language DeriveFunctor #-}

import Control.Comonad
import Data.Semigroup

data Store s a = Store { peek :: s -> a, pos :: s } deriving Functor

instance Comonad (Store s) where
  extract (Store f s) = f s
  duplicate (Store f s) = Store (Store f) s

instance (Semigroup s, Monoid s) => Applicative (Store s) where
  pure a = Store (const a) mempty
  Store f s <*> Store g t = Store (\m -> f m (g m)) (mappend s t)

instance Semigroup s => ComonadApply (Store s) where
  Store f s <@> Store g t = Store (\m -> f m (g m)) (s <> t)

instance (Semigroup s, Monoid s) => Monad (Store s) where
  return = pure
  m >>= k = Store (\s -> peek (k (peek m s)) s) (pos m `mappend` pos (k (peek m mempty)))

newtype State s a = State { runState :: s -> (a, s) } deriving Functor

instance Applicative (State s) where
  pure a = State $ \s -> (a, s)
  State mf <*> State ma = State $ \s -> case mf s of
    (f, s') -> case ma s' of
      (a, s'') -> (f a, s'')

instance Monad (State s) where
  return = pure
  State m >>= k = State $ \s -> case m s of
    (a, s') -> runState (k a) s'

instance Monoid s => Comonad (State s) where
  extract m = fst $ runState m mempty
  duplicate m = State $ \s -> (State $ \t -> runState m (mappend s t), snd $ runState m s)

instance Monoid s => ComonadApply (State s) where
  (<@>) = (<*>)