using and returning multiple STUArrays

I've been working out how to create and use multiple STUArrays in an ST computation. The specific scenarios are:

create multiple arrays but return only one of them
create multiple arrays but return none of them
create multiple arrays and return multiple arrays

I have answers to (1) and (2), but not (3).

First some imports so we know where everything is coming from:

import Control.Monad.ST (ST,runST)
import Data.Array.Base (unsafeFreezeSTUArray)
import Data.Array.ST (STUArray)
import Data.Array.Unboxed (UArray)
import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef)
import Data.Array.MArray (getBounds, newArray, readArray, writeArray, newListArray)
import Data.Array.ST (runSTUArray)

For (1), one trick is to define a new data type and constructor function:

data ArrayPair s = AP (STUArray s Int Int) (STUArray s Int Bool)

newAP n = do
  a1 <- newArray (1,n) 0
  a2 <- newArray (1,n) False
  return $ AP a1 a2

And here's how to return just one of the arrays:

foo :: UArray Int Int
foo = runSTUArray $ do
        AP ints bools <- newAP 10
        writeArray ints 1 42
        writeArray bools 1 True
        -- do stuff with ints and bools
        return ints

For (2), you can add an STRef to the data structure, end the computation with a readSTRef and run it with runST:

data WorkState s = WS (STUArray s Int Int)
                      (STUArray s Int Bool)
                      (STRef s Char)

newWS = do
  ints <- newArray (1,10) 0
  bools <- newArray (1,20) False
  char <- newSTRef 'X'
  return $ WS ints bools char

bar :: Char
bar = runST $ do
  WS ints bools char <- newWS
  writeArray ints 3 36
  writeArray bools 5 True
  writeSTRef char 'Z'
  -- ...
  readSTRef char

Comments on this approach for cases (1) and (2)? And what about case (3)?

baz :: (UArray Int Int, UArray Int Bool)
baz = runST??? $ do
  AP ints bools <- newAP
  ...
  return (ints,bools)  -- ???

Solution

You have 2 options:

Use freeze. This requires copying arrays, but you don't need to use any unsafe functions:

baz :: (UArray Int Int, UArray Int Bool)
baz = runST $ do
  AP ints bools <- newAP 12
  liftM2 (,) (freeze ints) (freeze bools)

Create your variant of runSTUArray for two arrays using unsafeFreezeSTUArray, knowing that the implementation is actually safe (because there will be no reference to the original mutable arrays left).

runSTUArray2 :: (Ix i1, Ix i2)
           => (forall s . (ST s (STUArray s i1 e1, STUArray s i2 e2)))
           -> (UArray i1 e1, UArray i2 e2)
runSTUArray2 st = runST $ do
    (a1, a2) <- st
    liftM2 (,) (unsafeFreezeSTUArray a1) (unsafeFreezeSTUArray a2)

baz' :: (UArray Int Int, UArray Int Bool)
baz' = runSTUArray2 $ do
  AP ints bools <- newAP 12
  return (ints, bools)

(Perhaps this approach could be even somehow generalized using Generics to allow to return any data structure containing ST arrays.)