{-# LANGUAGE CPP #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE Trustworthy #-}
{-# OPTIONS_GHC -Wall -funbox-strict-fields -fno-warn-orphans -fno-warn-type-defaults -O2 #-}
#ifdef ST_HACK
{-# OPTIONS_GHC -fno-full-laziness #-}
#endif
--------------------------------------------------------------------------------
-- |
-- Copyright   : (c) Edward Kmett 2015
-- License     : BSD-style
-- Maintainer  : Edward Kmett <ekmett@gmail.com>
-- Portability : non-portable
--
-- This module suppose a Word64-based array-mapped PATRICIA Trie.
--
-- The most significant nybble is isolated by using techniques based on
-- <https://www.fpcomplete.com/user/edwardk/revisiting-matrix-multiplication/part-4>
-- but modified to work nybble-by-nybble rather than bit-by-bit.
--
--------------------------------------------------------------------------------
module Data.Discrimination.Internal.WordMap
  ( WordMap
  , singleton
  , empty
  , insert
  , lookup
  , member
  , fromList
  ) where

import Control.Applicative hiding (empty)
import Control.DeepSeq
import Control.Monad.ST hiding (runST)
import Data.Bits
import Data.Discrimination.Internal.SmallArray
import Data.Foldable
import Data.Functor
import Data.Monoid
import Data.Traversable
import Data.Word
import qualified GHC.Exts as Exts
import Prelude hiding (lookup, length, foldr)
import GHC.Types
import GHC.ST

type Key = Word64
type Mask = Word16
type Offset = Int

ptrEq :: a -> a -> Bool
ptrEq :: forall a. a -> a -> Bool
ptrEq a
x a
y = Int# -> Bool
isTrue# (a -> a -> Int#
forall a b. a -> b -> Int#
Exts.reallyUnsafePtrEquality# a
x a
y Int# -> Int# -> Int#
Exts.==# Int#
1#)
{-# INLINEABLE ptrEq #-}

ptrNeq :: a -> a -> Bool
ptrNeq :: forall a. a -> a -> Bool
ptrNeq a
x a
y = Int# -> Bool
isTrue# (a -> a -> Int#
forall a b. a -> b -> Int#
Exts.reallyUnsafePtrEquality# a
x a
y Int# -> Int# -> Int#
Exts./=# Int#
1#)
{-# INLINEABLE ptrNeq #-}

data WordMap v
  = Full !Key !Offset !(SmallArray (WordMap v))
  | Node !Key !Offset !Mask !(SmallArray (WordMap v))
  | Tip  !Key v
  | Nil
  deriving Int -> WordMap v -> ShowS
[WordMap v] -> ShowS
WordMap v -> String
(Int -> WordMap v -> ShowS)
-> (WordMap v -> String)
-> ([WordMap v] -> ShowS)
-> Show (WordMap v)
forall v. Show v => Int -> WordMap v -> ShowS
forall v. Show v => [WordMap v] -> ShowS
forall v. Show v => WordMap v -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall v. Show v => Int -> WordMap v -> ShowS
showsPrec :: Int -> WordMap v -> ShowS
$cshow :: forall v. Show v => WordMap v -> String
show :: WordMap v -> String
$cshowList :: forall v. Show v => [WordMap v] -> ShowS
showList :: [WordMap v] -> ShowS
Show

node :: Key -> Offset -> Mask -> SmallArray (WordMap v) -> WordMap v
node :: forall v.
Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
node Key
k Int
o Word16
0xffff SmallArray (WordMap v)
a = Key -> Int -> SmallArray (WordMap v) -> WordMap v
forall v. Key -> Int -> SmallArray (WordMap v) -> WordMap v
Full Key
k Int
o SmallArray (WordMap v)
a
node Key
k Int
o Word16
m SmallArray (WordMap v)
a      = Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
forall v.
Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
Node Key
k Int
o Word16
m SmallArray (WordMap v)
a
{-# INLINE node #-}

instance NFData v => NFData (WordMap v) where
  rnf :: WordMap v -> ()
rnf (Full Key
_ Int
_ SmallArray (WordMap v)
a)   = SmallArray (WordMap v) -> ()
forall a. NFData a => a -> ()
rnf SmallArray (WordMap v)
a
  rnf (Node Key
_ Int
_ Word16
_ SmallArray (WordMap v)
a) = SmallArray (WordMap v) -> ()
forall a. NFData a => a -> ()
rnf SmallArray (WordMap v)
a
  rnf (Tip Key
_ v
v) = v -> ()
forall a. NFData a => a -> ()
rnf v
v
  rnf WordMap v
Nil = ()

instance Functor WordMap where
  fmap :: forall a b. (a -> b) -> WordMap a -> WordMap b
fmap a -> b
f = WordMap a -> WordMap b
go where
    go :: WordMap a -> WordMap b
go (Full Key
k Int
o SmallArray (WordMap a)
a) = Key -> Int -> SmallArray (WordMap b) -> WordMap b
forall v. Key -> Int -> SmallArray (WordMap v) -> WordMap v
Full Key
k Int
o ((WordMap a -> WordMap b)
-> SmallArray (WordMap a) -> SmallArray (WordMap b)
forall a b. (a -> b) -> SmallArray a -> SmallArray b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap WordMap a -> WordMap b
go SmallArray (WordMap a)
a)
    go (Node Key
k Int
o Word16
m SmallArray (WordMap a)
a) = Key -> Int -> Word16 -> SmallArray (WordMap b) -> WordMap b
forall v.
Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
Node Key
k Int
o Word16
m ((WordMap a -> WordMap b)
-> SmallArray (WordMap a) -> SmallArray (WordMap b)
forall a b. (a -> b) -> SmallArray a -> SmallArray b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap WordMap a -> WordMap b
go SmallArray (WordMap a)
a)
    go (Tip Key
k a
v) = Key -> b -> WordMap b
forall v. Key -> v -> WordMap v
Tip Key
k (a -> b
f a
v)
    go WordMap a
Nil = WordMap b
forall v. WordMap v
Nil
  {-# INLINEABLE fmap #-}

instance Foldable WordMap where
  foldMap :: forall m a. Monoid m => (a -> m) -> WordMap a -> m
foldMap a -> m
f = WordMap a -> m
go where
    go :: WordMap a -> m
go (Full Key
_ Int
_ SmallArray (WordMap a)
a) = (WordMap a -> m) -> SmallArray (WordMap a) -> m
forall m a. Monoid m => (a -> m) -> SmallArray a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap WordMap a -> m
go SmallArray (WordMap a)
a
    go (Node Key
_ Int
_ Word16
_ SmallArray (WordMap a)
a) = (WordMap a -> m) -> SmallArray (WordMap a) -> m
forall m a. Monoid m => (a -> m) -> SmallArray a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap WordMap a -> m
go SmallArray (WordMap a)
a
    go (Tip Key
_ a
v) = a -> m
f a
v
    go WordMap a
Nil = m
forall a. Monoid a => a
mempty
  {-# INLINEABLE foldMap #-}

instance Traversable WordMap where
  traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> WordMap a -> f (WordMap b)
traverse a -> f b
f = WordMap a -> f (WordMap b)
go where
    go :: WordMap a -> f (WordMap b)
go (Full Key
k Int
o SmallArray (WordMap a)
a) = Key -> Int -> SmallArray (WordMap b) -> WordMap b
forall v. Key -> Int -> SmallArray (WordMap v) -> WordMap v
Full Key
k Int
o (SmallArray (WordMap b) -> WordMap b)
-> f (SmallArray (WordMap b)) -> f (WordMap b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (WordMap a -> f (WordMap b))
-> SmallArray (WordMap a) -> f (SmallArray (WordMap b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> SmallArray a -> f (SmallArray b)
traverse WordMap a -> f (WordMap b)
go SmallArray (WordMap a)
a
    go (Node Key
k Int
o Word16
m SmallArray (WordMap a)
a) = Key -> Int -> Word16 -> SmallArray (WordMap b) -> WordMap b
forall v.
Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
Node Key
k Int
o Word16
m (SmallArray (WordMap b) -> WordMap b)
-> f (SmallArray (WordMap b)) -> f (WordMap b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (WordMap a -> f (WordMap b))
-> SmallArray (WordMap a) -> f (SmallArray (WordMap b))
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> SmallArray a -> f (SmallArray b)
traverse WordMap a -> f (WordMap b)
go SmallArray (WordMap a)
a
    go (Tip Key
k a
v) = Key -> b -> WordMap b
forall v. Key -> v -> WordMap v
Tip Key
k (b -> WordMap b) -> f b -> f (WordMap b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
v
    go WordMap a
Nil = WordMap b -> f (WordMap b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure WordMap b
forall v. WordMap v
Nil
  {-# INLINEABLE traverse #-}

-- Note: 'level 0' will return a negative shift, don't use it
level :: Key -> Int
level :: Key -> Int
level Key
w = Int
60 Int -> Int -> Int
forall a. Num a => a -> a -> a
- (Key -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros Key
w Int -> Int -> Int
forall a. Bits a => a -> a -> a
.&. Int
0x7c)
{-# INLINE level #-}

maskBit :: Key -> Offset -> Int
maskBit :: Key -> Int -> Int
maskBit Key
k Int
o = Key -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Key -> Int -> Key
forall a. Bits a => a -> Int -> a
unsafeShiftR Key
k Int
o Key -> Key -> Key
forall a. Bits a => a -> a -> a
.&. Key
0xf)
{-# INLINE maskBit #-}

mask :: Key -> Offset -> Word16
mask :: Key -> Int -> Word16
mask Key
k Int
o = Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
unsafeShiftL Word16
1 (Key -> Int -> Int
maskBit Key
k Int
o)
{-# INLINE mask #-}

-- offset :: Int -> Word16 -> Int
-- offset k w = popCount $ w .&. (unsafeShiftL 1 k - 1)
-- {-# INLINE offset #-}

fork :: Int -> Key -> WordMap v -> Key -> WordMap v -> WordMap v
fork :: forall v. Int -> Key -> WordMap v -> Key -> WordMap v -> WordMap v
fork Int
o Key
k WordMap v
n Key
ok WordMap v
on = Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
forall v.
Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
Node (Key
k Key -> Key -> Key
forall a. Bits a => a -> a -> a
.&. Key -> Int -> Key
forall a. Bits a => a -> Int -> a
unsafeShiftL Key
0xfffffffffffffff0 Int
o) Int
o (Key -> Int -> Word16
mask Key
k Int
o Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Key -> Int -> Word16
mask Key
ok Int
o) (SmallArray (WordMap v) -> WordMap v)
-> SmallArray (WordMap v) -> WordMap v
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (SmallArray (WordMap v))) -> SmallArray (WordMap v)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray (WordMap v)))
 -> SmallArray (WordMap v))
-> (forall s. ST s (SmallArray (WordMap v)))
-> SmallArray (WordMap v)
forall a b. (a -> b) -> a -> b
$ do
  arr <- Int
-> WordMap v
-> ST s (SmallMutableArray (PrimState (ST s)) (WordMap v))
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
2 WordMap v
n
  writeSmallArray arr (fromEnum (k < ok)) on
  unsafeFreezeSmallArray arr

insert :: Key -> v -> WordMap v -> WordMap v
insert :: forall v. Key -> v -> WordMap v -> WordMap v
insert !Key
k v
v WordMap v
xs0 = WordMap v -> WordMap v
go WordMap v
xs0 where
  go :: WordMap v -> WordMap v
go on :: WordMap v
on@(Full Key
ok Int
n SmallArray (WordMap v)
as)
    | Key
wd Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
0xf = Int -> Key -> WordMap v -> Key -> WordMap v -> WordMap v
forall v. Int -> Key -> WordMap v -> Key -> WordMap v -> WordMap v
fork (Key -> Int
level Key
okk) Key
k (Key -> v -> WordMap v
forall v. Key -> v -> WordMap v
Tip Key
k v
v) Key
ok WordMap v
on
    | !WordMap v
oz <- SmallArray (WordMap v) -> Int -> WordMap v
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray (WordMap v)
as Int
d
    , !WordMap v
z <- WordMap v -> WordMap v
go WordMap v
oz
    , WordMap v -> WordMap v -> Bool
forall a. a -> a -> Bool
ptrNeq WordMap v
z WordMap v
oz = Key -> Int -> SmallArray (WordMap v) -> WordMap v
forall v. Key -> Int -> SmallArray (WordMap v) -> WordMap v
Full Key
ok Int
n (Int
-> WordMap v -> SmallArray (WordMap v) -> SmallArray (WordMap v)
forall a. Int -> a -> SmallArray a -> SmallArray a
update16 Int
d WordMap v
z SmallArray (WordMap v)
as)
    | Bool
otherwise = WordMap v
on
    where
      okk :: Key
okk = Key -> Key -> Key
forall a. Bits a => a -> a -> a
xor Key
ok Key
k
      wd :: Key
wd  = Key -> Int -> Key
forall a. Bits a => a -> Int -> a
unsafeShiftR Key
okk Int
n
      d :: Int
d   = Key -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
wd
  go on :: WordMap v
on@(Node Key
ok Int
n Word16
m SmallArray (WordMap v)
as)
    | Key
wd Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
> Key
0xf = Int -> Key -> WordMap v -> Key -> WordMap v -> WordMap v
forall v. Int -> Key -> WordMap v -> Key -> WordMap v -> WordMap v
fork (Key -> Int
level Key
okk) Key
k (Key -> v -> WordMap v
forall v. Key -> v -> WordMap v
Tip Key
k v
v) Key
ok WordMap v
on
    | Word16
m Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. Word16
b Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
0 = Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
forall v.
Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
node Key
ok Int
n (Word16
m Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.|. Word16
b) (Int
-> WordMap v -> SmallArray (WordMap v) -> SmallArray (WordMap v)
forall a. Int -> a -> SmallArray a -> SmallArray a
insertSmallArray Int
odm (Key -> v -> WordMap v
forall v. Key -> v -> WordMap v
Tip Key
k v
v) SmallArray (WordMap v)
as)
    | !WordMap v
oz <- SmallArray (WordMap v) -> Int -> WordMap v
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray (WordMap v)
as Int
odm
    , !WordMap v
z <- WordMap v -> WordMap v
go WordMap v
oz
    , WordMap v -> WordMap v -> Bool
forall a. a -> a -> Bool
ptrNeq WordMap v
z WordMap v
oz = Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
forall v.
Key -> Int -> Word16 -> SmallArray (WordMap v) -> WordMap v
Node Key
ok Int
n Word16
m (Int
-> WordMap v -> SmallArray (WordMap v) -> SmallArray (WordMap v)
forall a. Int -> a -> SmallArray a -> SmallArray a
updateSmallArray Int
odm WordMap v
z SmallArray (WordMap v)
as)
    | Bool
otherwise = WordMap v
on
    where
      okk :: Key
okk = Key -> Key -> Key
forall a. Bits a => a -> a -> a
xor Key
ok Key
k
      wd :: Key
wd  = Key -> Int -> Key
forall a. Bits a => a -> Int -> a
unsafeShiftR Key
okk Int
n
      d :: Int
d   = Key -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
wd
      b :: Word16
b   = Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
unsafeShiftL Word16
1 Int
d
      odm :: Int
odm = Word16 -> Int
forall a. Bits a => a -> Int
popCount (Word16 -> Int) -> Word16 -> Int
forall a b. (a -> b) -> a -> b
$ Word16
m Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. (Word16
b Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
- Word16
1)
  go on :: WordMap v
on@(Tip Key
ok v
ov)
    | Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
/= Key
ok    = Int -> Key -> WordMap v -> Key -> WordMap v -> WordMap v
forall v. Int -> Key -> WordMap v -> Key -> WordMap v -> WordMap v
fork (Key -> Int
level (Key -> Key -> Key
forall a. Bits a => a -> a -> a
xor Key
ok Key
k)) Key
k (Key -> v -> WordMap v
forall v. Key -> v -> WordMap v
Tip Key
k v
v) Key
ok WordMap v
on
    | v -> v -> Bool
forall a. a -> a -> Bool
ptrEq v
v v
ov = WordMap v
on
    | Bool
otherwise  = Key -> v -> WordMap v
forall v. Key -> v -> WordMap v
Tip Key
k v
v
  go WordMap v
Nil = Key -> v -> WordMap v
forall v. Key -> v -> WordMap v
Tip Key
k v
v
{-# INLINEABLE insert #-}


lookup :: Key -> WordMap v -> Maybe v
lookup :: forall v. Key -> WordMap v -> Maybe v
lookup !Key
k (Full Key
ok Int
o SmallArray (WordMap v)
a)
  | Key
z <- Key -> Int -> Key
forall a. Bits a => a -> Int -> a
unsafeShiftR (Key -> Key -> Key
forall a. Bits a => a -> a -> a
xor Key
k Key
ok) Int
o, Key
z Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
<= Key
0xf = Key -> WordMap v -> Maybe v
forall v. Key -> WordMap v -> Maybe v
lookup Key
k (WordMap v -> Maybe v) -> WordMap v -> Maybe v
forall a b. (a -> b) -> a -> b
$ SmallArray (WordMap v) -> Int -> WordMap v
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray (WordMap v)
a (Key -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
z)
  | Bool
otherwise = Maybe v
forall a. Maybe a
Nothing
lookup Key
k (Node Key
ok Int
o Word16
m SmallArray (WordMap v)
a)
  | Key
z Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
<= Key
0xf Bool -> Bool -> Bool
&& Word16
m Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. Word16
b Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word16
0 = Key -> WordMap v -> Maybe v
forall v. Key -> WordMap v -> Maybe v
lookup Key
k (SmallArray (WordMap v) -> Int -> WordMap v
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray (WordMap v)
a (Word16 -> Int
forall a. Bits a => a -> Int
popCount (Word16
m Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. (Word16
b Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
- Word16
1))))
  | Bool
otherwise = Maybe v
forall a. Maybe a
Nothing
  where
    z :: Key
z = Key -> Int -> Key
forall a. Bits a => a -> Int -> a
unsafeShiftR (Key -> Key -> Key
forall a. Bits a => a -> a -> a
xor Key
k Key
ok) Int
o
    b :: Word16
b = Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
unsafeShiftL Word16
1 (Key -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
z)
lookup Key
k (Tip Key
ok v
ov)
  | Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
ok   = v -> Maybe v
forall a. a -> Maybe a
Just v
ov
  | Bool
otherwise = Maybe v
forall a. Maybe a
Nothing
lookup Key
_ WordMap v
Nil = Maybe v
forall a. Maybe a
Nothing
{-# INLINEABLE lookup #-}

member :: Key -> WordMap v -> Bool
member :: forall v. Key -> WordMap v -> Bool
member !Key
k (Full Key
ok Int
o SmallArray (WordMap v)
a)
  | Key
z <- Key -> Int -> Key
forall a. Bits a => a -> Int -> a
unsafeShiftR (Key -> Key -> Key
forall a. Bits a => a -> a -> a
xor Key
k Key
ok) Int
o = Key
z Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
<= Key
0xf Bool -> Bool -> Bool
&& Key -> WordMap v -> Bool
forall v. Key -> WordMap v -> Bool
member Key
k (SmallArray (WordMap v) -> Int -> WordMap v
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray (WordMap v)
a (Key -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
z))
member Key
k (Node Key
ok Int
o Word16
m SmallArray (WordMap v)
a)
  | Key
z <- Key -> Int -> Key
forall a. Bits a => a -> Int -> a
unsafeShiftR (Key -> Key -> Key
forall a. Bits a => a -> a -> a
xor Key
k Key
ok) Int
o
  = Key
z Key -> Key -> Bool
forall a. Ord a => a -> a -> Bool
<= Key
0xf Bool -> Bool -> Bool
&& let b :: Word16
b = Word16 -> Int -> Word16
forall a. Bits a => a -> Int -> a
unsafeShiftL Word16
1 (Key -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Key
z) in
    Word16
m Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. Word16
b Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word16
0 Bool -> Bool -> Bool
&& Key -> WordMap v -> Bool
forall v. Key -> WordMap v -> Bool
member Key
k (SmallArray (WordMap v) -> Int -> WordMap v
forall a. SmallArray a -> Int -> a
indexSmallArray SmallArray (WordMap v)
a (Word16 -> Int
forall a. Bits a => a -> Int
popCount (Word16
m Word16 -> Word16 -> Word16
forall a. Bits a => a -> a -> a
.&. (Word16
b Word16 -> Word16 -> Word16
forall a. Num a => a -> a -> a
- Word16
1))))
member Key
k (Tip Key
ok v
_) = Key
k Key -> Key -> Bool
forall a. Eq a => a -> a -> Bool
== Key
ok
member Key
_ WordMap v
Nil = Bool
False
{-# INLINEABLE member #-}

updateSmallArray :: Int -> a -> SmallArray a -> SmallArray a
updateSmallArray :: forall a. Int -> a -> SmallArray a -> SmallArray a
updateSmallArray !Int
k a
a SmallArray a
i = (forall s. ST s (SmallArray a)) -> SmallArray a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray a)) -> SmallArray a)
-> (forall s. ST s (SmallArray a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n = SmallArray a -> Int
forall a. SmallArray a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SmallArray a
i
  o <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
n a
forall a. HasCallStack => a
undefined
  copySmallArray o 0 i 0 n
  writeSmallArray o k a
  unsafeFreezeSmallArray o
{-# INLINEABLE updateSmallArray #-}

update16 :: Int -> a -> SmallArray a -> SmallArray a
update16 :: forall a. Int -> a -> SmallArray a -> SmallArray a
update16 !Int
k a
a SmallArray a
i = (forall s. ST s (SmallArray a)) -> SmallArray a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray a)) -> SmallArray a)
-> (forall s. ST s (SmallArray a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
  o <- SmallArray a -> ST s (SmallMutableArray s a)
forall a s. SmallArray a -> ST s (SmallMutableArray s a)
clone16 SmallArray a
i
  writeSmallArray o k a
  unsafeFreezeSmallArray o
{-# INLINEABLE update16 #-}

insertSmallArray :: Int -> a -> SmallArray a -> SmallArray a
insertSmallArray :: forall a. Int -> a -> SmallArray a -> SmallArray a
insertSmallArray !Int
k a
a SmallArray a
i = (forall s. ST s (SmallArray a)) -> SmallArray a
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (SmallArray a)) -> SmallArray a)
-> (forall s. ST s (SmallArray a)) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ do
  let n :: Int
n = SmallArray a -> Int
forall a. SmallArray a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length SmallArray a
i
  o <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) a
a
  copySmallArray  o 0 i 0 k
  copySmallArray  o (k+1) i k (n-k)
  unsafeFreezeSmallArray o
{-# INLINEABLE insertSmallArray #-}

clone16 :: SmallArray a -> ST s (SmallMutableArray s a)
clone16 :: forall a s. SmallArray a -> ST s (SmallMutableArray s a)
clone16 SmallArray a
i = do
  o <- Int -> a -> ST s (SmallMutableArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
16 a
forall a. HasCallStack => a
undefined
  indexSmallArrayM i 0 >>= writeSmallArray o 0
  indexSmallArrayM i 1 >>= writeSmallArray o 1
  indexSmallArrayM i 2 >>= writeSmallArray o 2
  indexSmallArrayM i 3 >>= writeSmallArray o 3
  indexSmallArrayM i 4 >>= writeSmallArray o 4
  indexSmallArrayM i 5 >>= writeSmallArray o 5
  indexSmallArrayM i 6 >>= writeSmallArray o 6
  indexSmallArrayM i 7 >>= writeSmallArray o 7
  indexSmallArrayM i 8 >>= writeSmallArray o 8
  indexSmallArrayM i 9 >>= writeSmallArray o 9
  indexSmallArrayM i 10 >>= writeSmallArray o 10
  indexSmallArrayM i 11 >>= writeSmallArray o 11
  indexSmallArrayM i 12 >>= writeSmallArray o 12
  indexSmallArrayM i 13 >>= writeSmallArray o 13
  indexSmallArrayM i 14 >>= writeSmallArray o 14
  indexSmallArrayM i 15 >>= writeSmallArray o 15
  return o
{-# INLINE clone16 #-}

-- | Build a singleton WordMap
singleton :: Key -> v -> WordMap v
singleton :: forall v. Key -> v -> WordMap v
singleton !Key
k v
v = Key -> v -> WordMap v
forall v. Key -> v -> WordMap v
Tip Key
k v
v
{-# INLINE singleton #-}

fromList :: [(Word64,v)] -> WordMap v
fromList :: forall v. [(Key, v)] -> WordMap v
fromList [(Key, v)]
xs = (WordMap v -> (Key, v) -> WordMap v)
-> WordMap v -> [(Key, v)] -> WordMap v
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\WordMap v
r (Key
k,v
v) -> Key -> v -> WordMap v -> WordMap v
forall v. Key -> v -> WordMap v -> WordMap v
insert Key
k v
v WordMap v
r) WordMap v
forall v. WordMap v
Nil [(Key, v)]
xs
{-# INLINE fromList #-}

empty :: WordMap a
empty :: forall v. WordMap v
empty = WordMap a
forall v. WordMap v
Nil
{-# INLINE empty #-}