Why can't i give an explicit type signature to a local let bound ST action

Question

I have the following simplified program that works fine:

{-# LANGUAGE Rank2Types #-}
module Temp where

import Control.Monad.ST
import Control.Monad
import Data.STRef

mystery :: Int -> Int
mystery start =
  let
    run :: ST s Int
    run = do
      count <- newSTRef (1::Int)
      loop count start
      readSTRef count
  in
    runST run

loop :: STRef s Int -> Int -> ST s ()
loop cnt n = when (n /= 1) 
                  (modifySTRef cnt succ >> 
                     if n `mod` 2 == 0 
                       then  loop cnt (n `div` 2) 
                       else  loop cnt (n * 3 + 1))

I move the loop definition inside the do block to be able to use the counter created in run like so:

mystery :: Int -> Int
mystery start =
  let
    run :: ST s Int
    run = do
      count <- newSTRef (1::Int)
      let
        loop :: Int -> ST s ()
        loop n = when (n /= 1) 
                      (modifySTRef count succ >> 
                        if n `mod` 2 == 0 
                          then  loop (n `div` 2) 
                          else  loop (n * 3 + 1))
      loop start
      readSTRef count
  in
    runST run

This gives me the following error:

Couldn't match type ‘s1’ with ‘s’
    ‘s1’ is a rigid type variable bound by
      the type signature for:
        loop :: forall s1. Int -> ST s1 ()
      at ...
    ‘s’ is a rigid type variable bound by
      the type signature for:
        run :: forall s. ST s Int
      at ...
    Expected type: ST s1 ()
    Actual type: ST s ()
    In the expression:
     ...

I understand that s isn't allowed to escape, but as far as i can see it won't? furthermore, when i remove the type signature for loop the issue goes away. I guess that indicates that they signature is incorrect somehow, but it's the same as before except without the counter and i have no idea what it should be.

Renaming s to match or not match the s mentioned in run makes no difference.

Answer 1

First, let's rename the type variables just so they're easier to talk about, and remove parts of the program that don't matter for this error:

mystery :: Int
mystery = runST run
  where run :: ST s Int
        run = do
          ref <- newSTRef 1
          let read :: ST t Int
              read = readSTRef ref
          read

This exhibits the same behavior, and commenting out the type signature for read fixes it, as before.

Now, let's ask: what is the type of ref? The type of newSTRef is a -> ST s (STRef s a), so ref :: STRef s Int, where s is the same as the s in run.

And what is the type of readSTRef ref? Well, readSTRef :: STRef s a -> ST s a. So, readSTRef ref :: ST s Int, where s is again the one in the definition of run. You give it a type signature that claims it works for any t, but it only works for the specific s in run, because it uses a ref from that transaction.

It is impossible to write the type for my read or your loop without turning on a language extension to allow you to refer to the s type variable that's already in scope. With ScopedTypeVariables, you can write:

{-# LANGUAGE ScopedTypeVariables #-}

import Control.Monad.ST
import Data.STRef

mystery :: Int
mystery = runST run
  where run :: forall s. ST s Int
        run = do
          ref <- newSTRef 1
          let read :: ST s Int
              read = readSTRef ref
          read

Using the forall brings s into scope explicitly, so that you can refer to it. Now the inner type signature's s actually refers to the outer one, rather than being a new, shadowing type variable. That's how you promise the type system that you'll only use this read function from inside the transaction owning the ref it uses.

Your original program, with the top-level loop, works for a similar reason. Instead of capturing an STRef (and thus its s implicitly), it declares a type which uses the same s for both the ref and the transaction. It works in any transaction, provided it's given a ref from that transaction.