random documentation and performance tweaks

src/TimerWheel.hs

@@ -294,10 +294,10 @@ type TimerBucket =
 
 timerBucketDelete :: Timestamp -> TimerId -> TimerBucket -> Maybe TimerBucket
 timerBucketDelete (coerce @Timestamp @Word64 -> timestamp) timerId bucket =
-  case WordMap.lookup timestamp bucket of
+  case WordMap.lookupExpectingHit timestamp bucket of
     Nothing -> Nothing
     Just (Timers1 timer)
-      | timerId == getTimerId timer -> Just $! WordMap.delete timestamp bucket
+      | timerId == getTimerId timer -> Just $! WordMap.deleteExpectingHit timestamp bucket
       | otherwise -> Nothing
     Just (TimersN timers0) ->
       case timersDelete timerId timers0 of
@@ -379,7 +379,7 @@ atomicExtractExpiredTimersFromBucket buckets index (coerce @Timestamp @Word64 ->
     loop :: Atomics.Ticket TimerBucket -> IO TimerBucket
     loop ticket = do
       let WordMap.Pair expired bucket1 = WordMap.splitL now (Atomics.peekTicket ticket)
-      if WordMap.null expired
+      if WordMap.isEmpty expired
         then pure WordMap.empty
         else do
           (success, ticket1) <- Atomics.casArrayElem buckets index ticket bucket1

src/wordmap/TimerWheel/Internal/WordMap.hs

@@ -3,21 +3,28 @@
 -- TODO: proper code attribution
 module TimerWheel.Internal.WordMap
   ( WordMap (..),
-    delete,
+
+    -- ** Construction
     empty,
+
+    -- ** Querying
+    isEmpty,
+    lookupExpectingHit,
+
+    -- ** Modification
     insert,
-    lookup,
-    null,
+    upsert,
+    deleteExpectingHit,
     Pop (..),
     pop,
     splitL,
-    upsert,
 
     -- * Strict pair
     Pair (..),
 
     -- * Internals
-    mask,
+    Mask (..),
+    prefixof,
   )
 where
 
@@ -27,113 +34,75 @@ import Prelude hiding (lookup, null)
 
 data WordMap a
   = Bin {-# UNPACK #-} !Prefix {-# UNPACK #-} !Mask !(WordMap1 a) !(WordMap1 a)
-  | Tip {-# UNPACK #-} !Word64 a
+  | Tip {-# UNPACK #-} !Word64 !a
   | Nil
 
--- A non-empty word map.
+-- A non-empty wordmap.
 type WordMap1 a = WordMap a
 
 type Prefix = Word64
 
-type Mask = Word64
+newtype Mask = Mask Word64
 
-bin :: Prefix -> Mask -> WordMap a -> WordMap a -> WordMap a
-bin _ _ l Nil = l
-bin _ _ Nil r = r
-bin p m l r = Bin p m l r
-{-# INLINE bin #-}
+empty :: WordMap a
+empty =
+  Nil
+{-# INLINE empty #-}
 
-binCheckLeft :: Prefix -> Mask -> WordMap a -> WordMap1 a -> WordMap1 a
-binCheckLeft _ _ Nil r = r
-binCheckLeft p m l r = Bin p m l r
-{-# INLINE binCheckLeft #-}
+isEmpty :: WordMap a -> Bool
+isEmpty = \case
+  Nil -> True
+  _ -> False
+{-# INLINE isEmpty #-}
 
-binCheckRight :: Prefix -> Mask -> WordMap1 a -> WordMap a -> WordMap1 a
-binCheckRight _ _ l Nil = l
-binCheckRight p m l r = Bin p m l r
-{-# INLINE binCheckRight #-}
+-- | Look up an element in a wordmap, expecting it to be there.
+lookupExpectingHit :: forall a. Word64 -> WordMap a -> Maybe a
+lookupExpectingHit !k =
+  go
+  where
+    go :: WordMap a -> Maybe a
+    go = \case
+      -- Ignore prefix because we expect to find the key
+      Bin _ m l r -> go (if leftside k m then l else r)
+      Tip kx x -> if k == kx then Just x else Nothing
+      Nil -> Nothing
 
-delete :: Word64 -> WordMap a -> WordMap a
-delete !k t =
+-- | Delete an element from a wordmap, expecting it to be there.
+deleteExpectingHit :: Word64 -> WordMap a -> WordMap a
+deleteExpectingHit !k t =
   case t of
     Bin p m l r
-      | nomatch k p m -> t
-      | zero k m -> binCheckLeft p m (delete k l) r
-      | otherwise -> binCheckRight p m l (delete k r)
+      -- Ignore prefix because we expect to find the key
+      | leftside k m -> binl p m (deleteExpectingHit k l) r
+      | otherwise -> binr p m l (deleteExpectingHit k r)
     Tip kx _
       | k == kx -> Nil
       | otherwise -> t
     Nil -> Nil
 
-empty :: WordMap a
-empty =
-  Nil
-{-# INLINE empty #-}
-
-insert :: Word64 -> a -> WordMap a -> WordMap a
+insert :: Word64 -> a -> WordMap a -> WordMap1 a
 insert !k !x t =
   case t of
     Bin p m l r
-      | nomatch k p m -> link k (Tip k x) p t
-      | zero k m -> Bin p m (insert k x l) r
+      | outsider k p m -> link k (Tip k x) p t
+      | leftside k m -> Bin p m (insert k x l) r
       | otherwise -> Bin p m l (insert k x r)
     Tip ky _
       | k == ky -> Tip k x
       | otherwise -> link k (Tip k x) ky t
     Nil -> Tip k x
 
-link :: Prefix -> WordMap a -> Prefix -> WordMap a -> WordMap a
-link p1 t1 p2 = linkWithMask (branchMask p1 p2) p1 t1
-{-# INLINE link #-}
-
--- `linkWithMask` is useful when the `branchMask` has already been computed
-linkWithMask :: Mask -> Prefix -> WordMap a -> WordMap a -> WordMap a
-linkWithMask m p0 t1 t2
-  | zero p0 m = Bin p m t1 t2
-  | otherwise = Bin p m t2 t1
-  where
-    p = mask p0 m
-{-# INLINE linkWithMask #-}
-
-lookup :: forall a. Word64 -> WordMap a -> Maybe a
-lookup !k =
-  go
-  where
-    go :: WordMap a -> Maybe a
-    go = \case
-      Bin _ m l r -> go (if zero k m then l else r)
-      Tip kx x -> if k == kx then Just x else Nothing
-      Nil -> Nothing
-
-null :: WordMap a -> Bool
-null = \case
-  Nil -> True
-  _ -> False
-{-# INLINE null #-}
-
-singleton :: Word64 -> a -> WordMap a
-singleton k !x = Tip k x
-{-# INLINE singleton #-}
-
--- @splitL k xs@ splits @xs@ into @ys@ and @zs@, where every value in @ys@ is less than or equal to @k@, and every value
--- in @zs@ is greater than @k@.
-splitL :: forall a. Word64 -> WordMap a -> Pair (WordMap a) (WordMap a)
-splitL k =
-  go
-  where
-    go :: WordMap a -> Pair (WordMap a) (WordMap a)
-    go = \case
-      t@(Bin p m l r)
-        | nomatch k p m ->
-            if k > p
-              then Pair t Nil
-              else Pair Nil t
-        | zero k m -> mapPairR (\lr -> bin p m lr r) (go l)
-        | otherwise -> mapPairL (\rl -> bin p m l rl) (go r)
-      t@(Tip k2 _)
-        | k >= k2 -> Pair t Nil
-        | otherwise -> Pair Nil t
-      Nil -> Pair Nil Nil
+upsert :: Word64 -> a -> (a -> a) -> WordMap a -> WordMap a
+upsert !k !x f t =
+  case t of
+    Bin p m l r
+      | outsider k p m -> link k (Tip k x) p t
+      | leftside k m -> Bin p m (upsert k x f l) r
+      | otherwise -> Bin p m l (upsert k x f r)
+    Tip ky y
+      | k == ky -> Tip k (f y)
+      | otherwise -> link k (Tip k x) ky t
+    Nil -> Tip k x
 
 data Pop a
   = PopNada
@@ -142,67 +111,109 @@ data Pop a
 pop :: WordMap a -> Pop a
 pop = \case
   Nil -> PopNada
-  Bin p m l r0 | m < 0 ->
-    case pop1 r0 of
-      PopAlgo k x r1 -> PopAlgo k x (binCheckRight p m l r1)
-      PopNada -> undefined
   t -> pop1 t
 {-# INLINE pop #-}
 
 pop1 :: WordMap1 a -> Pop a
 pop1 = \case
-  Bin p m l0 r ->
-    case pop1 l0 of
-      PopAlgo k x l1 -> PopAlgo k x (binCheckLeft p m l1 r)
+  Bin p m l r ->
+    case pop1 l of
+      PopAlgo k x l1 -> PopAlgo k x (binl p m l1 r)
       PopNada -> undefined
   Tip k x -> PopAlgo k x Nil
   Nil -> undefined
 
-upsert :: Word64 -> a -> (a -> a) -> WordMap a -> WordMap a
-upsert !k x f = \case
-  t@(Bin p m l r)
-    | nomatch k p m -> link k (singleton k x) p t
-    | zero k m -> Bin p m (upsert k x f l) r
-    | otherwise -> Bin p m l (upsert k x f r)
-  t@(Tip ky y)
-    | k == ky -> Tip k $! f y
-    | otherwise -> link k (singleton k x) ky t
-  Nil -> singleton k x
+-- @splitL k xs@ splits @xs@ into @ys@ and @zs@, where every value in @ys@ is less than or equal to @k@, and every value
+-- in @zs@ is greater than @k@.
+splitL :: forall a. Word64 -> WordMap a -> Pair (WordMap a) (WordMap a)
+splitL k =
+  go
+  where
+    go :: WordMap a -> Pair (WordMap a) (WordMap a)
+    go t =
+      case t of
+        Bin p m l r
+          | outsider k p m ->
+              if k > p
+                then Pair t Nil
+                else Pair Nil t
+          | leftside k m -> mapPairR (\lr -> bin p m lr r) (go l)
+          | otherwise -> mapPairL (bin p m l) (go r)
+        Tip k2 _
+          | k >= k2 -> Pair t Nil
+          | otherwise -> Pair Nil t
+        Nil -> Pair Nil Nil
+
+------------------------------------------------------------------------------------------------------------------------
+-- Low-level constructors
+
+-- | Combine two wordmaps with a Bin constructor if they are both non-empty. Otherwise, just return one that's not known
+-- to be empty.
+bin :: Prefix -> Mask -> WordMap a -> WordMap a -> WordMap a
+bin _ _ l Nil = l
+bin _ _ Nil r = r
+bin p m l r = Bin p m l r
+{-# INLINE bin #-}
+
+-- | Like 'bin', but for when the right wordmap is known to be non-empty.
+binl :: Prefix -> Mask -> WordMap a -> WordMap1 a -> WordMap1 a
+binl _ _ Nil r = r
+binl p m l r = Bin p m l r
+{-# INLINE binl #-}
+
+-- | Like 'bin', but for when the left wordmap is known to be non-empty.
+binr :: Prefix -> Mask -> WordMap1 a -> WordMap a -> WordMap1 a
+binr _ _ l Nil = l
+binr p m l r = Bin p m l r
+{-# INLINE binr #-}
+
+link :: Prefix -> WordMap a -> Prefix -> WordMap a -> WordMap a
+link p1 t1 p2 t2
+  | leftside p1 m = Bin p m t1 t2
+  | otherwise = Bin p m t2 t1
+  where
+    m = highestBitMask (p1 `xor` p2)
+    p = prefixof p1 m
+{-# INLINE link #-}
 
 ------------------------------------------------------------------------------------------------------------------------
 -- Bit twiddling
 
--- | Should this key follow the left subtree of a 'Bin' with switching bit @m@?
--- N.B., the answer is only valid when @match i p m@ is true.
-zero :: Word64 -> Mask -> Bool
-zero i m =
-  i .&. m == 0
-{-# INLINE zero #-}
-
--- | Does the key @i@ differ from the prefix @p@ before getting to the switching bit @m@?
-nomatch :: Word64 -> Prefix -> Mask -> Bool
-nomatch i p m =
-  mask i m /= p
-{-# INLINE nomatch #-}
-
--- | The prefix of key @i@ up to (but not including) the switching bit @m@.
-mask :: Word64 -> Mask -> Prefix
-mask i m =
-  i .&. ((-m) `xor` m)
-{-# INLINE mask #-}
-
--- | The first switching bit where the two prefixes disagree.
-branchMask :: Prefix -> Prefix -> Mask
-branchMask p1 p2 =
-  highestBitMask (p1 `xor` p2)
-{-# INLINE branchMask #-}
+-- @leftside w m@ returns whether key @k@ belongs of the left side of a bin with mask @m@.
+--
+-- Precondition: @w@ is not an "outsider".
+leftside :: Word64 -> Mask -> Bool
+leftside w (Mask m) =
+  w .&. m == 0
+{-# INLINE leftside #-}
+
+-- @outsider w p m@ returns whether @w@ is an "outsider" of a bin with prefix @p@ and mask @m@; that is, it is known
+-- not to be on either side.
+outsider :: Word64 -> Prefix -> Mask -> Bool
+outsider w p m =
+  prefixof w m /= p
+{-# INLINE outsider #-}
+
+-- As you can see from this:
+--
+--        m = 00000000 00000000 00000000 00010000 00000000 00000000 00000000 00000000
+--     -m   = 11111111 11111111 11111111 11110000 00000000 00000000 00000000 00000000
+-- xor -m m = 11111111 11111111 11111111 11100000 00000000 00000000 00000000 00000000
+--
+-- @prefixof w m@ retains only the prefix bits of word @w@, per mask @m@.
+prefixof :: Word64 -> Mask -> Prefix
+prefixof w (Mask m) =
+  w .&. ((-m) `xor` m)
+{-# INLINE prefixof #-}
 
 -- | Return a word where only the highest bit is set.
-highestBitMask :: Word64 -> Word64
-highestBitMask w = unsafeShiftL 1 (63 - countLeadingZeros w)
+highestBitMask :: Word64 -> Mask
+highestBitMask w = Mask (unsafeShiftL 1 (63 - countLeadingZeros w))
 {-# INLINE highestBitMask #-}
 
--- A strict pair
+------------------------------------------------------------------------------------------------------------------------
+-- Strict pair stuff
+
 data Pair a b
   = Pair !a !b