-- | This module provides an efficient implementation of maps from 'Int' keys to values.
-- It is a thin wrapper around 'Data.HashCons.WordMap', converting 'Int' keys to 'Word'.

module Data.HashCons.IntMap
  ( IntMap
  , empty
  , singleton
  , singletonNonEmpty
  , insert
  , insertNonEmpty
  , lookup
  , lookupNonEmpty
  , map
  , mapWithKey
  , traverseWithKey
  , union
  , unionNonEmpty
  , intersection
  , intersectionNonEmpty
  , difference
  , differenceNonEmpty
  , member
  ) where

import Prelude hiding (lookup, map)
import Data.HashCons.WordMap (WordMap, NonEmptyWordMap)
import qualified Data.HashCons.WordMap as WordMap
import Data.Word
import Data.Hashable

newtype IntMap a = IntMap { forall a. IntMap a -> WordMap a
unWordMap :: WordMap a }

newtype NonEmptyIntMap a = NonEmptyIntMap { forall a. NonEmptyIntMap a -> NonEmptyWordMap a
unNonEmptyWordMap :: NonEmptyWordMap a }

-- | /O(1)/. Create an empty map.
empty :: IntMap a
empty :: forall a. IntMap a
empty = WordMap a -> IntMap a
forall a. WordMap a -> IntMap a
IntMap WordMap a
forall a. WordMap a
WordMap.empty

-- | /O(1)/. Create a map with a single key-value pair.
singleton :: Hashable a => Int -> a -> IntMap a
singleton :: forall a. Hashable a => Int -> a -> IntMap a
singleton Int
k a
v = WordMap a -> IntMap a
forall a. WordMap a -> IntMap a
IntMap (WordMap a -> IntMap a) -> WordMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ Word -> a -> WordMap a
forall a. Hashable a => Word -> a -> WordMap a
WordMap.singleton (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) a
v

-- | /O(1)/. Create a map with a single key-value pair.
singletonNonEmpty :: Hashable a => Int -> a -> NonEmptyIntMap a
singletonNonEmpty :: forall a. Hashable a => Int -> a -> NonEmptyIntMap a
singletonNonEmpty Int
k a
v = NonEmptyWordMap a -> NonEmptyIntMap a
forall a. NonEmptyWordMap a -> NonEmptyIntMap a
NonEmptyIntMap (NonEmptyWordMap a -> NonEmptyIntMap a)
-> NonEmptyWordMap a -> NonEmptyIntMap a
forall a b. (a -> b) -> a -> b
$ Word -> a -> NonEmptyWordMap a
forall a. Hashable a => Word -> a -> NonEmptyWordMap a
WordMap.singletonNonEmpty (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) a
v

-- | /O(log n)/. Insert a key-value pair into the map.
insert :: Hashable a => Int -> a -> IntMap a -> IntMap a
insert :: forall a. Hashable a => Int -> a -> IntMap a -> IntMap a
insert Int
k a
v (IntMap WordMap a
m) = WordMap a -> IntMap a
forall a. WordMap a -> IntMap a
IntMap (WordMap a -> IntMap a) -> WordMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ Word -> a -> WordMap a -> WordMap a
forall a. Hashable a => Word -> a -> WordMap a -> WordMap a
WordMap.insert (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) a
v WordMap a
m

-- | /O(log n)/. Insert a key-value pair into the map.
insertNonEmpty :: Hashable a => Int -> a -> NonEmptyIntMap a -> NonEmptyIntMap a
insertNonEmpty :: forall a.
Hashable a =>
Int -> a -> NonEmptyIntMap a -> NonEmptyIntMap a
insertNonEmpty Int
k a
v (NonEmptyIntMap NonEmptyWordMap a
m) = NonEmptyWordMap a -> NonEmptyIntMap a
forall a. NonEmptyWordMap a -> NonEmptyIntMap a
NonEmptyIntMap (NonEmptyWordMap a -> NonEmptyIntMap a)
-> NonEmptyWordMap a -> NonEmptyIntMap a
forall a b. (a -> b) -> a -> b
$ Word -> a -> NonEmptyWordMap a -> NonEmptyWordMap a
forall a.
Hashable a =>
Word -> a -> NonEmptyWordMap a -> NonEmptyWordMap a
WordMap.insertNonEmpty (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) a
v NonEmptyWordMap a
m

