# Efficent lens operation for Data.Map et al.

roconnor at theorem.ca roconnor at theorem.ca
Wed Jan 18 17:06:36 CET 2012

```On Tue, 17 Jan 2012, Johan Tibell wrote:

> Try it out on the Criterion benchmarks in the source repo. I'd be
> curious to see the results.

Okay, I never intended to try and replace the existing operations with a
lens based one, but now you've got me curious to try.

Here is my fastLens implementation:

fastLens :: Key -> IntMap a -> (IntMap a -> (a -> IntMap a) -> a -> c) -> (IntMap a -> (a -> IntMap a) -> c) -> c
fastLens k = k `seq` go
where
go t@(Bin p m l r) c1 c2
| nomatch k p m = c2 t (\x -> join k (Tip k x) p t)
| zero k m      = go l (updateContL c1) (updateContL c2)
| otherwise     = go r (updateContR c1) (updateContR c2)
where
updateContL c dl il = c (bin p m dl r) (\x -> bin p m (il x) r)
updateContR c dr ir = c (bin p m l dr) (\x -> bin p m r (ir x))
go t@(Tip ky y) c1 c2
| k == ky   = c1 Nil (Tip ky) y
| otherwise = c2 t (\x -> join k (Tip k x) ky t)
go Nil _ c2   = c2 Nil (Tip k)

lookup :: Key -> IntMap a -> Maybe a
lookup k t = fastLens k t (\_ _ -> Just) (\_ _ -> Nothing)

insert :: Key -> a -> IntMap a -> IntMap a
insert k x t = fastLens k t (\_ it _ -> it x) (\_ it -> it x)

delete :: Key -> IntMap a -> IntMap a
delete k t = fastLens k t (\dt _ _ -> dt) (\dt _ -> dt)

update ::  (a -> Maybe a) -> Key -> IntMap a -> IntMap a
update f k t = fastLens k t (\dt it x -> maybe dt it (f x)) (\dt _ -> dt)

alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a
alter f k t = fastLens k t (\dt it x -> maybe dt it (f (Just x)))
(\dt it   -> maybe dt it (f Nothing))

lens :: Key -> IntMap a -> (Maybe a -> IntMap a, Maybe a)
lens k t = fastLens k t (\dt it x -> (maybe dt it,Just x))
(\dt it -> (maybe dt it,Nothing))