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

diff --git a/src/edu/berkeley/sbp/util/HashMapBag.java b/src/edu/berkeley/sbp/util/HashMapBag.java
index 6f5585d..210290d 100644
--- a/src/edu/berkeley/sbp/util/HashMapBag.java
+++ b/src/edu/berkeley/sbp/util/HashMapBag.java
@@ -1,3 +1,5 @@
+// Copyright 2006-2007 all rights reserved; see LICENSE file for BSD-style license
+
 package edu.berkeley.sbp.util;
 import java.util.*;
 
@@ -9,7 +11,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 +28,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 +45,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;
 }