Revisiting Polymorphic STUArrays with Constraint Kinds

Question

I want to implement a dynamic programming algorithm polymorphic in the score type; here's a simplified 1D version with no boundary conditions:

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

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

dynamicProgrammingSTU
  :: forall e i . (
    IArray UArray e,
    forall s. MArray (STUArray s) e (ST s),
    Ix i
  )
  => (forall m . Monad m => (i -> m e) -> (i -> m e))
  -> (i, i)
  -> (i -> e)
dynamicProgrammingSTU prog bnds = (arr !) where
  arr :: UArray i e
  arr = runSTUArray resultArrayST

  resultArrayST :: forall s . ST s (STUArray s i e)
  resultArrayST = do
    marr <- newArray_ bnds
    forM_ (range bnds) $ \i -> do
      result <- prog (readArray marr) i
      writeArray marr i result
    return marr

The constraint doesn't work;

    Could not deduce (MArray (STUArray s) e (ST s))
      arising from a use of `newArray_'
    from the context (IArray UArray e,
                      forall s. MArray (STUArray s) e (ST s),
                      Ix i)
      bound by the type signature for
                 dynamicProgrammingSTU :: (IArray UArray e,
                                           forall s. MArray (STUArray s) e (ST s
), Ix i) =>
                                          (forall (m :: * -> *). Monad m => (i -
> m e) -> i -> m e)
                                          -> (i, i) -> i -> e
      at example2.hs:(17,1)-(27,15)
    Possible fix:
      add (MArray (STUArray s) e (ST s)) to the context of
        the type signature for resultArrayST :: ST s (STUArray s i e)
        or the type signature for
             dynamicProgrammingSTU :: (IArray UArray e,
                                       forall s. MArray (STUArray s) e (ST s), I
x i) =>
                                      (forall (m :: * -> *). Monad m => (i -> m
e) -> i -> m e)
                                      -> (i, i) -> i -> e
      or add an instance declaration for (MArray (STUArray s) e (ST s))
    In a stmt of a 'do' block: marr <- newArray_ bnds
    In the expression:
      do { marr <- newArray_ bnds;
           forM_ (range bnds) $ \ i -> do { ... };
           return marr }
    In an equation for `resultArrayST':
        resultArrayST
          = do { marr <- newArray_ bnds;
                 forM_ (range bnds) $ \ i -> ...;
                 return marr }
Failed, modules loaded: none.

To summarize, Could not deduce (MArray (STUArray s) e (ST s)) from the context forall s. MArray (STUArray s) e (ST s i). Note that adding the constraint to resultArrayST just pushes the problem to runSTUArray.

I currently know of four flawed solutions:

1. Avoiding the problem with boxed STArrays or simply non-monadic Arrays, perhaps using seq and bang patterns to ease the resulting memory problems.
2. Breaking the type system with unsafeFreeze and unsafePerformIO, for which the damning constraint MArray IOUArray e IO works fine.
3. This solution to a similar problem using a typeclass and writing instances for every 'unboxable' type.
4. This one using GHC rewrite rules to pick a different function for each type (and a generic STArray version).

However, I'm asking this question in the hopes that modern language extensions like ConstraintKinds can allow me to express my original code's intent of forall s. MArray (STUArray s) e (ST s).

Empress

ghc-7.6.1 says Malformed predicate `forall s. MArray (STUArray s) e (ST s)', which to me makes more sense. – Censor

If prog function is monadic just for performance, I think your p. 1 (pure computations with probable bang patterns) would be the least of evils. – Lunik

Answer 1

Given the legendary helpfulness of the Haskell community, the lack of an answer at this point is a strong indication that there's no good solution in the current type system.

I've already outlined the flawed solutions in the question, so I'll just post a complete version of my example. This is basically what I used to solve most alignment problems on Rosalind:

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

import Control.Applicative
import Control.Monad
import Control.Monad.ST
import Data.Maybe

import Data.Array.ST
import Data.Array.Unboxed

class IArray UArray e => Unboxable e where
  newSTUArray_ :: forall s i. Ix i => (i, i) -> ST s (STUArray s i e)
  readSTUArray :: forall s i. Ix i => STUArray s i e -> i -> ST s e
  writeSTUArray :: forall s i. Ix i => STUArray s i e -> i -> e -> ST s ()


instance Unboxable Bool where 
  newSTUArray_ = newArray_
  readSTUArray = readArray
  writeSTUArray = writeArray

instance Unboxable Double where 
  newSTUArray_ = newArray_
  readSTUArray = readArray
  writeSTUArray = writeArray
{-
Same for Char, Float, (Int|Word)(|8|16|32|64)...
-}

{-# INLINE dynamicProgramming2DSTU #-}
dynamicProgramming2DSTU
  :: forall e i j . (
    Unboxable e,
    Ix i,
    Ix j,
    Enum i,
    Enum j
  )
  => (forall m . (Monad m, Applicative m) => (i -> j -> m e) -> (i -> j -> m e))
  -> (i -> j -> Maybe e)
  -> (i, i)
  -> (j, j)
  -> (i -> j -> e)
dynamicProgramming2DSTU program boundaryConditions (xl, xh) (yl, yh) = arrayLookup where
  arrayLookup :: i -> j -> e
  arrayLookup xi yj = fromMaybe (resultArray ! (xi, yj)) $ boundaryConditions xi yj

  arrB :: ((i, j), (i, j))
  arrB = ((xl, yl), (xh, yh))

  resultArray :: UArray (i, j) e
  resultArray = runSTUArray resultArrayST

  resultArrayST :: forall s. ST s (STUArray s (i, j) e)
  resultArrayST = do
    arr <- newSTUArray_ arrB
    let acc xi yj = maybe (readSTUArray arr (xi, yj)) return $ boundaryConditions xi yj

    forM_ [xl..xh] $ \xi -> do
      forM_ [yl..yh] $ \yj -> do
        result <- program acc xi yj
        writeSTUArray arr (xi, yj) result

    return arr

Empress