Implement lens-compatible `at` function
[packages/containers.git] / benchmarks / Map.hs
1 {-# LANGUAGE BangPatterns #-}
2 module Main where
3
4 import Control.Applicative (Const(Const, getConst), pure)
5 import Control.DeepSeq
6 import Control.Exception (evaluate)
7 import Control.Monad.Trans (liftIO)
8 import Criterion.Main
9 import Data.Functor.Identity (Identity(runIdentity))
10 import Data.List (foldl')
11 import qualified Data.Map as M
12 import Data.Maybe (fromMaybe)
13 import Prelude hiding (lookup)
14
15 main = do
16 let m = M.fromAscList elems :: M.Map Int Int
17 m_even = M.fromAscList elems_even :: M.Map Int Int
18 m_odd = M.fromAscList elems_odd :: M.Map Int Int
19 evaluate $ rnf [m, m_even, m_odd]
20 defaultMain
21 [ bench "lookup absent" $ whnf (lookup evens) m_odd
22 , bench "lookup present" $ whnf (lookup evens) m_even
23 , bench "at lookup absent" $ whnf (atLookup evens) m_odd
24 , bench "at lookup present" $ whnf (atLookup evens) m_even
25 , bench "atLens lookup absent" $ whnf (atLensLookup evens) m_odd
26 , bench "atLens lookup present" $ whnf (atLensLookup evens) m_even
27 , bench "insert absent" $ whnf (ins elems_even) m_odd
28 , bench "insert present" $ whnf (ins elems_even) m_even
29 , bench "at insert absent" $ whnf (atIns elems_even) m_odd
30 , bench "at insert present" $ whnf (atIns elems_even) m_even
31 , bench "atLens insert absent" $ whnf (atLensIns elems_even) m_odd
32 , bench "atLens insert present" $ whnf (atLensIns elems_even) m_even
33 , bench "insertWith absent" $ whnf (insWith elems_even) m_odd
34 , bench "insertWith present" $ whnf (insWith elems_even) m_even
35 , bench "insertWith' absent" $ whnf (insWith' elems_even) m_odd
36 , bench "insertWith' present" $ whnf (insWith' elems_even) m_even
37 , bench "insertWithKey absent" $ whnf (insWithKey elems_even) m_odd
38 , bench "insertWithKey present" $ whnf (insWithKey elems_even) m_even
39 , bench "insertWithKey' absent" $ whnf (insWithKey' elems_even) m_odd
40 , bench "insertWithKey' present" $ whnf (insWithKey' elems_even) m_even
41 , bench "insertLookupWithKey absent" $ whnf (insLookupWithKey elems_even) m_odd
42 , bench "insertLookupWithKey present" $ whnf (insLookupWithKey elems_even) m_even
43 , bench "insertLookupWithKey' absent" $ whnf (insLookupWithKey' elems_even) m_odd
44 , bench "insertLookupWithKey' present" $ whnf (insLookupWithKey' elems_even) m_even
45 , bench "map" $ whnf (M.map (+ 1)) m
46 , bench "mapWithKey" $ whnf (M.mapWithKey (+)) m
47 , bench "foldlWithKey" $ whnf (ins elems) m
48 -- , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sum 0) m
49 , bench "foldrWithKey" $ whnf (M.foldrWithKey consPair []) m
50 , bench "delete absent" $ whnf (del evens) m_odd
51 , bench "delete present" $ whnf (del evens) m
52 , bench "update absent" $ whnf (upd Just evens) m_odd
53 , bench "update present" $ whnf (upd Just evens) m_even
54 , bench "update delete" $ whnf (upd (const Nothing) evens) m
55 , bench "updateLookupWithKey absent" $ whnf (upd' Just evens) m_odd
56 , bench "updateLookupWithKey present" $ whnf (upd' Just evens) m_even
57 , bench "updateLookupWithKey delete" $ whnf (upd' (const Nothing) evens) m
58 , bench "alter absent" $ whnf (alt id evens) m_odd
59 , bench "alter insert" $ whnf (alt (const (Just 1)) evens) m_odd
60 , bench "alter update" $ whnf (alt id evens) m_even
61 , bench "alter delete" $ whnf (alt (const Nothing) evens) m
62 , bench "at alter absent" $ whnf (atAlt id evens) m_odd
63 , bench "at alter insert" $ whnf (atAlt (const (Just 1)) evens) m_odd
64 , bench "at alter update" $ whnf (atAlt id evens) m_even
65 , bench "at alter delete" $ whnf (atAlt (const Nothing) evens) m
66 , bench "atLens alter absent" $ whnf (atLensAlt id evens) m_odd
67 , bench "atLens alter insert" $ whnf (atLensAlt (const (Just 1)) evens) m_odd
68 , bench "atLens alter update" $ whnf (atLensAlt id evens) m_even
69 , bench "atLens alter delete" $ whnf (atLensAlt (const Nothing) evens) m
70 , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m
71 , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m
72 , bench "lookupIndex" $ whnf (lookupIndex keys) m
73 , bench "union" $ whnf (M.union m_even) m_odd
74 , bench "difference" $ whnf (M.difference m) m_even
75 , bench "intersection" $ whnf (M.intersection m) m_even
76 , bench "split" $ whnf (M.split (bound `div` 2)) m
77 , bench "fromList" $ whnf M.fromList elems
78 , bench "fromList-desc" $ whnf M.fromList (reverse elems)
79 , bench "fromAscList" $ whnf M.fromAscList elems
80 , bench "fromDistinctAscList" $ whnf M.fromDistinctAscList elems
81 ]
82 where
83 bound = 2^12
84 elems = zip keys values
85 elems_even = zip evens evens
86 elems_odd = zip odds odds
87 keys = [1..bound]
88 evens = [2,4..bound]
89 odds = [1,3..bound]
90 values = [1..bound]
91 sum k v1 v2 = k + v1 + v2
92 consPair k v xs = (k, v) : xs
93
94 add3 :: Int -> Int -> Int -> Int
95 add3 x y z = x + y + z
96 {-# INLINE add3 #-}
97
98 lookup :: [Int] -> M.Map Int Int -> Int
99 lookup xs m = foldl' (\n k -> fromMaybe n (M.lookup k m)) 0 xs
100
101 atLookup :: [Int] -> M.Map Int Int -> Int
102 atLookup xs m = foldl' (\n k -> fromMaybe n (getConst (M.at k Const m))) 0 xs
103
104 atLensLookup :: [Int] -> M.Map Int Int -> Int
105 atLensLookup xs m = foldl' (\n k -> fromMaybe n (getConst (atLens k Const m))) 0 xs
106
107 lookupIndex :: [Int] -> M.Map Int Int -> Int
108 lookupIndex xs m = foldl' (\n k -> fromMaybe n (M.lookupIndex k m)) 0 xs
109
110 ins :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
111 ins xs m = foldl' (\m (k, v) -> M.insert k v m) m xs
112
113 atIns :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
114 atIns xs m = foldl' (\m (k, v) -> runIdentity (M.at k (\_ -> pure (Just v)) m)) m xs
115
116 atLensIns :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
117 atLensIns xs m = foldl' (\m (k, v) -> runIdentity (atLens k (\_ -> pure (Just v)) m)) m xs
118
119 insWith :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
120 insWith xs m = foldl' (\m (k, v) -> M.insertWith (+) k v m) m xs
121
122 insWithKey :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
123 insWithKey xs m = foldl' (\m (k, v) -> M.insertWithKey add3 k v m) m xs
124
125 insWith' :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
126 insWith' xs m = foldl' (\m (k, v) -> M.insertWith' (+) k v m) m xs
127
128 insWithKey' :: [(Int, Int)] -> M.Map Int Int -> M.Map Int Int
129 insWithKey' xs m = foldl' (\m (k, v) -> M.insertWithKey' add3 k v m) m xs
130
131 data PairS a b = PS !a !b
132
133 insLookupWithKey :: [(Int, Int)] -> M.Map Int Int -> (Int, M.Map Int Int)
134 insLookupWithKey xs m = let !(PS a b) = foldl' f (PS 0 m) xs in (a, b)
135 where
136 f (PS n m) (k, v) = let !(n', m') = M.insertLookupWithKey add3 k v m
137 in PS (fromMaybe 0 n' + n) m'
138
139 insLookupWithKey' :: [(Int, Int)] -> M.Map Int Int -> (Int, M.Map Int Int)
140 insLookupWithKey' xs m = let !(PS a b) = foldl' f (PS 0 m) xs in (a, b)
141 where
142 f (PS n m) (k, v) = let !(n', m') = M.insertLookupWithKey' add3 k v m
143 in PS (fromMaybe 0 n' + n) m'
144
145 del :: [Int] -> M.Map Int Int -> M.Map Int Int
146 del xs m = foldl' (\m k -> M.delete k m) m xs
147
148 upd :: (Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
149 upd f xs m = foldl' (\m k -> M.update f k m) m xs
150
151 upd' :: (Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
152 upd' f xs m = foldl' (\m k -> snd $ M.updateLookupWithKey (\_ a -> f a) k m) m xs
153
154 alt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
155 alt f xs m = foldl' (\m k -> M.alter f k m) m xs
156
157 atAlt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
158 atAlt f xs m = foldl' (\m k -> runIdentity (M.at k (pure . f) m)) m xs
159
160 atLensAlt :: (Maybe Int -> Maybe Int) -> [Int] -> M.Map Int Int -> M.Map Int Int
161 atLensAlt f xs m = foldl' (\m k -> runIdentity (atLens k (pure . f) m)) m xs
162
163 -- implementation from Control.Lens.At for comparison
164 atLens :: (Functor f, Ord k) =>
165 k -> (Maybe a -> f (Maybe a)) -> M.Map k a -> f (M.Map k a)
166 atLens k f m = (`fmap` f mx) $ \ mx' ->
167 case mx' of
168 Just x' -> M.insert k x' m
169 Nothing ->
170 case mx of
171 Nothing -> m
172 Just x -> M.delete k m
173 where mx = M.lookup k m
174 {-# INLINE atLens #-}
175
176 maybeDel :: Int -> Maybe Int
177 maybeDel n | n `mod` 3 == 0 = Nothing
178 | otherwise = Just n