`FreeT f (StateT s m) a -> FreeT f m (a, s)`, Church-encoded https://stackoverflow.com/questions/59440787/converting-this-freet-explicitly-recursive-data-type-function-to-work-on-ft-c

RunStateFree.hs

{-# LANGUAGE BangPatterns, RankNTypes, DeriveFunctor #-}

import Control.Monad.Trans.Free (FreeT(..), FreeF(..), foldFreeT)
import Control.Monad.Trans.Free.Church
import Control.Monad.State

import Control.Applicative
import Data.Functor.Identity

import Test.QuickCheck

-- The original "state-hoist" operation, defined on FreeT
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)

-- An equivalent operation, on the Church-encoded variant FT
runStateF
    :: (Functor f, Monad m)
    => s
    -> FT f (StateT s m) a
    -> FT f m (a, s)
runStateF s0 (FT run) = FT $ \return0 handle0 ->
  -- Read Implementation notes below.
  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

{- = Implementation notes

  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 variants

    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 the continuation gets the state from before the call to 'handle0'
    --                                  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'.

  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 have to go out of your way to do
  something morally wrong with those continuations.

-}

{-
  Nevertheless, these are some subtle invariants.
  How to go about testing that function?
  (You can also verify it formally but I'm too lazy to do that.)

  We will test that 'runStateFree' and 'runStateF' are equivalent, in the sense of
  this equation between functions 'FreeT f (StateT s m) a -> FreeT f m (a, s)',
  for any 's0 :: s':

    fromFT . runStateF s0 . toFT  =  runStateFree s0

  ... by generating random inputs and testing that the results are equal.
  Note that we will generate inputs of type 'FreeT f (StateT s m) a', before
  converting them to 'FT f (StateT s m) a'. The effectiveness of that test
  relies on the assumption that the 'FT' computations we care about (satisfying
  the "invariant" that, to be clear, is not made formal here at all) can all be
  represented in 'FreeT'. I'll let you convince yourself of it.

  Then it can also be argued that, by parametricity in f, m, a, s, it is sufficient
  to test the above equation only using:

  -- (Pretend Int is Integer)

  f = (Int, Int -> _)  -- "most general" functor (in some sense), StoreF below
  m = Free StoreF      -- "most general" monad
  a = s = Int          -- "most general" type

  They are "most general" in the sense that they can be mapped surjectively
  (the mappings are allowed to be partial) into any other functor/monad/type
  (modulo the assumption that the type (->) represents only *computable*
  functions). The point is that correctness on this one "most general" domain
  can be transported to correctness on all other domains using parametricity
  (free theorems).

  Below, we implement a generator of 'FreeT f (StateT s m) a' for those
  instantiations (aka. M1 Int, see 'genM1'), and an equality test of 'FreeT f m (a, s)'
  (aka. M2 (Int, Int)), see 'eqM2').

  That's what we need to then define the QuickCheck property 'prop_runStateF'
  embodying the equation above.

  In 'eqM2', we also throw in some 'classify' to see that we are
  making some non-trivial comparisons.

-}

data StoreF a = StoreF Int (Int -> a)
  deriving Functor

type M0 = FreeT StoreF Identity  -- Free StoreF
type M1 = FreeT StoreF (StateT Int M0)
type M2 = FreeT StoreF M0

-- Generate trees of type 'FreeT f m a',
-- given generators for nodes 'm', 'f', and leaves 'a'.
genFreeT :: (forall a. Gen a -> Gen (m a)) -> (forall a. Gen a -> Gen (f a)) -> Gen a -> Gen (FreeT f m a)
genFreeT genM genF genA = go where
  go =
    sized $ \n ->
      FreeT <$> genM
        (if n == 0 then
          Pure <$> genA
         else resize (n-1) $ frequency
          [ (3, Free <$> genF go)
          , (1, Pure <$> genA)
          ])

