How to implement index-core style indexed state monad?

Question

I'm trying to understand indexed monads in the index-core style. I have become stuck in a paradox which is that I can't understand the principles until I have constructed a few examples, and I can't construct examples until I understand the principles.

I am trying to construct an indexed state monad. So far my intuition tells me it should be something like this

type a :* b = forall i. (a i, b i)
newtype IState f a i = IState { runIState :: f i -> (a :* f) }

and that I can recover the "restricted" state monad by setting f = Identity and choosing a appropriately:

type IState' s s' a = IState Identity (a := s') s

But I'm feeling rather lost. Can someone confirm that I'm on the right lines?

I'm using a similar question on the indexed continuation monad as a guide, but I don't think it's really close enough.

Answer 1

We can replicate Gabriel's argument from the indexed Cont answer linked. If the standard index state monad is

State s s' a = s -> (a, s')

then we generalize it in stages. First by using Identity to reflect the concrete types s and s' as fibers in the indexed type space Identity.

State s s' a = s          -> (a, s')
             ~ Identity s -> (a, Identity s')

then by generalizing the value type a to an indexed type as well over the "destination" index, the type s'.

             ~ Identity s -> (a   , Identity s')
             ~ Identity s -> (a s', Identity s')

and then by using an existential type to erase the destination index. We'll recover it later.

data a :* b = forall i . P (a i) (b i)

  Identity s -> (a s', Identity s') 
~ Identity s -> P a Identity

Finally, we'll note that Identity is the indexed type of the state space and a the indexed type of the value space, so we can write IState as

newtype IState s -- indexed state space
               a -- indexed value space
               i -- index
  = IState { runIState :: s i -> a :* s }
--   State { runState  :: s   -> (a, s) }        for comparison

---

Why use an existentially quantified pair instead of a universally quantified one? A first nudge comes from the fact that the index associated with a is occurring positively in IState while it appeared negatively in ICont. A second hint comes from writing returnI. If we use the universally quantified version and try to write returnI

newtype IState' s a i = IState' { runIState' :: s i -> (forall i . (a i, s i)) }

returnI :: a i -> IState' s a i
returnI a = IState' (\si -> (?forget a, ?forget si))

we would need to have this forget function which forgets all of the information about the index.

However, if we instead use the existentially quantified pair then it's up to the constructor of that returning pair, i.e. the implementor of the IState value, to pick the index. This lets us instantiate IFunctor and IMonad

instance IFunctor (IState s) where
  -- fmapI :: (a :-> b) -> IState s a :-> IState s b
  fmapI phi (IState go) = IState $ \si -> 
    case go si of 
      P ax sx -> P (phi ax) sx

instance IMonad (IState s) where
  -- returnI :: a :-> IState s a
  return ai = IState (\si -> P ai si)

  -- bindI :: (a :-> IState s b) -> (IState s a :-> IState s b)
  bindI f m = IState $ \s ->
    case runIState m s of
      P ax sx -> runIState (f ax) sx

The only downside to the use of this existential pair is that it's... pretty hard to actually use. For instance, we'd really like to be able to use the "pointed" indexed type constructor (:=) in order to fix a known index and project out of the existential pair.

one :: (a := i :* b) -> a
two :: (a := i :* b) -> b i

Unfortunately, Haskell isn't smart enough to coerce the existential even when we know what it it is, so the second of these projections has an unsavory implementation

one :: (a := i :* b) -> a
one (P (V a) _) = a

two :: (a := i :* b) -> b i
two (P _ s) = unsafeCoerce s

---

Finally, the proof is in the pudding. We can use IState to implement the standard complement of stateful effects we're used to seeing.

-- Higher order unit type
data U1 a = U1

put :: s i -> IState s U1 j
put s = IState (\_ -> P U1 s)

get :: IState s s i
get = IState (\s -> P s s)

and use them to implement some general higher order combinators such as modify (which needs an explicit type signature, but you can compute that from the implementation manually with some thought)

modify :: (s :-> s) -> IState s U1 i
modify f = get ?>= put . f

However, additionally to these, we have other ways of representing the combinators that are more explicit about the indexing due to restriction via (:=). This can be useful in order to pass around more information about the indexing.

put' :: s i1 -> IState s (() := i1) i
put' s = IState (\_ -> P (V ()) s)

get' :: IState s (s i := i) i
get' = IState (\s -> P (V s) s)

modify' :: (s -> s) -> IState (s := j) (() := j) i
modify' f = get >>= put' . V . f

modify'' :: (s i -> s k) -> IState s (() := k) i
modify'' f = get' >>= put' . f

Finally, we can use all this to implement an example. For instance, we can build the indexed type over file handle states not that it's tremendously useful.

data Open
data Closed
data Any a

data St where
  So :: Int -> St Open
  Sc ::        St Closed
  Sa :: a   -> St (Any a)

getAny :: St (Any a) -> a
getAny (Sa a) = a

Then we can build

open :: String -> File Closed Open ()
open name = put' (SOpen $ getHandle name) where
  getHandle = length

close :: File Open Closed ()
close = put' SClosed

getHandle :: File Open Open Int
getHandle = do
  SOpen i <- get'
  return i

putA :: a -> File x (Any a) ()
putA a = put' (SAny a)

where

open "foo" >> close                -- typechecks
open "foo" >> close >> close       -- fails
open "foo" >> getHandler >> close  -- typechecks
open "foo" >> close >> getHandler  -- fails

and things like

> one $ runIState (do putA 4
                      sa <- get'
                      return (getAny sa)) Sc
4
> one $ runIState (do putA ()
                      sa <- get'
                      return (getAny sa)) Sc
()
> one $ runIState (do putA 4
                      putA ()
                      sa <- get'
                      return (getAny sa)) Sc
()

all work.