Merge pull request #74 from purescript/strmap-map-with-key

Add StrMap.mapWithKey

Author: garyb

src/Data/StrMap.js

@@ -42,6 +42,17 @@ exports._fmapStrMap = function (m0, f) {
   return m;
 };
 
+// jshint maxparams: 2
+exports._mapWithKey = function (m0, f) {
+  var m = {};
+  for (var k in m0) {
+    if (m0.hasOwnProperty(k)) {
+      m[k] = f(k)(m0[k]);
+    }
+  }
+  return m;
+};
+
 // jshint maxparams: 1
 exports._foldM = function (bind) {
   return function (f) {

src/Data/StrMap.purs

@@ -21,6 +21,7 @@ module Data.StrMap
   , member
   , alter
   , update
+  , mapWithKey
   , keys
   , values
   , union
@@ -232,6 +233,12 @@ union m = mutate (\s -> foldM SM.poke s m)
 unions :: forall a. L.List (StrMap a) -> StrMap a
 unions = foldl union empty
 
+foreign import _mapWithKey :: forall a b. Fn2 (StrMap a) (String -> a -> b) (StrMap b)
+
+-- | Apply a function of two arguments to each key/value pair, producing a new map
+mapWithKey :: forall a b. (String -> a -> b) -> StrMap a -> StrMap b
+mapWithKey f m = runFn2 _mapWithKey m f
+
 instance semigroupStrMap :: (Semigroup a) => Semigroup (StrMap a) where
   append m1 m2 = mutate (\s1 -> foldM (\s2 k v2 -> SM.poke s2 k (runFn4 _lookup v2 (\v1 -> v1 <> v2) k m2)) s1 m1) m2
 

test/Test/Data/Map.purs

@@ -301,3 +301,10 @@ mapTests = do
   quickCheck $ \(TestMap m) -> case M.findMax (smallKeyToNumberMap m) of
     Nothing -> M.isEmpty m
     Just { key: k, value: v } -> M.lookup k m == Just v && all (_ <= k) (M.keys m)
+
+  log "mapWithKey is correct"
+  quickCheck $ \(TestMap m :: TestMap String Int) -> let
+    f k v = k <> show v
+    resultViaMapWithKey = m # M.mapWithKey f
+    resultViaLists = m # M.toList # map (\(Tuple k v) → Tuple k (f k v)) # M.fromFoldable
+    in resultViaMapWithKey === resultViaLists

test/Test/Data/StrMap.purs

@@ -17,7 +17,7 @@ import Data.Tuple (Tuple(..), fst)
 
 import Partial.Unsafe (unsafePartial)
 
-import Test.QuickCheck ((<?>), quickCheck, quickCheck')
+import Test.QuickCheck ((<?>), quickCheck, quickCheck', (===))
 import Test.QuickCheck.Arbitrary (class Arbitrary, arbitrary)
 
 newtype TestStrMap v = TestStrMap (M.StrMap v)
@@ -153,6 +153,13 @@ strMapTests = do
   log "fromFoldable = zip keys values"
   quickCheck $ \(TestStrMap m) -> M.toList m == zipWith Tuple (fromFoldable $ M.keys m) (M.values m :: List Int)
 
+  log "mapWithKey is correct"
+  quickCheck $ \(TestStrMap m :: TestStrMap Int) -> let
+    f k v = k <> show v
+    resultViaMapWithKey = m # M.mapWithKey f
+    resultViaLists = m # M.toList # map (\(Tuple k v) → Tuple k (f k v)) # M.fromFoldable
+    in resultViaMapWithKey === resultViaLists
+
   log "Bug #63: accidental observable mutation in foldMap"
   quickCheck \(TestStrMap m) ->
     let lhs = go m