subseq and rsubseq support for sorted maps and sets

5 files changed, 147 insertions, 5 deletions

diff --git a/src/boot.clj b/src/boot.clj
index 01ae3168..704a0628 100644
--- a/src/boot.clj
+++ b/src/boot.clj
@@ -2340,3 +2340,41 @@ not-every? (comp not every?))
                       [exprs nil])]
       `(binding [*math-context* (java.math.MathContext. ~precision ~@rm)]
          ~@body))) 
+
+(defn bound-fn
+  {:private true}
+  [#^clojure.lang.Sorted sc test key]
+  (fn [e]
+    (test (.. sc comparator (compare (. sc entryKey e) key)) 0)))
+  
+(defn subseq
+  "sc must be a sorted collection, test(s) one of <, <=, > or
+  >=. Returns a seq of those entries with keys ek for
+  which (test (.. sc comparator (compare ek key)) 0) is true"  
+  ([#^clojure.lang.Sorted sc test key]
+   (let [include (bound-fn sc test key)]
+     (if (#{> >=} test)
+       (when-let [e :as s] (. sc seqFrom key true)
+         (if (include e) s (rest s)))
+       (take-while include (. sc seq true)))))     
+  ([#^clojure.lang.Sorted sc start-test start-key end-test end-key]
+   (when-let [e :as s] (. sc seqFrom start-key true)
+     (take-while (bound-fn sc end-test end-key)
+                 (if ((bound-fn sc start-test start-key) e) s (rest s))))))
+
+(defn rsubseq
+  "sc must be a sorted collection, test(s) one of <, <=, > or
+  >=. Returns a reverse seq of those entries with keys ek for
+  which (test (.. sc comparator (compare ek key)) 0) is true"
+  ([#^clojure.lang.Sorted sc test key]
+   (let [include (bound-fn sc test key)]
+     (if (#{< <=} test)
+       (when-let [e :as s] (. sc seqFrom key false)
+         (if (include e) s (rest s)))
+       (take-while include (. sc seq false)))))     
+  ([#^clojure.lang.Sorted sc start-test start-key end-test end-key]
+   (when-let [e :as s] (. sc seqFrom end-key false)
+     (take-while (bound-fn sc start-test start-key)
+                 (if ((bound-fn sc end-test end-key) e) s (rest s))))))
+
+
diff --git a/src/jvm/clojure/lang/PersistentTreeMap.java b/src/jvm/clojure/lang/PersistentTreeMap.java
index 4326a8e4..fb9df3ad 100644
--- a/src/jvm/clojure/lang/PersistentTreeMap.java
+++ b/src/jvm/clojure/lang/PersistentTreeMap.java
@@ -22,7 +22,7 @@ import java.util.*;
  * See Okasaki, Kahrs, Larsen et al
  */
 
-public class PersistentTreeMap extends APersistentMap implements Reversible{
+public class PersistentTreeMap extends APersistentMap implements Reversible, Sorted{
 
 public final Comparator comp;
 public final Node tree;
@@ -31,7 +31,7 @@ public final int _count;
 final static public PersistentTreeMap EMPTY = new PersistentTreeMap();
 
 public PersistentTreeMap(){
-	this(null);
+	this(RT.DEFAULT_COMPARATOR);
 }
 
 public PersistentTreeMap withMeta(IPersistentMap meta){
@@ -132,6 +132,59 @@ public ISeq rseq() throws Exception{
 	return null;
 }
 
+public Comparator comparator(){
+	return comp;
+}
+
+public Object entryKey(Object entry){
+	return ((IMapEntry) entry).key();
+}
+
+public ISeq seq(boolean ascending){
+	if(_count > 0)
+		return Seq.create(tree, ascending, _count);
+	return null;
+}
+
+public ISeq seqFrom(Object key, boolean ascending){
+	if(_count > 0)
+		{
+		ISeq stack = null;
+		Node t = tree;
+		while(t != null)
+			{
+			int c = doCompare(key, t.key);
+			if(c == 0)
+				{
+				stack = RT.cons(t, stack);
+				return new Seq(stack, ascending);
+				}
+			else if(ascending)
+				{
+				if(c < 0)
+					{
+					stack = RT.cons(t, stack);
+					t = t.left();
+					}
+				else
+					t = t.right();
+				}
+			else
+				{
+				if(c > 0)
+					{
+					stack = RT.cons(t, stack);
+					t = t.right();
+					}
+				else
+					t = t.left();
+				}
+			}
+		if(stack != null)
+			return new Seq(stack, ascending);
+		}
+	return null;
+}
 
 public NodeIterator iterator(){
 	return new NodeIterator(tree, true);
@@ -230,9 +283,9 @@ public Node entryAt(Object key){
 }
 
 public int doCompare(Object k1, Object k2){
-	if(comp != null)
+//	if(comp != null)
 		return comp.compare(k1, k2);
-	return ((Comparable) k1).compareTo(k2);
+//	return ((Comparable) k1).compareTo(k2);
 }
 
 Node add(Node t, Object key, Object val, Box found){
diff --git a/src/jvm/clojure/lang/PersistentTreeSet.java b/src/jvm/clojure/lang/PersistentTreeSet.java
index ae96940d..8ffdb0d9 100644
--- a/src/jvm/clojure/lang/PersistentTreeSet.java
+++ b/src/jvm/clojure/lang/PersistentTreeSet.java
@@ -14,8 +14,9 @@ package clojure.lang;
 
 import java.util.List;
 import java.util.Iterator;
+import java.util.Comparator;
 
-public class PersistentTreeSet extends APersistentSet implements Reversible{
+public class PersistentTreeSet extends APersistentSet implements Reversible, Sorted{
 static public final PersistentTreeSet EMPTY = new PersistentTreeSet(null, PersistentTreeMap.EMPTY);
 
 public static PersistentTreeSet create(Object... init){
@@ -69,4 +70,23 @@ public ISeq rseq() throws Exception{
 public PersistentTreeSet withMeta(IPersistentMap meta){
 	return new PersistentTreeSet(meta, impl);
 }
+
+public Comparator comparator(){
+	return ((Sorted)impl).comparator();
+}
+
+public Object entryKey(Object entry){
+	return entry;
+}
+
+public ISeq seq(boolean ascending){
+	PersistentTreeMap m = (PersistentTreeMap) impl;
+	return RT.keys(m.seq(ascending));
+}
+
+public ISeq seqFrom(Object key, boolean ascending){
+	PersistentTreeMap m = (PersistentTreeMap) impl;
+	return RT.keys(m.seqFrom(key,ascending));
+}
+
 }
diff --git a/src/jvm/clojure/lang/RT.java b/src/jvm/clojure/lang/RT.java
index 0a7543d4..99df0982 100644
--- a/src/jvm/clojure/lang/RT.java
+++ b/src/jvm/clojure/lang/RT.java
@@ -189,6 +189,12 @@ public static List<String> processCommandLine(String[] args){
 //static Var NS_IMPORTS = Var.intern(CLOJURE_NS,Symbol.create("*ns-imports*"), IMPORTS);
 //static Var NS_REFERS = Var.intern(CLOJURE_NS,Symbol.create("*ns-refers*"), REFERS);
 static public final Object[] EMPTY_ARRAY = new Object[]{};
+static public final Comparator DEFAULT_COMPARATOR = new Comparator(){
+	public int compare(Object o1, Object o2){
+		return ((Comparable) o1).compareTo(o2);
+	}
+};
+
 //static public final Character[] chars;
 static AtomicInteger id = new AtomicInteger(1);
 
diff --git a/src/jvm/clojure/lang/Sorted.java b/src/jvm/clojure/lang/Sorted.java
new file mode 100644
index 00000000..28df9825
--- /dev/null
+++ b/src/jvm/clojure/lang/Sorted.java
@@ -0,0 +1,25 @@
+/**
+ *   Copyright (c) Rich Hickey. All rights reserved.
+ *   The use and distribution terms for this software are covered by the
+ *   Common Public License 1.0 (http://opensource.org/licenses/cpl.php)
+ *   which can be found in the file CPL.TXT 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.
+ **/
+
+/* rich Apr 15, 2008 */
+
+package clojure.lang;
+
+import java.util.Comparator;
+
+public interface Sorted{
+Comparator comparator();
+
+Object entryKey(Object entry);
+
+ISeq seq(boolean ascending);
+
+ISeq seqFrom(Object key, boolean ascending);
+}