Can I save the result of localSeqUnlift and call it some time later?

[ruifengx]

Update: in the end I settled on using a Day convolution for the staging, as proposed in "Breadth-First Traversal Via Staging". More specifically, I ended up with Staged es = Day (Eff (A : B : C : es)) (Eff (D : E : es)). It requires manually lifting the effects, but I think it is the best I can achieve for now.

---

I am trying to implement a custom effect to support multi-stage computation. In terms of mtl-style monad transformers, the effect is exactly WriterT w m (ReaderT w m a). The ReaderT delays access to w, so all tell operations in the WriterT stage are first accumulated, and the accumulated result is provided as the environment to the ReaderT stage.

Below is an experimental implementation (simplified, the actual effect would have w fixed and have more specific effect operations):

data Staged (w :: Type) :: Effect where
  Tell :: w -> Staged w m ()
  Ask :: (w -> m ()) -> Staged w m ()

type instance DispatchOf (Staged w) = Dynamic

runStaged :: forall w es a. Monoid w => Eff (Staged w : es) a -> Eff es a
runStaged = reinterpret finalise exec
  where
    finalise :: Eff (Writer (w, w -> Ap (Eff es) ()) : es) a -> Eff es a
    finalise m = do
      (x, (w, k)) <- runWriter m
      x <$ getAp (k w)
    exec
      :: LocalEnv ls (Writer (w, w -> Ap (Eff es) ()) : es)
      -> Staged w (Eff ls) r
      -> Eff (Writer (w, w -> Ap (Eff es) ()) : es) r
    exec _ (Tell x) = tell @(w, w -> Ap (Eff es) ()) (x, const mempty)
    exec env (Ask k) = do
      let lk w = Ap (localSeqUnlift env (\unlift -> unlift (k w)))
      tell @(w, w -> Ap (Eff es) ()) (mempty, lk)

The only problem I see is that I save the result of localSeqUnlift via the handler effect Writer and call it at the end of runStaged in finalise. Since the function has local in its name, I have a feeling that the unlifting is only valid within some scope, but I do not know whether (1) the scope is limited to the interpreter exec, or (2) the scope also extends to cover finalise, or even (3) there is no scope, it is valid as long as the program type-checks.

I have checked that the above code snippet compiles and runs (the example is completely meaningless; it is just a test program):

test :: (Staged String :> es, Writer [String] :> es) => Eff es Int
test = do
  send (Tell @String "x")
  send (Ask \(s :: String) -> tell [s])
  send (Tell @String "y")
  pure 0

>>> runPureEff (runWriter @[String] (runStaged @String test))
(0,["xy"])

But I want to make sure my use of localSeqUnlift is sound.

[arybczak]

But I want to make sure my use of localSeqUnlift is sound.

Not quite. It will break as soon as you try to use properly local effects, because then your code will attempt to run them after they went out of scope.

test :: (Staged String :> es, Writer [String] :> es, Reader String :> es) => Eff es Int
test = do
  send (Tell @String "x")
  send (Ask $ \(s :: String) -> ask >>= \z -> tell [s, z])
  send (Tell @String "y")
  pure 0

>>> runPureEff (runWriter @[String] (runStaged @String $ runReader @String "z" test))
*** Exception: version (4) /= storageVersion (0)
If you're attempting to run an unlifting function outside of the scope of effects it captures, have a look at UnliftStrategy (SeqForkUnlift).

HasCallStack backtrace:
  error, called at src/Effectful/Internal/Env.hs:438:5 in effectful-core-2.6.0.0-inplace:Effectful.Internal.Env
  getLocation, called at src/Effectful/Internal/Env.hs:390:14 in effectful-core-2.6.0.0-inplace:Effectful.Internal.Env
  getEnv, called at src/Effectful/Internal/Monad.hs:677:35 in effectful-core-2.6.0.0-inplace:Effectful.Internal.Monad
  getStaticRep, called at src/Effectful/Reader/Static.hs:53:15 in effectful-core-2.6.0.0-inplace:Effectful.Reader.Static
  ask, called at tests/Main.hs:52:33 in effectful-2.6.0.0-inplace-test:Main

Btw, just switching to SeqForkUnlift in your case won't help anyway because of this line:

let lk w = Ap (localSeqUnlift env (\unlift -> unlift (k w)))

You're delaying creation of the unlifting function until finalise, where the environment is no longer valid.

This "works":

runStaged :: forall w es a. Monoid w => Eff (Staged w : es) a -> Eff es a
runStaged = reinterpret finalise exec
  where
    finalise :: Eff (Writer (w, w -> Ap (Eff es) ()) : es) a -> Eff es a
    finalise m = do
      (x, (w, k)) <- runWriter m
      x <$ getAp (k w)
    exec
      :: LocalEnv ls (Writer (w, w -> Ap (Eff es) ()) : es)
      -> Staged w (Eff ls) r
      -> Eff (Writer (w, w -> Ap (Eff es) ()) : es) r
    exec _ (Tell x) = tell @(w, w -> Ap (Eff es) ()) (x, const mempty)
    exec env (Ask k) = do
      lk <- raise $ localUnlift env SeqForkUnlift $ \unlift -> do
        pure $ \w -> Ap (unlift (k w))
      tell @(w, w -> Ap (Eff es) ()) (mempty, lk)

but it doesn't produce the results you want:

>>> runPureEff (runWriter @[String] (runStaged @String $ runReader @String "z" test))
(0,[])

I don't know exactly what you're trying to do, but this doesn't look like the way to do it.

[ruifengx]

I want to avoid appealing to laziness and tying the knot between Writer and Reader. What I want is precisely my runStaged function, but properly encapsulated as a single effect so that I do not carry around that horrible Writer (w, w -> Ap (Eff es) ()). Besides being lengthy and ugly, it also has the problem that it refers to the "tail effect stack" es, so it is not compositional and works poorly with (:>).

The user program of this effect does two things interleaved. It generates w values, which are accumulated. It also does other things which depends on the final accumulated w value. Usually in Haskell, we do it by tying the knot and rely on laziness:

runStaged :: MonadFix m => WriterT w (ReaderT w m) a -> m a
runStaged m = do
  rec (a, w) <- runReaderT (runWriterT m) w
  pure a

But I want to avoid laziness here, because it is very easy to accidentally tell something that depends on a previous ask, so instead, I think this is more appropriate:

runStaged :: Monad m => WriterT (w, w -> Ap m ()) m a
runStaged m = do
  (x, (w, k)) <- runWriterT m
  x <$ getAp k w

You see, inside the ReaderT, we have no access to the WriterT, so the two stages are cleanly separated. All actions that depends on the final w is captured and accumulated alongside the w, and they are batch performed at the end of runStaged.

What I want is to implement this writer-reader composition as an Effect, but of course as you have shown, my version would not work. I do not in particular care about the exact signature of the effect operations, so if you know any alternative that works properly, I am happy to adopt it.

[ruifengx]

I have a feeling that the ability to subsume is what makes runStaged problematic. If I capture the actions with effects es, but later subsume the top-level Staged effect, the remaining effects might not support running the captured actions any more. Maybe the only way to make it work is to parametrise Staged with es to pin it to some specific position in the effect stack and make it clear what effects can be used.

[ruifengx]

I guess the answer to my original question is "no, one cannot save the result of localSeqUnlift and call it later". Since I do not see any obvious way to replicate the monad transformer using effectful, I will stick to mtl for now for this specific task.