-- SPDX-License-Identifier: MPL-2.0 AND BSD-3-Clause

-- This Source Code Form is subject to the terms of the Mozilla Public
-- License, v. 2.0. If a copy of the MPL was not distributed with this
-- file, You can obtain one at https://mozilla.org/MPL/2.0/.

{- |
Copyright   :  (c) 2016 Allele Dev; 2017 Ixperta Solutions s.r.o.; 2017 Alexis King; 2024 Sayo Koyoneda
License     :  MPL-2.0 (see the LICENSE file) AND BSD-3-Clause (see the LICENSE.BSD3 file)
Maintainer  :  ymdfield@outlook.jp

This module provides an ad-hoc specialized version of
 "Control.Monad.Hefty.Interpret" to accelerate interpretations that have a
single state type @s@, especially for effects like t'Data.Effect.State.State' or
 [@Writer@]("Data.Effect.Writer").
-}
module Control.Monad.Hefty.Interpret.State where

import Control.Effect (type (~>))
import Control.Monad.Hefty.Interpret (qApp)
import Control.Monad.Hefty.Types (Eff (Op, Val), sendUnionBy, sendUnionHBy)
import Data.Effect.OpenUnion.Internal (IsSuffixOf)
import Data.Effect.OpenUnion.Internal.FO (Union, prj, weakens, (!+), type (<|))
import Data.Effect.OpenUnion.Internal.HO (UnionH, hfmapUnion, nilH)
import Data.Kind (Type)

-- | An ad-hoc stateful version of t'Control.Monad.Hefty.Types.Interpreter' for performance.
type StateInterpreter s e m (ans :: Type) = forall x. e x -> s -> (s -> x -> m ans) -> m ans

-- | An ad-hoc stateful version of t'Control.Monad.Hefty.Types.Elaborator' for performance.
type StateElaborator s e m ans = StateInterpreter s (e m) m ans

-- * Interpretation functions

