0 | module Control.Monad.Indexed.State
 1 |
 2 | import Control.Functor.Indexed
 3 | import Control.Applicative.Indexed
 4 | import Control.Monad.Indexed
 5 | import Control.Monad.Indexed.Do
 6 |
 7 | ||| Support for wrapping an IndexedMonad in another stateful
 8 | ||| Monad.
 9 | ||| The order of the parameters may seem odd compared to other
10 | ||| StateT types, but the ordering is such that an IndexedMonad 
11 | ||| type can be elevated to a TransitionIndexedMonad type without 
12 | ||| much refactoring.
13 | public export
14 | record IndexedStateT stateType x y (m : Type -> x -> y -> Type) a (i : x) (j : y) where
15 |   constructor ST
16 |   runStateT' : stateType -> m (stateType, a) i j
17 |
18 | public export
19 | %inline
20 | runStateT : stateType -> IndexedStateT stateType x y m a i j -> m (stateType, a) i j
21 | runStateT s act = runStateT' act s
22 |
23 | public export
24 | %inline
25 | evalStateT : IndexedFunctor x y m => stateType -> IndexedStateT stateType x y m a i j -> m a i j
26 | evalStateT s = map snd . runStateT s
27 |
28 | public export
29 | %inline
30 | execStateT : IndexedFunctor x y m => stateType -> IndexedStateT stateType x y m a i j -> m stateType i j
31 | execStateT s = map fst . runStateT s
32 |
33 | public export
34 | %inline
35 | mapStateT : ({0 i : x} -> {0 j : y} -> m (s, a) i j -> n (s, b) i j) -> IndexedStateT s x y m a i j -> IndexedStateT s x y n b i j
36 | mapStateT f m = ST $ f . runStateT' m
37 |
38 | public export
39 | lift : IndexedFunctor x y m => m a i j -> IndexedStateT stateType x y m a i j
40 | lift z = ST $ \st => map (st,) z
41 |
42 | implementation IndexedFunctor x y f => IndexedFunctor x y (IndexedStateT stateType x y f) where
43 |   map f' (ST runStateT') = ST $ \st => map (mapSnd f') $ runStateT' st
44 |
45 | public export
46 | IndexedMonad z m => IndexedApplicative z (IndexedStateT stateType z z m) where
47 |   pure x = ST $ \st => pure (st, x)
48 |
49 |   (ST ff) `ap` (ST fx) = ST $ \st =>
50 |                           ff st >>>= \(st', f) =>
51 |                             fx st' >>>= \(st'', x) =>
52 |                               pure (st'', f x)
53 |
54 | public export
55 | IndexedMonad z m => IndexedMonad z (IndexedStateT stateType z z m) where
56 |   bind (ST f) g = ST $ \st =>
57 |                f st >>>= \(st', x) =>
58 |                  runStateT st' (g x)
59 |
60 | public export
61 | get : IndexedApplicative z m => IndexedStateT stateType z z m stateType i i
62 | get = ST $ \st => pure (st, st)
63 |
64 | public export
65 | gets : IndexedApplicative z m => (stateType -> a) -> IndexedStateT stateType z z m a i i
66 | gets f = ST $ \st => pure (st, f st)
67 |
68 | public export
69 | put : IndexedApplicative z m => stateType -> IndexedStateT stateType z z m () i i
70 | put st = ST $ \_ => pure (st, ())
71 |
72 | public export
73 | modify : IndexedMonad z m => (stateType -> stateType) -> IndexedStateT stateType z z m () i i
74 | modify f = do
75 |   x <- get
76 |   put (f x)
77 |
78 |