haskellmarrayconstraint-kindsstuarray

Revisiting Polymorphic STUArrays with Constraint Kinds


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).


Solution

  • 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