X-Git-Url: http://git.megacz.com/?a=blobdiff_plain;f=src%2Fedu%2Fberkeley%2Fsbp%2Futil%2FHashMapBag.java;h=73e2ccecbe2eca25d097b9b72d5af8af4fa1309c;hb=3473141880127c6aa52d094d51739e5eb67f1703;hp=6f5585d4d696cb272e9481ce8785fe6922ac45bf;hpb=842f3c9b981b35721bb50d49e85c11085b2040a3;p=sbp.git

diff --git a/src/edu/berkeley/sbp/util/HashMapBag.java b/src/edu/berkeley/sbp/util/HashMapBag.java
index 6f5585d..73e2cce 100644
--- a/src/edu/berkeley/sbp/util/HashMapBag.java
+++ b/src/edu/berkeley/sbp/util/HashMapBag.java
@@ -9,7 +9,9 @@ public final class HashMapBag<K,V> implements MapBag<K,V> {
     public void add(K k, V v) {
         HashSet<V> hs = hm.get(k);
         if (hs==null) hm.put(k, hs = new HashSet<V>());
+        size -= hs.size();
         hs.add(v);
+        size += hs.size();
     }
 
     public void addAll(K k, Iterable<V> iv) {
@@ -24,7 +26,15 @@ public final class HashMapBag<K,V> implements MapBag<K,V> {
 
     public void remove(K k, V v) {
         if (hm.get(k)==null) return;
-        hm.get(k).remove(v);
+        HashSet<V> hs = hm.get(k);
+        if (hs==null) return;
+        size -= hs.size();
+        hs.remove(v);
+        size += hs.size();
+    }
+
+    public void removeAll(K k, Iterable<V> iv) {
+        for(V v : iv) remove(k, v);
     }
 
     public void clear() { hm.clear(); }
@@ -33,5 +43,14 @@ public final class HashMapBag<K,V> implements MapBag<K,V> {
         return hm.get(k)!=null && hm.get(k).contains(v);
     }
 
+    public void addAll(HashMapBag<K,V> hmb) {
+        for(K k : hmb) addAll(k, hmb.getAll(k));
+    }
+    public void removeAll(HashMapBag<K,V> hmb) {
+        for(K k : hmb) removeAll(k, hmb.getAll(k));
+    }
+
     public Iterator<K> iterator() { return hm.keySet().iterator(); }
+    public int size() { return size; }
+    private int size = 0;
 }