Result of a continuation and running code after resuming - Polysemy

fun show(p : (int, list<int>)) : string
  show-tuple(p, show, show)

effect state<s>
  ctl get() : s
  ctl put(s : s) : ()

fun comp() : <state<int>> int
  val x = get()
  put(x + 1 : int)
  val y = get()
  put(y + y : int)
  get()

fun logState(init : s, action : () -> <state<s>> a) : <div> (a, list<s>)
  var s := init
  with handler
    return(x) -> (x, [])
    ctl get() -> resume(s)
    ctl put(s') -> {s := s'; val (x, ss) = resume(()); (x, Cons(s', ss))}
  action()

> logState(2, comp)
check  : interactive
check  : interactive
linking: interactive
created: .koka/v2.4.0/gcc-debug/interactive

(6,[3,6])

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module AfterResume where

import Control.Algebra (Algebra (..), Has, send, type (:+:) (..))
import Control.Carrier.Lift (Lift, runM, sendIO)
import Data.Kind (Type)
import Data.Type.Equality ()
import Debug.Trace (traceM)

data After (m :: Type -> Type) k where
  After :: String -> After m ()

after :: Has After sig m => String -> m ()
after s = send (After s)

newtype AfterC m a = AfterC {runAfterC :: forall b. (a -> m b) -> m b}
  deriving (Functor)

instance Applicative m => Applicative (AfterC m) where
  pure a = AfterC ($ a)
  AfterC f <*> AfterC v = AfterC $ \c -> f $ \g -> v (c . g)

instance Monad m => Monad (AfterC m) where
  return x = AfterC ($ x)
  m >>= k = AfterC $ \c -> runAfterC m (\x -> runAfterC (k x) c)

evalAfterC :: (Monad m) => AfterC m r -> m r
evalAfterC m = runAfterC m return

lift :: (Monad m) => m a -> AfterC m a
lift m = AfterC (m >>=)

resetT :: (Monad m) => AfterC m r -> AfterC m r
resetT = lift . evalAfterC

shiftT :: Monad m => (forall b. (a -> m b) -> AfterC m b) -> AfterC m a
shiftT f = AfterC (evalAfterC . f)

instance (Has (Lift IO) sig m, Algebra sig m) => Algebra (After :+: sig) (AfterC m) where
  alg hdl sig ctx = case sig of
    L (After s) -> shiftT (\k -> lift $ do
        res <- k (() <$ ctx)
        sendIO (putStrLn s)
        return res
      )
    R other -> lift $ alg (evalAfterC . hdl) other ctx

test :: (Has After sig m, Has (Lift IO) sig m) => m ()
test = do
  traceM "0"
  sendIO (putStrLn "a")
  after "a"
  traceM "1"
  sendIO (putStrLn "b")
  after "b"
  traceM "2"
  return ()

runTest :: IO ()
runTest = do
  runM . evalAfterC $ test

0
a
1
b
2
b
a