genStoreF :: Gen a -> Gen (StoreF a)
genStoreF ga = StoreF <$> arbitrary <*> liftArbitrary ga

genM0 :: Gen a -> Gen (M0 a)
genM0 = genFreeT (fmap Identity) genStoreF

genStateT :: CoArbitrary s => (forall a. Gen a -> Gen (m a)) -> Gen s -> Gen a -> Gen (StateT s m a)
genStateT genM genS genA = StateT <$> liftArbitrary (genM (liftA2 (,) genA genS))

genM1 :: Gen a -> Gen (M1 a)
genM1 = genFreeT (genStateT genM0 arbitrary) genStoreF

-- Type of comparison functions between functors.
type EqF f = (forall x. (x -> x -> Property) -> (f x -> f x -> Property))

-- Compare trees of type 'FreeT f m a',
-- given comparisons between nodes 'm', 'f', and leaves 'a'.
--
-- The comparisons may be partial: they are QuickCheck properties that may
-- still succeed even if the trees are different; this is to avoid looping
-- on infinite trees. But:
--
-- - if the property fails, then the trees are definitely different;
-- - if the trees are different, then there exists some execution of the property
--   which will fail.
--
-- Implementation:
-- Generate a path randomly, and check deterministically that the
-- restrictions of the two trees on that path are equal.
eqFreeT ::
  (Arbitrary e, Show e) =>
  EqF m -> (e -> EqF f) -> EqF (FreeT f m)
eqFreeT eqM eqF eqA t1 t2 =
  forAll (getNonEmpty <$> arbitrary) $ \es -> eqFreeT' eqM eqF eqA es t1 t2

eqFreeT' ::
  EqF m -> (e -> EqF f) -> (a -> a -> Property) ->
  [e] -> FreeT f m a -> FreeT f m a -> Property
eqFreeT' eqM eqF eqA = go 0 where
  go !i es (FreeT t1) (FreeT t2) = eqM (eqFreeF (eqEs i es) eqA) t1 t2
  eqEs i [] = \_ _ -> collect i (property True)
  eqEs i (e : es) = eqF e (go (i+1) es)

eqFreeF ::
  (f b -> f b -> Property) -> (a -> a -> Property) ->
  FreeF f a b -> FreeF f a b -> Property
eqFreeF eqF _ (Free u1) (Free u2) = eqF u1 u2
eqFreeF _ eqA (Pure a1) (Pure a2) = eqA a1 a2
eqFreeF _ _ _ _ = property False

eqStateT ::
  (Arbitrary s, Show s) =>
  EqF m -> (s -> s -> Property) ->
  EqF (StateT s m)
eqStateT eqM eqS eqA (StateT u1) (StateT u2) =
  property $ \s -> eqM (eq2 eqA eqS) (u1 s) (u2 s)

eq2 :: (a -> a -> Property) -> EqF ((,) a)
eq2 eqA eqB (a1, b1) (a2, b2) = eqA a1 a2 .&&. eqB b1 b2

eqStoreF :: Int -> EqF StoreF
eqStoreF i eqA (StoreF s1 f1) (StoreF s2 f2) =
  s1 === s2 .&&. eqA (f1 i) (f2 i)

type A = Int
type S = Int

eqIdentity :: EqF Identity
eqIdentity f (Identity a) (Identity b) = f a b

eqM0 :: EqF M0
eqM0 = eqFreeT eqIdentity eqStoreF

genM1_ :: Gen (M1 Int)
genM1_ = genM1 arbitrary

eqM2 :: EqF M2
eqM2 = eqFreeT eqM0 eqStoreF

prop_runStateF :: M1 A -> S -> Property
prop_runStateF t0 s0 =
  let f =? g = eqM2 (===) (f t0) (g t0) in
  (fromFT . runStateF s0 . toFT) =? runStateFree s0

prop_runStateF_ :: Property
prop_runStateF_ = forAllBlind genM1_ prop_runStateF

main :: IO ()
main = quickCheck prop_runStateF_