makes HAMT iteration faster

vavr-io · May 25, 2016 · e1235db · e1235db
1 parent 49a9aa6
commit e1235db
Show file tree

Hide file tree

Showing 2 changed files with 105 additions and 21 deletions.
diff --git a/javaslang/src/main/java/javaslang/collection/HashArrayMappedTrie.java b/javaslang/src/main/java/javaslang/collection/HashArrayMappedTrie.java
@@ -39,10 +39,10 @@ static <K, V> HashArrayMappedTrie<K, V> empty() {
 
     HashArrayMappedTrie<K, V> remove(K key);
 
-    // this is a javaslang.collection.Iterator!
     @Override
     Iterator<Tuple2<K, V>> iterator();
 
+    Iterator<K> keysIterator();
 }
 
 interface HashArrayMappedTrieModule {
@@ -51,6 +51,87 @@ enum Action {
         PUT, REMOVE
     }
 
+    class It<K, V> extends AbstractIterator<LeafNode<K, V>> {
+
+        private final static int MAX_LEVELS = Integer.SIZE / AbstractNode.SIZE + 1;
+
+        private final int total;
+        private final Object[] nodes = new Object[MAX_LEVELS];
+        private final int[] indexes = new int[MAX_LEVELS];
+
+        private int level;
+        private int ptr = 0;
+
+        It(AbstractNode<K, V> root) {
+            total = root.size();
+            level = downstairs(nodes, indexes, root, 0);
+        }
+
+        @Override
+        public boolean hasNext() {
+            return ptr < total;
+        }
+
+        @SuppressWarnings("unchecked")
+        @Override
+        protected LeafNode<K, V> getNext() {
+            Object node = nodes[level];
+            while (!(node instanceof LeafNode)) {
+                node = findNextLeaf();
+            }
+            ptr++;
+            if (node instanceof LeafList) {
+                LeafList<K, V> leaf = (LeafList<K, V>) node;
+                nodes[level] = leaf.tail;
+                return leaf;
+            } else {
+                nodes[level] = EmptyNode.instance();
+                return (LeafSingleton<K, V>) node;
+            }
+        }
+
+        @SuppressWarnings("unchecked")
+        private Object findNextLeaf() {
+            AbstractNode<K, V> node = null;
+            while (level > 0) {
+                level--;
+                indexes[level]++;
+                node = getChild((AbstractNode<K, V>) nodes[level], indexes[level]);
+                if (node != null) {
+                    break;
+                }
+            }
+            level = downstairs(nodes, indexes, node, level + 1);
+            return nodes[level];
+        }
+
+        private static <K, V> int downstairs(Object[] nodes, int[] indexes, AbstractNode<K, V> root, int level) {
+            while (true) {
+                nodes[level] = root;
+                indexes[level] = 0;
+                root = getChild(root, 0);
+                if (root == null) {
+                    break;
+                } else {
+                    level++;
+                }
+            }
+            return level;
+        }
+
+        @SuppressWarnings("unchecked")
+        private static <K, V> AbstractNode<K, V> getChild(AbstractNode<K, V> node, int index) {
+            if (node instanceof IndexedNode) {
+                Object[] subNodes = ((IndexedNode<K, V>) node).subNodes;
+                return index < subNodes.length ? (AbstractNode<K, V>) subNodes[index] : null;
+            } else if (node instanceof ArrayNode) {
+                ArrayNode<K, V> arrayNode = (ArrayNode<K, V>) node;
+                return index < 32 ? (AbstractNode<K, V>) arrayNode.subNodes[index] : null;
+            }
+            return null;
+        }
+    }
+
     /**
      * An abstract base class for nodes of a HAMT.
      *
@@ -101,6 +182,20 @@ static Object[] insert(Object[] arr, int index, Object newElem) {
 
         abstract AbstractNode<K, V> modify(int shift, int keyHashCode, K key, V value, Action action);
 
+        Iterator<LeafNode<K, V>> nodes() {
+            return new It<>(this);
+        }
+
+        @Override
+        public Iterator<Tuple2<K, V>> iterator() {
+            return nodes().map(node -> Tuple.of(node.key(), node.value()));
+        }
+
+        @Override
+        public Iterator<K> keysIterator() {
+            return nodes().map(LeafNode::key);
+        }
+
         @Override
         public Option<V> get(K key) {
             return lookup(0, Objects.hashCode(key), key);
@@ -184,7 +279,7 @@ public int size() {
         }
 
         @Override
-        public Iterator<Tuple2<K, V>> iterator() {
+        public Iterator<LeafNode<K, V>> nodes() {
             return Iterator.empty();
         }
 
@@ -282,9 +377,8 @@ public int size() {
         }
 
         @Override
-        public Iterator<Tuple2<K, V>> iterator() {
-            Tuple2<K, V> tuple = Tuple.of(key, value);
-            return Iterator.of(tuple);
+        public Iterator<LeafNode<K, V>> nodes() {
+            return Iterator.of(this);
         }
 
         @Override
@@ -398,8 +492,8 @@ public int size() {
         }
 
         @Override
-        public Iterator<Tuple2<K, V>> iterator() {
-            return new AbstractIterator<Tuple2<K, V>>() {
+        public Iterator<LeafNode<K, V>> nodes() {
+            return new AbstractIterator<LeafNode<K, V>>() {
                 LeafNode<K, V> node = LeafList.this;
 
                 @Override
@@ -408,14 +502,14 @@ public boolean hasNext() {
                 }
 
                 @Override
-                public Tuple2<K, V> getNext() {
-                    Tuple2<K, V> tuple = Tuple.of(node.key(), node.value());
+                public LeafNode<K, V> getNext() {
+                    LeafNode<K, V> result = node;
                     if (node instanceof LeafSingleton) {
                         node = null;
                     } else {
                         node = ((LeafList<K, V>) node).tail;
                     }
-                    return tuple;
+                    return result;
                 }
             };
         }
@@ -541,11 +635,6 @@ public int size() {
             return size;
         }
 
-        @Override
-        public Iterator<Tuple2<K, V>> iterator() {
-            return Iterator.concat(Array.wrap(subNodes));
-        }
-
         @Override
         public int hashCode() {
             return Objects.hash(subNodes);
@@ -626,11 +715,6 @@ public int size() {
             return size;
         }
 
-        @Override
-        public Iterator<Tuple2<K, V>> iterator() {
-            return Iterator.concat(Array.wrap(subNodes));
-        }
-
         @Override
         public int hashCode() {
             return Objects.hash(subNodes);

diff --git a/javaslang/src/main/java/javaslang/collection/HashSet.java b/javaslang/src/main/java/javaslang/collection/HashSet.java
@@ -619,7 +619,7 @@ public boolean isTraversableAgain() {
 
     @Override
     public Iterator<T> iterator() {
-        return tree.iterator().map(t -> t._1);
+        return tree.keysIterator();
     }
 
     @Override