added reduce-kv

Author: richhickey

clojure.iml

@@ -1,5 +1,10 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <module org.jetbrains.idea.maven.project.MavenProjectsManager.isMavenModule="true" type="JAVA_MODULE" version="4">
+  <component name="FacetManager">
+    <facet type="Clojure" name="Clojure">
+      <configuration />
+    </facet>
+  </component>
   <component name="NewModuleRootManager" LANGUAGE_LEVEL="JDK_1_5" inherit-compiler-output="false">
     <output url="file://$MODULE_DIR$/target/classes" />
     <output-test url="file://$MODULE_DIR$/target/test-classes" />

src/clj/clojure/core.clj

@@ -6032,6 +6032,44 @@
      (let [s (seq coll)]
        (clojure.core.protocols/internal-reduce s f val))))
 
+(extend-protocol clojure.core.protocols/IKVReduce
+ ;;slow path default
+ clojure.lang.IPersistentMap
+ (kv-reduce 
+  [amap f init]
+  (reduce (fn [ret [k v]] (f ret k v)) init amap))
+
+ clojure.lang.PersistentHashMap
+ (kv-reduce 
+  [amap f init]
+  (.kvreduce amap f init))
+
+ clojure.lang.PersistentArrayMap
+ (kv-reduce 
+  [amap f init]
+  (.kvreduce amap f init))
+
+ clojure.lang.PersistentTreeMap
+ (kv-reduce 
+  [amap f init]
+  (.kvreduce amap f init))
+
+ clojure.lang.PersistentVector
+ (kv-reduce 
+  [vec f init]
+  (.kvreduce vec f init)))
+
+(defn reduce-kv
+  "Reduces an associative collection. f should be a function of 3
+  arguments. Returns the result of applying f to init, the first key
+  and the first value in coll, then applying f to that result and the
+  2nd key and value, etc. If coll contains no entries, returns init
+  and f is not called. Note that reduce-kv is supported on vectors,
+  where the keys will be the ordinals."  
+  {:added "1.4"}
+  ([f init coll]
+     (clojure.core.protocols/kv-reduce coll f init)))
+
 (defn into
   "Returns a new coll consisting of to-coll with all of the items of
   from-coll conjoined."

src/clj/clojure/core/protocols.clj

@@ -8,6 +8,8 @@
 
 (ns clojure.core.protocols)
 
+(set! *warn-on-reflection* true)
+
 (defprotocol InternalReduce
   "Protocol for concrete seq types that can reduce themselves
    faster than first/next recursion. Called by clojure.core/reduce."
@@ -92,3 +94,5 @@
 
 (emit-array-impls int long float double byte char boolean)
 
+(defprotocol IKVReduce
+  (kv-reduce [amap f init]))

src/jvm/clojure/lang/PersistentArrayMap.java

@@ -279,6 +279,13 @@ public void remove(){
 
 }
 
+public Object kvreduce(IFn f, Object init){
+    for(int i=0;i < array.length;i+=2){
+        init = f.invoke(init, array[i], array[i+1]);
+        }
+    return init;
+}
+
 public ITransientMap asTransient(){
 	return new TransientArrayMap(array);
 }

src/jvm/clojure/lang/PersistentHashMap.java

@@ -178,6 +178,14 @@ public Iterator iterator(){
 	return new SeqIterator(seq());
 }
 
+public Object kvreduce(IFn f, Object init){
+    init = hasNull?f.invoke(init,null,nullValue):init;
+    if(root != null){
+        return root.kvreduce(f,init);
+    }
+    return init;
+}
+
 public int count(){
 	return count;
 }
@@ -311,6 +319,9 @@ static interface INode extends Serializable {
 	INode assoc(AtomicReference<Thread> edit, int shift, int hash, Object key, Object val, Box addedLeaf);
 
 	INode without(AtomicReference<Thread> edit, int shift, int hash, Object key, Box removedLeaf);
+
+    public Object kvreduce(IFn f, Object init);
+
 }
 
 final static class ArrayNode implements INode{
@@ -371,6 +382,16 @@ public ISeq nodeSeq(){
 		return Seq.create(array);
 	}
 
+    public Object kvreduce(IFn f, Object init){
+        for(INode node : array)
+            {
+            if(node != null)
+                init = node.kvreduce(f,init);
+            }
+        return init;
+    }
+
+
 	private ArrayNode ensureEditable(AtomicReference<Thread> edit){
 		if(this.edit == edit)
 			return this;
@@ -600,6 +621,11 @@ public ISeq nodeSeq(){
 		return NodeSeq.create(array);
 	}
 
+    public Object kvreduce(IFn f, Object init){
+         return NodeSeq.kvreduce(array,f,init);
+    }
+
+
 	private BitmapIndexedNode ensureEditable(AtomicReference<Thread> edit){
 		if(this.edit == edit)
 			return this;
@@ -784,6 +810,10 @@ public ISeq nodeSeq(){
 		return NodeSeq.create(array);
 	}
 
+    public Object kvreduce(IFn f, Object init){
+         return NodeSeq.kvreduce(array,f,init);
+    }
+
 	public int findIndex(Object key){
 		for(int i = 0; i < 2*count; i+=2)
 			{
@@ -1017,6 +1047,21 @@ static ISeq create(Object[] array) {
 		return create(array, 0, null);
 	}
 
+    static public Object kvreduce(Object[] array, IFn f, Object init){
+         for(int i=0;i<array.length;i+=2)
+             {
+             if(array[i] != null)
+                 init = f.invoke(init, array[i], array[i+1]);
+             else
+                 {
+                 INode node = (INode) array[i+1];
+                 if(node != null)
+                     init = node.kvreduce(f,init);
+                 }
+             }
+        return init;
+    }
+
 	private static ISeq create(Object[] array, int i, ISeq s) {
 		if(s != null)
 			return new NodeSeq(null, array, i, s);

src/jvm/clojure/lang/PersistentTreeMap.java

@@ -205,6 +205,13 @@ public NodeIterator iterator(){
 	return new NodeIterator(tree, true);
 }
 
+public Object kvreduce(IFn f, Object init){
+    if(tree != null)
+        return tree.kvreduce(f,init);
+    return init;
+}
+
+
 public NodeIterator reverseIterator(){
 	return new NodeIterator(tree, false);
 }
@@ -529,6 +536,16 @@ Node balanceRight(Node parent){
 
 	abstract Node replace(Object key, Object val, Node left, Node right);
 
+    public Object kvreduce(IFn f, Object init){
+        init = f.invoke(init, key(), val());
+        if(left() != null)
+            init = left().kvreduce(f, init);
+        if(right() != null)
+            init = right().kvreduce(f, init);
+        return init;
+    }
+
+
 }
 
 static class Black extends Node{

src/jvm/clojure/lang/PersistentVector.java

@@ -232,6 +232,17 @@ public ISeq seq(){
 	return chunkedSeq();
 }
 
+public Object kvreduce(IFn f, Object init){
+    int step = 0;
+    for(int i=0;i<cnt;i+=step){
+        Object[] array = arrayFor(i);
+        for(int j =0;j<array.length;++j)
+            init = f.invoke(init,j+i,array[j]);
+        step = array.length;
+    }
+    return init;
+}
+
 static public final class ChunkedSeq extends ASeq implements IChunkedSeq{
 
 	public final PersistentVector vec;