Add filter, etc.

bench/Main.purs

@@ -11,7 +11,7 @@ import Data.Array (range)
 import Data.Array as Array
 import Data.Array.ST as STA
 import Data.Foldable (foldMap, foldl, foldr)
-import Data.FoldableWithIndex (foldMapWithIndex, foldrWithIndex, forWithIndex_)
+import Data.FoldableWithIndex (foldMapWithIndex, foldlWithIndex, foldrWithIndex, forWithIndex_)
 import Data.HashMap (HashMap)
 import Data.HashMap as HM
 import Data.HashSet as HS
@@ -49,16 +49,21 @@ insertOM a = OM.fromFoldable a
 
 main :: Effect Unit
 main = do
+  let is100 = is 100
+  let hmIs100 = insertHM is100
+  let omIs100 = insertOM is100
+  let iKeys100 = range 1 100
+
   let is10000 = is 10000
   let is20000 = is 10000 <> is 10000
   let iKeys10000 = range 1 10000
   let hmIs10000 = insertHM is10000
   let omIs10000 = insertOM is10000
 
-  let is100 = is 100
-  let hmIs100 = insertHM is100
-  let omIs100 = insertOM is100
-  let iKeys100 = range 1 100
+  -- string keys
+  let si10000 = si 10000
+  let hmSi10000 = insertHM si10000
+  let omSi10000 = insertOM si10000
 
   log "HM singleton"
   benchWith 1000000 \_ -> HM.singleton 5 42
@@ -185,6 +190,29 @@ main = do
   log "OM map Just and sequence 10000"
   bench \_ -> sequence $ Just <$> omIs10000
 
+
+  log ""
+  log "Filtering"
+  log "---------"
+
+  log "HM filterWithKey key even"
+  bench \_ -> HM.filterWithKey (\k v -> k `mod` 2 == 0) hmIs10000
+
+  log "HM filterWithKey (naive) key even"
+  bench \_ -> (\f -> foldlWithIndex (\k m v -> if f k v then HM.insert k v m else m) HM.empty) (\k v -> k `mod` 2 == 0) hmIs10000
+
+  log "OM filterWithKey key even"
+  bench \_ -> OM.filterWithKey (\k v -> k `mod` 2 == 0) omIs10000
+
+  log "HM filterWithKey value even"
+  bench \_ -> HM.filterWithKey (\k v -> v `mod` 2 == 0) hmSi10000
+
+  log "HM filterWithKey (naive) value even"
+  bench \_ -> (\f -> foldlWithIndex (\k m v -> if f k v then HM.insert k v m else m) HM.empty) (\k v -> v `mod` 2 == 0) hmSi10000
+
+  log "OM filterWithKey value even"
+  bench \_ -> OM.filterWithKey (\k v -> v `mod` 2 == 0) omSi10000
+
   log ""
   log "UnionWith"
   log "---------"

src/Data/HashMap.js

@@ -335,10 +335,40 @@ MapNode.prototype.intersectionWith = function (eq, hash, f, that, shift) {
     return new MapNode(datamap, nodemap, data.concat(nodes.reverse()));
 }
 
+MapNode.prototype.filterWithKey = function filterWithKey(f) {
+    var datamap = 0;
+    var nodemap = 0;
+    var data = [];
+    var nodes = [];
+    for (var i = 0; i < 32; i++) {
+        var bit = 1 << i;
+        if ((this.datamap & bit) !== 0) {
+            var dataIndex = index(this.datamap, bit);
+            var k = this.content[dataIndex * 2];
+            var v = this.content[dataIndex * 2 + 1];
+            if (f(k)(v)) {
+                datamap |= bit;
+                data.push(k, v);
+            }
+        } else if ((this.nodemap & bit) !== 0) {
+            var nodeIndex = index(this.nodemap, bit);
+            var node = this.content[this.content.length - nodeIndex - 1].filterWithKey(f);
+            if (isEmpty(node)) continue;
+            if (node.isSingleton()) {
+                datamap |= bit;
+                data.push(node.content[0], node.content[1]);
+            } else {
+                nodemap |= bit;
+                nodes.push(node);
+            }
+        }
+    }
+    return new MapNode(datamap, nodemap, data.concat(nodes.reverse()));
+}
 
-// This builds an n-ary curried function that all values and all
+// This builds an n-ary curried function that takes all values and all
 // subnodes as arguments and places them in a copy of the hashmap
-// preserving the keys, datamap, and nodemap.  Basically, a (Hashmap k
+// preserving the keys, datamap, and nodemap. Basically, a (Hashmap k
 // v) with s key-value pairs and t nodes turns into a function:
 //
 // k_0 -> .. -> k_s -> HashMap_0 k v -> .. -> HashMap_t k v -> HashMap k v
@@ -573,6 +603,26 @@ Collision.prototype.intersectionWith = function (eq, hash, f, that, shift) {
     return new Collision(keys, values);
 }
 
+Collision.prototype.filterWithKey = function collisionFilterWithKey(f) {
+    var keys = [];
+    var values = [];
+    for (var i = 0; i < this.keys.length; i++) {
+        var k = this.keys[i];
+        var v = this.values[i];
+        if (f(k)(v)) {
+            keys.push(k);
+            values.push(v);
+        }
+    }
+    if (keys.length === 0) return empty;
+    // This is a bit dodgy. We return a fake MapNode (wrong datamap
+    // (WHICH CANNOT BE 0, OTHERWISE isEmpty THINKS IT'S EMPTY!) and
+    // nodemap), but it's okay, because we will immediately
+    // deconstruct it in MapNode's filterWithKey.
+    if (keys.length === 1) return new MapNode(1, 0, [keys[0], values[0]]);
+    return new Collision(keys, values);
+}
+
 function mask(keyHash, shift) {
     return 1 << ((keyHash >>> shift) & 31);
 }