-- | /O(log n)/. Lookup the value at a key in the map.
lookup :: Int -> IntMap a -> Maybe a
lookup :: forall a. Int -> IntMap a -> Maybe a
lookup Int
k (IntMap WordMap a
m) = Word -> WordMap a -> Maybe a
forall a. Word -> WordMap a -> Maybe a
WordMap.lookup (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) WordMap a
m

-- | /O(log n)/. Lookup the value at a key in the map.
lookupNonEmpty :: Int -> NonEmptyIntMap a -> Maybe a
lookupNonEmpty :: forall a. Int -> NonEmptyIntMap a -> Maybe a
lookupNonEmpty Int
k (NonEmptyIntMap NonEmptyWordMap a
m) = Word -> NonEmptyWordMap a -> Maybe a
forall a. Word -> NonEmptyWordMap a -> Maybe a
WordMap.lookupNonEmpty (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) NonEmptyWordMap a
m

-- | /O(n)/. Map a function over all values in the map.
map :: Hashable b => (a -> b) -> IntMap a -> IntMap b
map :: forall b a. Hashable b => (a -> b) -> IntMap a -> IntMap b
map a -> b
f (IntMap WordMap a
m) = WordMap b -> IntMap b
forall a. WordMap a -> IntMap a
IntMap (WordMap b -> IntMap b) -> WordMap b -> IntMap b
forall a b. (a -> b) -> a -> b
$ (a -> b) -> WordMap a -> WordMap b
forall b a. Hashable b => (a -> b) -> WordMap a -> WordMap b
WordMap.map a -> b
f WordMap a
m

-- | /O(n)/. Map a function over all key-value pairs in the map.
mapWithKey :: Hashable b => (Int -> a -> b) -> IntMap a -> IntMap b
mapWithKey :: forall b a. Hashable b => (Int -> a -> b) -> IntMap a -> IntMap b
mapWithKey Int -> a -> b
f (IntMap WordMap a
m) = WordMap b -> IntMap b
forall a. WordMap a -> IntMap a
IntMap (WordMap b -> IntMap b) -> WordMap b -> IntMap b
forall a b. (a -> b) -> a -> b
$ (Word -> a -> b) -> WordMap a -> WordMap b
forall b a.
Hashable b =>
(Word -> a -> b) -> WordMap a -> WordMap b
WordMap.mapWithKey (\Word
k a
v -> Int -> a -> b
f (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
k) a
v) WordMap a
m

-- | /O(n)/. Traverse the map with a function that accesses keys.
traverseWithKey :: (Applicative t, Hashable b) => (Int -> a -> t b) -> IntMap a -> t (IntMap b)
traverseWithKey :: forall (t :: * -> *) b a.
(Applicative t, Hashable b) =>
(Int -> a -> t b) -> IntMap a -> t (IntMap b)
traverseWithKey Int -> a -> t b
f (IntMap WordMap a
m) = WordMap b -> IntMap b
forall a. WordMap a -> IntMap a
IntMap (WordMap b -> IntMap b) -> t (WordMap b) -> t (IntMap b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Word -> a -> t b) -> WordMap a -> t (WordMap b)
forall (t :: * -> *) b a.
(Applicative t, Hashable b) =>
(Word -> a -> t b) -> WordMap a -> t (WordMap b)
WordMap.traverseWithKey (\Word
k a
v -> Int -> a -> t b
f (Word -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
k) a
v) WordMap a
m

-- | /O(n)/. The union of two maps. If a key occurs in both maps, the value from the left map is used.
union :: Hashable a => IntMap a -> IntMap a -> IntMap a
union :: forall a. Hashable a => IntMap a -> IntMap a -> IntMap a
union (IntMap WordMap a
m1) (IntMap WordMap a
m2) = WordMap a -> IntMap a
forall a. WordMap a -> IntMap a
IntMap (WordMap a -> IntMap a) -> WordMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ WordMap a -> WordMap a -> WordMap a
forall a. Hashable a => WordMap a -> WordMap a -> WordMap a
WordMap.union WordMap a
m1 WordMap a
m2

