STUArray with polymorphic type

Question

I want to implement an algorithm using the ST monad and STUArrays, and I want it to be able to work with both Float and Double data.

I'll demonstrate on a simpler example problem: calculating a memoized scanl (+) 0 (I know it can be solved without STUArray, just using as example).

{-# LANGUAGE FlexibleContexts, ScopedTypeVariables #-}

import Control.Monad
import Control.Monad.ST
import Data.Array.Unboxed
import Data.Array.ST

accumST :: forall a. (IArray UArray a, Num a) => [a] -> Int -> a
accumST vals = (!) . runSTUArray $ do
  arr <- newArray (0, length vals) 0 :: ST s (STUArray s Int a)
  forM_ (zip vals [1 .. length vals]) $ \(val, i) ->
    readArray arr (i - 1)
    >>= writeArray arr i . (+ val)
  return arr

This fails with:

Could not deduce (MArray (STUArray s) a (ST s)) from the context ()
  arising from a use of 'newArray'
Possible fix:
  add (MArray (STUArray s) a (ST s)) to the context of
    an expression type signature
  or add an instance declaration for (MArray (STUArray s) a (ST s))

I can't apply the suggested "Possible fix". Because I need to add something like (forall s. MArray (STUArray s) a (ST s)) to the context, but afaik that's impossible..

Wraf

Answer 1

Unforunately, you can't currently create a context that requires that an unboxed array be available for a specific type. Quantified Constraints aren't allowed. However, you can still accomplish what you're trying to do, (with the added advantage of having type-specific code versions.) For Longer functions, you could try to split out common expressions so that the repeated code is as small as possible.

{-# LANGUAGE FlexibleContexts, ScopedTypeVariables #-}
module AccumST where 

import Control.Monad
import Control.Monad.ST
import Data.Array.Unboxed
import Data.Array.ST
import Data.Array.IArray

-- General one valid for all instances of Num.
-- accumST :: forall a. (IArray UArray a, Num a) => [a] -> Int -> a
accumST :: forall a. (IArray UArray a, Num a) => [a] -> Int -> a
accumST vals = (!) . runSTArray $ do
  arr <- newArray (0, length vals) 0 :: (Num a) => ST s (STArray s Int a)
  forM_ (zip vals [1 .. length vals]) $ \(val, i) ->
    readArray arr (i - 1)
    >>= writeArray arr i . (+ val)
  return arr

accumSTFloat vals = (!) . runSTUArray $ do
  arr <- newArray (0, length vals) 0 :: ST s (STUArray s Int Float)
  forM_ (zip vals [1 .. length vals]) $ \(val, i) ->
    readArray arr (i - 1)
    >>= writeArray arr i . (+ val)
  return arr

accumSTDouble vals = (!) . runSTUArray $ do
  arr <- newArray (0, length vals) 0 :: ST s (STUArray s Int Double)
  forM_ (zip vals [1 .. length vals]) $ \(val, i) ->
    readArray arr (i - 1)
    >>= writeArray arr i . (+ val)
  return arr

{-# RULES "accumST/Float" accumST = accumSTFloat #-}
{-# RULES "accumST/Double" accumST = accumSTDouble #-}

The Generic Unboxed version (which doesn't work) would have a type constraint like the following:

accumSTU :: forall a. (IArray UArray a, Num a, 
    forall s. MArray (STUArray s) a (ST s)) => [a] -> Int -> a

You could simplify as follows:

-- accumST :: forall a. (IArray UArray a, Num a) => [a] -> Int -> a
accumST :: forall a. (IArray UArray a, Num a) => [a] -> Int -> a
accumST vals = (!) . runSTArray $ do
  arr <- newArray (0, length vals) 0 :: (Num a) => ST s (STArray s Int a)
  accumST_inner vals arr

accumST_inner vals arr = do
  forM_ (zip vals [1 .. length vals]) $ \(val, i) ->
    readArray arr (i - 1)
    >>= writeArray arr i . (+ val)
  return arr

accumSTFloat vals = (!) . runSTUArray $ do
  arr <- newArray (0, length vals) 0 :: ST s (STUArray s Int Float)
  accumST_inner vals arr

accumSTDouble vals = (!) . runSTUArray $ do
  arr <- newArray (0, length vals) 0 :: ST s (STUArray s Int Double)
  accumST_inner vals arr

{-# RULES "accumST/Float" accumST = accumSTFloat #-}
{-# RULES "accumST/Double" accumST = accumSTDouble #-}

Ulani

The rules only fire if compiled with optimizations enabled. – Ulani

I ended up using a different workaround for now - see answer below – Wraf

Answer 2

So here's the workaround I'm going with for now - creating a new typeclass for types for which (forall s. MArray (STUArray s) a (ST s)):

class IArray UArray a => Unboxed a where
  newSTUArray :: Ix i => (i, i) -> a -> ST s (STUArray s i a)
  readSTUArray :: Ix i => STUArray s i a -> i -> ST s a
  writeSTUArray :: Ix i => STUArray s i a -> i -> a -> ST s ()

instance Unboxed Float where
  newSTUArray = newArray
  readSTUArray = readArray
  writeSTUArray = writeArray

instance Unboxed Double where
  newSTUArray = newArray
  readSTUArray = readArray
  writeSTUArray = writeArray

While I'm not perfectly satisfied with this, I prefer it on rules because:

• rules depend on optimizations
• rules are not really supposed to change the algorithm (?). where in this case they would as boxed arrays have very different behaviour regarding lazyness etc.

Wraf