Add support for sorted-set and sorted-map types

README.md

@@ -51,6 +51,11 @@ nippy/stress-data
  :keyword      :keyword
  :ns-keyword   ::keyword
 
+ :queue        (-> (PersistentQueue/EMPTY) (conj :a :b :c :d :e :f :g))
+ :queue-empty  (PersistentQueue/EMPTY)
+ :sorted-set   (sorted-set 1 2 3 4 5)
+ :sorted-map   (sorted-map :b 2 :a 1 :d 4 :c 3)
+
  :list         (list 1 2 3 4 5 (list 6 7 8 (list 9 10)))
  :list-quoted  '(1 2 3 4 5 (6 7 8 (9 10)))
  :list-empty   (list)
@@ -61,8 +66,6 @@ nippy/stress-data
  :set          #{1 2 3 4 5 #{6 7 8 #{9 10}}}
  :set-empty    #{}
  :meta         (with-meta {:a :A} {:metakey :metaval})
- :queue        (-> (PersistentQueue/EMPTY) (conj :a :b :c :d :e :f :g))
- :queue-empty  (PersistentQueue/EMPTY)
  :coll         (repeatedly 1000 rand)
 
  :byte         (byte 16)

src/taoensso/nippy.clj

@@ -5,9 +5,9 @@
   (:require [taoensso.nippy.utils :as utils])
   (:import  [java.io DataInputStream DataOutputStream ByteArrayOutputStream
              ByteArrayInputStream]
-            [clojure.lang IPersistentList IPersistentVector IPersistentMap
-             IPersistentSet PersistentQueue IPersistentCollection Keyword
-             BigInt Ratio]))
+            [clojure.lang Keyword BigInt Ratio PersistentQueue PersistentTreeMap
+             PersistentTreeSet IPersistentList IPersistentVector IPersistentMap
+             IPersistentSet IPersistentCollection]))
 
 ;;;; Define type IDs
 
@@ -30,6 +30,8 @@
 (def ^:const id-meta       (int 25))
 (def ^:const id-queue      (int 26))
 (def ^:const id-map        (int 27))
+(def ^:const id-sorted-set (int 28))
+(def ^:const id-sorted-map (int 29))
 
 (def ^:const id-byte       (int 40))
 (def ^:const id-short      (int 41))
@@ -91,13 +93,21 @@
        ~@body)))
 
 (defmacro coll-freezer
-  "Helper to extend Freezable protocol to simple collection types."
+  "Extends Freezable to simple collection types."
   [type id & body]
-  `(freezer
-    ~type ~id
+  `(freezer ~type ~id
     (.writeInt ~'s (count ~'x)) ; Encode collection length
     (doseq [i# ~'x] (freeze-to-stream!* ~'s i#))))
 
+(defmacro kv-freezer
+  "Extends Freezable to key-value collection types."
+  [type id & body]
+  `(freezer ~type ~id
+    (.writeInt ~'s (* 2 (count ~'x))) ; Encode num kvs
+    (doseq [[k# v#] ~'x]
+      (freeze-to-stream!* ~'s k#)
+      (freeze-to-stream!* ~'s v#))))
+
 (freezer (Class/forName "[B") id-bytes   (write-bytes! s x))
 (freezer nil                  id-nil)
 (freezer Boolean              id-boolean (.writeBoolean s x))
@@ -110,15 +120,14 @@
 
 (declare freeze-to-stream!*)
 
+(coll-freezer PersistentQueue       id-queue)
+(coll-freezer PersistentTreeSet     id-sorted-set)
+(kv-freezer   PersistentTreeMap     id-sorted-map)
+
 (coll-freezer IPersistentList       id-list)
 (coll-freezer IPersistentVector     id-vector)
-(freezer      IPersistentMap        id-map
-              (.writeInt s (* 2 (count x))) ; Encode num kvs
-              (doseq [[k v] x]
-                (freeze-to-stream!* s k)
-                (freeze-to-stream!* s v)))
 (coll-freezer IPersistentSet        id-set)
-(coll-freezer PersistentQueue       id-queue)
+(kv-freezer   IPersistentMap        id-map)
 (coll-freezer IPersistentCollection id-coll) ; Must be LAST collection freezer!
 
 (freezer Byte       id-byte    (.writeByte s x))
@@ -171,12 +180,12 @@
 (declare thaw-from-stream!*)
 
 (defn coll-thaw!
-  "Helper to thaw simple collection types."
+  "Thaws simple collection types."
   [^DataInputStream s]
   (repeatedly (.readInt s) (partial thaw-from-stream!* s)))
 
-(defn coll-thaw-pairs!
-  "Helper to thaw pair-based collection types (e.g. hash maps)."
+(defn coll-thaw-kvs!
+  "Thaws key-value collection types."
   [^DataInputStream s]
   (repeatedly (/ (.readInt s) 2)
               (fn [] [(thaw-from-stream!* s) (thaw-from-stream!* s)])))
@@ -196,12 +205,15 @@
      id-string  (String. (read-bytes! s) "UTF-8")
      id-keyword (keyword (.readUTF s))
 
+     id-queue      (into (PersistentQueue/EMPTY) (coll-thaw! s))
+     id-sorted-set (into (sorted-set) (coll-thaw! s))
+     id-sorted-map (into (sorted-map) (coll-thaw-kvs! s))
+
      id-list    (apply list (coll-thaw! s)) ; TODO OOMs for big colls
      id-vector  (into  [] (coll-thaw! s))
      id-set     (into #{} (coll-thaw! s))
-     id-map     (into  {} (coll-thaw-pairs! s))
+     id-map     (into  {} (coll-thaw-kvs! s))
      id-coll    (doall (coll-thaw! s))
-     id-queue   (into  (PersistentQueue/EMPTY) (coll-thaw! s))
 
      id-meta (let [m (thaw-from-stream!* s)] (with-meta (thaw-from-stream!* s) m))
 
@@ -252,9 +264,7 @@
   (let [support-tagged-literals?
         (utils/version-sufficient? (clojure-version) "1.4.0")]
 
-    {;; Breaks reader, roundtrip equality
-     :bytes        (byte-array [(byte 1) (byte 2) (byte 3)])
-
+    {:bytes        (byte-array [(byte 1) (byte 2) (byte 3)])
      :nil          nil
      :boolean      true
 
@@ -264,6 +274,11 @@
      :keyword      :keyword
      :keyword-ns   ::keyword
 
+     :queue        (-> (PersistentQueue/EMPTY) (conj :a :b :c :d :e :f :g))
+     :queue-empty  (PersistentQueue/EMPTY)
+     :sorted-set   (sorted-set 1 2 3 4 5)
+     :sorted-map   (sorted-map :b 2 :a 1 :d 4 :c 3)
+
      :list         (list 1 2 3 4 5 (list 6 7 8 (list 9 10)))
      :list-quoted  '(1 2 3 4 5 (6 7 8 (9 10)))
      :list-empty   (list)
@@ -275,10 +290,6 @@
      :set-empty    #{}
      :meta         (with-meta {:a :A} {:metakey :metaval})
 
-     ;; Breaks reader
-     :queue        (-> (PersistentQueue/EMPTY) (conj :a :b :c :d :e :f :g))
-     :queue-empty  (PersistentQueue/EMPTY)
-
      :coll         (repeatedly 1000 rand)
 
      :byte         (byte 16)