-- | /O(n)/. The union of two maps. If a key occurs in both maps, the value from the left map is used.
unionNonEmpty :: Hashable a => NonEmptyIntMap a -> IntMap a -> NonEmptyIntMap a
unionNonEmpty :: forall a.
Hashable a =>
NonEmptyIntMap a -> IntMap a -> NonEmptyIntMap a
unionNonEmpty (NonEmptyIntMap NonEmptyWordMap a
m1) (IntMap WordMap a
m2) = NonEmptyWordMap a -> NonEmptyIntMap a
forall a. NonEmptyWordMap a -> NonEmptyIntMap a
NonEmptyIntMap (NonEmptyWordMap a -> NonEmptyIntMap a)
-> NonEmptyWordMap a -> NonEmptyIntMap a
forall a b. (a -> b) -> a -> b
$ NonEmptyWordMap a -> WordMap a -> NonEmptyWordMap a
forall a.
Hashable a =>
NonEmptyWordMap a -> WordMap a -> NonEmptyWordMap a
WordMap.unionNonEmpty NonEmptyWordMap a
m1 WordMap a
m2

-- | /O(n)/. Intersection of two maps. Only keys present in both maps are included.
intersection :: Hashable a => IntMap a -> IntMap a -> IntMap a
intersection :: forall a. Hashable a => IntMap a -> IntMap a -> IntMap a
intersection (IntMap WordMap a
m1) (IntMap WordMap a
m2) = WordMap a -> IntMap a
forall a. WordMap a -> IntMap a
IntMap (WordMap a -> IntMap a) -> WordMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ WordMap a -> WordMap a -> WordMap a
forall a. Hashable a => WordMap a -> WordMap a -> WordMap a
WordMap.intersection WordMap a
m1 WordMap a
m2

-- | /O(n)/. Intersection of two maps. Only keys present in both maps are included.
intersectionNonEmpty :: Hashable a => NonEmptyIntMap a -> IntMap a -> IntMap a
intersectionNonEmpty :: forall a. Hashable a => NonEmptyIntMap a -> IntMap a -> IntMap a
intersectionNonEmpty (NonEmptyIntMap NonEmptyWordMap a
m1) (IntMap WordMap a
m2) = WordMap a -> IntMap a
forall a. WordMap a -> IntMap a
IntMap (WordMap a -> IntMap a) -> WordMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ NonEmptyWordMap a -> WordMap a -> WordMap a
forall a. Hashable a => NonEmptyWordMap a -> WordMap a -> WordMap a
WordMap.intersectionNonEmpty NonEmptyWordMap a
m1 WordMap a
m2

-- | /O(n)/. Difference of two maps. Keys in the first map but not in the second are included.
difference :: Hashable a => IntMap a -> IntMap a -> IntMap a
difference :: forall a. Hashable a => IntMap a -> IntMap a -> IntMap a
difference (IntMap WordMap a
m1) (IntMap WordMap a
m2) = WordMap a -> IntMap a
forall a. WordMap a -> IntMap a
IntMap (WordMap a -> IntMap a) -> WordMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ WordMap a -> WordMap a -> WordMap a
forall a. Hashable a => WordMap a -> WordMap a -> WordMap a
WordMap.difference WordMap a
m1 WordMap a
m2

-- | /O(n)/. Intersection of two maps. Only keys present in both maps are included.
differenceNonEmpty :: Hashable a => NonEmptyIntMap a -> IntMap a -> IntMap a
differenceNonEmpty :: forall a. Hashable a => NonEmptyIntMap a -> IntMap a -> IntMap a
differenceNonEmpty (NonEmptyIntMap NonEmptyWordMap a
m1) (IntMap WordMap a
m2) = WordMap a -> IntMap a
forall a. WordMap a -> IntMap a
IntMap (WordMap a -> IntMap a) -> WordMap a -> IntMap a
forall a b. (a -> b) -> a -> b
$ NonEmptyWordMap a -> WordMap a -> WordMap a
forall a. Hashable a => NonEmptyWordMap a -> WordMap a -> WordMap a
WordMap.differenceNonEmpty NonEmptyWordMap a
m1 WordMap a
m2

-- | Check if a key is present in the map.
member :: Int -> IntMap a -> Bool
member :: forall a. Int -> IntMap a -> Bool
member Int
k (IntMap WordMap a
m) = Word -> WordMap a -> Bool
forall a. Word -> WordMap a -> Bool
WordMap.member (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) WordMap a
m

-- | Check if a key is present in the map.
memberNonEmpty :: Int -> NonEmptyIntMap a -> Bool
memberNonEmpty :: forall a. Int -> NonEmptyIntMap a -> Bool
memberNonEmpty Int
k (NonEmptyIntMap NonEmptyWordMap a
m) = Word -> NonEmptyWordMap a -> Bool
forall a. Word -> NonEmptyWordMap a -> Bool
WordMap.memberNonEmpty (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
k) NonEmptyWordMap a
m