{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE GADTs        #-}
{-# LANGUAGE MagicHash    #-}
{-# LANGUAGE PolyKinds    #-}
{-# LANGUAGE TypeFamilies #-}

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

TypeRepMap implementation based on the optimal Vector.
-}

module Data.TypeRep.OptimalVector
       ( -- * Map type
         TypeRepMap (..)

         -- 'TypeRepMap' interface
       , empty
       , insert
       , lookup
       , size

         -- * Helpful testing functions
       , TF (..)
       , fromList
       ) where

import Prelude hiding (lookup)

import Control.Arrow ((&&&))
import Control.DeepSeq
import Data.Kind (Type)
import Data.Proxy (Proxy (..))
import Data.Typeable (Typeable, typeRep, typeRepFingerprint)
import Data.Word (Word64)
import GHC.Base (Any, Int (..), Int#, uncheckedIShiftRA#, (+#), (-#), (<#))
import GHC.Exts (inline, sortWith)
import GHC.Fingerprint (Fingerprint (..))
import Unsafe.Coerce (unsafeCoerce)

import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as Unboxed


data TypeRepMap (f :: k -> Type) = TypeRepMap
    { TypeRepMap f -> Vector Word64
fingerprintAs :: Unboxed.Vector Word64
    , TypeRepMap f -> Vector Word64
fingerprintBs :: Unboxed.Vector Word64
    , TypeRepMap f -> Vector Any
anys          :: V.Vector Any
    }

instance NFData (TypeRepMap f) where
  rnf :: TypeRepMap f -> ()
rnf TypeRepMap f
x = TypeRepMap f
x TypeRepMap f -> () -> ()
`seq` ()

fromAny :: Any -> f a
fromAny :: Any -> f a
fromAny = Any -> f a
forall a b. a -> b
unsafeCoerce

-- | Empty structure.
empty :: TypeRepMap f
empty :: TypeRepMap f
empty = Vector Word64 -> Vector Word64 -> Vector Any -> TypeRepMap f
forall k (f :: k -> *).
Vector Word64 -> Vector Word64 -> Vector Any -> TypeRepMap f
TypeRepMap Vector Word64
forall a. Monoid a => a
mempty Vector Word64
forall a. Monoid a => a
mempty Vector Any
forall a. Monoid a => a
mempty

-- | Inserts the value with its type as a key.
insert :: forall a f . a -> TypeRepMap f -> TypeRepMap f
insert :: a -> TypeRepMap f -> TypeRepMap f
insert = a -> TypeRepMap f -> TypeRepMap f
forall a. HasCallStack => a
undefined

-- | Looks up the value at the type.
-- >>> let x = lookup $ insert (11 :: Int) empty
-- >>> x :: Maybe Int
-- Just 11
-- >>> x :: Maybe ()
-- Nothing
lookup :: forall a f . Typeable a => TypeRepMap f -> Maybe (f a)
lookup :: TypeRepMap f -> Maybe (f a)
lookup TypeRepMap f
tVect =  Any -> f a
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny (Any -> f a) -> (Int -> Any) -> Int -> f a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TypeRepMap f -> Vector Any
forall k (f :: k -> *). TypeRepMap f -> Vector Any
anys TypeRepMap f
tVect Vector Any -> Int -> Any
forall a. Vector a -> Int -> a
V.!)
            (Int -> f a) -> Maybe Int -> Maybe (f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Fingerprint -> Vector Word64 -> Vector Word64 -> Maybe Int
binarySearch (TypeRep -> Fingerprint
typeRepFingerprint (TypeRep -> Fingerprint) -> TypeRep -> Fingerprint
forall a b. (a -> b) -> a -> b
$ Proxy a -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy a -> TypeRep) -> Proxy a -> TypeRep
forall a b. (a -> b) -> a -> b
$ Proxy a
forall k (t :: k). Proxy t
Proxy @a)
                             (TypeRepMap f -> Vector Word64
forall k (f :: k -> *). TypeRepMap f -> Vector Word64
fingerprintAs TypeRepMap f
tVect)
                             (TypeRepMap f -> Vector Word64
forall k (f :: k -> *). TypeRepMap f -> Vector Word64
fingerprintBs TypeRepMap f
tVect)

-- | Returns the size of the 'TypeRepMap'.
size :: TypeRepMap f -> Int
size :: TypeRepMap f -> Int
size = Vector Word64 -> Int
forall a. Unbox a => Vector a -> Int
Unboxed.length (Vector Word64 -> Int)
-> (TypeRepMap f -> Vector Word64) -> TypeRepMap f -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRepMap f -> Vector Word64
forall k (f :: k -> *). TypeRepMap f -> Vector Word64
fingerprintAs

-- | Returns the index is found.
binarySearch :: Fingerprint -> Unboxed.Vector Word64 -> Unboxed.Vector Word64 -> Maybe Int
binarySearch :: Fingerprint -> Vector Word64 -> Vector Word64 -> Maybe Int
binarySearch (Fingerprint Word64
a Word64
b) Vector Word64
fpAs Vector Word64
fpBs =
    let
      !(I# Int#
len) = Vector Word64 -> Int
forall a. Unbox a => Vector a -> Int
Unboxed.length Vector Word64
fpAs
      checkfpBs :: Int# -> Maybe Int
      checkfpBs :: Int# -> Maybe Int
checkfpBs Int#
i =
        case Int#
i Int# -> Int# -> Int#
<# Int#
len of
          Int#
0# -> Maybe Int
forall a. Maybe a
Nothing
          Int#
_ | Word64
a Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector Word64 -> Int -> Word64
forall a. Unbox a => Vector a -> Int -> a
Unboxed.unsafeIndex Vector Word64
fpAs (Int# -> Int
I# Int#
i) -> Maybe Int
forall a. Maybe a
Nothing
            | Word64
b Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Vector Word64 -> Int -> Word64
forall a. Unbox a => Vector a -> Int -> a
Unboxed.unsafeIndex Vector Word64
fpBs (Int# -> Int
I# Int#
i) -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# Int#
i)
            | Bool
otherwise -> Int# -> Maybe Int
checkfpBs (Int#
i Int# -> Int# -> Int#
+# Int#
1#)
    in
      Maybe Int -> Maybe Int
forall a. a -> a
inline (Int# -> Maybe Int
checkfpBs (Int# -> Int# -> Int#
binSearchHelp (Int#
-1#) Int#
len))
  where
    binSearchHelp :: Int# -> Int# -> Int#
    binSearchHelp :: Int# -> Int# -> Int#
binSearchHelp Int#
l Int#
r = case Int#
l Int# -> Int# -> Int#
<# (Int#
r Int# -> Int# -> Int#
-# Int#
1#) of
        Int#
0# -> Int#
r
        Int#
_  ->
            let m :: Int#
m = Int# -> Int# -> Int#
uncheckedIShiftRA# (Int#
l Int# -> Int# -> Int#
+# Int#
r) Int#
1# in
            if Vector Word64 -> Int -> Word64
forall a. Unbox a => Vector a -> Int -> a
Unboxed.unsafeIndex Vector Word64
fpAs (Int# -> Int
I# Int#
m) Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
a
                then Int# -> Int# -> Int#
binSearchHelp Int#
m Int#
r
                else Int# -> Int# -> Int#
binSearchHelp Int#
l Int#
m

----------------------------------------------------------------------------
-- Functions for testing and benchmarking
----------------------------------------------------------------------------

data TF f where
  TF :: Typeable a => f a -> TF f

fromF :: f a -> Proxy a
fromF :: f a -> Proxy a
fromF f a
_ = Proxy a
forall k (t :: k). Proxy t
Proxy

fromList :: forall f . [TF f] -> TypeRepMap f
fromList :: [TF f] -> TypeRepMap f
fromList [TF f]
tfs = Vector Word64 -> Vector Word64 -> Vector Any -> TypeRepMap f
forall k (f :: k -> *).
Vector Word64 -> Vector Word64 -> Vector Any -> TypeRepMap f
TypeRepMap ([Word64] -> Vector Word64
forall a. Unbox a => [a] -> Vector a
Unboxed.fromList [Word64]
fpAs) ([Word64] -> Vector Word64
forall a. Unbox a => [a] -> Vector a
Unboxed.fromList [Word64]
fpBs) ([Any] -> Vector Any
forall a. [a] -> Vector a
V.fromList [Any]
ans)
  where
    ([Word64]
fpAs, [Word64]
fpBs) = [(Word64, Word64)] -> ([Word64], [Word64])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Word64, Word64)] -> ([Word64], [Word64]))
-> [(Word64, Word64)] -> ([Word64], [Word64])
forall a b. (a -> b) -> a -> b
$ (Fingerprint -> (Word64, Word64))
-> [Fingerprint] -> [(Word64, Word64)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Fingerprint Word64
a Word64
b) -> (Word64
a, Word64
b)) [Fingerprint]
fps
    ([Fingerprint]
fps, [Any]
ans) = [(Fingerprint, Any)] -> ([Fingerprint], [Any])
forall a b. [(a, b)] -> ([a], [b])
unzip ([(Fingerprint, Any)] -> ([Fingerprint], [Any]))
-> [(Fingerprint, Any)] -> ([Fingerprint], [Any])
forall a b. (a -> b) -> a -> b
$ ((Fingerprint, Any) -> Fingerprint)
-> [(Fingerprint, Any)] -> [(Fingerprint, Any)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortWith (Fingerprint, Any) -> Fingerprint
forall a b. (a, b) -> a
fst ([(Fingerprint, Any)] -> [(Fingerprint, Any)])
-> [(Fingerprint, Any)] -> [(Fingerprint, Any)]
forall a b. (a -> b) -> a -> b
$ (TF f -> (Fingerprint, Any)) -> [TF f] -> [(Fingerprint, Any)]
forall a b. (a -> b) -> [a] -> [b]
map (TF f -> Fingerprint
fp (TF f -> Fingerprint)
-> (TF f -> Any) -> TF f -> (Fingerprint, Any)
forall (a :: * -> * -> *) b c c'.
Arrow a =>
a b c -> a b c' -> a b (c, c')
&&& TF f -> Any
an) [TF f]
tfs

    fp :: TF f -> Fingerprint
    fp :: TF f -> Fingerprint
fp (TF f a
x) = TypeRep -> Fingerprint
typeRepFingerprint (TypeRep -> Fingerprint) -> TypeRep -> Fingerprint
forall a b. (a -> b) -> a -> b
$ Proxy a -> TypeRep
forall k (proxy :: k -> *) (a :: k).
Typeable a =>
proxy a -> TypeRep
typeRep (Proxy a -> TypeRep) -> Proxy a -> TypeRep
forall a b. (a -> b) -> a -> b
$ f a -> Proxy a
forall k (f :: k -> *) (a :: k). f a -> Proxy a
fromF f a
x

    an :: TF f -> Any
    an :: TF f -> Any
an (TF f a
x) = f a -> Any
forall a b. a -> b
unsafeCoerce f a
x