Converting this FreeT (explicitly recursive data type) function to work on FT (church encoding)

Question

I'm using the FreeT type from the free library to write this function which "runs" an underlying StateT:

runStateFree
    :: (Functor f, Monad m)
    => s
    -> FreeT f (StateT s m) a
    -> FreeT f m (a, s)
runStateFree s0 (FreeT x) = FreeT $ do
    flip fmap (runStateT x s0) $ \(r, s1) -> case r of
      Pure y -> Pure (y, s1)
      Free z -> Free (runStateFree s1 <$> z)

However, I'm trying to convert it to work on FT, the church-encoded version, instead:

runStateF
    :: (Functor f, Monad m)
    => s
    -> FT f (StateT s m) a
    -> FT f m (a, s)
runStateF s0 (FT x) = FT $ \ka kf -> ...

but I'm not quite having the same luck. Every sort of combination of things I get seems to not quite work out. The closest I've gotten is

runStateF s0 (FT x) = FT $ \ka kf ->
    ka =<< runStateT (x pure (\n -> _ . kf (_ . n)) s0

But the type of the first hole is m r -> StateT s m r and the type the second hole is StateT s m r -> m r...which means we necessarily lose the state in the process.

I know that all FreeT functions are possible to write with FT. Is there a nice way to write this that doesn't involve round-tripping through FreeT (that is, in a way that requires explicitly matching on Pure and Free)? (I've tried manually inlining things but I don't know how to deal with the recursion using different ss in the definition of runStateFree). Or maybe this is one of those cases where the explicit recursive data type is necessarily more performant than the church (mu) encoding?

Answer 1

Here's the definition. There are no tricks in the implementation itself. Don't think and make it type check. Yes, at least one of these fmap is morally questionable, but the difficulty is actually to convince ourselves it does the Right thing.

runStateF
    :: (Functor f, Monad m)
    => s
    -> FT f (StateT s m) a
    -> FT f m (a, s)
runStateF s0 (FT run) = FT $ \return0 handle0 ->
  let returnS a = StateT (\s -> fmap (\r -> (r, s)) (return0 (a, s)))
      handleS k e = StateT (\s -> fmap (\r -> (r, s)) (handle0 (\x -> evalStateT (k x) s) e))
  in evalStateT (run returnS handleS) s0

We have two stateless functions (i.e., plain m)

return0 :: a -> m r
handle0 :: forall x. (x -> m r) -> f x -> m r

and we must wrap them in two stateful (StateT s m) variants with the signatures below. The comments that follow give some details about what is going on in the definition of handleS.

returnS :: a -> StateT s m r
handleS :: forall x. (x -> StateT s m r) -> f x -> StateT s m r

-- 1.                                               --    ^   grab the current state 's' here
-- 2.                                               --      ^ call handle0 to produce that 'm'
-- 3.                             ^ here we will have to provide some state 's': pass the current state we just grabbed.
--                                  The idea is that 'handle0' is stateless in handling 'f x',
--                                  so it is fine for this continuation (x -> StateT s m r) to get the state from before the call to 'handle0'

There is an apparently dubious use of fmap in handleS, but it is valid as long as run never looks at the states produced by handleS. It is almost immediately thrown away by one of the evalStateT.

In theory, there exist terms of type FT f (StateT s m) a which break that invariant. In practice, that almost certainly doesn't occur; you would really have to go out of your way to do something morally wrong with those continuations.

In the following complete gist, I also show how to test with QuickCheck that it is indeed equivalent to your initial version using FreeT, with concrete evidence that the above invariant holds:

https://gist.github.com/Lysxia/a0afa3ca2ea9e39b400cde25b5012d18

Answer 2

I'd say that no, as even something as simple as cutoff converts to FreeT:

cutoff :: (Functor f, Monad m) => Integer -> FT f m a -> FT f m (Maybe a)
cutoff n = toFT . FreeT.cutoff n . fromFT

---

In general, you're probably looking at:

improve :: Functor f => (forall m. MonadFree f m => m a) -> Free f a

Improve the asymptotic performance of code that builds a free monad with only binds and returns by using F behind the scenes.

I.e. you'll construct Free efficiently, but then do whatever you need to do with Free (maybe again, by improveing).