Implement lens-compatible at function

Akin to `alter` but allows an arbitrary Functor.
Add benchmarks for `at`
Add tests for `at`
Add `at` from Lens to benchmarks for comparison

Data/Map/Base.hs

@@ -137,6 +137,7 @@ module Data.Map.Base (
     , updateWithKey
     , updateLookupWithKey
     , alter
+    , at
 
     -- * Combine
 
@@ -932,6 +933,27 @@ alter = go
 {-# INLINE alter #-}
 #endif
 
+at :: (Functor f, Ord k) =>
+      k -> (Maybe a -> f (Maybe a)) -> Map k a -> f (Map k a)
+at = go
+  where
+    STRICT_1_OF_3(go)
+    go k f Tip = (`fmap` f Nothing) $ \ mx -> case mx of
+               Nothing -> Tip
+               Just x  -> singleton k x
+
+    go k f (Bin sx kx x l r) = case compare k kx of
+               LT -> (\ m -> balance kx x m r) `fmap` go k f l
+               GT -> (\ m -> balance kx x l m) `fmap` go k f r
+               EQ -> (`fmap` f (Just x)) $ \ mx' -> case mx' of
+                       Just x' -> Bin sx kx x' l r
+                       Nothing -> glue l r
+#if __GLASGOW_HASKELL__ >= 700
+{-# INLINABLE at #-}
+#else
+{-# INLINE at #-}
+#endif
+
 {--------------------------------------------------------------------
   Indexing
 --------------------------------------------------------------------}

Data/Map/Lazy.hs

@@ -96,6 +96,7 @@ module Data.Map.Lazy (
     , updateWithKey
     , updateLookupWithKey
     , alter
+    , at
 
     -- * Combine
 

Data/Map/Strict.hs

@@ -104,6 +104,7 @@ module Data.Map.Strict
     , updateWithKey
     , updateLookupWithKey
     , alter
+    , at
 
     -- * Combine
 

benchmarks/Map.hs

@@ -1,10 +1,12 @@
 {-# LANGUAGE BangPatterns #-}
 module Main where
 
+import Control.Applicative (Const(Const, getConst), pure)
 import Control.DeepSeq
 import Control.Exception (evaluate)
 import Control.Monad.Trans (liftIO)
 import Criterion.Main
+import Data.Functor.Identity (Identity(runIdentity))
 import Data.List (foldl')
 import qualified Data.Map as M
 import Data.Maybe (fromMaybe)
@@ -18,8 +20,16 @@ main = do
     defaultMain
         [ bench "lookup absent" $ whnf (lookup evens) m_odd
         , bench "lookup present" $ whnf (lookup evens) m_even
+        , bench "at lookup absent" $ whnf (atLookup evens) m_odd
+        , bench "at lookup present" $ whnf (atLookup evens) m_even
+        , bench "atLens lookup absent" $ whnf (atLensLookup evens) m_odd
+        , bench "atLens lookup present" $ whnf (atLensLookup evens) m_even
         , bench "insert absent" $ whnf (ins elems_even) m_odd
         , bench "insert present" $ whnf (ins elems_even) m_even
+        , bench "at insert absent" $ whnf (atIns elems_even) m_odd
+        , bench "at insert present" $ whnf (atIns elems_even) m_even
+        , bench "atLens insert absent" $ whnf (atLensIns elems_even) m_odd
+        , bench "atLens insert present" $ whnf (atLensIns elems_even) m_even
         , bench "insertWith absent" $ whnf (insWith elems_even) m_odd
         , bench "insertWith present" $ whnf (insWith elems_even) m_even
         , bench "insertWith' absent" $ whnf (insWith' elems_even) m_odd
@@ -49,6 +59,14 @@ main = do
         , bench "alter insert"  $ whnf (alt (const (Just 1)) evens) m_odd
         , bench "alter update"  $ whnf (alt id evens) m_even
         , bench "alter delete"  $ whnf (alt (const Nothing) evens) m
+        , bench "at alter absent" $ whnf (atAlt id evens) m_odd
+        , bench "at alter insert" $ whnf (atAlt (const (Just 1)) evens) m_odd
+        , bench "at alter update" $ whnf (atAlt id evens) m_even
+        , bench "at alter delete" $ whnf (atAlt (const Nothing) evens) m
+        , bench "atLens alter absent" $ whnf (atLensAlt id evens) m_odd
+        , bench "atLens alter insert" $ whnf (atLensAlt (const (Just 1)) evens) m_odd
+        , bench "atLens alter update" $ whnf (atLensAlt id evens) m_even
+        , bench "atLens alter delete" $ whnf (atLensAlt (const Nothing) evens) m
         , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m
         , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m
         , bench "lookupIndex" $ whnf (lookupIndex keys) m
@@ -80,12 +98,24 @@ add3 x y z = x + y + z
 lookup :: [Int] -> M.Map Int Int -> Int
 lookup xs m = foldl' (\n k -> fromMaybe n (M.lookup k m)) 0 xs
 
+atLookup :: [Int] -> M.Map Int Int -> Int
+atLookup xs m = foldl' (\n k -> fromMaybe n (getConst (M.at k Const m))) 0 xs
+
+atLensLookup :: [Int] -> M.Map Int Int -> Int
+atLensLookup xs m = foldl' (\n k -> fromMaybe n (getConst (atLens k Const m))) 0 xs
+
 lookupIndex :: [Int] -> M.Map Int Int -> Int
 lookupIndex xs m = foldl' (\n k -> fromMaybe n (M.lookupIndex k m)) 0 xs
 
 ins :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
 ins xs m = foldl' (\m (k, v) -> M.insert k v m) m xs
 
+atIns :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
+atIns xs m = foldl' (\m (k, v) -> runIdentity (M.at k (\_ -> pure (Just v)) m)) m xs
+
+atLensIns :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
+atLensIns xs m = foldl' (\m (k, v) -> runIdentity (atLens k (\_ -> pure (Just v)) m)) m xs
+
 insWith :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
 insWith xs m = foldl' (\m (k, v) -> M.insertWith (+) k v m) m xs
 
@@ -124,6 +154,25 @@ upd' f xs m = foldl' (\m k -> snd $ M.updateLookupWithKey (\_ a -> f a) k m) m x
 alt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
 alt f xs m = foldl' (\m k -> M.alter f k m) m xs
 
+atAlt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
+atAlt f xs m = foldl' (\m k -> runIdentity (M.at k (pure . f) m)) m xs
+
+atLensAlt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
+atLensAlt f xs m = foldl' (\m k -> runIdentity (atLens k (pure . f) m)) m xs
+
+-- implementation from Control.Lens.At for comparison
+atLens :: (Functor f, Ord k) =>
+          k -> (Maybe a -> f (Maybe a)) -> M.Map k a -> f (M.Map k a)
+atLens k f m = (`fmap` f mx) $ \ mx' ->
+  case mx' of
+    Just x' -> M.insert k x' m
+    Nothing ->
+      case mx of
+        Nothing -> m
+        Just x  -> M.delete k m
+  where mx = M.lookup k m
+{-# INLINE atLens #-}
+
 maybeDel :: Int -> Maybe Int
 maybeDel n | n `mod` 3 == 0 = Nothing
            | otherwise      = Just n

containers.cabal

@@ -85,6 +85,8 @@ Test-suite map-lazy-properties
     extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types
 
     build-depends:
+        -- only needed for base < 4.8 to get Identity
+        transformers,
         HUnit,
         QuickCheck,
         test-framework,
@@ -103,6 +105,8 @@ Test-suite map-strict-properties
     extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types
 
     build-depends:
+        -- only needed for base < 4.8 to get Identity
+        transformers,
         HUnit,
         QuickCheck,
         test-framework,

tests/map-properties.hs

@@ -6,6 +6,8 @@ import Data.Map.Strict as Data.Map
 import Data.Map.Lazy as Data.Map
 #endif
 
+import Control.Applicative (Const(Const, getConst), pure)
+import Data.Functor.Identity (Identity(runIdentity))
 import Data.Monoid
 import Data.Maybe hiding (mapMaybe)
 import qualified Data.Maybe as Maybe (mapMaybe)
@@ -54,6 +56,7 @@ main = defaultMain
          , testCase "updateWithKey" test_updateWithKey
          , testCase "updateLookupWithKey" test_updateLookupWithKey
          , testCase "alter" test_alter
+         , testCase "at" test_at
          , testCase "union" test_union
          , testCase "mappend" test_mappend
          , testCase "unionWith" test_unionWith
@@ -405,6 +408,28 @@ test_alter = do
     f _ = Nothing
     g _ = Just "c"
 
+test_at :: Assertion
+test_at = do
+    employeeCurrency "John" @?= Just "Euro"
+    employeeCurrency "Pete" @?= Nothing
+    atAlter f 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a")]
+    atAlter f 5 (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "b"
+    atAlter g 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a"), (7, "c")]
+    atAlter g 5 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "c")]
+  where
+    atAlter f k m = runIdentity (at k (pure . f) m)
+    atLookup k m = getConst (at k Const m)
+    f _ = Nothing
+    g _ = Just "c"
+    employeeDept = fromList([("John","Sales"), ("Bob","IT")])
+    deptCountry = fromList([("IT","USA"), ("Sales","France")])
+    countryCurrency = fromList([("USA", "Dollar"), ("France", "Euro")])
+    employeeCurrency :: String -> Maybe String
+    employeeCurrency name = do
+        dept <- atLookup name employeeDept
+        country <- atLookup dept deptCountry
+        atLookup country countryCurrency
+
 ----------------------------------------------------------------
 -- Combine