Fused in-place updates of mutable values

PipelineMutation.hs

{-# OPTIONS -O #-}
-- Perform in-place updates on mutable data.
import Control.Applicative
import Control.Monad
import Data.IORef
import Data.Monoid
import Debug.Trace
import System.IO.Unsafe

-- Presumably we can duplicate the values we want to mutate. Here, we
-- use IORef as a proxy for some typically abstract type used with the
-- FFI.
clone :: IORef a -> IO (IORef a)
clone original = trace "cloning" $ 
                 readIORef original >>= newIORef

withClone :: (IORef a -> IO ()) -> IORef a -> IO (IORef a)
withClone f = clone >=> (\x -> f x >> return x)

-- An operation mutates a reference.
newtype Op a = Op (IORef a -> IO ())

(<>) :: Monoid m => m -> m -> m
(<>) = mappend

-- We use the 'Monoid' instance to compose 'Op's.
instance Monoid (Op a) where
  mempty = Op . const $ return ()
  Op f `mappend` Op g = Op $ \x -> g x >> f x

-- We want to present a functional interface, so we don't actually
-- mutate our data, but clone it, mutate the clone, then return the
-- clone.
operate :: Op a -> IORef a -> IORef a
operate (Op f) = unsafePerformIO . withClone f

op1 :: Num a => Op a
op1 = Op $ flip modifyIORef (+ 1)

op2 :: Num a => Op a
op2 = Op $ flip modifyIORef (* 2)

-- If we manually compose operations before evaluating them, we can
-- get away with a single clone for a stack of operations.

test1 :: IO Int
test1 = newIORef 3 >>= readIORef . operate (op2 <> op1)

-- Alternately, we can push 'operate' into the operations themselves
-- to present a pure interface at the expense of excessive cloning.

type PureOp a = IORef a -> IORef a

op1' :: Num a => PureOp a
op1' = operate op1

op2' :: Num a => PureOp a
op2' = operate op2

-- Now we can use standard function composition, but we create a new
-- clone for each operation.
test2 :: IO Int
test2 = newIORef 3 >>= readIORef . op2' . op1'

-- An improvement is to fuse operations with a rewrite rule.

operate' :: Op a -> IORef a -> IORef a
operate' (Op f) = unsafePerformIO . withClone f
{-# NOINLINE operate' #-}

{-# RULES "operate/fuse"
  forall f g x. operate' f (operate' g x) = operate' (f <> g) x #-}

op1'', op2'' :: Num a => PureOp a
op1'' = operate' op1
op2'' = operate' op2
{-# INLINE op1'' #-}
{-# INLINE op2'' #-}

-- Now standard function composition can be automatically rewritten by
-- the compiler to fuse multiple operations under a single clone.

test3 :: IO Int
test3 = newIORef 3 >>= readIORef . op2'' . op1''

main = do putStr "Manual cloning: " >> test1 >>= print
          putStr "Pure interface: " >> test2 >>= print
          putStr "Rewritten pure: " >> test3 >>= print