Using a monadic rank-2 type

Here's the code:

{-# LANGUAGE RankNTypes, FlexibleContexts, ScopedTypeVariables #-}

module Foo where

import Data.Vector.Generic.Mutable as M
import Data.Vector.Generic as V
import Control.Monad.ST
import Control.Monad.Primitive
import Control.Monad

data DimFun v s r = 
  DimFun {dim::Int, func :: v (PrimState s) r -> s ()}

runFun :: (Vector v r) => 
  (forall s . (PrimMonad s) => DimFun (Mutable v) s r) -> v r -> v r
runFun t x = runST $ do
  y <- thaw x
  evalFun t y
  unsafeFreeze y

evalFun :: (PrimMonad s, MVector v r) => DimFun v s r -> v (PrimState s) r -> s ()
evalFun (DimFun dim f) y | dim == M.length y = f y

fm :: (MVector v r, PrimMonad s, Num r, Monad m) => m (DimFun v s r)
fm = error ""

f :: forall v r m . (Vector v r, Num r, Monad m) => m (v r -> v r)
f = liftM runFun $ (fm :: forall s . (PrimMonad s) => m (DimFun (Mutable v) s r))

This results in errors:

Couldn't match type ‘DimFun (Mutable v) s0 r’
              with ‘forall (s :: * -> *). PrimMonad s => DimFun (Mutable v) s r’
Expected type: DimFun (Mutable v) s0 r -> v r -> v r
  Actual type: (forall (s :: * -> *).
                PrimMonad s =>
                DimFun (Mutable v) s r)
               -> v r -> v r
Relevant bindings include
  f :: m (v r -> v r) (bound at Testing/Foo.hs:36:1)
In the first argument of ‘liftM’, namely ‘runFun’
In the expression: liftM runFun

However, I'm not sure how to fix or diagnose the problem. It might be as simple as a well-place (and well-written) type signature.

While trying to figure out what was going on, I write a non-monadic version (useless to me), but it compiles:

gm :: (MVector v r, PrimMonad s, Num r) => DimFun v s r
gm = error ""

g :: forall v r m . (Vector v r, Num r) => v r -> v r
g = runFun (gm :: forall s . (PrimMonad s) => DimFun (Mutable v) s r)

This makes me thing the error above is related to the this question where there is no place for the dictionary to go, but that's really just a stab in the dark.

One solution is to move the PrimMonad constraint inside the DimFun datatype.

data DimFun v r = DimFun 
   { dim  :: Int
   , func :: forall s . PrimMonad s => v (PrimState s) r -> s ()
   }

The rest of your code compiles as-is, removing the s parameter from DimFun:

runFun :: Vector v r => DimFun (Mutable v) r -> v r -> v r
runFun = ...

evalFun :: (PrimMonad s, MVector v r) => DimFun v r -> v (PrimState s) r -> s () 
evalFun = ...

fm :: (MVector v r, Num r, Monad m) => m (DimFun v r)
fm = ...

f :: (Vector v r, Num r, Monad m) => m (v r -> v r)
f = liftM runFun fm

Moving the class constraint into the datatype may seem scary to you, but in reality, you already had the class constraint there anyways. PrimState is an associated type family of PrimMonad, so in order to produce or consume a v (PrimState s) r, you need the PrimMonad constraint.

If you want to avoid it nevertheless, you will have to change the type of something. To see why the function you have is illtyped, consider the following (which requires ImpredictiveTypes):

fm :: (MVector v r, PrimMonad s, Num r, Monad m) => m (DimFun v s r)
fm = error ""

g :: (Vector v r, Monad m) 
  => m (forall s . PrimMonad s => DimFun (Mutable v) s r) -> m (v r -> v r)
g = liftM runFun

It should be cleared why g fm is illtyped: g expects something where the forall s . PrimMonad s => is inside the m, which is not the case for fm. You will have to write a function of type:

fm' :: (MVector v r, Monad m, Num r) => m (forall s . PrimMonad s => DimFun v s r)
fm' = error ""

f :: forall v r m . (Vector v r, Num r, Monad m) => m (v r -> v r)
f = g fm'