# Implement “at” lens for Data.Map

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)
#define DEFINE_ALTERF_FALLBACK 1
#endif


...

alterF :: (Functor f, Ord k) =>
(Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a)
#if DEFINE_ALTERF_FALLBACK
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
#endif
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 #-}
#else
{-# INLINE alterF #-}
#endif

#if DEFINE_ALTERF_FALLBACK
alterFCutoff :: Int
#if WORD_SIZE_IN_BITS == 32
alterFCutoff = 55744454
#else
alterFCutoff = case wordSize of
30 -> 17637893
31 -> 31356255
32 -> 55744454
x -> (4^(x*2-2)) quot (3^(x*2-2))  -- Unlikely
#endif
#endif


...

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
where
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)

• It's been over a year. Want to review your own code? – Zeta Nov 5 '17 at 17:48