{-# LANGUAGE Strict #-}
-- | A minimal, strict map keyed by 'Word64' values (adaptation of
-- 'Data.IntMap.Strict').
--
-- This module is intended for internal use only, and may change without warning
-- in subsequent releases.
module Effectful.Internal.Utils.Word64Map
  ( Word64Map
  , empty
  , lookup
  , insert
  , delete
  , updateLookupWithKey
  ) where

import Data.Bits
import Data.Word
import Prelude hiding (lookup)

-- | A map of 'Word64' keys to values of type @a@.
data Word64Map a
  = Bin Prefix (Word64Map a) (Word64Map a)
  | Tip Word64 a
  | Nil

-- | A @Prefix@ represents some prefix of high-order bits of a @Word64@.
newtype Prefix = Prefix Word64

unPrefix :: Prefix -> Word64
unPrefix :: Prefix -> Word64
unPrefix (Prefix Word64
p) = Word64
p

----------------------------------------

-- | The empty map.
empty :: Word64Map a
empty :: forall a. Word64Map a
empty = Word64Map a
forall a. Word64Map a
Nil

-- | Look up the value at a key in the map.
lookup :: Word64 -> Word64Map a -> Maybe a
lookup :: forall a. Word64 -> Word64Map a -> Maybe a
lookup Word64
k = Word64Map a -> Maybe a
go
  where
    go :: Word64Map a -> Maybe a
go (Bin Prefix
p Word64Map a
l Word64Map a
r) | Word64 -> Prefix -> Bool
left Word64
k Prefix
p  = Word64Map a -> Maybe a
go Word64Map a
l
                   | Bool
otherwise = Word64Map a -> Maybe a
go Word64Map a
r
    go (Tip Word64
kx a
x) | Word64
k Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
kx   = a -> Maybe a
forall a. a -> Maybe a
Just a
x
                  | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing
    go Word64Map a
Nil = Maybe a
forall a. Maybe a
Nothing

-- | Insert a new key/value pair in the map. If the key is already present, the
-- associated value is replaced with the supplied one.
--
-- The value is evaluated to WHNF when it is inserted into the map.
insert :: Word64 -> a -> Word64Map a -> Word64Map a
insert :: forall a. Word64 -> a -> Word64Map a -> Word64Map a
insert Word64
k a
x = Word64Map a -> Word64Map a
go
  where
    go :: Word64Map a -> Word64Map a
go t :: Word64Map a
t@(Bin Prefix
p Word64Map a
l Word64Map a
r)
      | Word64 -> Prefix -> Bool
nomatch Word64
k Prefix
p = Word64 -> Word64Map a -> Prefix -> Word64Map a -> Word64Map a
forall a.
Word64 -> Word64Map a -> Prefix -> Word64Map a -> Word64Map a
linkKey Word64
k (Word64 -> a -> Word64Map a
forall a. Word64 -> a -> Word64Map a
Tip Word64
k a
x) Prefix
p Word64Map a
t
      | Word64 -> Prefix -> Bool
left Word64
k Prefix
p    = Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
Bin Prefix
p (Word64Map a -> Word64Map a
go Word64Map a
l) Word64Map a
r
      | Bool
otherwise   = Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
Bin Prefix
p Word64Map a
l (Word64Map a -> Word64Map a
go Word64Map a
r)
    go t :: Word64Map a
t@(Tip Word64
ky a
_)
      | Word64
k Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
ky     = Word64 -> a -> Word64Map a
forall a. Word64 -> a -> Word64Map a
Tip Word64
k a
x
      | Bool
otherwise   = Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
forall a.
Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
link Word64
k (Word64 -> a -> Word64Map a
forall a. Word64 -> a -> Word64Map a
Tip Word64
k a
x) Word64
ky Word64Map a
t
    go Word64Map a
Nil = Word64 -> a -> Word64Map a
forall a. Word64 -> a -> Word64Map a
Tip Word64
k a
x

-- | Delete a key and its value from the map. When the key is not a member of
-- the map, the original map is returned.
delete :: Word64 -> Word64Map a -> Word64Map a
delete :: forall a. Word64 -> Word64Map a -> Word64Map a
delete Word64
k = Word64Map a -> Word64Map a
go
  where
    go :: Word64Map a -> Word64Map a
go t :: Word64Map a
t@(Bin Prefix
p Word64Map a
l Word64Map a
r)
      | Word64 -> Prefix -> Bool
nomatch Word64
k Prefix
p = Word64Map a
t
      | Word64 -> Prefix -> Bool
left Word64
k Prefix
p    = Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
binCheckLeft Prefix
p (Word64Map a -> Word64Map a
go Word64Map a
l) Word64Map a
r
      | Bool
otherwise   = Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
binCheckRight Prefix
p Word64Map a
l (Word64Map a -> Word64Map a
go Word64Map a
r)
    go t :: Word64Map a
t@(Tip Word64
ky a
_)
      | Word64
k Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
ky     = Word64Map a
forall a. Word64Map a
Nil
      | Bool
otherwise   = Word64Map a
t
    go Word64Map a
Nil = Word64Map a
forall a. Word64Map a
Nil

-- | Look up and update the value at a key in the map. The function returns the
-- original value, if it exists, and the updated map.
--
-- The updated value is evaluated to WHNF when it is inserted into the map.
updateLookupWithKey
  :: (Word64 -> a -> Maybe a)
  -> Word64
  -> Word64Map a
  -> (Maybe a, Word64Map a)
updateLookupWithKey :: forall a.
(Word64 -> a -> Maybe a)
-> Word64 -> Word64Map a -> (Maybe a, Word64Map a)
updateLookupWithKey Word64 -> a -> Maybe a
f Word64
k = Word64Map a -> (Maybe a, Word64Map a)
go
  where
    go :: Word64Map a -> (Maybe a, Word64Map a)
go t :: Word64Map a
t@(Bin Prefix
p Word64Map a
l Word64Map a
r)
      | Word64 -> Prefix -> Bool
nomatch Word64
k Prefix
p = (Maybe a
forall a. Maybe a
Nothing, Word64Map a
t)
      | Word64 -> Prefix -> Bool
left Word64
k Prefix
p    = let (Maybe a
found, Word64Map a
l') = Word64Map a -> (Maybe a, Word64Map a)
go Word64Map a
l in (Maybe a
found, Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
binCheckLeft Prefix
p Word64Map a
l' Word64Map a
r)
      | Bool
otherwise   = let (Maybe a
found, Word64Map a
r') = Word64Map a -> (Maybe a, Word64Map a)
go Word64Map a
r in (Maybe a
found, Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
binCheckRight Prefix
p Word64Map a
l Word64Map a
r')
    go t :: Word64Map a
t@(Tip Word64
ky a
y)
      | Word64
k Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== Word64
ky     = case Word64 -> a -> Maybe a
f Word64
ky a
y of
          Just a
y' -> (a -> Maybe a
forall a. a -> Maybe a
Just a
y, Word64 -> a -> Word64Map a
forall a. Word64 -> a -> Word64Map a
Tip Word64
ky a
y')
          Maybe a
Nothing -> (a -> Maybe a
forall a. a -> Maybe a
Just a
y, Word64Map a
forall a. Word64Map a
Nil)
      | Bool
otherwise   = (Maybe a
forall a. Maybe a
Nothing, Word64Map a
t)
    go Word64Map a
Nil = (Maybe a
forall a. Maybe a
Nothing, Word64Map a
forall a. Word64Map a
Nil)

----------------------------------------
-- Internal helpers

-- | Whether the @Word64@ does not start with the given @Prefix@.
--
-- A @Word64@ starts with a @Prefix@ if it shares the high bits with the
-- internal @Word64@ value of the @Prefix@ up to the mask bit.
--
-- @nomatch@ is usually used to determine whether a key belongs in a @Bin@,
-- since all keys in a @Bin@ share a @Prefix@.
nomatch :: Word64 -> Prefix -> Bool
nomatch :: Word64 -> Prefix -> Bool
nomatch Word64
i (Prefix Word64
p) = (Word64
i Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
p) Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. Word64
prefixMask Word64 -> Word64 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word64
0
  where
    prefixMask :: Word64
prefixMask = Word64
p Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (-Word64
p)

-- | Whether the @Word64@ is to the left of the split created by a @Bin@ with
-- this @Prefix@.
--
-- This does not imply that the @Word64@ belongs in this @Bin@. That fact is
-- usually determined first using @nomatch@.
left :: Word64 -> Prefix -> Bool
left :: Word64 -> Prefix -> Bool
left Word64
i Prefix
p = Word64
i Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Prefix -> Word64
unPrefix Prefix
p

-- | Link two @Word64Map@s. The maps must not be empty. The @Prefix@es of the
-- two maps must be different. @k1@ must share the prefix of @t1@. @p2@ must be
-- the prefix of @t2@.
linkKey :: Word64 -> Word64Map a -> Prefix -> Word64Map a -> Word64Map a
linkKey :: forall a.
Word64 -> Word64Map a -> Prefix -> Word64Map a -> Word64Map a
linkKey Word64
k1 Word64Map a
t1 Prefix
p2 Word64Map a
t2 = Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
forall a.
Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
link Word64
k1 Word64Map a
t1 (Prefix -> Word64
unPrefix Prefix
p2) Word64Map a
t2

-- | Link two @Word64Map@s. The maps must not be empty. The @Prefix@es of the
-- two maps must be different. @k1@ must share the prefix of @t1@ and @k2@ must
-- share the prefix of @t2@.
link :: Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
link :: forall a.
Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
link Word64
k1 Word64Map a
t1 Word64
k2 Word64Map a
t2 = Word64
-> Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
forall a.
Word64
-> Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
linkWithMask (Word64 -> Word64 -> Word64
branchMask Word64
k1 Word64
k2) Word64
k1 Word64Map a
t1 Word64
k2 Word64Map a
t2

-- `linkWithMask` is useful when the `branchMask` has already been computed
linkWithMask :: Word64 -> Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
linkWithMask :: forall a.
Word64
-> Word64 -> Word64Map a -> Word64 -> Word64Map a -> Word64Map a
linkWithMask Word64
m Word64
k1 Word64Map a
t1 Word64
k2 Word64Map a
t2
  | Word64
k1 Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
k2   = Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
Bin Prefix
p Word64Map a
t1 Word64Map a
t2
  | Bool
otherwise = Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
Bin Prefix
p Word64Map a
t2 Word64Map a
t1
  where
    p :: Prefix
p = Word64 -> Prefix
Prefix (Word64 -> Word64 -> Word64
mask Word64
k1 Word64
m Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.|. Word64
m)

-- | The prefix of key @i@ up to (but not including) the switching bit @m@.
mask :: Word64 -> Word64 -> Word64
mask :: Word64 -> Word64 -> Word64
mask Word64
i Word64
m = Word64
i Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
.&. (Word64
m Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` (-Word64
m))

-- | The first switching bit where the two prefixes disagree.
--
-- Precondition for defined behavior: p1 /= p2.
branchMask :: Word64 -> Word64 -> Word64
branchMask :: Word64 -> Word64 -> Word64
branchMask Word64
k1 Word64
k2 =
  Word64 -> Int -> Word64
forall a. Bits a => a -> Int -> a
unsafeShiftL Word64
1 (Word64 -> Int
forall b. FiniteBits b => b -> Int
finiteBitSize (Word64
0 :: Word64) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- Word64 -> Int
forall b. FiniteBits b => b -> Int
countLeadingZeros (Word64
k1 Word64 -> Word64 -> Word64
forall a. Bits a => a -> a -> a
`xor` Word64
k2))

-- | Smart constructor that collapses an empty left subtree.
binCheckLeft :: Prefix -> Word64Map a -> Word64Map a -> Word64Map a
binCheckLeft :: forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
binCheckLeft Prefix
_ Word64Map a
Nil Word64Map a
r = Word64Map a
r
binCheckLeft Prefix
p Word64Map a
l   Word64Map a
r = Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
Bin Prefix
p Word64Map a
l Word64Map a
r

-- | Smart constructor that collapses an empty right subtree.
binCheckRight :: Prefix -> Word64Map a -> Word64Map a -> Word64Map a
binCheckRight :: forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
binCheckRight Prefix
_ Word64Map a
l Word64Map a
Nil = Word64Map a
l
binCheckRight Prefix
p Word64Map a
l   Word64Map a
r = Prefix -> Word64Map a -> Word64Map a -> Word64Map a
forall a. Prefix -> Word64Map a -> Word64Map a -> Word64Map a
Bin Prefix
p Word64Map a
l Word64Map a
r