Direct iterators for PersistentHashMap, APersistentSet, PersistentQueue, and PersistentStructMap, and records. Added new IMapIterable interface for key and val iterators.

Signed-off-by: Stuart Halloway <[email protected]>

Author: puredanger

src/clj/clojure/core_deftype.clj

@@ -215,7 +215,8 @@
                                               (clojure.lang.MapEntry. k# v#))))
                    `(seq [this#] (seq (concat [~@(map #(list `new `clojure.lang.MapEntry (keyword %) %) base-fields)] 
                                               ~'__extmap)))
-                   `(iterator [this#] (clojure.lang.SeqIterator. (.seq this#)))
+                   `(iterator [~gs]
+                        (clojure.lang.RecordIterator. ~gs [~@(map keyword base-fields)] (RT/iter ~'__extmap)))
                    `(assoc [this# k# ~gs]
                      (condp identical? k#
                        ~@(mapcat (fn [fld]

src/jvm/clojure/lang/APersistentMap.java

@@ -194,6 +194,23 @@ public ValSeq withMeta(IPersistentMap meta){
 	}
 }
 
+static final IFn MAKE_ENTRY = new AFn() {
+    public Object invoke(Object key, Object val) {
+        return new MapEntry(key, val);
+    }
+};
+
+static final IFn MAKE_KEY = new AFn() {
+    public Object invoke(Object key, Object val) {
+        return key;
+    }
+};
+
+static final IFn MAKE_VAL = new AFn() {
+    public Object invoke(Object key, Object val) {
+        return val;
+    }
+};
 
 public Object invoke(Object arg1) {
 	return valAt(arg1);

src/jvm/clojure/lang/APersistentSet.java

@@ -170,7 +170,23 @@ public boolean isEmpty(){
 }
 
 public Iterator iterator(){
-	return new SeqIterator(this);
+    if(impl instanceof IMapIterable)
+        return ((IMapIterable)impl).keyIterator();
+    else return new Iterator() {
+        private final Iterator iter = impl.iterator();
+
+        public boolean hasNext() {
+            return iter.hasNext();
+        }
+
+        public Object next() {
+            return ((IMapEntry)iter.next()).key();
+        }
+
+        public void remove() {
+            throw new UnsupportedOperationException();
+        }
+    };
 }
 
 }

src/jvm/clojure/lang/IMapIterable.java

@@ -0,0 +1,26 @@
+/**
+ *   Copyright (c) Rich Hickey. All rights reserved.
+ *   The use and distribution terms for this software are covered by the
+ *   Eclipse Public License 1.0 (http://opensource.org/licenses/eclipse-1.0.php)
+ *   which can be found in the file epl-v10.html at the root of this distribution.
+ *   By using this software in any fashion, you are agreeing to be bound by
+ * 	 the terms of this license.
+ *   You must not remove this notice, or any other, from this software.
+ **/
+
+/* Alex Miller Dec 3, 2014 */
+
+package clojure.lang;
+
+import java.util.Iterator;
+
+/**
+ * Indicate a map can provide more efficient key and val iterators.
+ */
+public interface IMapIterable {
+
+Iterator keyIterator();
+
+Iterator valIterator();
+
+}

src/jvm/clojure/lang/PersistentArrayMap.java

@@ -26,7 +26,7 @@
  * null keys and values are ok, but you won't be able to distinguish a null value via valAt - use contains/entryAt
  */
 
-public class PersistentArrayMap extends APersistentMap implements IObj, IEditableCollection {
+public class PersistentArrayMap extends APersistentMap implements IObj, IEditableCollection, IMapIterable {
 
 final Object[] array;
 static final int HASHTABLE_THRESHOLD = 16;
@@ -281,7 +281,15 @@ static boolean equalKey(Object k1, Object k2){
 }
 
 public Iterator iterator(){
-	return new Iter(array);
+	return new Iter(array,APersistentMap.MAKE_ENTRY);
+}
+
+public Iterator keyIterator(){
+    return new Iter(array,APersistentMap.MAKE_KEY);
+}
+
+public Iterator valIterator() {
+    return new Iter(array,APersistentMap.MAKE_VAL);
 }
 
 public ISeq seq(){
@@ -329,18 +337,20 @@ public Obj withMeta(IPersistentMap meta){
 }
 
 static class Iter implements Iterator{
+    IFn f;
 	Object[] array;
 	int i;
 
 	//for iterator
-	Iter(Object[] array){
-		this(array, -2);
+	Iter(Object[] array, IFn f){
+		this(array, -2, f);
 	}
 
 	//for entryAt
-	Iter(Object[] array, int i){
+	Iter(Object[] array, int i, IFn f){
 		this.array = array;
 		this.i = i;
+        this.f = f;
 	}
 
 	public boolean hasNext(){
@@ -349,7 +359,7 @@ public boolean hasNext(){
 
 	public Object next(){
 		i += 2;
-		return new MapEntry(array[i],array[i+1]);
+		return f.invoke(array[i],array[i+1]);
 	}
 
 	public void remove(){

src/jvm/clojure/lang/PersistentHashMap.java

@@ -25,7 +25,7 @@
  Any errors are my own
  */
 
-public class PersistentHashMap extends APersistentMap implements IEditableCollection, IObj {
+public class PersistentHashMap extends APersistentMap implements IEditableCollection, IObj, IMapIterable {
 
 final int count;
 final INode root;
@@ -173,8 +173,59 @@ public IPersistentMap without(Object key){
 	return new PersistentHashMap(meta(), count - 1, newroot, hasNull, nullValue); 
 }
 
+static final Iterator EMPTY_ITER = new Iterator(){
+    public boolean hasNext(){
+        return false;
+    }
+
+    public Object next(){
+        throw new NoSuchElementException();
+    }
+
+    public void remove(){
+        throw new UnsupportedOperationException();
+    }
+};
+
+private Iterator iterator(final IFn f){
+    final Iterator rootIter = (root == null) ? EMPTY_ITER : root.iterator(f);
+    if(hasNull) {
+        return new Iterator() {
+            private boolean seen = false;
+            public boolean hasNext() {
+                if (!seen)
+                    return true;
+                else
+                    return rootIter.hasNext();
+            }
+
+            public Object next(){
+                if (!seen) {
+                    seen = true;
+                    return f.invoke(null, nullValue);
+                } else
+                    return rootIter.next();
+            }
+
+            public void remove(){
+                throw new UnsupportedOperationException();
+            }
+        };
+    }
+    else
+        return rootIter;
+}
+
 public Iterator iterator(){
-	return new SeqIterator(this);
+    return iterator(APersistentMap.MAKE_ENTRY);
+}
+
+public Iterator keyIterator(){
+    return iterator(APersistentMap.MAKE_KEY);
+}
+
+public Iterator valIterator(){
+    return iterator(APersistentMap.MAKE_VAL);
 }
 
 public Object kvreduce(IFn f, Object init){
@@ -340,6 +391,9 @@ static interface INode extends Serializable {
     public Object kvreduce(IFn f, Object init);
 
 	Object fold(IFn combinef, IFn reducef, IFn fjtask, IFn fjfork, IFn fjjoin);
+
+    // returns the result of (f [k v]) for each iterated element
+    Iterator iterator(IFn f);
 }
 
 final static class ArrayNode implements INode{
@@ -400,6 +454,10 @@ public ISeq nodeSeq(){
 		return Seq.create(array);
 	}
 
+    public Iterator iterator(IFn f){
+        return new Iter(array, f);
+    }
+
     public Object kvreduce(IFn f, Object init){
         for(INode node : array){
             if(node != null){
@@ -563,6 +621,49 @@ public ISeq next() {
 		}
 
 	}
+
+    static class Iter implements Iterator {
+        private final INode[] array;
+        private final IFn f;
+        private int i = 0;
+        private Iterator nestedIter;
+
+        private Iter(INode[] array, IFn f){
+            this.array = array;
+            this.f = f;
+        }
+
+        public boolean hasNext(){
+            while(true)
+            {
+                if(nestedIter != null)
+                    if(nestedIter.hasNext())
+                        return true;
+                    else
+                        nestedIter = null;
+
+                if(i < array.length)
+                {
+                    INode node = array[i++];
+                    if (node != null)
+                        nestedIter = node.iterator(f);
+                }
+                else
+                    return false;
+            }
+        }
+
+        public Object next(){
+            if(hasNext())
+                return nestedIter.next();
+            else
+                throw new NoSuchElementException();
+        }
+
+        public void remove(){
+            throw new UnsupportedOperationException();
+        }
+    }
 }
 
 final static class BitmapIndexedNode implements INode{
@@ -687,6 +788,10 @@ public ISeq nodeSeq(){
 		return NodeSeq.create(array);
 	}
 
+    public Iterator iterator(IFn f){
+        return new NodeIter(array, f);
+    }
+
     public Object kvreduce(IFn f, Object init){
          return NodeSeq.kvreduce(array,f,init);
     }
@@ -879,6 +984,10 @@ public ISeq nodeSeq(){
 		return NodeSeq.create(array);
 	}
 
+    public Iterator iterator(IFn f){
+        return new NodeIter(array, f);
+    }
+
     public Object kvreduce(IFn f, Object init){
          return NodeSeq.kvreduce(array,f,init);
     }
@@ -1107,6 +1216,73 @@ private static int bitpos(int hash, int shift){
 	return 1 << mask(hash, shift);
 }
 
+static final class NodeIter implements Iterator {
+    private static final Object NULL = new Object();
+    final Object[] array;
+    final IFn f;
+    private int i = 0;
+    private Object nextEntry = NULL;
+    private Iterator nextIter;
+
+    NodeIter(Object[] array, IFn f){
+        this.array = array;
+        this.f = f;
+    }
+
+    private boolean advance(){
+        while (i<array.length)
+        {
+            Object key = array[i];
+            Object nodeOrVal = array[i+1];
+            i += 2;
+            if (key != null)
+            {
+                nextEntry = f.invoke(key, nodeOrVal);
+                return true;
+            }
+            else if(nodeOrVal != null)
+            {
+                Iterator iter = ((INode) nodeOrVal).iterator(f);
+                if(iter != null && iter.hasNext())
+                {
+                    nextIter = iter;
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
+    public boolean hasNext(){
+        if (nextEntry != NULL || nextIter != null)
+            return true;
+        return advance();
+    }
+
+    public Object next(){
+        Object ret = nextEntry;
+        if(ret != NULL)
+        {
+            nextEntry = NULL;
+            return ret;
+        }
+        else if(nextIter != null)
+        {
+            ret = nextIter.next();
+            if(! nextIter.hasNext())
+                nextIter = null;
+            return ret;
+        }
+        else if(advance())
+            return next();
+        throw new NoSuchElementException();
+    }
+
+    public void remove(){
+        throw new UnsupportedOperationException();
+    }
+}
+
 static final class NodeSeq extends ASeq {
 	final Object[] array;
 	final int i;