diffarray.hs

{- cabal:
default-language: GHC2021
build-depends: base ^>= 4.17
             , array

ghc-options: -threaded
             -rtsopts
             -with-rtsopts=-N
-}
{-# LANGUAGE CPP, Strict, UnliftedDatatypes, StandaloneKindSignatures, MagicHash, UnliftedNewtypes #-}
{-# LANGUAGE UnboxedTuples #-}
{-# OPTIONS_GHC -O2 -ddump-simpl -dsuppress-all -ddump-to-file -ddump-rule-firings #-}
{-# OPTIONS_GHC -Wno-unused-imports #-}

import Data.Array.IArray
import qualified Data.Array as A

------------------------------------------------------------------------
-- Imports.

import Data.Array.Base
import Data.Array.IO (IOArray, IOUArray)

import Data.Int (Int16, Int32, Int64, Int8)
import Data.Word (Word, Word16, Word32, Word64, Word8)
import Foreign.Ptr (FunPtr, Ptr)
import Foreign.StablePtr (StablePtr)

import Control.Concurrent.MVar (MVar, newMVar, putMVar, readMVar, takeMVar)
import Control.Exception (evaluate)
import System.IO.Unsafe (unsafePerformIO, unsafeDupablePerformIO)
import GHC.Base (UnliftedType, MutVar#, RealWorld, IO (IO), newMutVar#, unIO, readMutVar#, writeMutVar#)

newtype UnliftedIORef a = UnliftedIORef (MutVar# RealWorld a)

newUnliftedIORef :: a -> (UnliftedIORef a -> IO s) -> IO s
newUnliftedIORef x f = IO $ \s ->
    case newMutVar# x s of
        (# s', var #) -> unIO (f (UnliftedIORef var)) s'

readUnliftedIORef :: UnliftedIORef a -> IO a
readUnliftedIORef (UnliftedIORef var) = IO $ \s -> readMutVar# var s

writeUnliftedIORef :: UnliftedIORef a -> a -> IO ()
writeUnliftedIORef (UnliftedIORef var) x = IO $ \s ->
    case writeMutVar# var x s of
        s' -> (# s', () #)

------------------------------------------------------------------------
-- Diff array types.

{- | An arbitrary 'MArray' type living in the 'IO' monad can be converted
to a diff array.
-}
data IOToDiffArray a i e
    = DiffArray
        { varDiffArray :: UnliftedIORef (DiffArrayData a i e)
        , bounds' :: {-# UNPACK #-} (i, i)
        , numElems :: {-# UNPACK #-} Int}

-- Internal representation: either a mutable array, or a link to
-- another diff array patched with a list of index+element pairs.
data DiffArrayData a i e
    = Current (a i e)
    | Diff (IOToDiffArray a i e) [(Int, e)]

-- | Fully polymorphic lazy boxed diff array.
type DiffArray = IOToDiffArray IOArray

{- | Strict unboxed diff array, working only for elements
of primitive types but more compact and usually faster than 'DiffArray'.
-}
type DiffUArray = IOToDiffArray IOUArray

-- Having 'MArray a e IO' in instance context would require
-- -XUndecidableInstances, so each instance is separate here.

------------------------------------------------------------------------
-- Showing DiffArrays

instance (Ix ix, Show ix, Show e) => Show (DiffArray ix e) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Bool) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Char) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Int) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Word) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Float) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Double) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Int8) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Int16) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Int32) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Int64) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Word8) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Word16) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Word32) where
    showsPrec = showsIArray

instance (Ix ix, Show ix) => Show (DiffUArray ix Word64) where
    showsPrec = showsIArray

------------------------------------------------------------------------
-- Boring instances.

instance IArray (IOToDiffArray IOArray) e where
    {-# INLINE bounds #-}
    bounds = boundsDiffArray
    {-# INLINE numElements #-}
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafeDupablePerformIO $ newDiffArray lu ies
    {-# INLINE unsafeAt #-}
    unsafeAt a i = unsafeDupablePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafeDupablePerformIO $ a `replaceDiffArray1` ies

instance IArray (IOToDiffArray IOUArray) Bool where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Char where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Int where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Word where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) (Ptr a) where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) (FunPtr a) where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Float where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Double where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) (StablePtr a) where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Int8 where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Int16 where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Int32 where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Int64 where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Word8 where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Word16 where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Word32 where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

