-- | 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 {
    forall a. Array a -> Int
arrayStart    :: {-# UNPACK #-} !Int,
    -- | The contents of the array.
    forall a. Array a -> SmallArray a
arrayContents :: {-# UNPACK #-} !(P.SmallArray a) }

arraySize :: Array a -> Int
arraySize :: forall a. Array a -> Int
arraySize = forall a. SmallArray a -> Int
P.sizeofSmallArray forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Array a -> SmallArray a
arrayContents

-- | Convert an array to a list of (index, value) pairs.
{-# INLINE toList #-}
toList :: Array a -> [(Int, a)]
toList :: forall a. Array a -> [(Int, a)]
toList Array a
arr =
  [ (Int
iforall a. Num a => a -> a -> a
+forall a. Array a -> Int
arrayStart Array a
arr, a
x)
  | Int
i <- [Int
0..forall a. Array a -> Int
arraySize Array a
arrforall a. Num a => a -> a -> a
-Int
1],
    let x :: a
x = forall a. SmallArray a -> Int -> a
P.indexSmallArray (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
"{" forall a. [a] -> [a] -> [a]
++
    forall a. [a] -> [[a]] -> [a]
intercalate String
", "
      [ forall a. Show a => a -> String
show Int
i forall a. [a] -> [a] -> [a]
++ String
"->" forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show a
x
      | (Int
i, a
x) <- forall a. Array a -> [(Int, a)]
toList Array a
arr ] forall a. [a] -> [a] -> [a]
++
    String
"}"

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

{-# INLINE singleton #-}
-- | Create an array with one element.
singleton :: Default a => Int -> a -> Array a
singleton :: forall a. Default a => Int -> a -> Array a
singleton Int
i a
x = forall a. Default a => Int -> a -> Array a -> Array a
update Int
i a
x forall a. Array a
newArray

-- | Index into an array. O(1) time.
{-# INLINE (!) #-}
(!) :: Default a => Array a -> Int -> a
Array a
arr ! :: forall a. Default a => Array a -> Int -> a
! Int
n = forall a. a -> Int -> Array a -> a
getWithDefault 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 :: forall a. a -> Int -> Array a -> a
getWithDefault a
def Int
n Array a
arr
  | forall a. Array a -> Int
arrayStart Array a
arr forall a. Ord a => a -> a -> Bool
<= Int
n Bool -> Bool -> Bool
&& Int
n forall a. Ord a => a -> a -> Bool
< forall a. Array a -> Int
arrayStart Array a
arr forall a. Num a => a -> a -> a
+ forall a. Array a -> Int
arraySize Array a
arr =
    forall a. SmallArray a -> Int -> a
P.indexSmallArray (forall a. Array a -> SmallArray a
arrayContents Array a
arr) (Int
n forall a. Num a => a -> a -> a
- forall a. Array a -> Int
arrayStart Array a
arr)
  | Bool
otherwise = a
def

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

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