-- | An array of key-value pairs, where the keys are integers.
-- Can be accessed both as a map ("give me the value corresponding
-- to the key 2") and as an array ("give me the 3rd key-value pair").
-- Array-like access is fast; everything else is slow.
module Data.Numbered(
  Numbered,
  empty, fromList, singleton, toList, size, (!),
  lookup, put, modify, filter, delete) where

import Prelude hiding (filter, lookup)
import qualified Data.List as List
import Data.Primitive.ByteArray
import Data.Primitive.SmallArray
import Data.Int
import Data.Maybe

-- | An array of key-value pairs.
data Numbered a =
  Numbered
    {-# UNPACK #-} !ByteArray
    {-# UNPACK #-} !(SmallArray a)

instance Show a => Show (Numbered a) where show :: Numbered a -> String
show = forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Numbered a -> [(Int, a)]
toList

-- | An empty array.
empty :: Numbered a
empty :: forall a. Numbered a
empty = forall a. [(Int, a)] -> Numbered a
fromList []

-- | A singleton array.
singleton :: Int -> a -> Numbered a
singleton :: forall a. Int -> a -> Numbered a
singleton Int
i a
x = forall a. [(Int, a)] -> Numbered a
fromList [(Int
i, a
x)]

-- | Convert a list of pairs to an array.
-- Duplicate keys are allowed.
fromList :: [(Int, a)] -> Numbered a
fromList :: forall a. [(Int, a)] -> Numbered a
fromList [(Int, a)]
xs =
  forall a. ByteArray -> SmallArray a -> Numbered a
Numbered
    (forall a. Prim a => [a] -> ByteArray
byteArrayFromList (forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Int, a)]
xs :: [Int32]))
    (forall a. [a] -> SmallArray a
smallArrayFromList (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Int, a)]
xs))

-- | Convert an array to a list.
toList :: Numbered a -> [(Int, a)]
toList :: forall a. Numbered a -> [(Int, a)]
toList Numbered a
num =
  [Numbered a
num forall a. Numbered a -> Int -> (Int, a)
! Int
i | Int
i <- [Int
0..forall a. Numbered a -> Int
size Numbered a
numforall a. Num a => a -> a -> a
-Int
1]]

-- | Get the number of key-value pairs in an array. O(1) time.
size :: Numbered a -> Int
size :: forall a. Numbered a -> Int
size (Numbered ByteArray
_ SmallArray a
elems) = forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
elems

-- | Index into the array. O(1) time.
(!) :: Numbered a -> Int -> (Int, a)
Numbered ByteArray
idxs SmallArray a
elems ! :: forall a. Numbered a -> Int -> (Int, a)
! Int
i =
  (forall a b. (Integral a, Num b) => a -> b
fromIntegral (forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
idxs Int
i :: Int32),
   forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
elems Int
i)

-- | Look up the value associated with a particular key.
-- If the key occurs multiple times, any of its values
-- may be returned. O(n) time.
lookup :: Int -> Numbered a -> Maybe a
lookup :: forall a. Int -> Numbered a -> Maybe a
lookup Int
i Numbered a
num =
  forall a b. Eq a => a -> [(a, b)] -> Maybe b
List.lookup Int
i (forall a. Numbered a -> [(Int, a)]
toList Numbered a
num)

-- | Associate a value with a key. Any existing occurences
-- of the key are removed. O(n) time.
put :: Int -> a -> Numbered a -> Numbered a
put :: forall a. Int -> a -> Numbered a -> Numbered a
put Int
i a
x Numbered a
num =
  forall a. [(Int, a)] -> Numbered a
fromList forall a b. (a -> b) -> a -> b
$ [(Int, a)]
lt forall a. [a] -> [a] -> [a]
++ [(Int
i, a
x)] forall a. [a] -> [a] -> [a]
++ [(Int, a)]
gt
  where
    xs :: [(Int, a)]
xs = forall a. Numbered a -> [(Int, a)]
toList Numbered a
num
    lt :: [(Int, a)]
lt = forall a. (a -> Bool) -> [a] -> [a]
List.filter ((forall a. Ord a => a -> a -> Bool
< Int
i) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Int, a)]
xs
    gt :: [(Int, a)]
gt = forall a. (a -> Bool) -> [a] -> [a]
List.filter ((forall a. Ord a => a -> a -> Bool
> Int
i) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(Int, a)]
xs

-- | Remove a given key. O(n) time.
delete :: Int -> Numbered a -> Numbered a
delete :: forall a. Int -> Numbered a -> Numbered a
delete Int
i = forall a. [(Int, a)] -> Numbered a
fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
List.filter ((forall a. Eq a => a -> a -> Bool
/= Int
i) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Numbered a -> [(Int, a)]
toList

-- | Modify the value associated with a given key.
-- If the key occurs multiple times, one of its values
-- is chosen and the others deleted. In the call
-- @modify k def f@, if the key @k@ is not present,
-- the entry @k -> f def@ is added. O(n) time.
modify :: Int -> a -> (a -> a) -> Numbered a -> Numbered a
modify :: forall a. Int -> a -> (a -> a) -> Numbered a -> Numbered a
modify Int
i a
def a -> a
f Numbered a
num =
  forall a. Int -> a -> Numbered a -> Numbered a
put Int
i (a -> a
f (forall a. a -> Maybe a -> a
fromMaybe a
def (forall a. Int -> Numbered a -> Maybe a
lookup Int
i Numbered a
num))) Numbered a
num

-- | Keep only keys satisfying a predicate.
filter :: (a -> Bool) -> Numbered a -> Numbered a
filter :: forall a. (a -> Bool) -> Numbered a -> Numbered a
filter a -> Bool
p = forall a. [(Int, a)] -> Numbered a
fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
List.filter (a -> Bool
p forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> b
snd) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Numbered a -> [(Int, a)]
toList