Haskell function that works with STUArray

Question

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?

Answer 1

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.

Answer 2

No, you're confused; that's not possible. Think of STUArray s i e as being a pointer to the beginning of a block of memory. You have to pass that pointer around to anything that needs to modify that block of memory; you can't just conjure it up out of thin air.

But you don't need to return it. Presumably the caller already has the pointer.

Answer 3

You can use the freeze and thaw functions to convert to and from a UArray.

However either this will incur a performance penalty, or you need to use the "unsafe" variants. Since you're already doing unsafe writing, that's probably OK.