Having ST(U)Arrays in a data structure?

Question

What do I have to do to make GHC accept this code:

{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances #-}

module STTest where

import Data.Array.ST
import Control.Monad.ST.Strict as S
import Control.Monad.ST.Lazy as L

-- ST monad arrays (unboxed in actual code)
type Arr s a = STArray s Int a

-- representing some algorithm that works on these STArrays
data ArrGen s a = ArrGen (a -> S.ST s (Arr s a)) (Arr s a -> S.ST s ())

-- class for some "generator"
class Generator g a where
  gen :: g -> a -> [a]

instance Generator (ArrGen s a) a where
  gen (ArrGen create apply) s = L.runST $ do
    a <- strictToLazyST $ create s -- DOES NOT WORK
    strictToLazyST $ apply a >> getElems a

The error I get is the following:

Couldn't match type `s' with `s1'
  `s' is a rigid type variable bound by
     the instance declaration at STTest.hs:20:28
  `s1' is a rigid type variable bound by
     a type expected by the context: L.ST s1 [a] at STTest.hs:21:33

However, this works fine:

data Dummy
create' :: a -> S.ST s (Arr s a)
create' = undefined
apply' :: Arr s a -> S.ST s [a]
apply' = undefined

instance Generator Dummy a where
  gen _ s = L.runST $ do
    a <- strictToLazyST $ create' s
    strictToLazyST $ apply' a >> getElems a

Why does it work with the second and not the first? And what can I do with the data declaration to make it work? Or can I add some sort of "forall" on the instance declaration?

The above is just a minimal test program. I actually loop the apply forever to create an infinite Stream of the output values. (So I can't just merge the two steps together.) And I really want to be able to instantiate once for the ArrGen data type and then make a variety of values of it using these STArray algorithms.

EDIT:

Didn't think to put the forall inside the functions to ArrGen (I put it on the overall type). Though now I have the a problem of getting it to work on STUArray. Like if I use the following:

class (Integral a, Bits a, forall s. MArray (STUArray s) a (S.ST s)) => HasSTU a
type AC a = (HasSTU a) => forall s. a -> S.ST s (STUArray s Int a)
type AU a = (HasSTU a) => forall s. STUArray s Int a -> S.ST s ()
type TX a = (HasSTU a) => a -> a -- or without the context
data ArrayGen a = AG (AC a) (AU a) (TX a)

Then this fails:

instance (HasSTU a) => Generator (ArrayGen a) a [a] where
  gens (AG c u p) s = fmap (fmap p) $ L.runST $ do
    ar <- strictToLazyST $ (c s)
    streamM $ strictToLazyST $ u ar >> getElems ar -- can't use getElems here!

streamM :: (Applicative f) => f a -> f (Stream a))
streamM = Cons <$> a <*> streamM a

It complains:

Could not deduce (MArray (STUArray s) a (S.ST s))
  arising from a use of `getElems'
from the context (HasSTU a)

Even though the context (HasSTU a) says (in my mind) that there is an (MArray (STUArray s) a (S.ST s)) context for all s, it doesn't seem to think so. I tried to fix it by changing the (AU a) type:

type AU a = (HasSTU a) => forall s. STUArray s Int a -> S.ST s [a]

And it seems to type check, but I am unable to actually use it. Similarly if I change to:

class (Integral a, Bits a, forall s. MArray (STUArray s) a (S.ST s)) => HasSTU s a
type AC a = (forall s. HasSTU s a) => a -> S.ST s (STUArray s Int a)
...
instance (forall s. HasSTU s a) => Generator (ArrayGen a) a [a] where
  ...

instance forall s. HasSTU s Word32 -- !!!

But then when I try to run something:

Could not deduce (forall s. HasSTU s Word32)

I hate this s! Why? I have an instance for all s! And I am really lost as to where I should put my foralls and what's really going on.

Answer 1

The problem is that runST requires a forall s. ST s t argument, but your type fixes s, so a use of create and apply in the monadic action makes it unsuitable for runST.

It does not seem to me that your use case forbids giving ArrGen polymorphic (in s) arguments, so

{-# LANGUAGE MultiParamTypeClasses, FlexibleInstances, RankNTypes #-}

module STTest where

import Data.Array.ST
import Control.Monad.ST.Strict as S
import Control.Monad.ST.Lazy as L

-- ST monad arrays (unboxed in actual code)
type Arr s a = STArray s Int a

-- representing some algorithm that works on these STArrays
data ArrGen a = ArrGen (forall s. a -> S.ST s (Arr s a)) (forall s. Arr s a -> S.ST s ())

-- class for some "generator"
class Generator g a where
  gen :: g -> a -> [a]

instance Generator (ArrGen a) a where
  gen (ArrGen create apply) s = L.runST $ do
    a <- strictToLazyST $ create s -- DOES NOT WORK
    strictToLazyST $ apply a >> getElems a

making the components polymorphic works (at least in the sense that it compiles, your use case may forbid this approach).

Why does it work with the second and not the first?

Because there, the s was not fixed, the computation is fully polymorphic in s, as required by runST.