@@ -0,0 +1,415 @@
1+module Ether.State
2+ (
3+ -- * MonadState class
4+ MonadState
5+ , get
6+ , put
7+ , state
8+ , modify
9+ , gets
10+ -- * The State monad
11+ , State
12+ , runState
13+ , evalState
14+ , execState
15+ -- * The StateT monad transformer
16+ , StateT
17+ , stateT
18+ , runStateT
19+ , evalStateT
20+ , execStateT
21+ -- * The State monad (lazy)
22+ , LazyState
23+ , runLazyState
24+ , evalLazyState
25+ , execLazyState
26+ -- * The StateT monad transformer (lazy)
27+ , LazyStateT
28+ , lazyStateT
29+ , runLazyStateT
30+ , evalLazyStateT
31+ , execLazyStateT
32+ -- * The State monad (flattened)
33+ , States
34+ , runStates
35+ -- * The StateT monad transformer (flattened)
36+ , StatesT
37+ , runStatesT
38+ -- * MonadState class (implicit)
39+ , MonadState'
40+ , get'
41+ , put'
42+ , state'
43+ , modify'
44+ , gets'
45+ -- * The State monad (implicit)
46+ , State'
47+ , runState'
48+ , evalState'
49+ , execState'
50+ -- * The StateT monad transformer (implicit)
51+ , StateT'
52+ , stateT'
53+ , runStateT'
54+ , evalStateT'
55+ , execStateT'
56+ -- * The State monad (lazy, implicit)
57+ , LazyState'
58+ , runLazyState'
59+ , evalLazyState'
60+ , execLazyState'
61+ -- * The StateT monad transformer (lazy, implicit)
62+ , LazyStateT'
63+ , lazyStateT'
64+ , runLazyStateT'
65+ , evalLazyStateT'
66+ , execLazyStateT'
67+ -- * Zoom
68+ , ZoomT
69+ , zoom
70+ -- * Internal labels
71+ , TAGGED
72+ , STATE
73+ , STATES
74+ , ZOOM
75+ ) where
76+
77+import qualified Control.Monad.State.Class as T
78+import qualified Control.Monad.Trans as Lift
79+import Control.Monad.Trans.Identity
80+import qualified Control.Monad.Trans.State.Lazy as T.Lazy
81+import qualified Control.Monad.Trans.State.Strict as T.Strict
82+import Data.Coerce
83+import Data.Functor.Identity
84+import Data.Kind
85+import Data.Proxy
86+import Data.Reflection
87+
88+import Ether.Internal
89+import Ether.TaggedTrans
90+
91+class Monad m => MonadState tag s m | m tag -> s where
92+
93+ {-# MINIMAL state | get, put #-}
94+
95+ -- | Return the state from the internals of the monad.
96+ get :: m s
97+ get = state @tag (\s -> (s, s))
98+
99+ -- | Replace the state inside the monad.
100+ put :: s -> m ()
101+ put s = state @tag (\_ -> ((), s))
102+
103+ -- | Embed a simple state action into the monad.
104+ state :: (s -> (a, s)) -> m a
105+ state f = do
106+ s <- get @tag
107+ let ~(a, s') = f s
108+ put @tag s'
109+ return a
110+
111+instance {-# OVERLAPPABLE #-}
112+ ( Lift.MonadTrans t
113+ , Monad (t m)
114+ , MonadState tag s m
115+ ) => MonadState tag s (t m)
116+ where
117+ get = Lift.lift (get @tag)
118+ put = Lift.lift . put @tag
119+ state = Lift.lift . state @tag
120+
121+instance {-# OVERLAPPABLE #-}
122+ ( Monad (trans m)
123+ , MonadState tag s (TaggedTrans effs trans m)
124+ ) => MonadState tag s (TaggedTrans (eff ': effs) trans (m :: Type -> Type))
125+ where
126+
127+ get =
128+ (coerce ::
129+ TaggedTrans effs trans m s ->
130+ TaggedTrans (eff ': effs) trans m s)
131+ (get @tag)
132+
133+ put =
134+ (coerce ::
135+ (s -> TaggedTrans effs trans m ()) ->
136+ (s -> TaggedTrans (eff ': effs) trans m ()))
137+ (put @tag)
138+
139+ state =
140+ (coerce :: forall a .
141+ ((s -> (a, s)) -> TaggedTrans effs trans m a) ->
142+ ((s -> (a, s)) -> TaggedTrans (eff ': effs) trans m a))
143+ (state @tag)
144+
145+-- | Modifies the state inside a state monad.
146+modify :: forall tag s m . MonadState tag s m => (s -> s) -> m ()
147+modify f = state @tag (\s -> ((), f s))
148+
149+-- | Gets specific component of the state, using a projection function supplied.
150+gets :: forall tag s m a . MonadState tag s m => (s -> a) -> m a
151+gets f = fmap f (get @tag)
152+
153+-- | Encode type-level information for 'StateT'.
154+data STATE
155+
156+type instance HandleSuper STATE s trans = ()
157+type instance HandleConstraint STATE s trans m =
158+ T.MonadState s (trans m)
159+
160+instance Handle STATE s (T.Strict.StateT s) where
161+ handling r = r
162+
163+instance Handle STATE s (T.Lazy.StateT s) where
164+ handling r = r
165+
166+instance
167+ ( Handle STATE s trans
168+ , Monad m, Monad (trans m)
169+ ) => MonadState tag s (TaggedTrans (TAGGED STATE tag) trans m)
170+ where
171+
172+ get =
173+ handling @STATE @s @trans @m $
174+ coerce (T.get @s @(trans m))
175+
176+ put =
177+ handling @STATE @s @trans @m $
178+ coerce (T.put @s @(trans m))
179+
180+ state =
181+ handling @STATE @s @trans @m $
182+ coerce (T.state @s @(trans m) @a) ::
183+ forall eff a . (s -> (a, s)) -> TaggedTrans eff trans m a
184+
185+instance
186+ ( HasLens tag payload s
187+ , Handle STATE payload trans
188+ , Monad m, Monad (trans m)
189+ ) => MonadState tag s (TaggedTrans (TAGGED STATE tag ': effs) trans m)
190+ where
191+
192+ get =
193+ handling @STATE @payload @trans @m $
194+ (coerce :: forall eff a .
195+ trans m a ->
196+ TaggedTrans eff trans m a)
197+ (T.gets (view (lensOf @tag @payload @s)))
198+
199+ put s =
200+ handling @STATE @payload @trans @m $
201+ (coerce :: forall eff a .
202+ trans m a ->
203+ TaggedTrans eff trans m a)
204+ (T.modify (over (lensOf @tag @payload @s) (const s)))
205+
206+ state f =
207+ handling @STATE @payload @trans @m $
208+ (coerce :: forall eff a .
209+ trans m a ->
210+ TaggedTrans eff trans m a)
211+ (T.state (lensOf @tag @payload @s f))
212+
213+-- | The parametrizable state monad.
214+--
215+-- Computations have access to a mutable state.
216+--
217+-- The 'return' function leaves the state unchanged, while '>>=' uses
218+-- the final state of the first computation as the initial state of the second.
219+type State tag r = StateT tag r Identity
220+
221+-- | The state monad transformer.
222+--
223+-- The 'return' function leaves the state unchanged, while '>>=' uses
224+-- the final state of the first computation as the initial state of the second.
225+type StateT tag s = TaggedTrans (TAGGED STATE tag) (T.Strict.StateT s)
226+
227+-- | Constructor for computations in the state monad transformer.
228+stateT :: forall tag s m a . (s -> m (a, s)) -> StateT tag s m a
229+stateT = coerce (T.Strict.StateT @s @m @a)
230+
231+-- | Runs a 'StateT' with the given initial state
232+-- and returns both the final value and the final state.
233+runStateT :: forall tag s m a . StateT tag s m a -> s -> m (a, s)
234+runStateT = coerce (T.Strict.runStateT @s @m @a)
235+
236+-- | Runs a 'StateT' with the given initial state
237+-- and returns the final value, discarding the final state.
238+evalStateT :: forall tag s m a . Monad m => StateT tag s m a -> s -> m a
239+evalStateT = coerce (T.Strict.evalStateT @m @s @a)
240+
241+-- | Runs a 'StateT' with the given initial state
242+-- and returns the final state, discarding the final value.
243+execStateT :: forall tag s m a . Monad m => StateT tag s m a -> s -> m s
244+execStateT = coerce (T.Strict.execStateT @m @s @a)
245+
246+-- | Runs a 'State' with the given initial state
247+-- and returns both the final value and the final state.
248+runState :: forall tag s a . State tag s a -> s -> (a, s)
249+runState = coerce (T.Strict.runState @s @a)
250+
251+-- | Runs a 'State' with the given initial state
252+-- and returns the final value, discarding the final state.
253+evalState :: forall tag s a . State tag s a -> s -> a
254+evalState = coerce (T.Strict.evalState @s @a)
255+
256+-- | Runs a 'State' with the given initial state
257+-- and returns the final state, discarding the final value.
258+execState :: forall tag s a . State tag s a -> s -> s
259+execState = coerce (T.Strict.execState @s @a)
260+
261+-- | The parametrizable state monad.
262+--
263+-- Computations have access to a mutable state.
264+--
265+-- The 'return' function leaves the state unchanged, while '>>=' uses
266+-- the final state of the first computation as the initial state of the second.
267+type LazyState tag r = LazyStateT tag r Identity
268+
269+-- | The state monad transformer.
270+--
271+-- The 'return' function leaves the state unchanged, while '>>=' uses
272+-- the final state of the first computation as the initial state of the second.
273+type LazyStateT tag s = TaggedTrans (TAGGED STATE tag) (T.Lazy.StateT s)
274+
275+-- | Constructor for computations in the state monad transformer.
276+lazyStateT :: forall tag s m a . (s -> m (a, s)) -> LazyStateT tag s m a
277+lazyStateT = coerce (T.Lazy.StateT @s @m @a)
278+
279+-- | Runs a 'StateT' with the given initial state
280+-- and returns both the final value and the final state.
281+runLazyStateT :: forall tag s m a . LazyStateT tag s m a -> s -> m (a, s)
282+runLazyStateT = coerce (T.Lazy.runStateT @s @m @a)
283+
284+-- | Runs a 'StateT' with the given initial state
285+-- and returns the final value, discarding the final state.
286+evalLazyStateT :: forall tag s m a . Monad m => LazyStateT tag s m a -> s -> m a
287+evalLazyStateT = coerce (T.Lazy.evalStateT @m @s @a)
288+
289+-- | Runs a 'StateT' with the given initial state
290+-- and returns the final state, discarding the final value.
291+execLazyStateT :: forall tag s m a . Monad m => LazyStateT tag s m a -> s -> m s
292+execLazyStateT = coerce (T.Lazy.execStateT @m @s @a)
293+
294+-- | Runs a 'State' with the given initial state
295+-- and returns both the final value and the final state.
296+runLazyState :: forall tag s a . LazyState tag s a -> s -> (a, s)
297+runLazyState = coerce (T.Lazy.runState @s @a)
298+
299+-- | Runs a 'State' with the given initial state
300+-- and returns the final value, discarding the final state.
301+evalLazyState :: forall tag s a . LazyState tag s a -> s -> a
302+evalLazyState = coerce (T.Lazy.evalState @s @a)
303+
304+-- | Runs a 'State' with the given initial state
305+-- and returns the final state, discarding the final value.
306+execLazyState :: forall tag s a . LazyState tag s a -> s -> s
307+execLazyState = coerce (T.Lazy.execState @s @a)
308+
309+type family STATES (ts :: HList xs) :: [Type] where
310+ STATES 'HNil = '[]
311+ STATES ('HCons t ts) = TAGGED STATE t ': STATES ts
312+
313+type StatesT s = TaggedTrans (STATES (Tags s)) (T.Strict.StateT s)
314+
315+type States s = StatesT s Identity
316+
317+runStatesT :: forall p m a . StatesT p m a -> p -> m (a, p)
318+runStatesT = coerce (T.Strict.runStateT @p @m @a)
319+
320+runStates :: forall p a . States p a -> p -> (a, p)
321+runStates = coerce (T.Strict.runState @p @a)
322+
323+type StateT' s = StateT s s
324+
325+stateT' :: (s -> m (a, s)) -> StateT' s m a
326+stateT' = stateT
327+
328+runStateT' :: StateT' s m a -> s -> m (a, s)
329+runStateT' = runStateT
330+
331+runState' :: State' s a -> s -> (a, s)
332+runState' = runState
333+
334+evalStateT' :: Monad m => StateT' s m a -> s -> m a
335+evalStateT' = evalStateT
336+
337+type State' s = State s s
338+
339+evalState' :: State' s a -> s -> a
340+evalState' = evalState
341+
342+execStateT' :: Monad m => StateT' s m a -> s -> m s
343+execStateT' = execStateT
344+
345+execState' :: State' s a -> s -> s
346+execState' = execState
347+
348+type LazyStateT' s = LazyStateT s s
349+
350+lazyStateT' :: (s -> m (a, s)) -> LazyStateT' s m a
351+lazyStateT' = lazyStateT
352+
353+runLazyStateT' :: LazyStateT' s m a -> s -> m (a, s)
354+runLazyStateT' = runLazyStateT
355+
356+runLazyState' :: LazyState' s a -> s -> (a, s)
357+runLazyState' = runLazyState
358+
359+evalLazyStateT' :: Monad m => LazyStateT' s m a -> s -> m a
360+evalLazyStateT' = evalLazyStateT
361+
362+type LazyState' s = LazyState s s
363+
364+evalLazyState' :: LazyState' s a -> s -> a
365+evalLazyState' = evalLazyState
366+
367+execLazyStateT' :: Monad m => LazyStateT' s m a -> s -> m s
368+execLazyStateT' = execLazyStateT
369+
370+execLazyState' :: LazyState' s a -> s -> s
371+execLazyState' = execLazyState
372+
373+type MonadState' s = MonadState s s
374+
375+get' :: forall s m . MonadState' s m => m s
376+get' = get @s
377+
378+gets' :: forall s m a . MonadState' s m => (s -> a) -> m a
379+gets' = gets @s
380+
381+put' :: forall s m . MonadState' s m => s -> m ()
382+put' = put @s
383+
384+state' :: forall s m a . MonadState' s m => (s -> (a, s)) -> m a
385+state' = state @s
386+
387+modify' :: forall s m . MonadState' s m => (s -> s) -> m ()
388+modify' = modify @s
389+
390+-- | Encode type-level information for 'zoom'.
391+data ZOOM t z
392+
393+type ZoomT t (z :: Type) = TaggedTrans (ZOOM t z) IdentityT
394+
395+-- | Zoom into a part of a state using a lens.
396+zoom
397+ :: forall tag sOuter sInner m a
398+ . Lens' sOuter sInner
399+ -> (forall z . Reifies z (ReifiedLens' sOuter sInner) => ZoomT tag z m a)
400+ -> m a
401+zoom l m = reify (Lens l) (\(_ :: Proxy z) -> coerce (m @z))
402+
403+instance
404+ ( MonadState tag sOuter m
405+ , Reifies z (ReifiedLens' sOuter sInner)
406+ , trans ~ IdentityT
407+ ) => MonadState tag sInner (TaggedTrans (ZOOM tag z) trans m)
408+ where
409+ state =
410+ (coerce :: forall eff r a .
411+ (r -> m a) ->
412+ (r -> TaggedTrans eff trans m a))
413+ (state @tag . l)
414+ where
415+ Lens l = reflect (Proxy :: Proxy z)