{-# OPTIONS_GHC -Wno-missing-export-lists #-}

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE BangPatterns          #-}
{-# LANGUAGE CPP                   #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE InstanceSigs          #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE RoleAnnotations       #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeInType            #-}
{-# LANGUAGE ViewPatterns          #-}

#if __GLASGOW_HASKELL__ >= 806
{-# LANGUAGE QuantifiedConstraints #-}
#endif

-- {-# OPTIONS_GHC -ddump-simpl -dsuppress-idinfo -dsuppress-coercions -dsuppress-type-applications -dsuppress-uniques -dsuppress-module-prefixes #-}

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

Internal API for 'TypeRepMap' and operations on it. The functions here do
not have any stability guarantees and can change between minor versions.

If you need to use this module for purposes other than tests,
create an issue.
-}

#include "MachDeps.h"

module Data.TypeRepMap.Internal where

import Prelude hiding (lookup)

import Control.DeepSeq
import Control.Monad.ST (ST, runST)
import Control.Monad.Zip (mzip)
import Data.Function (on)
import Data.Kind (Type)
import Data.List (intercalate, nubBy)
import Data.Maybe (fromMaybe)
import Data.Primitive.Array (Array, MutableArray, indexArray, mapArray', sizeofArray, thawArray,
                             unsafeFreezeArray, writeArray)
import Data.Primitive.PrimArray (MutablePrimArray, PrimArray, indexPrimArray, newPrimArray,
                                 primArrayFromListN, primArrayToList, sizeofPrimArray,
                                 unsafeFreezePrimArray, writePrimArray)
import Data.Semigroup (All (..), Semigroup (..))
import Data.Type.Equality (TestEquality (..), (:~:) (..))
import GHC.Base (Any, Int (..), Int#, (*#), (+#), (<#))
import GHC.Exts (IsList (..), inline, sortWith)
import GHC.Fingerprint (Fingerprint (..))
#if WORD_SIZE_IN_BITS >= 64
import GHC.Prim (eqWord#, ltWord#)
#else
import GHC.IntWord64 (eqWord64#, ltWord64#)
#define eqWord eqWord64
#define ltWord ltWord64
#endif
import GHC.Word (Word64 (..))
import Type.Reflection (SomeTypeRep (..), TypeRep, Typeable, typeRep, withTypeable)
import Type.Reflection.Unsafe (typeRepFingerprint)
import Unsafe.Coerce (unsafeCoerce)

import qualified GHC.Exts as GHC (fromList, toList)

{- |

'TypeRepMap' 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. In
addition to that, each value is wrapped in an interpretation @f@.

Here is an example of using 'Prelude.Maybe' as an interpretation, with a
comparison to 'Data.Map.Map':

@
 'Data.Map.Map' 'Prelude.String' ('Prelude.Maybe' 'Prelude.String')          'TypeRepMap' 'Prelude.Maybe'
---------------------------       ---------------------
 \"Int\"  -> Just \"5\"                 'Prelude.Int'  -> Just 5
 \"Bool\" -> Just \"True\"              'Prelude.Bool' -> Just 'Prelude.True'
 \"Char\" -> Nothing                  'Prelude.Char' -> Nothing
@

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

-}
type role TypeRepMap representational
data TypeRepMap (f :: k -> Type) =
  TypeRepMap
    { TypeRepMap f -> PrimArray Word64
fingerprintAs :: {-# UNPACK #-} !(PrimArray Word64) -- ^ first components of key fingerprints
    , TypeRepMap f -> PrimArray Word64
fingerprintBs :: {-# UNPACK #-} !(PrimArray Word64) -- ^ second components of key fingerprints
    , TypeRepMap f -> Array Any
trAnys        :: {-# UNPACK #-} !(Array Any)        -- ^ values stored in the map
    , TypeRepMap f -> Array Any
trKeys        :: {-# UNPACK #-} !(Array Any)        -- ^ typerep keys
    }
  -- ^ an unsafe constructor for 'TypeRepMap'

instance NFData (TypeRepMap f) where
   rnf :: TypeRepMap f -> ()
rnf TypeRepMap f
x = [SomeTypeRep] -> ()
forall a. NFData a => a -> ()
rnf (TypeRepMap f -> [SomeTypeRep]
forall k (f :: k -> *). TypeRepMap f -> [SomeTypeRep]
keys TypeRepMap f
x) () -> () -> ()
`seq` ()

-- | Shows only keys.
instance Show (TypeRepMap f) where
    show :: TypeRepMap f -> String
show TypeRepMap{PrimArray Word64
Array Any
trKeys :: Array Any
trAnys :: Array Any
fingerprintBs :: PrimArray Word64
fingerprintAs :: PrimArray Word64
trKeys :: forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys :: forall k (f :: k -> *). TypeRepMap f -> Array Any
fingerprintBs :: forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs :: forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
..} = String
"TypeRepMap [" String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
showKeys String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"]"
      where
        showKeys :: String
        showKeys :: String
showKeys = String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
intercalate String
", " ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ Array String -> [Item (Array String)]
forall l. IsList l => l -> [Item l]
toList (Array String -> [Item (Array String)])
-> Array String -> [Item (Array String)]
forall a b. (a -> b) -> a -> b
$ (Any -> String) -> Array Any -> Array String
forall a b. (a -> b) -> Array a -> Array b
mapArray' (TypeRep Any -> String
forall a. Show a => a -> String
show (TypeRep Any -> String) -> (Any -> TypeRep Any) -> Any -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any -> TypeRep Any
forall k (f :: k). Any -> TypeRep f
anyToTypeRep) Array Any
trKeys

-- | Uses 'union' to combine 'TypeRepMap's.
instance Semigroup (TypeRepMap f) where
    (<>) :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
    <> :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
(<>) = TypeRepMap f -> TypeRepMap f -> TypeRepMap f
forall k (f :: k -> *).
TypeRepMap f -> TypeRepMap f -> TypeRepMap f
union
    {-# INLINE (<>) #-}

instance Monoid (TypeRepMap f) where
    mempty :: TypeRepMap f
mempty = PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
forall k (f :: k -> *).
PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
TypeRepMap PrimArray Word64
forall a. Monoid a => a
mempty PrimArray Word64
forall a. Monoid a => a
mempty Array Any
forall a. Monoid a => a
mempty Array Any
forall a. Monoid a => a
mempty
    mappend :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
mappend = TypeRepMap f -> TypeRepMap f -> TypeRepMap f
forall a. Semigroup a => a -> a -> a
(<>)
    {-# INLINE mempty #-}
    {-# INLINE mappend #-}

#if __GLASGOW_HASKELL__ >= 806
instance (forall a. Typeable a => Eq (f a)) => Eq (TypeRepMap f) where
    TypeRepMap f
tm1 == :: TypeRepMap f -> TypeRepMap f -> Bool
== TypeRepMap f
tm2 = TypeRepMap f -> Int
forall k (f :: k -> *). TypeRepMap f -> Int
size TypeRepMap f
tm1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== TypeRepMap f -> Int
forall k (f :: k -> *). TypeRepMap f -> Int
size TypeRepMap f
tm2 Bool -> Bool -> Bool
&& Int -> Bool
go Int
0
      where
        go :: Int -> Bool
        go :: Int -> Bool
go Int
i
            | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== TypeRepMap f -> Int
forall k (f :: k -> *). TypeRepMap f -> Int
size TypeRepMap f
tm1 = Bool
True
            | Bool
otherwise = case TypeRep Any -> TypeRep Any -> Maybe (Any :~: Any)
forall k (f :: k -> *) (a :: k) (b :: k).
TestEquality f =>
f a -> f b -> Maybe (a :~: b)
testEquality TypeRep Any
forall k (x :: k). TypeRep x
tr1i TypeRep Any
forall k (x :: k). TypeRep x
tr2i of
                  Maybe (Any :~: Any)
Nothing   -> Bool
False
                  Just Any :~: Any
Refl -> TypeRep Any -> f Any -> f Any -> Bool
forall (x :: k). TypeRep x -> f x -> f x -> Bool
repEq TypeRep Any
forall k (x :: k). TypeRep x
tr1i (Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
tv1i) (Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
tv2i) Bool -> Bool -> Bool
&& Int -> Bool
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
          where
            tr1i :: TypeRep x
            tr1i :: TypeRep x
tr1i = Any -> TypeRep x
forall k (f :: k). Any -> TypeRep f
anyToTypeRep (Any -> TypeRep x) -> Any -> TypeRep x
forall a b. (a -> b) -> a -> b
$ Array Any -> Int -> Any
forall a. Array a -> Int -> a
indexArray (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trKeys TypeRepMap f
tm1) Int
i

            tr2i :: TypeRep y
            tr2i :: TypeRep y
tr2i = Any -> TypeRep y
forall k (f :: k). Any -> TypeRep f
anyToTypeRep (Any -> TypeRep y) -> Any -> TypeRep y
forall a b. (a -> b) -> a -> b
$ Array Any -> Int -> Any
forall a. Array a -> Int -> a
indexArray (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trKeys TypeRepMap f
tm2) Int
i

            tv1i, tv2i :: Any
            tv1i :: Any
tv1i = Array Any -> Int -> Any
forall a. Array a -> Int -> a
indexArray (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys TypeRepMap f
tm1) Int
i
            tv2i :: Any
tv2i = Array Any -> Int -> Any
forall a. Array a -> Int -> a
indexArray (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys TypeRepMap f
tm2) Int
i

            repEq :: TypeRep x -> f x -> f x -> Bool
            repEq :: TypeRep x -> f x -> f x -> Bool
repEq TypeRep x
tr = TypeRep x
-> (Typeable x => f x -> f x -> Bool) -> f x -> f x -> Bool
forall k (a :: k) r. TypeRep a -> (Typeable a => r) -> r
withTypeable TypeRep x
tr Typeable x => f x -> f x -> Bool
forall a. Eq a => a -> a -> Bool
(==)
#endif

-- | Returns the list of 'Fingerprint's from 'TypeRepMap'.
toFingerprints :: TypeRepMap f -> [Fingerprint]
toFingerprints :: TypeRepMap f -> [Fingerprint]
toFingerprints TypeRepMap{PrimArray Word64
Array Any
trKeys :: Array Any
trAnys :: Array Any
fingerprintBs :: PrimArray Word64
fingerprintAs :: PrimArray Word64
trKeys :: forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys :: forall k (f :: k -> *). TypeRepMap f -> Array Any
fingerprintBs :: forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs :: forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
..} =
    (Word64 -> Word64 -> Fingerprint)
-> [Word64] -> [Word64] -> [Fingerprint]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Word64 -> Word64 -> Fingerprint
Fingerprint (PrimArray Word64 -> [Item (PrimArray Word64)]
forall l. IsList l => l -> [Item l]
GHC.toList PrimArray Word64
fingerprintAs) (PrimArray Word64 -> [Item (PrimArray Word64)]
forall l. IsList l => l -> [Item l]
GHC.toList PrimArray Word64
fingerprintBs)

{- |

A 'TypeRepMap' with no values stored in it.

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

-}
empty :: TypeRepMap f
empty :: TypeRepMap f
empty = TypeRepMap f
forall a. Monoid a => a
mempty
{-# INLINE empty #-}

{- |

Construct a 'TypeRepMap' with a single element.

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

-}
one :: forall a f . Typeable a => f a -> TypeRepMap f
one :: f a -> TypeRepMap f
one f a
x = PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
forall k (f :: k -> *).
PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
TypeRepMap (Int -> [Word64] -> PrimArray Word64
forall a. Prim a => Int -> [a] -> PrimArray a
primArrayFromListN Int
1 [Word64
fa])
                   (Int -> [Word64] -> PrimArray Word64
forall a. Prim a => Int -> [a] -> PrimArray a
primArrayFromListN Int
1 [Word64
fb])
                   (Any -> Array Any
forall (f :: * -> *) a. Applicative f => a -> f a
pure @Array Any
v)
                   (Any -> Array Any
forall (f :: * -> *) a. Applicative f => a -> f a
pure @Array Any
k)
  where
    (Fingerprint Word64
fa Word64
fb, Any
v, Any
k) = (Typeable a => Fingerprint
forall k (a :: k). Typeable a => Fingerprint
calcFp @a, f a -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny f a
x, TypeRep a -> Any
forall a b. a -> b
unsafeCoerce (TypeRep a -> Any) -> TypeRep a -> Any
forall a b. (a -> b) -> a -> b
$ Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @a)
{-# INLINE one #-}

{- |

Insert a value into a 'TypeRepMap'.
TypeRepMap 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 :: f a) tm) == True

-}
insert :: forall a f . Typeable a => f a -> TypeRepMap f -> TypeRepMap f
insert :: f a -> TypeRepMap f -> TypeRepMap f
insert f a
x TypeRepMap f
m
  | TypeRepMap f -> Int
forall k (f :: k -> *). TypeRepMap f -> Int
size TypeRepMap f
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = f a -> TypeRepMap f
forall k (a :: k) (f :: k -> *). Typeable a => f a -> TypeRepMap f
one f a
x
  | Bool
otherwise = case Fingerprint -> PrimArray Word64 -> PrimArray Word64 -> Maybe Int
cachedBinarySearch (Typeable a => Fingerprint
forall k (a :: k). Typeable a => Fingerprint
typeFp @a) (TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs TypeRepMap f
m) (TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintBs TypeRepMap f
m) of
      Maybe Int
Nothing -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
forall k (f :: k -> *).
TypeRepMap f -> TypeRepMap f -> TypeRepMap f
union TypeRepMap f
m (TypeRepMap f -> TypeRepMap f) -> TypeRepMap f -> TypeRepMap f
forall a b. (a -> b) -> a -> b
$ f a -> TypeRepMap f
forall k (a :: k) (f :: k -> *). Typeable a => f a -> TypeRepMap f
one f a
x
      Just Int
i  -> TypeRepMap f
m {trAnys :: Array Any
trAnys = Int -> Array Any -> Array Any
changeAnyArr Int
i (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys TypeRepMap f
m)}
  where
    changeAnyArr :: Int -> Array Any -> Array Any
    changeAnyArr :: Int -> Array Any -> Array Any
changeAnyArr Int
i Array Any
trAs = (forall s. ST s (Array Any)) -> Array Any
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array Any)) -> Array Any)
-> (forall s. ST s (Array Any)) -> Array Any
forall a b. (a -> b) -> a -> b
$ do
      let n :: Int
n = Array Any -> Int
forall a. Array a -> Int
sizeofArray Array Any
trAs
      MutableArray s Any
mutArr <- Array Any
-> Int -> Int -> ST s (MutableArray (PrimState (ST s)) Any)
forall (m :: * -> *) a.
PrimMonad m =>
Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
thawArray Array Any
trAs Int
0 Int
n
      MutableArray (PrimState (ST s)) Any -> Int -> Any -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Any
MutableArray (PrimState (ST s)) Any
mutArr Int
i (Any -> ST s ()) -> Any -> ST s ()
forall a b. (a -> b) -> a -> b
$ f a -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny f a
x
      MutableArray (PrimState (ST s)) Any -> ST s (Array Any)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
unsafeFreezeArray MutableArray s Any
MutableArray (PrimState (ST s)) Any
mutArr
{-# INLINE insert #-}

-- Extract the kind of a type. We use it to work around lack of syntax for
-- inferred type variables (which are not subject to type applications).
type KindOf (a :: k) = k

type ArgKindOf (f :: k -> l) = k

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

TypeRepMap 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 (Just True) $ one (Just 'a')
>>> size tm
1
>>> member @Bool tm
False
>>> member @Char tm
True
-}
delete :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> TypeRepMap f
delete :: TypeRepMap f -> TypeRepMap f
delete TypeRepMap f
m
  -- Lookups are fast, so check if we even have the element first.
  | Bool -> Bool
not (TypeRepMap f -> Bool
forall k (a :: k) (f :: k -> *). Typeable a => TypeRepMap f -> Bool
member @a TypeRepMap f
m) = TypeRepMap f
m
  -- We know we have the element, If the map has exactly one element, we can return the empty map
  | TypeRepMap f -> Int
forall k (f :: k -> *). TypeRepMap f -> Int
size TypeRepMap f
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = TypeRepMap f
forall k (f :: k -> *). TypeRepMap f
empty
  -- Otherwise, filter out the element in linear time.
  | Bool
otherwise = [(Fingerprint, Any, Any)] -> TypeRepMap f
forall k (f :: k -> *). [(Fingerprint, Any, Any)] -> TypeRepMap f
fromSortedTriples ([(Fingerprint, Any, Any)] -> TypeRepMap f)
-> (TypeRepMap f -> [(Fingerprint, Any, Any)])
-> TypeRepMap f
-> TypeRepMap f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Fingerprint, Any, Any) -> Bool)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. (a -> Bool) -> [a] -> [a]
deleteFirst ((Fingerprint -> Fingerprint -> Bool
forall a. Eq a => a -> a -> Bool
== Typeable a => Fingerprint
forall k (a :: k). Typeable a => Fingerprint
typeFp @a) (Fingerprint -> Bool)
-> ((Fingerprint, Any, Any) -> Fingerprint)
-> (Fingerprint, Any, Any)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Fingerprint, Any, Any) -> Fingerprint
forall a b c. (a, b, c) -> a
fst3) ([(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)])
-> (TypeRepMap f -> [(Fingerprint, Any, Any)])
-> TypeRepMap f
-> [(Fingerprint, Any, Any)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRepMap f -> [(Fingerprint, Any, Any)]
forall k (f :: k -> *). TypeRepMap f -> [(Fingerprint, Any, Any)]
toSortedTriples (TypeRepMap f -> TypeRepMap f) -> TypeRepMap f -> TypeRepMap f
forall a b. (a -> b) -> a -> b
$ TypeRepMap f
m
{-# INLINE delete #-}

deleteFirst :: (a -> Bool) -> [a] -> [a]
deleteFirst :: (a -> Bool) -> [a] -> [a]
deleteFirst a -> Bool
_ []       = []
deleteFirst a -> Bool
p (a
x : [a]
xs) = if a -> Bool
p a
x then [a]
xs else a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: (a -> Bool) -> [a] -> [a]
forall a. (a -> Bool) -> [a] -> [a]
deleteFirst a -> Bool
p [a]
xs

{- |
Update a value at a specific key with the result of the provided function.
When the key is not a member of the map, the original map is returned.

>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "a"]
>>> lookup @String $ adjust (fmap (++ "ww")) trmap
Just (Identity "aww")
-}
adjust :: forall a f . Typeable a => (f a -> f a) -> TypeRepMap f -> TypeRepMap f
adjust :: (f a -> f a) -> TypeRepMap f -> TypeRepMap f
adjust f a -> f a
fun = (Maybe (f a) -> Maybe (f a)) -> TypeRepMap f -> TypeRepMap f
forall k (a :: k) (f :: k -> *).
Typeable a =>
(Maybe (f a) -> Maybe (f a)) -> TypeRepMap f -> TypeRepMap f
alter ((f a -> f a) -> Maybe (f a) -> Maybe (f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap f a -> f a
fun)
{-# 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.

>>> func = (\case Nothing -> Just (Identity "new"); Just (Identity s) -> Just (Identity (reverse s)))
>>> lookup @String $ alter func empty
Just (Identity "new")
>>> trmap = fromList @(TypeRepMap Identity) [WrapTypeable $ Identity "helllo"]
>>> lookup @String $ alter func trmap
>>> Just (Identity "olleh")
-}
alter :: forall a f . Typeable a => (Maybe (f a) -> Maybe (f a)) -> TypeRepMap f -> TypeRepMap f
alter :: (Maybe (f a) -> Maybe (f a)) -> TypeRepMap f -> TypeRepMap f
alter Maybe (f a) -> Maybe (f a)
fun TypeRepMap f
tr = case Fingerprint -> PrimArray Word64 -> PrimArray Word64 -> Maybe Int
cachedBinarySearch (Typeable a => Fingerprint
forall k (a :: k). Typeable a => Fingerprint
typeFp @a) (TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs TypeRepMap f
tr) (TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintBs TypeRepMap f
tr) of
    Maybe Int
Nothing ->
        case (Maybe (f a) -> Maybe (f a)
fun Maybe (f a)
forall a. Maybe a
Nothing) of
            Maybe (f a)
Nothing -> TypeRepMap f
tr
            Just f a
v  -> f a -> TypeRepMap f -> TypeRepMap f
forall k (a :: k) (f :: k -> *).
Typeable a =>
f a -> TypeRepMap f -> TypeRepMap f
insert f a
v TypeRepMap f
tr
    Just Int
i  ->
        case Maybe (f a) -> Maybe (f a)
fun (f a -> Maybe (f a)
forall a. a -> Maybe a
Just (f a -> Maybe (f a)) -> (Any -> f a) -> Any -> Maybe (f a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any -> f a
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny (Any -> Maybe (f a)) -> Any -> Maybe (f a)
forall a b. (a -> b) -> a -> b
$ Array Any -> Int -> Any
forall a. Array a -> Int -> a
indexArray (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys TypeRepMap f
tr) Int
i) of
            Maybe (f a)
Nothing -> TypeRepMap f -> TypeRepMap f
forall k (a :: k) (f :: k -> *).
Typeable a =>
TypeRepMap f -> TypeRepMap f
delete @a TypeRepMap f
tr
            Just f a
v  -> TypeRepMap f
tr{trAnys :: Array Any
trAnys = Int -> Any -> Array Any -> Array Any
replaceAnyAt Int
i (f a -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny f a
v) (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys TypeRepMap f
tr)}
  where
    replaceAnyAt :: Int -> Any -> Array Any -> Array Any
    replaceAnyAt :: Int -> Any -> Array Any -> Array Any
replaceAnyAt Int
i Any
v Array Any
trAs = (forall s. ST s (Array Any)) -> Array Any
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Array Any)) -> Array Any)
-> (forall s. ST s (Array Any)) -> Array Any
forall a b. (a -> b) -> a -> b
$ do
        let n :: Int
n = Array Any -> Int
forall a. Array a -> Int
sizeofArray Array Any
trAs
        MutableArray s Any
mutArr <- Array Any
-> Int -> Int -> ST s (MutableArray (PrimState (ST s)) Any)
forall (m :: * -> *) a.
PrimMonad m =>
Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
thawArray Array Any
trAs Int
0 Int
n
        MutableArray (PrimState (ST s)) Any -> Int -> Any -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s Any
MutableArray (PrimState (ST s)) Any
mutArr Int
i Any
v
        MutableArray (PrimState (ST s)) Any -> ST s (Array Any)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
unsafeFreezeArray MutableArray s Any
MutableArray (PrimState (ST s)) Any
mutArr
{-# INLINE alter #-}

{- | Map over the elements of a 'TypeRepMap'.

>>> tm = insert (Identity True) $ one (Identity 'a')
>>> lookup @Bool tm
Just (Identity True)
>>> lookup @Char tm
Just (Identity 'a')
>>> tm2 = hoist ((:[]) . runIdentity) tm
>>> lookup @Bool tm2
Just [True]
>>> lookup @Char tm2
Just "a"
-}
hoist :: (forall x. f x -> g x) -> TypeRepMap f -> TypeRepMap g
hoist :: (forall (x :: k). f x -> g x) -> TypeRepMap f -> TypeRepMap g
hoist forall (x :: k). f x -> g x
f (TypeRepMap PrimArray Word64
as PrimArray Word64
bs Array Any
ans Array Any
ks) = PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap g
forall k (f :: k -> *).
PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
TypeRepMap PrimArray Word64
as PrimArray Word64
bs ((Any -> Any) -> Array Any -> Array Any
forall a b. (a -> b) -> Array a -> Array b
mapArray' (g Any -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny (g Any -> Any) -> (Any -> g Any) -> Any -> Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f Any -> g Any
forall (x :: k). f x -> g x
f (f Any -> g Any) -> (Any -> f Any) -> Any -> g Any
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny) Array Any
ans) Array Any
ks
{-# INLINE hoist #-}

hoistA :: (Applicative t) => (forall x. f x -> t (g x)) -> TypeRepMap f -> t (TypeRepMap g)
hoistA :: (forall (x :: k). f x -> t (g x))
-> TypeRepMap f -> t (TypeRepMap g)
hoistA forall (x :: k). f x -> t (g x)
f (TypeRepMap PrimArray Word64
as PrimArray Word64
bs (Array Any -> [Item (Array Any)]
forall l. IsList l => l -> [Item l]
toList -> [Item (Array Any)]
ans) Array Any
ks) = (\[g Any]
l -> PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap g
forall k (f :: k -> *).
PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
TypeRepMap PrimArray Word64
as PrimArray Word64
bs ([Item (Array Any)] -> Array Any
forall l. IsList l => [Item l] -> l
fromList ([Item (Array Any)] -> Array Any)
-> [Item (Array Any)] -> Array Any
forall a b. (a -> b) -> a -> b
$ (g Any -> Any) -> [g Any] -> [Any]
forall a b. (a -> b) -> [a] -> [b]
map g Any -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny [g Any]
l) Array Any
ks)
    ([g Any] -> TypeRepMap g) -> t [g Any] -> t (TypeRepMap g)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Any -> t (g Any)) -> [Any] -> t [g Any]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (f Any -> t (g Any)
forall (x :: k). f x -> t (g x)
f (f Any -> t (g Any)) -> (Any -> f Any) -> Any -> t (g Any)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny) [Any]
[Item (Array Any)]
ans
{-# INLINE hoistA #-}

hoistWithKey :: forall f g. (forall x. Typeable x => f x -> g x) -> TypeRepMap f -> TypeRepMap g
hoistWithKey :: (forall (x :: k). Typeable x => f x -> g x)
-> TypeRepMap f -> TypeRepMap g
hoistWithKey forall (x :: k). Typeable x => f x -> g x
f (TypeRepMap PrimArray Word64
as PrimArray Word64
bs Array Any
ans Array Any
ks) = PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap g
forall k (f :: k -> *).
PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
TypeRepMap PrimArray Word64
as PrimArray Word64
bs Array Any
newAns Array Any
ks
  where
    newAns :: Array Any
newAns = ((Any, Any) -> Any) -> Array (Any, Any) -> Array Any
forall a b. (a -> b) -> Array a -> Array b
mapArray' (Any, Any) -> Any
mapAns (Array Any -> Array Any -> Array (Any, Any)
forall (m :: * -> *) a b. MonadZip m => m a -> m b -> m (a, b)
mzip Array Any
ans Array Any
ks)
    mapAns :: (Any, Any) -> Any
mapAns (Any
a, Any
k) = g Any -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny (g Any -> Any) -> g Any -> Any
forall a b. (a -> b) -> a -> b
$ TypeRep Any -> f Any -> g Any
forall (x :: k). TypeRep x -> f x -> g x
withTr (Any -> TypeRep Any
forall a b. a -> b
unsafeCoerce Any
k) (f Any -> g Any) -> f Any -> g Any
forall a b. (a -> b) -> a -> b
$ Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
a

    withTr :: forall x. TypeRep x -> f x -> g x
    withTr :: TypeRep x -> f x -> g x
withTr TypeRep x
t = TypeRep x -> (Typeable x => f x -> g x) -> f x -> g x
forall k (a :: k) r. TypeRep a -> (Typeable a => r) -> r
withTypeable TypeRep x
t Typeable x => f x -> g x
forall (x :: k). Typeable x => f x -> g x
f
{-# INLINE hoistWithKey #-}

-- | The union of two 'TypeRepMap's using a combining function for conflicting entries. @O(n + m)@
unionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
unionWith :: (forall (x :: k). Typeable x => f x -> f x -> f x)
-> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
unionWith forall (x :: k). Typeable x => f x -> f x -> f x
f TypeRepMap f
ma TypeRepMap f
mb = do
    [(Fingerprint, Any, Any)] -> TypeRepMap f
forall k (f :: k -> *). [(Fingerprint, Any, Any)] -> TypeRepMap f
fromSortedTriples ([(Fingerprint, Any, Any)] -> TypeRepMap f)
-> [(Fingerprint, Any, Any)] -> TypeRepMap f
forall a b. (a -> b) -> a -> b
$ [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps (TypeRepMap f -> [(Fingerprint, Any, Any)]
forall k (f :: k -> *). TypeRepMap f -> [(Fingerprint, Any, Any)]
toSortedTriples TypeRepMap f
ma) (TypeRepMap f -> [(Fingerprint, Any, Any)]
forall k (f :: k -> *). TypeRepMap f -> [(Fingerprint, Any, Any)]
toSortedTriples TypeRepMap f
mb)
  where
    f' :: forall x. TypeRep x -> f x -> f x -> f x
    f' :: TypeRep x -> f x -> f x -> f x
f' TypeRep x
tr = TypeRep x -> (Typeable x => f x -> f x -> f x) -> f x -> f x -> f x
forall k (a :: k) r. TypeRep a -> (Typeable a => r) -> r
withTypeable TypeRep x
tr Typeable x => f x -> f x -> f x
forall (x :: k). Typeable x => f x -> f x -> f x
f

    combine :: (Fingerprint, Any, Any) -> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
    combine :: (Fingerprint, Any, Any)
-> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
combine (Fingerprint
fp, Any
av, Any
ak) (Fingerprint
_, Any
bv, Any
_) = (Fingerprint
fp, f Any -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny (f Any -> Any) -> f Any -> Any
forall a b. (a -> b) -> a -> b
$ TypeRep Any -> f Any -> f Any -> f Any
forall (x :: k). TypeRep x -> f x -> f x -> f x
f' (Any -> TypeRep Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
ak) (Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
av) (Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
bv), Any
ak)

    -- Merges two typrepmaps into a sorted, dedup'd list of triples.
    -- Using 'toSortedTriples' allows us to assume the triples are sorted by fingerprint,
    -- Given O(n) performance from 'toSortedTriples', and given that we can merge-sort in
    -- O(n + m) time, then can '.fromSortedTriples' back into cachedBinarySearch order in O(n + m)
    -- that gives a total of O(n + m).
    mergeMaps :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
    -- We've addressed all elements from both maps
    mergeMaps :: [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps [(Fingerprint, Any, Any)]
as [] = [(Fingerprint, Any, Any)]
as
    mergeMaps [] [(Fingerprint, Any, Any)]
bs = [(Fingerprint, Any, Any)]
bs
    -- Merge
    mergeMaps (a :: (Fingerprint, Any, Any)
a@(Fingerprint
af, Any
_, Any
_) : [(Fingerprint, Any, Any)]
as) (b :: (Fingerprint, Any, Any)
b@(Fingerprint
bf, Any
_, Any
_) : [(Fingerprint, Any, Any)]
bs) =
        case Fingerprint -> Fingerprint -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Fingerprint
af Fingerprint
bf of
            -- Fingerprints are equal, union the elements using our function
            -- If the incoming maps were de-duped, there shouldn't be any other equivalent
            -- fingerprints
            Ordering
EQ -> (Fingerprint, Any, Any)
-> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
combine (Fingerprint, Any, Any)
a (Fingerprint, Any, Any)
b (Fingerprint, Any, Any)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. a -> [a] -> [a]
: [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps [(Fingerprint, Any, Any)]
as [(Fingerprint, Any, Any)]
bs
            -- First fingerprint must not be in the second map or we would have seen it by now
            -- Add it to the result as-is
            Ordering
LT -> (Fingerprint, Any, Any)
a (Fingerprint, Any, Any)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. a -> [a] -> [a]
: [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps [(Fingerprint, Any, Any)]
as ((Fingerprint, Any, Any)
b (Fingerprint, Any, Any)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. a -> [a] -> [a]
: [(Fingerprint, Any, Any)]
bs)
            -- Second fingerprint must not be in the first map or we would have seen it by now
            -- Add it to the result as-is
            Ordering
GT -> (Fingerprint, Any, Any)
b (Fingerprint, Any, Any)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. a -> [a] -> [a]
: [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps ((Fingerprint, Any, Any)
a(Fingerprint, Any, Any)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. a -> [a] -> [a]
:[(Fingerprint, Any, Any)]
as) [(Fingerprint, Any, Any)]
bs
{-# INLINE unionWith #-}

-- | The (left-biased) union of two 'TypeRepMap's in @O(n + m)@. It prefers the first map when
-- duplicate keys are encountered, i.e. @'union' == 'unionWith' const@.
union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
union :: TypeRepMap f -> TypeRepMap f -> TypeRepMap f
union = (forall (x :: k). Typeable x => f x -> f x -> f x)
-> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
forall k (f :: k -> *).
(forall (x :: k). Typeable x => f x -> f x -> f x)
-> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
unionWith forall (x :: k). Typeable x => f x -> f x -> f x
forall a b. a -> b -> a
const
{-# INLINE union #-}

-- | The 'intersection' of two 'TypeRepMap's using a combining function
--
-- @O(n + m)@
intersectionWith :: forall f. (forall x. Typeable x => f x -> f x -> f x) -> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
intersectionWith :: (forall (x :: k). Typeable x => f x -> f x -> f x)
-> TypeRepMap f -> TypeRepMap f -> TypeRepMap f
intersectionWith forall (x :: k). Typeable x => f x -> f x -> f x
f TypeRepMap f
ma TypeRepMap f
mb =
    [(Fingerprint, Any, Any)] -> TypeRepMap f
forall k (f :: k -> *). [(Fingerprint, Any, Any)] -> TypeRepMap f
fromSortedTriples ([(Fingerprint, Any, Any)] -> TypeRepMap f)
-> [(Fingerprint, Any, Any)] -> TypeRepMap f
forall a b. (a -> b) -> a -> b
$ [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps (TypeRepMap f -> [(Fingerprint, Any, Any)]
forall k (f :: k -> *). TypeRepMap f -> [(Fingerprint, Any, Any)]
toSortedTriples TypeRepMap f
ma) (TypeRepMap f -> [(Fingerprint, Any, Any)]
forall k (f :: k -> *). TypeRepMap f -> [(Fingerprint, Any, Any)]
toSortedTriples TypeRepMap f
mb)
  where
    f' :: forall x. TypeRep x -> f x -> f x -> f x
    f' :: TypeRep x -> f x -> f x -> f x
f' TypeRep x
tr = TypeRep x -> (Typeable x => f x -> f x -> f x) -> f x -> f x -> f x
forall k (a :: k) r. TypeRep a -> (Typeable a => r) -> r
withTypeable TypeRep x
tr Typeable x => f x -> f x -> f x
forall (x :: k). Typeable x => f x -> f x -> f x
f

    combine :: (Fingerprint, Any, Any) -> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
    combine :: (Fingerprint, Any, Any)
-> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
combine (Fingerprint
fp, Any
av, Any
ak) (Fingerprint
_, Any
bv, Any
_) = (Fingerprint
fp, f Any -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny (f Any -> Any) -> f Any -> Any
forall a b. (a -> b) -> a -> b
$ TypeRep Any -> f Any -> f Any -> f Any
forall (x :: k). TypeRep x -> f x -> f x -> f x
f' (Any -> TypeRep Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
ak) (Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
av) (Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
bv), Any
ak)

    -- Merges two typrepmaps into a sorted, dedup'd list of triples.
    mergeMaps :: [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
    -- If either list is empty, the intersection must be finished.
    mergeMaps :: [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps [(Fingerprint, Any, Any)]
_ [] = []
    mergeMaps [] [(Fingerprint, Any, Any)]
_ = []
    -- Merge the two maps considering one element at a time.
    mergeMaps (a :: (Fingerprint, Any, Any)
a@(Fingerprint
af, Any
_, Any
_) : [(Fingerprint, Any, Any)]
as) (b :: (Fingerprint, Any, Any)
b@(Fingerprint
bf, Any
_, Any
_) : [(Fingerprint, Any, Any)]
bs) =
        case Fingerprint -> Fingerprint -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Fingerprint
af Fingerprint
bf of
            -- Fingerprints are equal, union the elements using our function
            -- If the incoming maps were de-duped, there shouldn't be any other equivalent
            -- fingerprints
            Ordering
EQ -> (Fingerprint, Any, Any)
-> (Fingerprint, Any, Any) -> (Fingerprint, Any, Any)
combine (Fingerprint, Any, Any)
a (Fingerprint, Any, Any)
b (Fingerprint, Any, Any)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. a -> [a] -> [a]
: [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps [(Fingerprint, Any, Any)]
as [(Fingerprint, Any, Any)]
bs
            -- First fingerprint must not be in the second map or we would have seen it by now
            -- Skip it an move on
            Ordering
LT -> [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps [(Fingerprint, Any, Any)]
as ((Fingerprint, Any, Any)
b (Fingerprint, Any, Any)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. a -> [a] -> [a]
: [(Fingerprint, Any, Any)]
bs)
            -- Second fingerprint must not be in the first map or we would have seen it by now
            -- Skip it an move on
            Ordering
GT -> [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
mergeMaps ((Fingerprint, Any, Any)
a(Fingerprint, Any, Any)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. a -> [a] -> [a]
:[(Fingerprint, Any, Any)]
as) [(Fingerprint, Any, Any)]
bs
{-# INLINE intersectionWith #-}

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


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

>>> member @Char $ one (Identity 'a')
True
>>> member @Bool $ one (Identity 'a')
False
-}
member :: forall a (f :: KindOf a -> Type) . Typeable a => TypeRepMap f -> Bool
member :: TypeRepMap f -> Bool
member TypeRepMap f
tm = case TypeRepMap f -> Maybe (f a)
forall k (a :: k) (f :: k -> *).
Typeable a =>
TypeRepMap f -> Maybe (f a)
lookup @a TypeRepMap f
tm of
    Maybe (f a)
Nothing -> Bool
False
    Just f a
_  -> Bool
True
{-# INLINE member #-}

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

>>> x = lookup $ insert (Identity (11 :: Int)) empty
>>> x :: Maybe (Identity Int)
Just (Identity 11)
>>> x :: Maybe (Identity ())
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 -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys TypeRepMap f
tVect Array Any -> Int -> Any
forall a. Array a -> Int -> a
`indexArray`)
           (Int -> f a) -> Maybe Int -> Maybe (f a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Fingerprint -> PrimArray Word64 -> PrimArray Word64 -> Maybe Int
cachedBinarySearch (Typeable a => Fingerprint
forall k (a :: k). Typeable a => Fingerprint
typeFp @a)
                                  (TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs TypeRepMap f
tVect)
                                  (TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintBs TypeRepMap f
tVect)
{-# INLINE lookup #-}

-- | Get the amount of elements in a 'TypeRepMap'.
size :: TypeRepMap f -> Int
size :: TypeRepMap f -> Int
size = PrimArray Word64 -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray (PrimArray Word64 -> Int)
-> (TypeRepMap f -> PrimArray Word64) -> TypeRepMap f -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs
{-# INLINE size #-}

-- | Return the list of 'SomeTypeRep' from the keys.
keys :: TypeRepMap f -> [SomeTypeRep]
keys :: TypeRepMap f -> [SomeTypeRep]
keys = (forall (a :: k). TypeRep a -> SomeTypeRep)
-> TypeRepMap f -> [SomeTypeRep]
forall k (f :: k -> *) r.
(forall (a :: k). TypeRep a -> r) -> TypeRepMap f -> [r]
keysWith forall (a :: k). TypeRep a -> SomeTypeRep
forall k (a :: k). TypeRep a -> SomeTypeRep
SomeTypeRep
{-# INLINE keys #-}

-- | Return the list of keys by wrapping them with a user-provided function.
keysWith :: (forall (a :: ArgKindOf f). TypeRep a -> r) -> TypeRepMap f -> [r]
keysWith :: (forall (a :: ArgKindOf f). TypeRep a -> r) -> TypeRepMap f -> [r]
keysWith forall (a :: ArgKindOf f). TypeRep a -> r
f TypeRepMap{PrimArray Word64
Array Any
trKeys :: Array Any
trAnys :: Array Any
fingerprintBs :: PrimArray Word64
fingerprintAs :: PrimArray Word64
trKeys :: forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys :: forall k (f :: k -> *). TypeRepMap f -> Array Any
fingerprintBs :: forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs :: forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
..} = TypeRep Any -> r
forall (a :: ArgKindOf f). TypeRep a -> r
f (TypeRep Any -> r) -> (Any -> TypeRep Any) -> Any -> r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Any -> TypeRep Any
forall k (f :: k). Any -> TypeRep f
anyToTypeRep (Any -> r) -> [Any] -> [r]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Array Any -> [Item (Array Any)]
forall l. IsList l => l -> [Item l]
toList Array Any
trKeys
{-# INLINE keysWith #-}

-- | Return the list of key-value pairs by wrapping them with a user-provided function.
toListWith :: forall f r . (forall (a :: ArgKindOf f) . Typeable a => f a -> r) -> TypeRepMap f -> [r]
toListWith :: (forall (a :: ArgKindOf f). Typeable a => f a -> r)
-> TypeRepMap f -> [r]
toListWith forall (a :: ArgKindOf f). Typeable a => f a -> r
f = ((Fingerprint, Any, Any) -> r) -> [(Fingerprint, Any, Any)] -> [r]
forall a b. (a -> b) -> [a] -> [b]
map (Fingerprint, Any, Any) -> r
toF ([(Fingerprint, Any, Any)] -> [r])
-> (TypeRepMap f -> [(Fingerprint, Any, Any)])
-> TypeRepMap f
-> [r]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TypeRepMap f -> [(Fingerprint, Any, Any)]
forall k (f :: k -> *). TypeRepMap f -> [(Fingerprint, Any, Any)]
toTriples
  where
    withTypeRep :: TypeRep a -> f a -> r
    withTypeRep :: TypeRep a -> f a -> r
withTypeRep TypeRep a
tr f a
an = TypeRep a -> (Typeable a => r) -> r
forall k (a :: k) r. TypeRep a -> (Typeable a => r) -> r
withTypeable TypeRep a
tr ((Typeable a => r) -> r) -> (Typeable a => r) -> r
forall a b. (a -> b) -> a -> b
$ f a -> r
forall (a :: ArgKindOf f). Typeable a => f a -> r
f f a
an
    toF :: (Fingerprint, Any, Any) -> r
toF (Fingerprint
_, Any
an, Any
k) = TypeRep Any -> f Any -> r
forall (a :: ArgKindOf f). TypeRep a -> f a -> r
withTypeRep (Any -> TypeRep Any
forall a b. a -> b
unsafeCoerce Any
k) (Any -> f Any
forall k (f :: k -> *) (a :: k). Any -> f a
fromAny Any
an)

-- | Binary searched based on this article
-- http://bannalia.blogspot.com/2015/06/cache-friendly-binary-search.html
-- with modification for our two-vector search case.
cachedBinarySearch :: Fingerprint -> PrimArray Word64 -> PrimArray Word64 -> Maybe Int
cachedBinarySearch :: Fingerprint -> PrimArray Word64 -> PrimArray Word64 -> Maybe Int
cachedBinarySearch (Fingerprint (W64# Word#
a) (W64# Word#
b)) PrimArray Word64
fpAs PrimArray Word64
fpBs = Maybe Int -> Maybe Int
forall a. a -> a
inline (Int# -> Maybe Int
go Int#
0#)
  where
    go :: Int# -> Maybe Int
    go :: Int# -> Maybe Int
go Int#
i = case Int#
i Int# -> Int# -> Int#
<# Int#
len of
        Int#
0# -> Maybe Int
forall a. Maybe a
Nothing
        Int#
_  -> let !(W64# Word#
valA) = PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fpAs (Int# -> Int
I# Int#
i) in case Word#
a Word# -> Word# -> Int#
`ltWord#` Word#
valA of
            Int#
0#  -> case Word#
a Word# -> Word# -> Int#
`eqWord#` Word#
valA of
                Int#
0# -> Int# -> Maybe Int
go (Int#
2# Int# -> Int# -> Int#
*# Int#
i Int# -> Int# -> Int#
+# Int#
2#)
                Int#
_ -> let !(W64# Word#
valB) = PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fpBs (Int# -> Int
I# Int#
i) in case Word#
b Word# -> Word# -> Int#
`eqWord#` Word#
valB of
                    Int#
0# -> case Word#
b Word# -> Word# -> Int#
`ltWord#` Word#
valB of
                        Int#
0# -> Int# -> Maybe Int
go (Int#
2# Int# -> Int# -> Int#
*# Int#
i Int# -> Int# -> Int#
+# Int#
2#)
                        Int#
_  -> Int# -> Maybe Int
go (Int#
2# Int# -> Int# -> Int#
*# Int#
i Int# -> Int# -> Int#
+# Int#
1#)
                    Int#
_ -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int# -> Int
I# Int#
i)
            Int#
_ -> Int# -> Maybe Int
go (Int#
2# Int# -> Int# -> Int#
*# Int#
i Int# -> Int# -> Int#
+# Int#
1#)

    len :: Int#
    len :: Int#
len = let !(I# Int#
l) = PrimArray Word64 -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Word64
fpAs in Int#
l
{-# INLINE cachedBinarySearch #-}

----------------------------------------------------------------------------
-- Internal functions
----------------------------------------------------------------------------

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

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

anyToTypeRep :: Any -> TypeRep f
anyToTypeRep :: Any -> TypeRep f
anyToTypeRep = Any -> TypeRep f
forall a b. a -> b
unsafeCoerce

typeFp :: forall a . Typeable a => Fingerprint
typeFp :: Fingerprint
typeFp = TypeRep a -> Fingerprint
forall k (a :: k). TypeRep a -> Fingerprint
typeRepFingerprint (TypeRep a -> Fingerprint) -> TypeRep a -> Fingerprint
forall a b. (a -> b) -> a -> b
$ Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @a
{-# INLINE typeFp #-}

toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
toTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
toTriples TypeRepMap f
tm = [Fingerprint] -> [Any] -> [Any] -> [(Fingerprint, Any, Any)]
forall a b c. [a] -> [b] -> [c] -> [(a, b, c)]
zip3 (TypeRepMap f -> [Fingerprint]
forall k (f :: k -> *). TypeRepMap f -> [Fingerprint]
toFingerprints TypeRepMap f
tm) (Array Any -> [Item (Array Any)]
forall l. IsList l => l -> [Item l]
GHC.toList (Array Any -> [Item (Array Any)])
-> Array Any -> [Item (Array Any)]
forall a b. (a -> b) -> a -> b
$ TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys TypeRepMap f
tm) (Array Any -> [Item (Array Any)]
forall l. IsList l => l -> [Item l]
GHC.toList (Array Any -> [Item (Array Any)])
-> Array Any -> [Item (Array Any)]
forall a b. (a -> b) -> a -> b
$ TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trKeys TypeRepMap f
tm)

-- | Efficiently get sorted triples from a map in O(n) time
--
-- We assume the incoming TypeRepMap is already sorted into 'cachedBinarySearch' order using fromSortedList.
-- Then we can construct the index mapping from the "cached" ordering into monotonically
-- increasing order using 'generateOrderMapping' with the length of the TRM. This takes @O(n).
-- We then pull those indexes from the source TRM to get the sorted triples in a total of @O(n).
toSortedTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
toSortedTriples :: TypeRepMap f -> [(Fingerprint, Any, Any)]
toSortedTriples TypeRepMap f
tm = Int -> (Fingerprint, Any, Any)
trip (Int -> (Fingerprint, Any, Any))
-> [Int] -> [(Fingerprint, Any, Any)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
ordering
  where
    trip :: Int -> (Fingerprint, Any, Any)
trip Int
i = ( Word64 -> Word64 -> Fingerprint
Fingerprint (PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray (TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs TypeRepMap f
tm) Int
i) (PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray (TypeRepMap f -> PrimArray Word64
forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintBs TypeRepMap f
tm) Int
i)
             , Array Any -> Int -> Any
forall a. Array a -> Int -> a
indexArray (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys TypeRepMap f
tm) Int
i
             , Array Any -> Int -> Any
forall a. Array a -> Int -> a
indexArray (TypeRepMap f -> Array Any
forall k (f :: k -> *). TypeRepMap f -> Array Any
trKeys TypeRepMap f
tm) Int
i)
    ordering :: [ Int ]
    ordering :: [Int]
ordering = Int -> [Int]
generateOrderMapping (TypeRepMap f -> Int
forall k (f :: k -> *). TypeRepMap f -> Int
size TypeRepMap f
tm)

nubByFst :: (Eq a) => [(a, b, c)] -> [(a, b, c)]
nubByFst :: [(a, b, c)] -> [(a, b, c)]
nubByFst = ((a, b, c) -> (a, b, c) -> Bool) -> [(a, b, c)] -> [(a, b, c)]
forall a. (a -> a -> Bool) -> [a] -> [a]
nubBy (a -> a -> Bool
forall a. Eq a => a -> a -> Bool
(==) (a -> a -> Bool)
-> ((a, b, c) -> a) -> (a, b, c) -> (a, b, c) -> Bool
forall b c a. (b -> b -> c) -> (a -> b) -> a -> a -> c
`on` (a, b, c) -> a
forall a b c. (a, b, c) -> a
fst3)

fst3 :: (a, b, c) -> a
fst3 :: (a, b, c) -> a
fst3 (a
a, b
_, c
_) = a
a

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

-- | Existential wrapper around 'Typeable' indexed by @f@ type parameter.
-- Useful for 'TypeRepMap' structure creation form list of 'WrapTypeable's.
data WrapTypeable f where
    WrapTypeable :: Typeable a => f a -> WrapTypeable f

instance Show (WrapTypeable f) where
    show :: WrapTypeable f -> String
show (WrapTypeable (f a
_ :: f a)) = Fingerprint -> String
forall a. Show a => a -> String
show (Fingerprint -> String) -> Fingerprint -> String
forall a b. (a -> b) -> a -> b
$ Typeable a => Fingerprint
forall k (a :: k). Typeable a => Fingerprint
calcFp @a

{- |

prop> fromList . toList == 'id'

Creates 'TypeRepMap' from a list of 'WrapTypeable's.

>>> show $ fromList [WrapTypeable $ Identity True, WrapTypeable $ Identity 'a']
TypeRepMap [Bool, Char]


-}
instance IsList (TypeRepMap f) where
    type Item (TypeRepMap f) = WrapTypeable f

    fromList :: [WrapTypeable f] -> TypeRepMap f
    fromList :: [WrapTypeable f] -> TypeRepMap f
fromList = [(Fingerprint, Any, Any)] -> TypeRepMap f
forall k (f :: k -> *). [(Fingerprint, Any, Any)] -> TypeRepMap f
fromTriples ([(Fingerprint, Any, Any)] -> TypeRepMap f)
-> ([WrapTypeable f] -> [(Fingerprint, Any, Any)])
-> [WrapTypeable f]
-> TypeRepMap f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (WrapTypeable f -> (Fingerprint, Any, Any))
-> [WrapTypeable f] -> [(Fingerprint, Any, Any)]
forall a b. (a -> b) -> [a] -> [b]
map (\WrapTypeable f
x -> (WrapTypeable f -> Fingerprint
fp WrapTypeable f
x, WrapTypeable f -> Any
an WrapTypeable f
x, WrapTypeable f -> Any
k WrapTypeable f
x))
      where
        fp :: WrapTypeable f -> Fingerprint
        fp :: WrapTypeable f -> Fingerprint
fp (WrapTypeable (f a
_ :: f a)) = Typeable a => Fingerprint
forall k (a :: k). Typeable a => Fingerprint
calcFp @a

        an :: WrapTypeable f -> Any
        an :: WrapTypeable f -> Any
an (WrapTypeable f a
x) = f a -> Any
forall k (f :: k -> *) (a :: k). f a -> Any
toAny f a
x

        k :: WrapTypeable f -> Any
        k :: WrapTypeable f -> Any
k (WrapTypeable (f a
_ :: f a)) = TypeRep a -> Any
forall a b. a -> b
unsafeCoerce (TypeRep a -> Any) -> TypeRep a -> Any
forall a b. (a -> b) -> a -> b
$ Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @a

    toList :: TypeRepMap f -> [WrapTypeable f]
    toList :: TypeRepMap f -> [WrapTypeable f]
toList = (forall (a :: k). Typeable a => f a -> WrapTypeable f)
-> TypeRepMap f -> [WrapTypeable f]
forall k (f :: k -> *) r.
(forall (a :: k). Typeable a => f a -> r) -> TypeRepMap f -> [r]
toListWith forall (a :: k). Typeable a => f a -> WrapTypeable f
forall k (a :: k) (f :: k -> *).
Typeable a =>
f a -> WrapTypeable f
WrapTypeable

calcFp :: forall a . Typeable a => Fingerprint
calcFp :: Fingerprint
calcFp = TypeRep a -> Fingerprint
forall k (a :: k). TypeRep a -> Fingerprint
typeRepFingerprint (TypeRep a -> Fingerprint) -> TypeRep a -> Fingerprint
forall a b. (a -> b) -> a -> b
$ Typeable a => TypeRep a
forall k (a :: k). Typeable a => TypeRep a
typeRep @a

fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
fromTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
fromTriples = [(Fingerprint, Any, Any)] -> TypeRepMap f
forall k (f :: k -> *). [(Fingerprint, Any, Any)] -> TypeRepMap f
fromSortedTriples ([(Fingerprint, Any, Any)] -> TypeRepMap f)
-> ([(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)])
-> [(Fingerprint, Any, Any)]
-> TypeRepMap f
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Fingerprint, Any, Any) -> Fingerprint)
-> [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortWith (Fingerprint, Any, Any) -> Fingerprint
forall a b c. (a, b, c) -> a
fst3 ([(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)])
-> ([(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)])
-> [(Fingerprint, Any, Any)]
-> [(Fingerprint, Any, Any)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a b c. Eq a => [(a, b, c)] -> [(a, b, c)]
nubByFst

fromSortedTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
fromSortedTriples :: [(Fingerprint, Any, Any)] -> TypeRepMap f
fromSortedTriples [(Fingerprint, Any, Any)]
kvs = PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
forall k (f :: k -> *).
PrimArray Word64
-> PrimArray Word64 -> Array Any -> Array Any -> TypeRepMap f
TypeRepMap ([Item (PrimArray Word64)] -> PrimArray Word64
forall l. IsList l => [Item l] -> l
GHC.fromList [Word64]
[Item (PrimArray Word64)]
fpAs) ([Item (PrimArray Word64)] -> PrimArray Word64
forall l. IsList l => [Item l] -> l
GHC.fromList [Word64]
[Item (PrimArray Word64)]
fpBs) ([Item (Array Any)] -> Array Any
forall l. IsList l => [Item l] -> l
GHC.fromList [Any]
[Item (Array Any)]
ans) ([Item (Array Any)] -> Array Any
forall l. IsList l => [Item l] -> l
GHC.fromList [Any]
[Item (Array Any)]
ks)
  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 a b. (a -> b) -> [a] -> [b]
map (\(Fingerprint Word64
a Word64
b) -> (Word64
a, Word64
b)) [Fingerprint]
fps
    ([Fingerprint]
fps, [Any]
ans, [Any]
ks) = [(Fingerprint, Any, Any)] -> ([Fingerprint], [Any], [Any])
forall a b c. [(a, b, c)] -> ([a], [b], [c])
unzip3 ([(Fingerprint, Any, Any)] -> ([Fingerprint], [Any], [Any]))
-> [(Fingerprint, Any, Any)] -> ([Fingerprint], [Any], [Any])
forall a b. (a -> b) -> a -> b
$ [(Fingerprint, Any, Any)] -> [(Fingerprint, Any, Any)]
forall a. [a] -> [a]
fromSortedList [(Fingerprint, Any, Any)]
kvs

----------------------------------------------------------------------------
-- Tree-like conversion
----------------------------------------------------------------------------

fromSortedList :: forall a . [a] -> [a]
fromSortedList :: [a] -> [a]
fromSortedList [a]
l = (forall s. ST s [a]) -> [a]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [a]) -> [a]) -> (forall s. ST s [a]) -> [a]
forall a b. (a -> b) -> a -> b
$ do
    let n :: Int
n = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
l
    let arrOrigin :: Array a
arrOrigin = Int -> [Item (Array a)] -> Array a
forall l. IsList l => Int -> [Item l] -> l
fromListN Int
n [a]
[Item (Array a)]
l
    MutableArray s a
arrResult <- Array a -> Int -> Int -> ST s (MutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Array a -> Int -> Int -> m (MutableArray (PrimState m) a)
thawArray Array a
arrOrigin Int
0 Int
n
    Int -> MutableArray s a -> Array a -> ST s ()
forall s. Int -> MutableArray s a -> Array a -> ST s ()
go Int
n MutableArray s a
arrResult Array a
arrOrigin
    Array a -> [a]
forall l. IsList l => l -> [Item l]
toList (Array a -> [a]) -> ST s (Array a) -> ST s [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableArray (PrimState (ST s)) a -> ST s (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
unsafeFreezeArray MutableArray s a
MutableArray (PrimState (ST s)) a
arrResult
  where
    -- state monad could be used here, but it's another dependency
    go :: forall s . Int -> MutableArray s a -> Array a -> ST s ()
    go :: Int -> MutableArray s a -> Array a -> ST s ()
go Int
len MutableArray s a
result Array a
origin = () () -> ST s Int -> ST s ()
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ Int -> Int -> ST s Int
loop Int
0 Int
0
      where
        loop :: Int -> Int -> ST s Int
        loop :: Int -> Int -> ST s Int
loop Int
i Int
first =
            if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len
            then Int -> ST s Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
first
            else do
                Int
newFirst <- Int -> Int -> ST s Int
loop (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
first
                MutableArray (PrimState (ST s)) a -> Int -> a -> ST s ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
writeArray MutableArray s a
MutableArray (PrimState (ST s)) a
result Int
i (Array a -> Int -> a
forall a. Array a -> Int -> a
indexArray Array a
origin Int
newFirst)
                Int -> Int -> ST s Int
loop (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) (Int
newFirst Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

-- Returns a list of indexes which represents the "sorted" order of an array generated by
-- fromSortedList of the provided length.
-- I.e. fmap (fromSortedList [1, 2, 3, 4, 5, 6] !!) (generateOrderMapping 6) == [1, 2, 3, 4, 5, 6]
--
--     >>> generateOrderMapping 6
--     [3,1,4,0,5,2]
--
--     >>> generateOrderMapping 8
--     [7,3,1,4,0,5,2,6]
generateOrderMapping :: Int -> [Int]
generateOrderMapping :: Int -> [Int]
generateOrderMapping Int
len = (forall s. ST s [Int]) -> [Int]
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s [Int]) -> [Int])
-> (forall s. ST s [Int]) -> [Int]
forall a b. (a -> b) -> a -> b
$ do
    MutablePrimArray s Int
orderMappingArr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) Int)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
    Int
_ <- MutablePrimArray s Int -> Int -> Int -> ST s Int
forall s. MutablePrimArray s Int -> Int -> Int -> ST s Int
loop MutablePrimArray s Int
orderMappingArr Int
0 Int
0
    PrimArray Int -> [Int]
forall a. Prim a => PrimArray a -> [a]
primArrayToList (PrimArray Int -> [Int]) -> ST s (PrimArray Int) -> ST s [Int]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState (ST s)) Int -> ST s (PrimArray Int)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
orderMappingArr
  where
    loop :: MutablePrimArray s Int -> Int -> Int -> ST s Int
    loop :: MutablePrimArray s Int -> Int -> Int -> ST s Int
loop MutablePrimArray s Int
result Int
i Int
first =
        if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len
        then Int -> ST s Int
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
first
        else do
            Int
newFirst <- MutablePrimArray s Int -> Int -> Int -> ST s Int
forall s. MutablePrimArray s Int -> Int -> Int -> ST s Int
loop MutablePrimArray s Int
result (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
first
            MutablePrimArray (PrimState (ST s)) Int -> Int -> Int -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
writePrimArray MutablePrimArray s Int
MutablePrimArray (PrimState (ST s)) Int
result Int
newFirst Int
i
            MutablePrimArray s Int -> Int -> Int -> ST s Int
forall s. MutablePrimArray s Int -> Int -> Int -> ST s Int
loop MutablePrimArray s Int
result (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2) (Int
newFirst Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

----------------------------------------------------------------------------
--  Helper functions.
----------------------------------------------------------------------------

-- | Check that invariant of the structure holds.
-- The structure maintains the following invariant.
-- For each element @A@ at index @i@:
--
--   1. if there is an element @B@ at index @2*i+1@,
--      then @B < A@.
--
--   2. if there is an element @C@ at index @2*i+2@,
--      then @A < C@.
--
invariantCheck :: TypeRepMap f -> Bool
invariantCheck :: TypeRepMap f -> Bool
invariantCheck TypeRepMap{PrimArray Word64
Array Any
trKeys :: Array Any
trAnys :: Array Any
fingerprintBs :: PrimArray Word64
fingerprintAs :: PrimArray Word64
trKeys :: forall k (f :: k -> *). TypeRepMap f -> Array Any
trAnys :: forall k (f :: k -> *). TypeRepMap f -> Array Any
fingerprintBs :: forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
fingerprintAs :: forall k (f :: k -> *). TypeRepMap f -> PrimArray Word64
..} = All -> Bool
getAll (Int -> All
check Int
0)
  where
    lastMay :: [a] -> Maybe a
lastMay []     = Maybe a
forall a. Maybe a
Nothing
    lastMay [a
x]    = a -> Maybe a
forall a. a -> Maybe a
Just a
x
    lastMay (a
_:[a]
xs) = [a] -> Maybe a
lastMay [a]
xs
    sz :: Int
sz = PrimArray Word64 -> Int
forall a. Prim a => PrimArray a -> Int
sizeofPrimArray PrimArray Word64
fingerprintAs
    check :: Int -> All
check Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
sz = Bool -> All
All Bool
True
            | Bool
otherwise =
      let left :: Int
left = Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1
          right :: Int
right = Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2
          -- maximum value in the left branch
          leftMax :: Maybe (Word64, Word64)
leftMax =
               (Int -> (Word64, Word64)) -> Maybe Int -> Maybe (Word64, Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
j -> (PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintAs Int
j, PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintBs Int
j))
             (Maybe Int -> Maybe (Word64, Word64))
-> Maybe Int -> Maybe (Word64, Word64)
forall a b. (a -> b) -> a -> b
$ [Int] -> Maybe Int
forall a. [a] -> Maybe a
lastMay
             ([Int] -> Maybe Int) -> [Int] -> Maybe Int
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
sz)
             ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (\Int
j -> Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
2) Int
left
          -- minimum value in the right branch
          rightMin :: Maybe (Word64, Word64)
rightMin =
               (Int -> (Word64, Word64)) -> Maybe Int -> Maybe (Word64, Word64)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
j -> (PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintAs Int
j, PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintBs Int
j))
             (Maybe Int -> Maybe (Word64, Word64))
-> Maybe Int -> Maybe (Word64, Word64)
forall a b. (a -> b) -> a -> b
$ [Int] -> Maybe Int
forall a. [a] -> Maybe a
lastMay
             ([Int] -> Maybe Int) -> [Int] -> Maybe Int
forall a b. (a -> b) -> a -> b
$ (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<Int
sz)
             ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Int -> Int) -> Int -> [Int]
forall a. (a -> a) -> a -> [a]
iterate (\Int
j -> Int
jInt -> Int -> Int
forall a. Num a => a -> a -> a
*Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) Int
right
      in [All] -> All
forall a. Monoid a => [a] -> a
mconcat
          [ Bool -> All
All (Bool -> All) -> Bool -> All
forall a b. (a -> b) -> a -> b
$
            if Int
left Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
            then
              case PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintAs Int
i Word64 -> Word64 -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintAs Int
left of
                Ordering
LT -> Bool
False
                Ordering
EQ -> PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintBs Int
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintBs Int
left
                Ordering
GT -> Bool
True
            else Bool
True
         , Bool -> All
All (Bool -> All) -> Bool -> All
forall a b. (a -> b) -> a -> b
$
           if Int
right Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
sz
           then
              case PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintAs Int
i Word64 -> Word64 -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintAs Int
right of
                Ordering
LT -> Bool
True
                Ordering
EQ -> PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintBs Int
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= PrimArray Word64 -> Int -> Word64
forall a. Prim a => PrimArray a -> Int -> a
indexPrimArray PrimArray Word64
fingerprintBs Int
right
                Ordering
GT -> Bool
False
            else Bool
True
         , Bool -> All
All (Bool -> All) -> Bool -> All
forall a b. (a -> b) -> a -> b
$ Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
True (Maybe Bool -> Bool) -> Maybe Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (Word64, Word64) -> (Word64, Word64) -> Bool
forall a. Ord a => a -> a -> Bool
(<=) ((Word64, Word64) -> (Word64, Word64) -> Bool)
-> Maybe (Word64, Word64) -> Maybe ((Word64, Word64) -> Bool)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe (Word64, Word64)
leftMax Maybe ((Word64, Word64) -> Bool)
-> Maybe (Word64, Word64) -> Maybe Bool
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Maybe (Word64, Word64)
rightMin
         , Int -> All
check (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
         ]