Mixing IO w/ ST Monad - "type variable `s2' would escape its scope"

I decided to simplify my code to see what conditions produced the error. I start with a simple nested "s" similar to ST s (STArray s x y) like so:

{-# LANGUAGE RankNTypes #-}
import Control.Monad.ST
import Control.Applicative

data Foo s = Foo { foo::Bool }

newFoo :: ST s (Foo s)
newFoo = return $ Foo False

To test the state code, I run the following transformation:

changeFoo :: (forall s. ST s (Foo s)) -> ST s (Foo s)
changeFoo sf = do
   f <- sf
   let f' = Foo (not $ foo f)
   return f'

I would like to extract a value from the state while keeping the state, so the next step is to extract the Bool value:

splitChangeFoo :: (forall s. ST s (Foo s)) -> ST s (Bool,(Foo s))
splitChangeFoo sf = do
   f <- changeFoo sf
   let b = foo f
   return (b,f)

In order to extract that Bool I must use runST. My understanding is that this will create an additional state, which I specify by providing a forall s. in the return type:

extractFoo :: (forall s. ST s (Bool,(Foo s))) -> (forall s. (Bool,ST s ((Foo s))))
extractFoo sbf = (runST $ fst <$> sbf,snd <$> sbf)

The example above does compile without the final forall however in the circumstance I am trying to debug this is not possible. Regardless, in this case it compiles either way.

I am able to use the above code to chain multiple uses of the state together:

testFoo :: (Bool, ST s (Foo s))
testFoo = (b && b',sf')
   where
      (b,sf) = extractFoo $ splitChangeFoo newFoo
      (b',sf') = extractFoo $ splitChangeFoo sf

Now I try to throw IO into the mix and so I make use of the applicative fmap <$>. This does not compile: (NB. same problem if I use fmap or >>= return rather than <$>)

testBar :: IO (Bool, ST s (Foo s))
testBar = (\(b,sf) -> extractFoo $ splitChangeFoo sf) <$> testBar'
   where
      testBar' :: IO (Bool, ST s (Foo s))
      testBar' = return $ extractFoo $ splitChangeFoo newFoo

With the following error:

Couldn't match type `s0' with `s2'
  because type variable `s2' would escape its scope
This (rigid, skolem) type variable is bound by
  a type expected by the context: ST s2 (Foo s2)
The following variables have types that mention s0
  sf :: ST s0 (Foo s0) (bound at src\Tests.hs:132:16)
Expected type: ST s2 (Foo s2)
  Actual type: ST s0 (Foo s0)
In the first argument of `splitChangeFoo', namely `sf'
In the second argument of `($)', namely `splitChangeFoo sf'
In the expression: extractFoo $ splitChangeFoo sf

I'm aware from this SO question about ST function composition that not all possible uses of ST are accounted for by the compiler. To test this assumption I modified the above function to not use IO and simply pass the return value into a lambda:

testFooBar :: (Bool, ST s (Foo s))
testFooBar = (\(b,sf) -> extractFoo $ splitChangeFoo sf) testFooBar'
   where
      testFooBar' :: (Bool, ST s (Foo s))
      testFooBar' = extractFoo $ splitChangeFoo newFoo

Predictably this also does not compile with an identical error message.

This presents a challenge. I have a reasonable amount of IO that needs to interact with a set of deeply nested ST operations. It works perfectly fine for a single iteration, however I am not able to make further use of the return value.

Solution

I'm not sure how all of this works out in your actual code, and it seems to me that you're doing something more convoluted than strictly necessary, but as I can't say much definitive about that I'll only address the immediate problems in your question.

Why `testBar` doesn't work

In (\(b,sf) -> extractFoo $ splitChangeFoo sf) <$> testBar', we have <$>, which has type forall a b. (a -> b) -> IO a -> IO b. However, testBar has type forall s. IO (Bool, ST s (Foo s)). We can only apply fmap to testBar after the s there has been instantiated.

In your case, s is instantiated to a fresh variable s0, so inside the lambda sf also has that parameter. Unfortunately, s0 is distinct from the topmost s parameter (which we'd like to see in the return type), so sf is pretty much unusable now.

In the case of testFooBar, the issue is that lambda arguments are monomorphic by default (or else type inference would become undecidable), so we just have to annotate the type:

testFooBar :: (Bool, ST s (Foo s))
testFooBar =
  (\(x :: forall s. (Bool, (ST s (Foo s)))) -> extractFoo $ splitChangeFoo $ snd x) testFooBar'
   where
      testFooBar' :: (Bool, ST s (Foo s))
      testFooBar' = extractFoo $ splitChangeFoo newFoo

All we've done here is to introduce an unwieldy alternative for let and where expressions. But this has little significance, since in testFoo the fmap necessitates that the lambda argument has a fully instantiated type, so there we can't annotate our woes away.

We might try the following:

{-# LANGUAGE ScopedTypeVariables #-}

testBar :: forall s. IO (Bool, ST s (Foo s))
testBar = (\(b,sf) -> extractFoo $ splitChangeFoo sf) <$> testBar'
   where
      testBar' :: IO (Bool, ST s (Foo s))
      testBar' = return $ extractFoo $ splitChangeFoo newFoo

Now testBar has the right s parameter to begin with; since it's not generalized anymore, we can apply fmap to it right away. Inside the lambda, sf has type ST s (Foo s). But this also doesn't work out, since extractFoo expects a forall quantified argument, and sf is monomorphic. Now - if you try this version in GHCi - the error message we get is no longer about escaping skolems, but a good old-fashioned can't match type s with s2.

How to help it

Clearly, extract must return a forall s.-ed ST s (Foo s) somehow. Your original type is no good because GHC doesn't have polymorphic return types, so

(forall s. ST s (Bool, Foo s)) -> (forall s. (Bool, ST s (Foo s)))

actually means

forall s. (forall s. ST s (Bool, Foo s)) -> (Bool, ST s (Foo s))

since GHC floats out all forall-s to the top of the scope. This makes type inference easier and allows more terms to typecheck. Indeed, you can enter :t extract in GHCi and be presented with the floated-out type.

A possible solution would be enabling ImprediactiveTypes and having

extract :: (forall s. ST s (Bool, Foo s)) -> (Bool, forall s. ST s (Foo s))

But there's a reason why ImpredicativeTypes is rarely used: it's poorly integrated with the rest of the type checking machinery and has a tendency to fail for all but the simplest purposes. As an example, assuming extract with the modified type, the following example doesn't typecheck:

testBar :: IO (Bool, forall s. ST s (Foo s))
testBar = (\(b, sf) -> (b, sf)) <$> testBar'
  where
    testBar' :: IO (Bool, forall s. ST s (Foo s))
    testBar' = return $ extractFoo $ splitChangeFoo newFoo

But it does typecheck when we change (\(b, sf) -> (b, sf)) to (\x -> x)! So let's forget about impredicative types.

The usual solution is to just wrap the quantified types in a newtype. This is a lot less flexible than impredicative types, but at least it works.

{-# LANGUAGE RankNTypes #-}

import Control.Monad.ST
import Control.Applicative
import Control.Arrow

data Foo s = Foo { foo::Bool }
newtype ST' f = ST' {unST' :: forall s. ST s (f s)}

newFoo :: ST s (Foo s)
newFoo = return $ Foo False

splitChangeFoo :: ST s (Foo s) -> ST s (Bool, Foo s)
splitChangeFoo = fmap (\(Foo b) -> (b, Foo (not b)))

extract :: (forall s. ST s (Bool, Foo s)) -> (Bool, ST' Foo)
extract s = (runST $ fst <$> s, ST' $ snd <$> s)

testBar :: IO (Bool, ST s (Foo s))
testBar = (\(_, ST' s) -> second unST' $ extract $ splitChangeFoo s) <$>
          (return $ extract $ splitChangeFoo newFoo)

So, whenever you need to return quantified types, wrap it in a newtype, then unwrap as needed. You might want to stick to a common type argument scheme for you custom types (for example, always make the s parameter the last, so you can use a single ST' type for multiple to-be-wrapped types).

Mixing IO w/ ST Monad - "type variable `s2' would escape its scope"

Why testBar doesn't work

How to help it

Why `testBar` doesn't work