Haskell function that works with STUArray

I have a small Haskell function that is supposed to accept an STUArray, modify some of the elements, and then return the changed array. It will be called from another function working in the ST s (STUArray s Int Word32) monad. It part of a fast PBKDF2 function that I am trying to write. This function does SHA-1 padding for a fixed sized message (160-bits).

Here is my code:

padFixed :: STUArray s Int Word32 -> ST s (STUArray s Int Word32)
padFixed block = do
  unsafeWrite block 5 0x80000000
  unsafeWrite block 15 160
  return block

The array will contain the 20 bytes from a previous SHA-1 run, plus 44 bytes of zeros. It will add the required padding as per RFC 3174.

How can I rewrite it so is "takes" the array out of the monad, works on it, and then puts it back? The signature should be padFixed :: ST s (STUArray s Int Word32), without the block parameter.

Is this possible? I could not find any functions in the library that let me extract the array from the monad, but maybe I missed something.

Are there any good tutorials on the STArray?

Solution

No, it's not possible; ST doesn't have those semantics. The monad is ST s, and not ST s (STUArray s a). ST s is just a monad for keeping track of mutable state; which structures you choose to allocate and use inside a single ST region are up to you. If you have a bunch of computations which all operate on the same STUArray, you can use ReaderT:

type Hasher s = ReaderT (STUArray s Int Word32) (ST s)

padFixed :: Hasher ()
padFixed = do
  block <- ask
  unsafeWrite block 5  0x80000000
  unsafeWrite block 15 160

The Reader r monad is just a wrapper around r ->; a value of type Reader r a is just a function r -> a. This is essentially a way to compute a while having access to a value of type r. The ReaderT r monad transformer just allows you to provide access to a variable of type r to an arbitrary monadic computation; thus, ReaderT (STUArray s Int Word32) (ST s) is an ST s computation which has access to some array. Note that you don't need to return the array from padFixed; the monad bind will handle all of that.

This'll be a little bit of a pain to write, since we'll have to keep asking for the array. Luckily, we can write some combinators to handle this for us:

{-# LANGUAGE RankNTypes, GeneralizedNewtypeDeriving #-}

import Data.Word
import Control.Applicative
import Control.Monad.Reader
import Control.Monad.ST
import Data.Array.ST (STUArray, runSTUArray)
import qualified Data.Array.Base as A
import Data.Array.Unboxed (UArray)

newtype Hasher s a =
  Hasher { getHasher :: ReaderT (STUArray s Int Word32) (ST s) a }
  deriving (Functor, Applicative, Monad, MonadReader (A.STUArray s Int Word32))

hasherToST :: Hasher s () -> (Int,Int) -> ST s (STUArray s Int Word32)
hasherToST (Hasher r) bounds = do
  block <- A.newArray bounds 0
  runReaderT r block
  return block

runHasher :: (forall s. Hasher s ()) -> (Int,Int) -> UArray Int Word32
runHasher h bounds = runSTUArray $ hasherToST h bounds

-- Perhaps private to this module, perhaps not
liftST :: ST s a -> Hasher s a
liftST = Hasher . lift

----- We can lift the functions which act on an STUArray -----

getBounds :: Hasher s (Int,Int)
getBounds = liftST . A.getBounds =<< ask

-- I'd recommend against removing the `unsafe` from the name; this function
-- could segfault, after all.
unsafeReadBlock :: Int -> Hasher s Word32
unsafeReadBlock i = do
  block <- ask
  liftST $ A.unsafeRead block i

unsafeWriteBlock :: Int -> Word32 -> Hasher s ()
unsafeWriteBlock i x = do
  block <- ask
  liftST $ A.unsafeWrite block i x

----- And then, perhaps in a separate module: -----

padFixed :: Hasher s ()
padFixed = do
  unsafeWriteBlock 5  0x80000000
  unsafeWriteBlock 15 160

(Note that I couldn't inline hasherToST inside of runHasher, probably because of the higher-rank types blocking inference.)

Basically, we wrap the ReaderT (STUArray s Int Word32) (ST s) into a newtype instead of a type synonym, and lift some basic array primitives up to work on the always-available block. You don't even need to derive MonadReader for the Hasher type if you don't want, as long as you lift all the necessary functions. But once you've done this, your hashing code can talk about the array implicitly.