-- | Zero-indexed dynamic arrays, optimised for lookup.
-- Modification is slow. Uninitialised indices have a default value.
{-# LANGUAGE CPP #-}
module Data.DynamicArray where

#ifdef BOUNDS_CHECKS
import qualified Data.Primitive.SmallArray.Checked as P
#else
import qualified Data.Primitive.SmallArray as P
#endif
import Control.Monad.ST
import Data.List

-- | A type which has a default value.
class Default a where
  -- | The default value.
  def :: a

-- | An array.
data Array a =
  Array {
    -- | The size of the array.
    Array a -> Int
arraySize     :: {-# UNPACK #-} !Int,
    -- | The contents of the array.
    Array a -> SmallArray a
arrayContents :: {-# UNPACK #-} !(P.SmallArray a) }

-- | Convert an array to a list of (index, value) pairs.
{-# INLINE toList #-}
toList :: Array a -> [(Int, a)]
toList :: Array a -> [(Int, a)]
toList Array a
arr =
  [ (Int
i, a
x)
  | Int
i <- [Int
0..Array a -> Int
forall a. Array a -> Int
arraySize Array a
arrInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1],
    let x :: a
x = SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
P.indexSmallArray (Array a -> SmallArray a
forall a. Array a -> SmallArray a
arrayContents Array a
arr) Int
i ]

instance Show a => Show (Array a) where
  show :: Array a -> String
show Array a
arr =
    String
"{" String -> ShowS
forall a. [a] -> [a] -> [a]
++
    String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", "
      [ Int -> String
forall a. Show a => a -> String
show Int
i String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"->" String -> ShowS
forall a. [a] -> [a] -> [a]
++ a -> String
forall a. Show a => a -> String
show a
x
      | (Int
i, a
x) <- Array a -> [(Int, a)]
forall a. Array a -> [(Int, a)]
toList Array a
arr ] String -> ShowS
forall a. [a] -> [a] -> [a]
++
    String
"}"

-- | Create an empty array.
newArray :: Array a
newArray :: Array a
newArray = (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array a)) -> Array a)
-> (forall s. ST s (Array a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
  SmallMutableArray s a
marr <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
P.newSmallArray Int
0 a
forall a. HasCallStack => a
undefined
  SmallArray a
arr  <- SmallMutableArray (PrimState (ST s)) a -> ST s (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
P.unsafeFreezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr
  Array a -> ST s (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> SmallArray a -> Array a
forall a. Int -> SmallArray a -> Array a
Array Int
0 SmallArray a
arr)

-- | Index into an array. O(1) time.
{-# INLINE (!) #-}
(!) :: Default a => Array a -> Int -> a
Array a
arr ! :: Array a -> Int -> a
! Int
n = a -> Int -> Array a -> a
forall a. a -> Int -> Array a -> a
getWithDefault a
forall a. Default a => a
def Int
n Array a
arr

-- | Index into an array. O(1) time.
{-# INLINE getWithDefault #-}
getWithDefault :: a -> Int -> Array a -> a
getWithDefault :: a -> Int -> Array a -> a
getWithDefault a
def Int
n Array a
arr
  | Int
0 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n Bool -> Bool -> Bool
&& Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Array a -> Int
forall a. Array a -> Int
arraySize Array a
arr =
    SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
P.indexSmallArray (Array a -> SmallArray a
forall a. Array a -> SmallArray a
arrayContents Array a
arr) Int
n
  | Bool
otherwise = a
def

-- | Update the array. O(n) time.
{-# INLINE update #-}
update :: Default a => Int -> a -> Array a -> Array a
update :: Int -> a -> Array a -> Array a
update Int
n a
x Array a
arr = a -> Int -> a -> Array a -> Array a
forall a. a -> Int -> a -> Array a -> Array a
updateWithDefault a
forall a. Default a => a
def Int
n a
x Array a
arr

{-# INLINEABLE updateWithDefault #-}
updateWithDefault :: a -> Int -> a -> Array a -> Array a
updateWithDefault :: a -> Int -> a -> Array a -> Array a
updateWithDefault a
def Int
n a
x Array a
arr = (forall s. ST s (Array a)) -> Array a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array a)) -> Array a)
-> (forall s. ST s (Array a)) -> Array a
forall a b. (a -> b) -> a -> b
$ do
  let size :: Int
size = Array a -> Int
forall a. Array a -> Int
arraySize Array a
arr Int -> Int -> Int
forall a. Ord a => a -> a -> a
`max` (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
  SmallMutableArray s a
marr <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
P.newSmallArray Int
size a
def
  SmallMutableArray (PrimState (ST s)) a
-> Int -> SmallArray a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallArray a -> Int -> Int -> m ()
P.copySmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr Int
0 (Array a -> SmallArray a
forall a. Array a -> SmallArray a
arrayContents Array a
arr) Int
0 (Array a -> Int
forall a. Array a -> Int
arraySize Array a
arr)
  SmallMutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
P.writeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr Int
n (a -> ST s ()) -> a -> ST s ()
forall a b. (a -> b) -> a -> b
$! a
x
  SmallArray a
arr' <- SmallMutableArray (PrimState (ST s)) a -> ST s (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> m (SmallArray a)
P.unsafeFreezeSmallArray SmallMutableArray s a
SmallMutableArray (PrimState (ST s)) a
marr
  Array a -> ST s (Array a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> SmallArray a -> Array a
forall a. Int -> SmallArray a -> Array a
Array Int
size SmallArray a
arr')