CPS-ed State monad

cps.hs

{-# LANGUAGE RankNTypes #-}

module CPS where

import Control.Applicative

newtype StateCPS s a = StateCPS { runStateCPS :: forall r. s -> (a -> s -> r) -> r }

instance Functor (StateCPS s) where
  fmap f (StateCPS k) = StateCPS $ \s c -> k s (c . f)

instance Applicative (StateCPS s) where
  pure a = StateCPS $ \s k -> k a s

  StateCPS kf <*> StateCPS ka = StateCPS $ \s c -> kf s (\f s' -> ka s' (\a s'' -> c (f a) s''))

instance Monad (StateCPS s) where
  return a = StateCPS $ \s k -> k a s

  StateCPS k >>= f = StateCPS $ \s c -> k s (\a s' -> runStateCPS (f a) s' c)

get :: StateCPS s s
get = StateCPS $ \s k -> k s s

put :: s -> StateCPS s ()
put s = StateCPS $ \_ k -> k () s 


modify :: (s -> s) -> StateCPS s ()
modify f = do
  s <- get
  put (f s)

main = do
  runStateCPS go 1 (\a s -> print (a, s))
    where
      go = do
        x <- get
        modify (+1)
        return x