Permalink
Browse files

added reduce-kv

  • Loading branch information...
1 parent 3bd2af4 commit 68ec7ed0d0d88e346b8e069281a116883bc6494b @richhickey richhickey committed Apr 6, 2012
View
@@ -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" />
View
@@ -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."
@@ -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]))
@@ -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);
}
@@ -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 @@ void ensureEditable(){
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);
@@ -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{
@@ -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;

0 comments on commit 68ec7ed

Please sign in to comment.