instance IArray (IOToDiffArray IOUArray) Word64 where
    bounds = boundsDiffArray
    numElements = numElementsDiffArray
    unsafeArray lu ies = unsafePerformIO $ newDiffArray lu ies
    unsafeAt a i = unsafePerformIO $ a `readDiffArray` i
    unsafeReplace a ies = unsafePerformIO $ a `replaceDiffArray2` ies

------------------------------------------------------------------------
-- The important stuff.

-- diffArray ::
--     (MArray a e IO, Ix i) =>
--     (i, i) ->
--     Int ->
--     a i e ->
--     IO (IOToDiffArray a i e)
-- diffArray (l, u) n a = IO $ \s ->
--     newMutVar# s a)

newDiffArray ::
    (MArray a e IO, Ix i) =>
    (i, i) ->
    [(Int, e)] ->
    IO (IOToDiffArray a i e)
newDiffArray (l, u) ies = do
    a <- newArray_ (l, u)
    sequence_ [unsafeWrite a i e | (i, e) <- ies]
    n <- getNumElements a
    newUnliftedIORef (Current a) $ \var -> return (DiffArray var (l, u) n)

{-# INLINE readDiffArray #-}
readDiffArray ::
    (MArray a e IO, Ix i) =>
    IOToDiffArray a i e ->
    Int ->
    IO e
a `readDiffArray` i = do
    d <- readUnliftedIORef (varDiffArray a)
    case d of
        Current a' -> unsafeRead a' i
        Diff a' ies -> maybe (readDiffArray a' i) return (lookup i ies)

replaceDiffArray ::
    (MArray a e IO, Ix i) =>
    IOToDiffArray a i e ->
    [(Int, e)] ->
    IO (IOToDiffArray a i e)
