From c2a964932cda5e335e0582e73bac464b6871b24c Mon Sep 17 00:00:00 2001
From: Peter Taoussanis <ptaoussanis@gmail.com>
Date: Fri, 2 Aug 2013 14:41:59 +0700
Subject: [PATCH] Add `custom-freezer` macro for easier Freezable extension

---
 src/taoensso/nippy.clj             | 44 +++++++++++++++++++-----------
 test/taoensso/nippy/tests/main.clj | 12 ++++----
 2 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/src/taoensso/nippy.clj b/src/taoensso/nippy.clj
index 001991f..ce7c0cc 100644
--- a/src/taoensso/nippy.clj
+++ b/src/taoensso/nippy.clj
@@ -68,9 +68,9 @@
 (def ^:const id-old-keyword (int 12)) ; as of 2.0.0-alpha5, for str consistecy
 
 ;;;; Freezing
-(defprotocol Freezable (freeze-to-stream* [this ^DataOutputStream stream]))
+(defprotocol Freezable (freeze-to-stream* [this stream]))
 
-(defmacro ^:private write-id    [s id] `(.writeByte ~s ~id))
+(defmacro           write-id    [s id] `(.writeByte ~s ~id))
 (defmacro ^:private write-bytes [s ba]
   `(let [s# ~s ba# ~ba]
      (let [size# (alength ba#)]
@@ -84,7 +84,7 @@
   [s x]
   `(let [x# ~x s# ~s]
      (when-let [m# (meta x#)]
-       (write-id s# ~id-meta)
+       (write-id  s# ~id-meta)
        (freeze-to-stream* m# s#))
      (try (freeze-to-stream* x# s#)
           (catch java.lang.IllegalArgumentException _#
@@ -103,11 +103,27 @@
   "Helper to extend Freezable protocol."
   [type id & body]
   `(extend-type ~type
-     ~'Freezable
+     Freezable
      (~'freeze-to-stream* [~'x ~(with-meta 's {:tag 'DataOutputStream})]
        (write-id ~'s ~id)
        ~@body)))
 
+(defmacro custom-freezer
+  "Helper to extend Freezable protocol to custom types with id ∈[1, 128]:
+  (defrecord MyType [data])
+  (custom-freezer MyType 1 x s (.writeUTF s (:data x)))"
+  [type id x data-output-stream & body]
+  (assert (and (>= id 1) (<= id 128)))
+  `(extend-type ~type
+     Freezable
+     (~'freeze-to-stream* [~x ~(with-meta data-output-stream
+                                 {:tag 'DataOutputStream})]
+       (write-id ~data-output-stream ~(int (- id)))
+       ~@body)))
+
+(comment (defrecord MyType [data])
+         (custom-freezer MyType 1 x s (.writeUTF s (:data x))))
+
 (defmacro ^:private coll-freezer
   "Extends Freezable to simple collection types."
   [type id & body]
@@ -176,14 +192,8 @@
 
 (defn freeze
   "Serializes arg (any Clojure data type) to a byte array. Set :legacy-mode to
-  true to produce bytes readble by Nippy < 2.x.
-
-  For custom types extend the Clojure reader or Nippy's `Freezable` protocol:
-  (defrecord MyType [data]
-    nippy/Freezable
-    (freeze-to-stream* [x stream]
-      (.writeByte stream -1) ; Custom type id ∈ [-128, -1]
-      (.writeUTF  stream (:data x))))"
+  true to produce bytes readble by Nippy < 2.x. For custom types extend the
+  Clojure reader or see `custom-freezer`."
   ^bytes [x & [{:keys [print-dup? password compressor encryptor legacy-mode]
                 :or   {print-dup? true
                        compressor snappy-compressor
@@ -270,13 +280,14 @@
      id-old-keyword (keyword (.readUTF s))
 
      ;;; Custom types
-     (or (when-let [reader (get readers type-id)]
+     (or (when-let [reader (get readers (- type-id))]
            (try (reader s)
                 (catch Exception e
                   (throw (Exception. (str "Reader exception for custom type ID: "
-                                          type-id) e)))))
+                                          (- type-id)) e)))))
          (if (neg? type-id)
-           (throw (Exception. (str "No reader provided for custom type ID: " type-id)))
+           (throw (Exception. (str "No reader provided for custom type ID: "
+                                   (- type-id))))
            (throw (Exception. (str "Unknown type ID: " type-id))))))))
 
 (defn thaw-from-stream!
@@ -298,7 +309,8 @@
   Nippy.
 
   For custom `Freezable` types provide a `:readers` arg:
-  (thaw (freeze (MyType. \"Joe\")) {:readers {-1 (fn [stream] (.readUTF stream))}})
+  (thaw (freeze (MyType. \"Joe\"))
+    {:readers {1 (fn [^DataInputStream stream] (.readUTF stream))}})
 
   WARNING: Enabling `:read-eval?` can lead to security vulnerabilities unless
   you are sure you know what you're doing."
diff --git a/test/taoensso/nippy/tests/main.clj b/test/taoensso/nippy/tests/main.clj
index 794e9cb..6acfce8 100644
--- a/test/taoensso/nippy/tests/main.clj
+++ b/test/taoensso/nippy/tests/main.clj
@@ -1,7 +1,8 @@
 (ns taoensso.nippy.tests.main
   (:require [expectations   :as test :refer :all]
             [taoensso.nippy :as nippy :refer (freeze thaw)]
-            [taoensso.nippy.benchmarks :as benchmarks]))
+            [taoensso.nippy.benchmarks :as benchmarks])
+  (:import  [java.io DataInputStream DataOutputStream]))
 
 ;; Remove stuff from stress-data that breaks roundtrip equality
 (def test-data (dissoc nippy/stress-data :bytes))
@@ -31,14 +32,11 @@
       (thaw (org.iq80.snappy.Snappy/uncompress   xerial-ba  0 (alength xerial-ba))))))
 
 ;;; Custom types
-(defrecord MyType [data]
-  nippy/Freezable
-  (freeze-to-stream* [x s]
-    (.writeByte s -1)
-    (.writeUTF s (:data x))))
+(defrecord MyType [data])
+(nippy/custom-freezer MyType 1 x s (.writeUTF s (:data x)))
 
 (expect Exception (thaw (freeze (MyType. "Joe"))))
 (expect "Joe"     (thaw (freeze (MyType. "Joe"))
-                        {:readers {-1 (fn [s] (.readUTF s))}}))
+                        {:readers {1 (fn [^DataInputStream s] (.readUTF s))}}))
 
 (expect (benchmarks/bench {:reader? false})) ; Also tests :cached passwords
\ No newline at end of file