X-Git-Url: http://git.megacz.com/?p=sbp.git;a=blobdiff_plain;f=src%2Fedu%2Fberkeley%2Fsbp%2Futil%2FTopologicalBag.java;h=1b2711d4a1fade2ebb6ecfb82f278a62f2afc0bd;hp=9a18eef7fa98eda82d0bc51d6b33b8732748403b;hb=8c293278cc3351b872f691c0b866367747bad44a;hpb=3eb9fa560dee4f8c3b3af0460ff850d755676d81

diff --git a/src/edu/berkeley/sbp/util/TopologicalBag.java b/src/edu/berkeley/sbp/util/TopologicalBag.java
index 9a18eef..1b2711d 100644
--- a/src/edu/berkeley/sbp/util/TopologicalBag.java
+++ b/src/edu/berkeley/sbp/util/TopologicalBag.java
@@ -29,6 +29,8 @@ public class TopologicalBag<K,V> implements MapBag<Topology<K>,V>, VisitableMap<
 
     public void putAll(Topology<K> k, Iterable<V> vi) { if (vi!=null) for(V v : vi) put(k, v); }
     public void put(Topology<K> t, V v) {
+            // XXX hack
+            if (t instanceof Atom) t = ((Atom)t).underlying();
         for(Topology<K> ht : h.keySet()) {
             if (t.disjoint(ht)) continue;
             if (t.containsAll(ht) && ht.containsAll(t)) {
@@ -135,7 +137,7 @@ public class TopologicalBag<K,V> implements MapBag<Topology<K>,V>, VisitableMap<
         }
     }
 
-    public VisitableMap<K,V> optimize() {
+    public VisitableMap<K,V> optimize(final Functor<K,Integer> f) {
         ArrayList<Long> min_ = new ArrayList<Long>();
         ArrayList<Long> max_ = new ArrayList<Long>();
         ArrayList<Object[]> v_ = new ArrayList<Object[]>();
@@ -157,16 +159,14 @@ public class TopologicalBag<K,V> implements MapBag<Topology<K>,V>, VisitableMap<
         final Object[][]   v = new Object[size][]; v_.toArray(v);
         return new VisitableMap<K,V>() {
             public boolean contains(K k) {
-                IntegerMappable im = (IntegerMappable)k;
-                int asint = im.toInt();
+                int asint = f.invoke(k);
                 for(int i=0; i<size; i++)
                     if (min[i] <= asint && max[i] >= asint && v[i].length > 0)
                         return true;
                 return false;
             }
             public <B,C> void invoke(K k, Invokable<V,B,C> ivbc, B b, C c) {
-                IntegerMappable im = (IntegerMappable)k;
-                int asint = im.toInt();
+                int asint = f.invoke(k);
                 for(int i=0; i<size; i++) {
                     if (min[i] <= asint && max[i] >= asint) {
                         Object[] arr = v[i];