A monad for piecewise mutable state

Question

I know how ordinary State works (edit: apparently not!).

If i need to create an array, and it is inconvenient to create the entire array at once, I can create an STArray, populate it, and then freeze and return a normal immutable array to the user.

Now suppose I need to create two arrays of different types simultaneously.

More generally, I may want to create an arbitrary graph with mutable nodes, modify it node by node for a while like I would modify an STArray cell by cell, and then freeze the whole graph at once and return normal immutable data.

I don't want to resort to IOArrays or anything in the IO monad. What are my options?

Answer 1

Here are some options.

{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
module TestSTArray where

import Control.Monad.ST
import Data.Array.ST
import Data.Array
import Data.Array.MArray

-- not needed until later on    
import GHC.Arr (unsafeFreezeSTArray)

The basic, safe way is to freeze. This will cause a copy, though.

test2Safe :: (Array Int Char, Array Int Bool)
test2Safe = runST $ do
   a1 <- newArray (0,9) 'A' :: ST s (STArray s Int Char)
   a2 <- newArray (0,9) False :: ST s (STArray s Int Bool)
   writeArray a1 5 'B'
   x <- readArray a2 6
   writeArray a1 7 (if x then 'X' else 'Y')
   writeArray a2 5 True
   arr1 <- freeze a1
   arr2 <- freeze a2
   return (arr1, arr2)

More risky, but probably still safe is to leverage lower lever / unsafe GHC primitives and build an extended variant of the safe runSTArray. In this way we avoid the copy.

runSTArray2 :: (forall s. ST s (STArray s i1 e1, STArray s i2 e2))
            -> (Array i1 e1, Array i2 e2) 
runSTArray2 st = runST $ do
  (a1, a2) <- st
  (,) <$> unsafeFreezeSTArray a1 <*> unsafeFreezeSTArray a2

I believe the above use of unsafe stuff is actually safe, since we no longer use a1,a2 after the unsafe freeze, so no copy should be needed.

The above wrapper can be generalized to more arrays, of course. Arguably, a more general version should be put in the libraries.

Finally, we can exploit the auxiliary function:

test2LessSafe :: (Array Int Char, Array Int Bool)
test2LessSafe = runSTArray2 $ do
   a1 <- newArray (0,9) 'A' :: ST s (STArray s Int Char)
   a2 <- newArray (0,9) False :: ST s (STArray s Int Bool)
   writeArray a1 5 'B'
   x <- readArray a2 6
   writeArray a1 7 (if x then 'X' else 'Y')
   writeArray a2 5 True
   return (a1, a2)