@@ -0,0 +1,475 @@
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.Handler
89+import Ether.Internal
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+
118+ get = Lift.lift (get @tag)
119+ {-# INLINE get #-}
120+
121+ put = Lift.lift . put @tag
122+ {-# INLINE put #-}
123+
124+ state = Lift.lift . state @tag
125+ {-# INLINE state #-}
126+
127+instance {-# OVERLAPPABLE #-}
128+ ( Monad (trans m)
129+ , MonadState tag s (Handler effs trans m)
130+ ) => MonadState tag s (Handler (eff ': effs) trans (m :: Type -> Type))
131+ where
132+
133+ get =
134+ (coerce ::
135+ Handler effs trans m s ->
136+ Handler (eff ': effs) trans m s)
137+ (get @tag)
138+ {-# INLINE get #-}
139+
140+ put =
141+ (coerce ::
142+ (s -> Handler effs trans m ()) ->
143+ (s -> Handler (eff ': effs) trans m ()))
144+ (put @tag)
145+ {-# INLINE put #-}
146+
147+ state =
148+ (coerce :: forall a .
149+ ((s -> (a, s)) -> Handler effs trans m a) ->
150+ ((s -> (a, s)) -> Handler (eff ': effs) trans m a))
151+ (state @tag)
152+ {-# INLINE state #-}
153+
154+-- | Encode type-level information for 'StateT'.
155+data STATE
156+
157+-- | Modifies the state inside a state monad.
158+modify :: forall tag s m . MonadState tag s m => (s -> s) -> m ()
159+modify f = state @tag (\s -> ((), f s))
160+{-# INLINABLE modify #-}
161+
162+-- | Gets specific component of the state, using a projection function supplied.
163+gets :: forall tag s m a . MonadState tag s m => (s -> a) -> m a
164+gets f = fmap f (get @tag)
165+{-# INLINABLE gets #-}
166+
167+type instance HandleSuper STATE s trans = ()
168+type instance HandleConstraint STATE s trans m =
169+ T.MonadState s (trans m)
170+
171+instance Handle STATE s (T.Strict.StateT s) where
172+ handling r = r
173+ {-# INLINE handling #-}
174+
175+instance Handle STATE s (T.Lazy.StateT s) where
176+ handling r = r
177+ {-# INLINE handling #-}
178+
179+instance
180+ ( Handle STATE s trans
181+ , Monad m, Monad (trans m)
182+ ) => MonadState tag s (Handler (TAGGED STATE tag) trans m)
183+ where
184+
185+ get =
186+ handling @STATE @s @trans @m $
187+ coerce (T.get @s @(trans m))
188+ {-# INLINE get #-}
189+
190+ put =
191+ handling @STATE @s @trans @m $
192+ coerce (T.put @s @(trans m))
193+ {-# INLINE put #-}
194+
195+ state =
196+ handling @STATE @s @trans @m $
197+ coerce (T.state @s @(trans m) @a) ::
198+ forall eff a . (s -> (a, s)) -> Handler eff trans m a
199+ {-# INLINE state #-}
200+
201+instance
202+ ( HasLens tag payload s
203+ , Handle STATE payload trans
204+ , Monad m, Monad (trans m)
205+ ) => MonadState tag s (Handler (TAGGED STATE tag ': effs) trans m)
206+ where
207+
208+ get =
209+ handling @STATE @payload @trans @m $
210+ (coerce :: forall eff a .
211+ trans m a ->
212+ Handler eff trans m a)
213+ (T.gets (view (lensOf @tag @payload @s)))
214+ {-# INLINE get #-}
215+
216+ put s =
217+ handling @STATE @payload @trans @m $
218+ (coerce :: forall eff a .
219+ trans m a ->
220+ Handler eff trans m a)
221+ (T.modify (over (lensOf @tag @payload @s) (const s)))
222+ {-# INLINE put #-}
223+
224+ state f =
225+ handling @STATE @payload @trans @m $
226+ (coerce :: forall eff a .
227+ trans m a ->
228+ Handler eff trans m a)
229+ (T.state (lensOf @tag @payload @s f))
230+ {-# INLINE state #-}
231+
232+-- | The parametrizable state monad.
233+--
234+-- Computations have access to a mutable state.
235+--
236+-- The 'return' function leaves the state unchanged, while '>>=' uses
237+-- the final state of the first computation as the initial state of the second.
238+type State tag r = StateT tag r Identity
239+
240+-- | The state monad transformer.
241+--
242+-- The 'return' function leaves the state unchanged, while '>>=' uses
243+-- the final state of the first computation as the initial state of the second.
244+type StateT tag s = Handler (TAGGED STATE tag) (T.Strict.StateT s)
245+
246+-- | Constructor for computations in the state monad transformer.
247+stateT :: forall tag s m a . (s -> m (a, s)) -> StateT tag s m a
248+stateT = coerce (T.Strict.StateT @s @m @a)
249+{-# INLINE stateT #-}
250+
251+-- | Runs a 'StateT' with the given initial state
252+-- and returns both the final value and the final state.
253+runStateT :: forall tag s m a . StateT tag s m a -> s -> m (a, s)
254+runStateT = coerce (T.Strict.runStateT @s @m @a)
255+{-# INLINE runStateT #-}
256+
257+-- | Runs a 'StateT' with the given initial state
258+-- and returns the final value, discarding the final state.
259+evalStateT :: forall tag s m a . Monad m => StateT tag s m a -> s -> m a
260+evalStateT = coerce (T.Strict.evalStateT @m @s @a)
261+{-# INLINE evalStateT #-}
262+
263+-- | Runs a 'StateT' with the given initial state
264+-- and returns the final state, discarding the final value.
265+execStateT :: forall tag s m a . Monad m => StateT tag s m a -> s -> m s
266+execStateT = coerce (T.Strict.execStateT @m @s @a)
267+{-# INLINE execStateT #-}
268+
269+-- | Runs a 'State' with the given initial state
270+-- and returns both the final value and the final state.
271+runState :: forall tag s a . State tag s a -> s -> (a, s)
272+runState = coerce (T.Strict.runState @s @a)
273+{-# INLINE runState #-}
274+
275+-- | Runs a 'State' with the given initial state
276+-- and returns the final value, discarding the final state.
277+evalState :: forall tag s a . State tag s a -> s -> a
278+evalState = coerce (T.Strict.evalState @s @a)
279+{-# INLINE evalState #-}
280+
281+-- | Runs a 'State' with the given initial state
282+-- and returns the final state, discarding the final value.
283+execState :: forall tag s a . State tag s a -> s -> s
284+execState = coerce (T.Strict.execState @s @a)
285+{-# INLINE execState #-}
286+
287+-- | The parametrizable state monad.
288+--
289+-- Computations have access to a mutable state.
290+--
291+-- The 'return' function leaves the state unchanged, while '>>=' uses
292+-- the final state of the first computation as the initial state of the second.
293+type LazyState tag r = LazyStateT tag r Identity
294+
295+-- | The state monad transformer.
296+--
297+-- The 'return' function leaves the state unchanged, while '>>=' uses
298+-- the final state of the first computation as the initial state of the second.
299+type LazyStateT tag s = Handler (TAGGED STATE tag) (T.Lazy.StateT s)
300+
301+-- | Constructor for computations in the state monad transformer.
302+lazyStateT :: forall tag s m a . (s -> m (a, s)) -> LazyStateT tag s m a
303+lazyStateT = coerce (T.Lazy.StateT @s @m @a)
304+{-# INLINE lazyStateT #-}
305+
306+-- | Runs a 'StateT' with the given initial state
307+-- and returns both the final value and the final state.
308+runLazyStateT :: forall tag s m a . LazyStateT tag s m a -> s -> m (a, s)
309+runLazyStateT = coerce (T.Lazy.runStateT @s @m @a)
310+{-# INLINE runLazyStateT #-}
311+
312+-- | Runs a 'StateT' with the given initial state
313+-- and returns the final value, discarding the final state.
314+evalLazyStateT :: forall tag s m a . Monad m => LazyStateT tag s m a -> s -> m a
315+evalLazyStateT = coerce (T.Lazy.evalStateT @m @s @a)
316+{-# INLINE evalLazyStateT #-}
317+
318+-- | Runs a 'StateT' with the given initial state
319+-- and returns the final state, discarding the final value.
320+execLazyStateT :: forall tag s m a . Monad m => LazyStateT tag s m a -> s -> m s
321+execLazyStateT = coerce (T.Lazy.execStateT @m @s @a)
322+{-# INLINE execLazyStateT #-}
323+
324+-- | Runs a 'State' with the given initial state
325+-- and returns both the final value and the final state.
326+runLazyState :: forall tag s a . LazyState tag s a -> s -> (a, s)
327+runLazyState = coerce (T.Lazy.runState @s @a)
328+{-# INLINE runLazyState #-}
329+
330+-- | Runs a 'State' with the given initial state
331+-- and returns the final value, discarding the final state.
332+evalLazyState :: forall tag s a . LazyState tag s a -> s -> a
333+evalLazyState = coerce (T.Lazy.evalState @s @a)
334+{-# INLINE evalLazyState #-}
335+
336+-- | Runs a 'State' with the given initial state
337+-- and returns the final state, discarding the final value.
338+execLazyState :: forall tag s a . LazyState tag s a -> s -> s
339+execLazyState = coerce (T.Lazy.execState @s @a)
340+{-# INLINE execLazyState #-}
341+
342+type family STATES (ts :: HList xs) :: [Type] where
343+ STATES 'HNil = '[]
344+ STATES ('HCons t ts) = TAGGED STATE t ': STATES ts
345+
346+type StatesT s = Handler (STATES (Tags s)) (T.Strict.StateT s)
347+
348+type States s = StatesT s Identity
349+
350+runStatesT :: forall p m a . StatesT p m a -> p -> m (a, p)
351+runStatesT = coerce (T.Strict.runStateT @p @m @a)
352+{-# INLINE runStatesT #-}
353+
354+runStates :: forall p a . States p a -> p -> (a, p)
355+runStates = coerce (T.Strict.runState @p @a)
356+{-# INLINE runStates #-}
357+
358+type StateT' s = StateT s s
359+
360+stateT' :: (s -> m (a, s)) -> StateT' s m a
361+stateT' = stateT
362+{-# INLINE stateT' #-}
363+
364+runStateT' :: StateT' s m a -> s -> m (a, s)
365+runStateT' = runStateT
366+{-# INLINE runStateT' #-}
367+
368+runState' :: State' s a -> s -> (a, s)
369+runState' = runState
370+{-# INLINE runState' #-}
371+
372+evalStateT' :: Monad m => StateT' s m a -> s -> m a
373+evalStateT' = evalStateT
374+{-# INLINE evalStateT' #-}
375+
376+type State' s = State s s
377+
378+evalState' :: State' s a -> s -> a
379+evalState' = evalState
380+{-# INLINE evalState' #-}
381+
382+execStateT' :: Monad m => StateT' s m a -> s -> m s
383+execStateT' = execStateT
384+{-# INLINE execStateT' #-}
385+
386+execState' :: State' s a -> s -> s
387+execState' = execState
388+{-# INLINE execState' #-}
389+
390+type LazyStateT' s = LazyStateT s s
391+
392+lazyStateT' :: (s -> m (a, s)) -> LazyStateT' s m a
393+lazyStateT' = lazyStateT
394+{-# INLINE lazyStateT' #-}
395+
396+runLazyStateT' :: LazyStateT' s m a -> s -> m (a, s)
397+runLazyStateT' = runLazyStateT
398+{-# INLINE runLazyStateT' #-}
399+
400+runLazyState' :: LazyState' s a -> s -> (a, s)
401+runLazyState' = runLazyState
402+{-# INLINE runLazyState' #-}
403+
404+evalLazyStateT' :: Monad m => LazyStateT' s m a -> s -> m a
405+evalLazyStateT' = evalLazyStateT
406+{-# INLINE evalLazyStateT' #-}
407+
408+type LazyState' s = LazyState s s
409+
410+evalLazyState' :: LazyState' s a -> s -> a
411+evalLazyState' = evalLazyState
412+{-# INLINE evalLazyState' #-}
413+
414+execLazyStateT' :: Monad m => LazyStateT' s m a -> s -> m s
415+execLazyStateT' = execLazyStateT
416+{-# INLINE execLazyStateT' #-}
417+
418+execLazyState' :: LazyState' s a -> s -> s
419+execLazyState' = execLazyState
420+{-# INLINE execLazyState' #-}
421+
422+type MonadState' s = MonadState s s
423+
424+get' :: forall s m . MonadState' s m => m s
425+get' = get @s
426+{-# INLINE get' #-}
427+
428+gets' :: forall s m a . MonadState' s m => (s -> a) -> m a
429+gets' = gets @s
430+{-# INLINE gets' #-}
431+
432+put' :: forall s m . MonadState' s m => s -> m ()
433+put' = put @s
434+{-# INLINE put' #-}
435+
436+state' :: forall s m a . MonadState' s m => (s -> (a, s)) -> m a
437+state' = state @s
438+{-# INLINE state' #-}
439+
440+modify' :: forall s m . MonadState' s m => (s -> s) -> m ()
441+modify' = modify @s
442+{-# INLINE modify' #-}
443+
444+-- | Encode type-level information for 'zoom'.
445+data ZOOM t z
446+
447+type ZoomT t (z :: Type) = Handler (ZOOM t z) IdentityT
448+
449+newtype ReifiedLens s t a b = Lens (Lens s t a b)
450+
451+type ReifiedLens' s a = ReifiedLens s s a a
452+
453+-- | Zoom into a part of a state using a lens.
454+zoom
455+ :: forall tag sOuter sInner m a
456+ . Lens' sOuter sInner
457+ -> (forall z . Reifies z (ReifiedLens' sOuter sInner) => ZoomT tag z m a)
458+ -> m a
459+zoom l m = reify (Lens l) (\(_ :: Proxy z) -> coerce (m @z))
460+{-# INLINE zoom #-}
461+
462+instance
463+ ( MonadState tag sOuter m
464+ , Reifies z (ReifiedLens' sOuter sInner)
465+ , trans ~ IdentityT
466+ ) => MonadState tag sInner (Handler (ZOOM tag z) trans m)
467+ where
468+ state =
469+ (coerce :: forall eff r a .
470+ (r -> m a) ->
471+ (r -> Handler eff trans m a))
472+ (state @tag . l)
473+ where
474+ Lens l = reflect (Proxy :: Proxy z)
475+ {-# INLINE state #-}