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

data Numbered a =
  Numbered
    {-# UNPACK #-} !ByteArray
    {-# UNPACK #-} !(SmallArray a)

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

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

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

fromList :: [(Int, a)] -> Numbered a
fromList :: [(Int, a)] -> Numbered a
fromList [(Int, a)]
xs =
  ByteArray -> SmallArray a -> Numbered a
forall a. ByteArray -> SmallArray a -> Numbered a
Numbered
    ([Int32] -> ByteArray
forall a. Prim a => [a] -> ByteArray
byteArrayFromList (((Int, a) -> Int32) -> [(Int, a)] -> [Int32]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Int32) -> ((Int, a) -> Int) -> (Int, a) -> Int32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> Int
forall a b. (a, b) -> a
fst) [(Int, a)]
xs :: [Int32]))
    ([a] -> SmallArray a
forall a. [a] -> SmallArray a
smallArrayFromList (((Int, a) -> a) -> [(Int, a)] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map (Int, a) -> a
forall a b. (a, b) -> b
snd [(Int, a)]
xs))

toList :: Numbered a -> [(Int, a)]
toList :: Numbered a -> [(Int, a)]
toList Numbered a
num =
  [Numbered a
num Numbered a -> Int -> (Int, a)
forall a. Numbered a -> Int -> (Int, a)
! Int
i | Int
i <- [Int
0..Numbered a -> Int
forall a. Numbered a -> Int
size Numbered a
numInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]]

size :: Numbered a -> Int
size :: Numbered a -> Int
size (Numbered ByteArray
_ SmallArray a
elems) = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray SmallArray a
elems

(!) :: Numbered a -> Int -> (Int, a)
Numbered ByteArray
idxs SmallArray a
elems ! :: Numbered a -> Int -> (Int, a)
! Int
i =
  (Int32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteArray -> Int -> Int32
forall a. Prim a => ByteArray -> Int -> a
indexByteArray ByteArray
idxs Int
i :: Int32),
   SmallArray a -> Int -> a
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray a
elems Int
i)

lookup :: Int -> Numbered a -> Maybe a
lookup :: Int -> Numbered a -> Maybe a
lookup Int
i Numbered a
num =
  Int -> [(Int, a)] -> Maybe a
forall a b. Eq a => a -> [(a, b)] -> Maybe b
List.lookup Int
i (Numbered a -> [(Int, a)]
forall a. Numbered a -> [(Int, a)]
toList Numbered a
num)

put :: Int -> a -> Numbered a -> Numbered a
put :: Int -> a -> Numbered a -> Numbered a
put Int
i a
x Numbered a
num =
  [(Int, a)] -> Numbered a
forall a. [(Int, a)] -> Numbered a
fromList ([(Int, a)] -> Numbered a) -> [(Int, a)] -> Numbered a
forall a b. (a -> b) -> a -> b
$ [(Int, a)]
lt [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a] -> [a]
++ [(Int
i, a
x)] [(Int, a)] -> [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a] -> [a]
++ [(Int, a)]
gt
  where
    xs :: [(Int, a)]
xs = Numbered a -> [(Int, a)]
forall a. Numbered a -> [(Int, a)]
toList Numbered a
num
    lt :: [(Int, a)]
lt = ((Int, a) -> Bool) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> Bool) -> [a] -> [a]
List.filter ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i) (Int -> Bool) -> ((Int, a) -> Int) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> Int
forall a b. (a, b) -> a
fst) [(Int, a)]
xs
    gt :: [(Int, a)]
gt = ((Int, a) -> Bool) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> Bool) -> [a] -> [a]
List.filter ((Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
i) (Int -> Bool) -> ((Int, a) -> Int) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> Int
forall a b. (a, b) -> a
fst) [(Int, a)]
xs

delete :: Int -> Numbered a -> Numbered a
delete :: Int -> Numbered a -> Numbered a
delete Int
i = [(Int, a)] -> Numbered a
forall a. [(Int, a)] -> Numbered a
fromList ([(Int, a)] -> Numbered a)
-> (Numbered a -> [(Int, a)]) -> Numbered a -> Numbered a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, a) -> Bool) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> Bool) -> [a] -> [a]
List.filter ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
i) (Int -> Bool) -> ((Int, a) -> Int) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> Int
forall a b. (a, b) -> a
fst) ([(Int, a)] -> [(Int, a)])
-> (Numbered a -> [(Int, a)]) -> Numbered a -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Numbered a -> [(Int, a)]
forall a. Numbered a -> [(Int, a)]
toList

modify :: Int -> a -> (a -> a) -> Numbered a -> Numbered a
modify :: Int -> a -> (a -> a) -> Numbered a -> Numbered a
modify Int
i a
def a -> a
f Numbered a
num =
  Int -> a -> Numbered a -> Numbered a
forall a. Int -> a -> Numbered a -> Numbered a
put Int
i (a -> a
f (a -> Maybe a -> a
forall a. a -> Maybe a -> a
fromMaybe a
def (Int -> Numbered a -> Maybe a
forall a. Int -> Numbered a -> Maybe a
lookup Int
i Numbered a
num))) Numbered a
num

filter :: (a -> Bool) -> Numbered a -> Numbered a
filter :: (a -> Bool) -> Numbered a -> Numbered a
filter a -> Bool
p = [(Int, a)] -> Numbered a
forall a. [(Int, a)] -> Numbered a
fromList ([(Int, a)] -> Numbered a)
-> (Numbered a -> [(Int, a)]) -> Numbered a -> Numbered a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Int, a) -> Bool) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> Bool) -> [a] -> [a]
List.filter (a -> Bool
p (a -> Bool) -> ((Int, a) -> a) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> a
forall a b. (a, b) -> b
snd) ([(Int, a)] -> [(Int, a)])
-> (Numbered a -> [(Int, a)]) -> Numbered a -> [(Int, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Numbered a -> [(Int, a)]
forall a. Numbered a -> [(Int, a)]
toList