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

{- |
Copyright:  (c) 2017-2020 Kowainik
SPDX-License-Identifier: MPL-2.0
Maintainer: Kowainik <xrom.xkov@gmail.com>

'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
       , map
       , adjust

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

import Prelude hiding (lookup, map)

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

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 x :: a
x = TMap -> TMap
forall a b. Coercible a b => a -> b
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
forall a b. Coercible a b => a -> b
coerce a
x)
{-# INLINE one #-}

{- |

Insert a value into a 'TMap'.

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 x :: a
x = (TMap -> TMap) -> TMap -> TMap
forall a b. Coercible a b => a -> b
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
forall a b. Coercible a b => a -> b
coerce a
x)
{-# INLINE insert #-}

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

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 f :: 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 y1 :: Identity y
y1 y2 :: 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
forall a b. Coercible a b => a -> b
coerce Identity y
y1) (Identity y -> y
forall a b. Coercible a b => a -> b
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 #-}

{- | 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
forall a b. Coercible a b => a -> b
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 #-}

-- | 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 f :: 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 f :: 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 #-}

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