interpretStateBy
    :: forall s e ef ans a
     . s
    -> (s -> a -> Eff '[] ef ans)
    -> StateInterpreter s e (Eff '[] ef) ans
    -> Eff '[] (e ': ef) a
    -> Eff '[] ef ans
interpretStateBy :: forall s (e :: * -> *) (ef :: [* -> *]) ans a.
s
-> (s -> a -> Eff '[] ef ans)
-> StateInterpreter s e (Eff '[] ef) ans
-> Eff '[] (e : ef) a
-> Eff '[] ef ans
interpretStateBy = s
-> (s -> a -> Eff '[] ef ans)
-> StateInterpreter s e (Eff '[] ef) ans
-> Eff '[] (e : ef) a
-> Eff '[] ef ans
forall s (e :: * -> *) (ef' :: [* -> *]) (ef :: [* -> *]) ans a.
IsSuffixOf ef ef' =>
s
-> (s -> a -> Eff '[] ef' ans)
-> StateInterpreter s e (Eff '[] ef') ans
-> Eff '[] (e : ef) a
-> Eff '[] ef' ans
reinterpretStateBy
{-# INLINE interpretStateBy #-}

reinterpretStateBy
    :: forall s e ef' ef ans a
     . (ef `IsSuffixOf` ef')
    => s
    -> (s -> a -> Eff '[] ef' ans)
    -> StateInterpreter s e (Eff '[] ef') ans
    -> Eff '[] (e ': ef) a
    -> Eff '[] ef' ans
reinterpretStateBy :: forall s (e :: * -> *) (ef' :: [* -> *]) (ef :: [* -> *]) ans a.
IsSuffixOf ef ef' =>
s
-> (s -> a -> Eff '[] ef' ans)
-> StateInterpreter s e (Eff '[] ef') ans
-> Eff '[] (e : ef) a
-> Eff '[] ef' ans
reinterpretStateBy s
s0 s -> a -> Eff '[] ef' ans
ret StateInterpreter s e (Eff '[] ef') ans
hdl =
    s
-> (s -> a -> Eff '[] ef' ans)
-> StateInterpreter
     s (UnionH '[] (Eff '[] (e : ef))) (Eff '[] ef') ans
-> StateInterpreter s (Union (e : ef)) (Eff '[] ef') ans
-> Eff '[] (e : ef) a
-> Eff '[] ef' ans
forall s (eh :: [EffectH]) (ef :: [* -> *]) (m :: * -> *) ans a.
Monad m =>
s
-> (s -> a -> m ans)
-> StateInterpreter s (UnionH eh (Eff eh ef)) m ans
-> StateInterpreter s (Union ef) m ans
-> Eff eh ef a
-> m ans
iterStateAllEffHFBy s
s0 s -> a -> Eff '[] ef' ans
ret UnionH '[] (Eff '[] (e : ef)) x
-> s -> (s -> x -> Eff '[] ef' ans) -> Eff '[] ef' ans
StateInterpreter
  s (UnionH '[] (Eff '[] (e : ef))) (Eff '[] ef') ans
forall (f :: * -> *) a r. UnionH '[] f a -> r
nilH (e x -> s -> (s -> x -> Eff '[] ef' ans) -> Eff '[] ef' ans
StateInterpreter s e (Eff '[] ef') ans
hdl (e x -> s -> (s -> x -> Eff '[] ef' ans) -> Eff '[] ef' ans)
-> (Union ef x
    -> s -> (s -> x -> Eff '[] ef' ans) -> Eff '[] ef' ans)
-> Union (e : ef) x
-> s
-> (s -> x -> Eff '[] ef' ans)
-> Eff '[] ef' ans
forall (e :: * -> *) a r (es :: [* -> *]).
(e a -> r) -> (Union es a -> r) -> Union (e : es) a -> r
!+ \Union ef x
u s
s s -> x -> Eff '[] ef' ans
k -> (x -> Eff '[] ef' ans) -> Union ef' x -> Eff '[] ef' ans
forall a (eh :: [EffectH]) (ef :: [* -> *]) ans.
(a -> Eff eh ef ans) -> Union ef a -> Eff eh ef ans
sendUnionBy (s -> x -> Eff '[] ef' ans
k s
s) (Union ef x -> Union ef' x
forall (es :: [* -> *]) (es' :: [* -> *]) a.
IsSuffixOf es es' =>
Union es a -> Union es' a
weakens Union ef x
u))
{-# INLINE reinterpretStateBy #-}

interpretStateRecWith
    :: forall s e ef eh a
     . s
    -> (forall ans. StateInterpreter s e (Eff eh ef) ans)
    -> Eff eh (e ': ef) a
    -> Eff eh ef a
interpretStateRecWith :: forall s (e :: * -> *) (ef :: [* -> *]) (eh :: [EffectH]) a.
s
-> (forall ans x.
    e x -> s -> (s -> x -> Eff eh ef ans) -> Eff eh ef ans)
-> Eff eh (e : ef) a
-> Eff eh ef a
interpretStateRecWith = s
-> (forall ans x.
    e x -> s -> (s -> x -> Eff eh ef ans) -> Eff eh ef ans)
-> Eff eh (e : ef) a
-> Eff eh ef a
forall s (e :: * -> *) (ef' :: [* -> *]) (ef :: [* -> *])
       (eh :: [EffectH]) a.
IsSuffixOf ef ef' =>
s
-> (forall ans x.
    e x -> s -> (s -> x -> Eff eh ef' ans) -> Eff eh ef' ans)
-> Eff eh (e : ef) a
-> Eff eh ef' a
reinterpretStateRecWith
{-# INLINE interpretStateRecWith #-}

reinterpretStateRecWith
    :: forall s e ef' ef eh a
     . (ef `IsSuffixOf` ef')
    => s
    -> (forall ans. StateInterpreter s e (Eff eh ef') ans)
    -> Eff eh (e ': ef) a
    -> Eff eh ef' a
reinterpretStateRecWith :: forall s (e :: * -> *) (ef' :: [* -> *]) (ef :: [* -> *])
       (eh :: [EffectH]) a.
IsSuffixOf ef ef' =>
s
-> (forall ans x.
    e x -> s -> (s -> x -> Eff eh ef' ans) -> Eff eh ef' ans)
-> Eff eh (e : ef) a
-> Eff eh ef' a
reinterpretStateRecWith s
s0 forall ans x.
e x -> s -> (s -> x -> Eff eh ef' ans) -> Eff eh ef' ans
hdl = s -> Eff eh (e : ef) ~> Eff eh ef'
loop s
s0
  where
    loop :: s -> Eff eh (e ': ef) ~> Eff eh ef'
    loop :: s -> Eff eh (e : ef) ~> Eff eh ef'
loop s
s =
        s
-> (s -> x -> Eff eh ef' x)
-> StateInterpreter s (UnionH eh (Eff eh (e : ef))) (Eff eh ef') x
-> StateInterpreter s (Union (e : ef)) (Eff eh ef') x
-> Eff eh (e : ef) x
-> Eff eh ef' x
forall s (eh :: [EffectH]) (ef :: [* -> *]) (m :: * -> *) ans a.
Monad m =>
s
-> (s -> a -> m ans)
-> StateInterpreter s (UnionH eh (Eff eh ef)) m ans
-> StateInterpreter s (Union ef) m ans
-> Eff eh ef a
-> m ans
iterStateAllEffHFBy
            s
s
            ((x -> Eff eh ef' x) -> s -> x -> Eff eh ef' x
forall a b. a -> b -> a
const x -> Eff eh ef' x
forall a. a -> Eff eh ef' a
forall (f :: * -> *) a. Applicative f => a -> f a
pure)
            (\UnionH eh (Eff eh (e : ef)) x
u s
s' s -> x -> Eff eh ef' x
k -> (x -> Eff eh ef' x) -> UnionH eh (Eff eh ef') x -> Eff eh ef' x
forall a (eh :: [EffectH]) (ef :: [* -> *]) ans.
(a -> Eff eh ef ans) -> UnionH eh (Eff eh ef) a -> Eff eh ef ans
sendUnionHBy (s -> x -> Eff eh ef' x
k s
s') (UnionH eh (Eff eh ef') x -> Eff eh ef' x)
-> UnionH eh (Eff eh ef') x -> Eff eh ef' x
forall a b. (a -> b) -> a -> b
$ (Eff eh (e : ef) ~> Eff eh ef')
-> UnionH eh (Eff eh (e : ef)) x -> UnionH eh (Eff eh ef') x
forall (f :: * -> *) (g :: * -> *) (es :: [EffectH]) a.
(f ~> g) -> UnionH es f a -> UnionH es g a
hfmapUnion (s -> Eff eh (e : ef) ~> Eff eh ef'
loop s
s') UnionH eh (Eff eh (e : ef)) x
u)
            (e x -> s -> (s -> x -> Eff eh ef' x) -> Eff eh ef' x
forall ans x.
e x -> s -> (s -> x -> Eff eh ef' ans) -> Eff eh ef' ans
hdl (e x -> s -> (s -> x -> Eff eh ef' x) -> Eff eh ef' x)
-> (Union ef x -> s -> (s -> x -> Eff eh ef' x) -> Eff eh ef' x)
-> Union (e : ef) x
-> s
-> (s -> x -> Eff eh ef' x)
-> Eff eh ef' x
forall (e :: * -> *) a r (es :: [* -> *]).
(e a -> r) -> (Union es a -> r) -> Union (e : es) a -> r
!+ \Union ef x
u s
s' s -> x -> Eff eh ef' x
k -> (x -> Eff eh ef' x) -> Union ef' x -> Eff eh ef' x
forall a (eh :: [EffectH]) (ef :: [* -> *]) ans.
(a -> Eff eh ef ans) -> Union ef a -> Eff eh ef ans
sendUnionBy (s -> x -> Eff eh ef' x
k s
s') (Union ef x -> Union ef' x
forall (es :: [* -> *]) (es' :: [* -> *]) a.
IsSuffixOf es es' =>
Union es a -> Union es' a
weakens Union ef x
u))
{-# INLINE reinterpretStateRecWith #-}

-- * Interposition functions

interposeStateBy
    :: forall s e ef ans a
     . (e <| ef)
    => s
    -> (s -> a -> Eff '[] ef ans)
    -> StateInterpreter s e (Eff '[] ef) ans
    -> Eff '[] ef a
    -> Eff '[] ef ans
interposeStateBy :: forall s (e :: * -> *) (ef :: [* -> *]) ans a.
(e <| ef) =>
s
-> (s -> a -> Eff '[] ef ans)
-> StateInterpreter s e (Eff '[] ef) ans
-> Eff '[] ef a
-> Eff '[] ef ans
interposeStateBy s
s0 s -> a -> Eff '[] ef ans
ret StateInterpreter s e (Eff '[] ef) ans
f =
    s
-> (s -> a -> Eff '[] ef ans)
-> StateInterpreter s (UnionH '[] (Eff '[] ef)) (Eff '[] ef) ans
-> StateInterpreter s (Union ef) (Eff '[] ef) ans
-> Eff '[] ef a
-> Eff '[] ef ans
forall s (eh :: [EffectH]) (ef :: [* -> *]) (m :: * -> *) ans a.
Monad m =>
s
-> (s -> a -> m ans)
-> StateInterpreter s (UnionH eh (Eff eh ef)) m ans
-> StateInterpreter s (Union ef) m ans
-> Eff eh ef a
-> m ans
iterStateAllEffHFBy s
s0 s -> a -> Eff '[] ef ans
ret UnionH '[] (Eff '[] ef) x
-> s -> (s -> x -> Eff '[] ef ans) -> Eff '[] ef ans
StateInterpreter s (UnionH '[] (Eff '[] ef)) (Eff '[] ef) ans
forall (f :: * -> *) a r. UnionH '[] f a -> r
nilH \Union ef x
u s
s s -> x -> Eff '[] ef ans
k ->
        Eff '[] ef ans
-> (e x -> Eff '[] ef ans) -> Maybe (e x) -> Eff '[] ef ans
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ((x -> Eff '[] ef ans) -> Union ef x -> Eff '[] ef ans
forall a (eh :: [EffectH]) (ef :: [* -> *]) ans.
(a -> Eff eh ef ans) -> Union ef a -> Eff eh ef ans
sendUnionBy (s -> x -> Eff '[] ef ans
k s
s) Union ef x
u) (\e x
e -> e x -> s -> (s -> x -> Eff '[] ef ans) -> Eff '[] ef ans
StateInterpreter s e (Eff '[] ef) ans
f e x
e s
s s -> x -> Eff '[] ef ans
k) (forall (e :: * -> *) (es :: [* -> *]) a.
Member e es =>
Union es a -> Maybe (e a)
prj @e Union ef x
u)
{-# INLINE interposeStateBy #-}

-- * Transformation to monads

iterStateAllEffHFBy
    :: forall s eh ef m ans a
     . (Monad m)
    => s
    -> (s -> a -> m ans)
    -> StateInterpreter s (UnionH eh (Eff eh ef)) m ans
    -> StateInterpreter s (Union ef) m ans
    -> Eff eh ef a
    -> m ans
iterStateAllEffHFBy :: forall s (eh :: [EffectH]) (ef :: [* -> *]) (m :: * -> *) ans a.
Monad m =>
s
-> (s -> a -> m ans)
-> StateInterpreter s (UnionH eh (Eff eh ef)) m ans
-> StateInterpreter s (Union ef) m ans
-> Eff eh ef a
-> m ans
iterStateAllEffHFBy s
s0 s -> a -> m ans
ret StateInterpreter s (UnionH eh (Eff eh ef)) m ans
fh StateInterpreter s (Union ef) m ans
ff = s -> Eff eh ef a -> m ans
loop s
s0
  where
    loop :: s -> Eff eh ef a -> m ans
loop s
s = \case
        Val a
x -> s -> a -> m ans
ret s
s a
x
        Op Either (UnionH eh (Eff eh ef) x) (Union ef x)
u FTCQueue (Eff eh ef) x a
q -> (UnionH eh (Eff eh ef) x -> (s -> x -> m ans) -> m ans)
-> (Union ef x -> (s -> x -> m ans) -> m ans)
-> Either (UnionH eh (Eff eh ef) x) (Union ef x)
-> (s -> x -> m ans)
-> m ans
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (UnionH eh (Eff eh ef) x -> s -> (s -> x -> m ans) -> m ans
StateInterpreter s (UnionH eh (Eff eh ef)) m ans
`fh` s
s) (Union ef x -> s -> (s -> x -> m ans) -> m ans
StateInterpreter s (Union ef) m ans
`ff` s
s) Either (UnionH eh (Eff eh ef) x) (Union ef x)
u s -> x -> m ans
k
          where
            k :: s -> x -> m ans
k s
s' = s -> Eff eh ef a -> m ans
loop s
s' (Eff eh ef a -> m ans) -> (x -> Eff eh ef a) -> x -> m ans
forall b c a. (b -> c) -> (a -> b) -> a -> c
. FTCQueue (Eff eh ef) x a -> x -> Eff eh ef a
forall (eh :: [EffectH]) (ef :: [* -> *]) a b.
FTCQueue (Eff eh ef) a b -> a -> Eff eh ef b
qApp FTCQueue (Eff eh ef) x a
q
{-# INLINE iterStateAllEffHFBy #-}

-- TODO: add other pattern functions.