a `replaceDiffArray` ies = do
    d <- readUnliftedIORef (varDiffArray a)
    case d of
        Current a' -> case ies of
            [] -> do
                -- We don't do the copy when there is nothing to change
                -- and this is the current version. But see below.
                writeUnliftedIORef (varDiffArray a) d
                return a
            _ : _ -> do
                diff <-
                    sequence
                        [ do e <- unsafeRead a' i; return (i, e)
                        | (i, _) <- ies
                        ]
                sequence_ [unsafeWrite a' i e | (i, e) <- ies]
                newUnliftedIORef (Current a') $ \var' -> do
                    writeUnliftedIORef (varDiffArray a) (Diff (DiffArray var' (bounds' a) (numElems a)) diff)
                    return (DiffArray var' (bounds' a) (numElems a))
        Diff _ _ -> do
            -- We still do the copy when there is nothing to change
            -- but this is not the current version. So you can use
            -- 'a // []' to make sure that the resulting array has
            -- fast element access.
            writeUnliftedIORef (varDiffArray a) d
            a' <- thawDiffArray a
            -- thawDiffArray gives a fresh array which we can
            -- safely mutate.
            sequence_ [unsafeWrite a' i e | (i, e) <- ies]
            newUnliftedIORef (Current a') $ \var' ->
                return (DiffArray var' (bounds' a) (numElems a))

-- The elements of the diff list might recursively reference the
-- array, so we must seq them before taking the MVar to avoid
-- deadlock.
replaceDiffArray1 ::
    (MArray a e IO, Ix i) =>
    IOToDiffArray a i e ->
    [(Int, e)] ->
    IO (IOToDiffArray a i e)
a `replaceDiffArray1` ies = do
    mapM_ (evaluate . fst) ies
    a `replaceDiffArray` ies

-- If the array contains unboxed elements, then the elements of the
-- diff list may also recursively reference the array from inside
-- replaceDiffArray, so we must seq them too.
replaceDiffArray2 ::
    (MArray a e IO, Ix i) =>
    IOToDiffArray a i e ->
    [(Int, e)] ->
    IO (IOToDiffArray a i e)
arr `replaceDiffArray2` ies = do
    mapM_ (\(a, b) -> evaluate a >> evaluate b) ies
    arr `replaceDiffArray` ies

{-# INLINE boundsDiffArray #-}
boundsDiffArray ::
    (MArray a e IO, Ix ix) =>
    IOToDiffArray a ix e ->
    (ix, ix)
boundsDiffArray = bounds'

{-# INLINE numElementsDiffArray #-}
numElementsDiffArray ::
    (MArray a e IO, Ix ix) =>
    IOToDiffArray a ix e ->
    Int
numElementsDiffArray = numElems

freezeDiffArray ::
    (MArray a e IO, Ix ix) =>
    a ix e ->
    IO (IOToDiffArray a ix e)
freezeDiffArray a = do
    (l, u) <- getBounds a
    n <- getNumElements a
    a' <- newArray_ (l, u)
    sequence_ [unsafeRead a i >>= unsafeWrite a' i | i <- [0 .. rangeSize (l, u) - 1]]
    newUnliftedIORef (Current a') $ \var ->
        return (DiffArray var (l, u) n)

{-# RULES
"freeze/DiffArray" freeze = freezeDiffArray
    #-}

-- unsafeFreezeDiffArray is really unsafe. Better don't use the old
-- array at all after freezing. The contents of the source array will
-- be changed when '//' is applied to the resulting array.

unsafeFreezeDiffArray ::
    (MArray a e IO, Ix ix) =>
    a ix e ->
    IO (IOToDiffArray a ix e)
unsafeFreezeDiffArray a = do
    b <- getBounds a
    n <- getNumElements a
    newUnliftedIORef (Current a) $ \var ->
        return (DiffArray var b n)

{-# RULES
"unsafeFreeze/DiffArray" unsafeFreeze = unsafeFreezeDiffArray
    #-}

thawDiffArray ::
    (MArray a e IO, Ix ix) =>
    IOToDiffArray a ix e ->
    IO (a ix e)
thawDiffArray a = do
    d <- readUnliftedIORef (varDiffArray a)
    case d of
        Current a' -> do
            (l, u) <- getBounds a'
            a'' <- newArray_ (l, u)
            sequence_ [unsafeRead a' i >>= unsafeWrite a'' i | i <- [0 .. rangeSize (l, u) - 1]]
            return a''
        Diff a' ies -> do
            a'' <- thawDiffArray a'
            sequence_ [unsafeWrite a'' i e | (i, e) <- ies]
            return a''

{-# RULES
"thaw/DiffArray" thaw = thawDiffArray
    #-}

-- unsafeThawDiffArray is really unsafe. Better don't use the old
-- array at all after thawing. The contents of the resulting array
-- will be changed when '//' is applied to the source array.
unsafeThawDiffArray ::
    (MArray a e IO, Ix ix) =>
    IOToDiffArray a ix e ->
    IO (a ix e)
unsafeThawDiffArray a = do
    d <- readUnliftedIORef (varDiffArray a)
    case d of
        Current a' -> return a'
        Diff a' ies -> do
            a'' <- unsafeThawDiffArray a'
            sequence_ [unsafeWrite a'' i e | (i, e) <- ies]
            return a''

{-# RULES
"unsafeThaw/DiffArray" unsafeThaw = unsafeThawDiffArray
    #-}

type Arr = DiffArray Int Int

maxPos = 10000
maxLoops = 100000

loop arr i c
    | i > maxPos = arr -- // [(maxPos, c)]
    | otherwise = loop arr (i + 1) $! arr ! i + c

loop2 i arr
    | i > maxLoops = arr
    | otherwise = loop2 (i + 1) $ loop arr 1 0

main = print $ loop2 1 arr ! maxPos

arr :: Arr
arr = array (1, maxPos) [(i, 1) | i <- [1 .. maxPos]]