I've been working on adding an operation like Data.Lens.At.at to Data.Map. This uses an extremely fast, extremely small, and extremely limited sort of queue of bits representing the path taken down a tree of bounded balance with delta = 3. If the word size is 64, or nearly so, then it's unlikely that the code will overflow the queue (which can hold up to 2 * bitSize (0 :: Word) - 2 bits). If the word size is closer to 32 bits, then I need to be a bit more careful about very large maps. I believe that I've calculated appropriate cutoffs based on the maximum tree depths described in the paper, but I'm not entirely confident. Furthermore, I don't really know how to test this code—it requires a setup with a small word size and a giant, badly balanced map.

#if !(WORD_SIZE_IN_BITS >= 61)


alterF :: (Functor f, Ord k) =>
      (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
alterF f !k m
-- It doesn't seem sensible to worry about overflowing the queue
-- if the word size is 61 or more. If I calculate it correctly,
-- that would take a map with nearly a quadrillion entries.
  | wordSize < 61 && size m >= alterFCutoff = alterFFallback f k m
alterF f !k m = case lookupTrace k m of
  TraceResult mv q -> (<$> f mv) $ \ fres ->
    case fres of
      Nothing -> case mv of
                   Nothing -> m
                   Just old -> deleteAlong old q m
      Just new -> case mv of
                   Nothing -> insertAlong q k new m
                   Just _ -> replaceAlong q new m

{-# INLINABLE alterF #-}
{-# INLINE alterF #-}

alterFCutoff :: Int
alterFCutoff = 55744454
alterFCutoff = case wordSize of
      30 -> 17637893
      31 -> 31356255
      32 -> 55744454
      x -> (4^(x*2-2)) `quot` (3^(x*2-2))  -- Unlikely


data TraceResult a = TraceResult (Maybe a) {-# UNPACK #-} !BitQueue

-- Look up a key and return a result indicating whether it was found
-- and what path was taken.
lookupTrace :: Ord k => k -> Map k a -> TraceResult a
lookupTrace = go emptyQB
    go :: Ord k => BitQueueB -> k -> Map k a -> TraceResult a
    go !q !_ Tip = TraceResult Nothing (buildQ q)
    go q k (Bin _ kx x l r) = case compare k kx of
      LT -> (go $! q `snocQB` False) k l
      GT -> (go $! q `snocQB` True) k r
      EQ -> TraceResult (Just x) (buildQ q)
  • \$\begingroup\$ It's been over a year. Want to review your own code? \$\endgroup\$ – Zeta Nov 5 '17 at 17:48

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.