{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE Rank2Types          #-}

{- |
Module                  : Data.TMap
Copyright               : (c) 2017-2022 Kowainik
SPDX-License-Identifier : MPL-2.0
Maintainer              : Kowainik <xrom.xkov@gmail.com>
Stability               : Stable
Portability             : Portable

'TMap' is a heterogeneous data structure similar in its essence to
'Data.Map.Map' with types as keys, where each value has the type of its key.

Here is an example of a 'TMap' with a comparison to 'Data.Map.Map':

@
 'Data.Map.Map' 'Prelude.String' 'Prelude.String'             'TMap'
--------------------     -----------------
 \"Int\"  -> \"5\"             'Prelude.Int'  -> 5
 \"Bool\" -> \"True\"          'Prelude.Bool' -> 'Prelude.True'
 \"Char\" -> \"\'x\'\"           'Prelude.Char' -> \'x\'
@

The runtime representation of 'TMap' is an array, not a tree. This makes
'lookup' significantly more efficient.
-}

module Data.TMap
       ( -- * Map type
         TMap

         -- * Construction
       , empty
       , one

         -- * Modification
       , insert
       , delete
       , unionWith
       , union
       , intersectionWith
       , intersection
       , map
       , adjust
       , alter

         -- * Query
       , lookup
       , member
       , size
       , keys
       , keysWith
       , toListWith
       ) where

import Prelude hiding (lookup, map)

import Data.Functor.Identity (Identity (..))
import Data.Typeable (Typeable)
import GHC.Exts (coerce)
import Type.Reflection (SomeTypeRep, TypeRep)

import qualified Data.TypeRepMap as F

-- | 'TMap' is a special case of 'F.TypeRepMap' when the interpretation is
-- 'Identity'.
type TMap = F.TypeRepMap Identity

{- |

A 'TMap' with no values stored in it.

prop> size empty == 0
prop> member @a empty == False

-}
empty :: TMap
empty :: TMap
empty = TMap
forall k (f :: k -> *). TypeRepMap f
F.empty
{-# INLINE empty #-}

{- |

Construct a 'TMap' with a single element.

prop> size (one x) == 1
prop> member @a (one (x :: a)) == True

-}
one :: forall a . Typeable a => a -> TMap
one :: a -> TMap
one a
x = TMap -> TMap
coerce (Typeable a => Identity a -> TMap
forall k (a :: k) (f :: k -> *). Typeable a => f a -> TypeRepMap f
F.one @a @Identity (Identity a -> TMap) -> Identity a -> TMap
forall a b. (a -> b) -> a -> b
$ a -> Identity a
coerce a
x)
{-# INLINE one #-}

{- |

Insert a value into a 'TMap'.
TMap optimizes for fast reads rather than inserts, as a trade-off inserts are @O(n)@.

prop> size (insert v tm) >= size tm
prop> member @a (insert (x :: a) tm) == True

-}
insert :: forall a . Typeable a => a -> TMap -> TMap
insert :: a -> TMap -> TMap
insert a
x = (TMap -> TMap) -> TMap -> TMap
coerce (Typeable a => Identity a -> TMap -> TMap
forall k (a :: k) (f :: k -> *).
Typeable a =>
f a -> TypeRepMap f -> TypeRepMap f
F.insert @a @Identity (Identity a -> TMap -> TMap) -> Identity a -> TMap -> TMap
forall a b. (a -> b) -> a -> b
$ a -> Identity a
coerce a
x)
{-# INLINE insert #-}

{- | Delete a value from a 'TMap'.

TMap optimizes for fast reads rather than modifications, as a trade-off deletes are @O(n)@,
with an @O(log(n))@ optimization for when the element is already missing.

prop> size (delete @a tm) <= size tm
prop> member @a (delete @a tm) == False

>>> tm = delete @Bool $ insert True $ one 'a'
>>> size tm
1
>>> member @Bool tm
False
>>> member @Char tm
True
-}
delete :: forall a . Typeable a => TMap -> TMap
delete :: TMap -> TMap
delete = Typeable a => TMap -> TMap
forall k (a :: k) (f :: k -> *).
Typeable a =>
TypeRepMap f -> TypeRepMap f
F.delete @a @Identity
{-# INLINE delete #-}

-- | The union of two 'TMap's using a combining function.
unionWith :: (forall x . Typeable x => x -> x -> x) -> TMap -> TMap -> TMap
unionWith :: (forall x. Typeable x => x -> x -> x) -> TMap -> TMap -> TMap
unionWith forall x. Typeable x => x -> x -> x
f = (forall x. Typeable x => Identity x -> Identity x -> Identity x)
-> TMap -> TMap -> TMap
forall k (f :: k -> *).
(forall (x :: k). Typeable x => f x -> f x -> f x)
-> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
F.unionWith forall x. Typeable x => Identity x -> Identity x -> Identity x
fId
  where
    fId :: forall y . Typeable y => Identity y -> Identity y -> Identity y
    fId :: Identity y -> Identity y -> Identity y
fId Identity y
y1 Identity y
y2 = y -> Identity y
forall a. a -> Identity a
Identity (y -> Identity y) -> y -> Identity y
forall a b. (a -> b) -> a -> b
$ y -> y -> y
forall x. Typeable x => x -> x -> x
f (Identity y -> y
coerce Identity y
y1) (Identity y -> y
coerce Identity y
y2)
{-# INLINE unionWith #-}

-- | The (left-biased) union of two 'TMap's. It prefers the first map when
-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
union :: TMap -> TMap -> TMap
union :: TMap -> TMap -> TMap
union = TMap -> TMap -> TMap
forall k (f :: k -> *).
TypeRepMap f -> TypeRepMap f -> TypeRepMap f
F.union
{-# INLINE union #-}

-- | The intersection of two 'TMap's using a combining function.
--
-- @O(n + m)@
intersectionWith :: (forall x . Typeable x => x -> x -> x) -> TMap -> TMap -> TMap
intersectionWith :: (forall x. Typeable x => x -> x -> x) -> TMap -> TMap -> TMap
intersectionWith forall x. Typeable x => x -> x -> x
f = (forall x. Typeable x => Identity x -> Identity x -> Identity x)
-> TMap -> TMap -> TMap
forall k (f :: k -> *).
(forall (x :: k). Typeable x => f x -> f x -> f x)
-> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
F.intersectionWith forall x. Typeable x => Identity x -> Identity x -> Identity x
fId
  where
    fId :: forall y . Typeable y => Identity y -> Identity y -> Identity y
    fId :: Identity y -> Identity y -> Identity y
fId Identity y
y1 Identity y
y2 = Identity y -> Identity y -> Identity y
forall x. Typeable x => x -> x -> x
f (Identity y -> Identity y
coerce Identity y
y1) (Identity y -> Identity y
coerce Identity y
y2)
{-# INLINE intersectionWith #-}

-- | The intersection of two 'TMap's.
-- It keeps all values from the first map whose keys are present in the second.
--
-- @O(n + m)@
intersection :: TMap -> TMap -> TMap
intersection :: TMap -> TMap -> TMap
intersection = TMap -> TMap -> TMap
forall k (f :: k -> *).
TypeRepMap f -> TypeRepMap f -> TypeRepMap f
F.intersection
{-# INLINE intersection #-}

{- | Lookup a value of the given type in a 'TMap'.

>>> x = lookup $ insert (11 :: Int) empty
>>> x :: Maybe Int
Just 11
>>> x :: Maybe ()
Nothing
-}
lookup :: forall a . Typeable a => TMap -> Maybe a
lookup :: TMap -> Maybe a
lookup = (TMap -> Maybe (Identity a)) -> TMap -> Maybe a
coerce (Typeable a => TMap -> Maybe (Identity a)
forall k (a :: k) (f :: k -> *).
Typeable a =>
TypeRepMap f -> Maybe (f a)
F.lookup @a @Identity)
{-# INLINE lookup #-}

{- | Check if a value of the given type is present in a 'TMap'.

>>> member @Char $ one 'a'
True
>>> member @Bool $ one 'a'
False
-}
member :: forall a . Typeable a => TMap -> Bool
member :: TMap -> Bool
member = Typeable a => TMap -> Bool
forall k (a :: k) (f :: k -> *). Typeable a => TypeRepMap f -> Bool
F.member @a @Identity
{-# INLINE member #-}

-- | Get the amount of elements in a 'TMap'.
size :: TMap -> Int
size :: TMap -> Int
size = TMap -> Int
forall k (f :: k -> *). TypeRepMap f -> Int
F.size
{-# INLINE size #-}

-- | Returns the list of 'SomeTypeRep's from keys.
keys :: TMap -> [SomeTypeRep]
keys :: TMap -> [SomeTypeRep]
keys = TMap -> [SomeTypeRep]
forall k (f :: k -> *). TypeRepMap f -> [SomeTypeRep]
F.keys
{-# INLINE keys #-}

-- | Return the list of keys by wrapping them with a user-provided function.
keysWith :: (forall a . TypeRep a -> r) -> TMap -> [r]
keysWith :: (forall a. TypeRep a -> r) -> TMap -> [r]
keysWith = (forall a. TypeRep a -> r) -> TMap -> [r]
forall k (f :: k -> *) r.
(forall (a :: k). TypeRep a -> r) -> TypeRepMap f -> [r]
F.keysWith
{-# INLINE keysWith #-}

-- | Return the list of key-value pairs by wrapping them with a user-provided function.
toListWith :: (forall a . Typeable a => a -> r) -> TMap -> [r]
toListWith :: (forall a. Typeable a => a -> r) -> TMap -> [r]
toListWith forall a. Typeable a => a -> r
f = (forall a. Typeable a => Identity a -> r) -> TMap -> [r]
forall k (f :: k -> *) r.
(forall (a :: k). Typeable a => f a -> r) -> TypeRepMap f -> [r]
F.toListWith (a -> r
forall a. Typeable a => a -> r
f (a -> r) -> (Identity a -> a) -> Identity a -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Identity a -> a
forall a. Identity a -> a
runIdentity)
{-# INLINE toListWith #-}

-- | Map a function over the values.
map :: (forall a . Typeable a => a -> a) -> TMap -> TMap
map :: (forall a. Typeable a => a -> a) -> TMap -> TMap
map forall a. Typeable a => a -> a
f = (forall x. Typeable x => Identity x -> Identity x) -> TMap -> TMap
forall k (f :: k -> *) (g :: k -> *).
(forall (x :: k). Typeable x => f x -> g x)
-> TypeRepMap f -> TypeRepMap g
F.hoistWithKey ((x -> x) -> Identity x -> Identity x
forall a. (a -> a) -> Identity a -> Identity a
liftToIdentity x -> x
forall a. Typeable a => a -> a
f)
{-# INLINE map #-}

-- | Update a value with the result of the provided function.
adjust :: Typeable a => (a -> a) -> TMap -> TMap
adjust :: (a -> a) -> TMap -> TMap
adjust a -> a
f = (Identity a -> Identity a) -> TMap -> TMap
forall k (a :: k) (f :: k -> *).
Typeable a =>
(f a -> f a) -> TypeRepMap f -> TypeRepMap f
F.adjust ((a -> a) -> Identity a -> Identity a
forall a. (a -> a) -> Identity a -> Identity a
liftToIdentity a -> a
f)
{-# INLINE adjust #-}

-- | Updates a value at a specific key, whether or not it exists.
--   This can be used to insert, delete, or update a value of a given type in the map.
alter :: Typeable a => (Maybe a -> Maybe a) -> TMap -> TMap
alter :: (Maybe a -> Maybe a) -> TMap -> TMap
alter Maybe a -> Maybe a
f = (Maybe (Identity a) -> Maybe (Identity a)) -> TMap -> TMap
forall k (a :: k) (f :: k -> *).
Typeable a =>
(Maybe (f a) -> Maybe (f a)) -> TypeRepMap f -> TypeRepMap f
F.alter ((Maybe a -> Maybe a) -> Maybe (Identity a) -> Maybe (Identity a)
forall a.
(Maybe a -> Maybe a) -> Maybe (Identity a) -> Maybe (Identity a)
liftF Maybe a -> Maybe a
f)
  where
    liftF :: forall a . (Maybe a -> Maybe a) -> Maybe (Identity a) -> Maybe (Identity a)
    liftF :: (Maybe a -> Maybe a) -> Maybe (Identity a) -> Maybe (Identity a)
liftF = (Maybe a -> Maybe a) -> Maybe (Identity a) -> Maybe (Identity a)
coerce
{-# INLINE alter #-}

liftToIdentity :: forall a . (a -> a) -> Identity a -> Identity a
liftToIdentity :: (a -> a) -> Identity a -> Identity a
liftToIdentity = (a -> a) -> Identity a -> Identity a
coerce