@@ -799,6 +849,12 @@ exports.hashPurs = function (vhash) {
     };
 };
 
+exports.filterWithKey = function (f) {
+    return function (m) {
+        return m.filterWithKey(f);
+    };
+};
+
 // TODO ideally this would use a mutable HashSet instead
 exports.nubHashPurs = function (eq) {
     return function (hash) {

src/Data/HashMap.purs

@@ -20,6 +20,10 @@ module Data.HashMap (
   update,
   alter,
 
+  filter,
+  filterWithKey,
+  filterKeys,
+
   fromFoldable,
   fromFoldableBy,
   toArrayBy,
@@ -258,6 +262,31 @@ intersectionWith = intersectionWithPurs eq hash
 difference :: forall k v w. Hashable k => HashMap k v -> HashMap k w -> HashMap k v
 difference l r = foldr delete l (keys r)
 
+-- | Remove key-value-pairs from a map for which the predicate does
+-- | not hold.
+-- |
+-- | ```PureScript
+-- | filter (const False) m == empty
+-- | filter (const True) m == m
+-- | ```
+filter :: forall k v. Hashable k => (v -> Boolean) -> HashMap k v -> HashMap k v
+filter f = filterWithKey (const f)
+
+-- | Remove key-value-pairs from a map for which the predicate does
+-- | not hold.
+-- |
+-- | Like `filter`, but the predicate takes both key and value.
+foreign import filterWithKey :: forall k v. (k -> v -> Boolean) -> HashMap k v -> HashMap k v
+
+-- | Remove all keys from the map for which the predicate does not
+-- | hold.
+-- |
+-- | `difference m1 m2 == filterKeys (\k -> member k m2) m1`
+filterKeys :: forall k v. Hashable k => (k -> Boolean) -> HashMap k v -> HashMap k v
+filterKeys f = filterWithKey (\k v -> f k)
+
+
+
 -- | Remove duplicates from an array.
 -- |
 -- | Like `nub` from `Data.Array`, but uses a `Hashable` constraint

src/Data/HashSet.purs

@@ -12,6 +12,8 @@ module Data.HashSet (
   member,
   delete,
 
+  filter,
+
   union,
   intersection,
   difference,
@@ -78,6 +80,13 @@ fromFoldable = foldr insert empty
 toArray :: forall a. HashSet a -> Array a
 toArray (HashSet m) = M.keys m
 
+-- | Remove all elements from the set for which the predicate does not
+-- | hold.
+-- |
+-- | `filter (const false) s == empty`
+filter :: forall a. (a -> Boolean) -> HashSet a -> HashSet a
+filter f (HashSet m) = HashSet (M.filterWithKey (\k v -> f k) m)
+
 -- | Union two sets.
 union :: forall a. Hashable a => HashSet a -> HashSet a -> HashSet a
 union (HashSet l) (HashSet r) = HashSet (M.unionWith const l r)

test/Main.purs

@@ -157,7 +157,7 @@ main = do
          "\nhmab: " <> show (HM.fromFoldable (a <> b)))
 
   log "union = repeated insertion"
-  quickCheck' 100000 $ \(a :: Array (Tuple CollidingInt String)) b ->
+  quickCheck $ \(a :: Array (Tuple CollidingInt String)) b ->
     let m = arbitraryHM a
         n = arbitraryHM b
     in HM.union m n  === foldrWithIndex HM.insert n m
@@ -166,6 +166,20 @@ main = do
         n = arbitraryHM b
     in HM.union n m === foldlWithIndex (\k m' v -> HM.insert k v m') m n
 
+  log "filterWithKey agrees with OrdMap"
+  quickCheck' 10000 $ \(a :: Array (Tuple Int Int)) f ->
+    Array.sort (HM.toArrayBy Tuple (HM.filterWithKey f (HM.fromFoldable a))) ==
+    Array.sort (OM.toUnfoldable (OM.filterWithKey f (OM.fromFoldable a)))
+    <?> (  "        a: " <> show a <>
+         "\n dbg befo: " <> show (HM.debugShow (HM.fromFoldable a)) <>
+         "\n dbg afte: " <> show (HM.debugShow (HM.filterWithKey f (HM.fromFoldable a))) <>
+         "\n filtered: " <> show (Array.sort (HM.toArrayBy Tuple (HM.filterWithKey f (HM.fromFoldable a)))) <>
+         "\n shouldbe: " <> show (Array.sort (OM.toUnfoldable (OM.filterWithKey f (OM.fromFoldable a)))))
+
+  log "filter empty"
+  quickCheck \(a :: Array (Tuple CollidingInt Int)) ->
+    HM.empty === HM.filter (const false) (HM.fromFoldable a)
+
   log "unionWith agrees with OrdMap"
   quickCheck' 10000 $ \(a :: Array (Tuple CollidingInt Int)) (b :: Array (Tuple CollidingInt Int)) c f ->
     Array.sort (HM.toArrayBy Tuple (HM.unionWith f (HM.fromFoldable (a <> c)) (HM.fromFoldable (b <> map (map (_ + 1)) c)))) ===