Misc performance improvements for CRAM reader

src/cljam/io/cram/codecs/rans4x8.clj

@@ -1,50 +1,58 @@
 (ns cljam.io.cram.codecs.rans4x8
   (:require [cljam.io.util.byte-buffer :as bb]
-            [cljam.io.cram.itf8 :as itf8]))
+            [cljam.io.cram.itf8 :as itf8])
+  (:import [java.util Arrays]))
 
-(defn- read-frequencies* [bb read-fn]
-  (loop [sym (long (bb/read-ubyte bb))
-         rle 0
-         freqs (transient {})]
-    (let [freqs' (assoc! freqs sym (read-fn bb))]
-      (if (pos? rle)
-        (recur (inc sym) (dec rle) freqs')
-        (let [sym' (long (bb/read-ubyte bb))
-              rle' (if (= sym' (inc sym))
-                     (long (bb/read-ubyte bb))
-                     rle)]
-          (if (zero? sym')
-            (persistent! freqs')
-            (recur sym' rle' freqs')))))))
+(def ^:private byte-array-type (type (byte-array 0)))
+(def ^:private int-array-type (type (int-array 0)))
 
-(defn- read-frequencies0 [bb]
-  (read-frequencies* bb itf8/decode-itf8))
+(defmacro ^:private read-frequencies* [bb init-expr read-expr]
+  `(let [ret# ~init-expr]
+     (loop [sym# (long (bb/read-ubyte ~bb))
+            rle# 0]
+       (aset ret# sym# ~read-expr)
+       (if (pos? rle#)
+         (recur (inc sym#) (dec rle#))
+         (let [sym'# (long (bb/read-ubyte ~bb))
+               rle'# (if (= sym'# (inc sym#))
+                       (long (bb/read-ubyte ~bb))
+                       rle#)]
+           (if (zero? sym'#)
+             ret#
+             (recur sym'# rle'#)))))))
 
-(defn- read-frequencies1 [bb]
-  (read-frequencies* bb read-frequencies0))
+(defn- read-frequencies0 ^ints [bb]
+  (read-frequencies* bb (int-array 256) (int (itf8/decode-itf8 bb))))
 
-(defn- cumulative-frequencies [freqs]
-  (loop [i 0
-         sum 0
-         cum-freqs (transient [])]
-    (if (< i 256)
-      (let [f (get freqs i 0)]
-        (recur (inc i) (+ sum (long f)) (conj! cum-freqs sum)))
-      (persistent! (conj! cum-freqs sum)))))
+(defn- read-frequencies1 ^"[[I" [bb]
+  (read-frequencies* bb ^"[[I" (make-array int-array-type 256) (read-frequencies0 bb)))
 
-(defn- lookup-symbol ^long [cum-freqs ^long f]
-  (loop [l 0
-         r (dec (count cum-freqs))]
-    (if (< l r)
-      (let [m (quot (+ l r) 2)
-            fm (long (nth cum-freqs m))]
-        (cond (and (<= fm f)
-                   (< f (long (nth cum-freqs (inc m)))))
-              m
+(def ^:private zero-int-array (int-array 256))
 
-              (< f fm) (recur l (dec m))
-              :else (recur (inc m) r)))
-      l)))
+(defn- cumulative-frequencies ^ints [^ints freqs]
+  (if (nil? freqs)
+    zero-int-array
+    (let [cum-freqs (int-array 256)]
+      (loop [i 0
+             sum 0]
+        (when (< i 256)
+          (let [f (aget freqs i)]
+            (aset cum-freqs i sum)
+            (recur (inc i) (+ sum (long f))))))
+      cum-freqs)))
+
+(defn- reverse-lookup-table ^bytes [^ints cum-freqs]
+  (let [arr (byte-array 4096)
+        n (alength cum-freqs)]
+    (loop [i 1, start 0]
+      (if (< i n)
+        (let [curr (aget cum-freqs i)]
+          (if (= start curr)
+            (recur (inc i) start)
+            (do (Arrays/fill arr start curr (byte (dec i)))
+                (recur (inc i) curr))))
+        (Arrays/fill arr start 4096 (byte 255))))
+    arr))
 
 (defn- advance-step ^long [^long c ^long f ^long state]
   (-> (* f (bit-shift-right state 12))
@@ -54,33 +62,35 @@
 (defn- renormalize-state ^long [bb ^long state]
   (loop [state state]
     (if (< state 0x800000)
-      (recur (+ (bit-shift-left state 8)
-                (long (bb/read-ubyte bb))))
+      (recur (bit-or (bit-shift-left state 8) (long (bb/read-ubyte bb))))
       state)))
 
 (defn- decode0 [bb ^long n-out]
   (let [freqs (read-frequencies0 bb)
         cum-freqs (cumulative-frequencies freqs)
+        table (reverse-lookup-table cum-freqs)
         states (bb/read-ints bb 4)
         out (byte-array n-out)]
     (dotimes [i n-out]
       (let [j (rem i 4)
             state (aget states j)
             f (bit-and state 0xfff)
-            sym (lookup-symbol cum-freqs f)
+            sym (bit-and (aget table f) 0xff)
             state' (->> state
-                        (advance-step (nth cum-freqs sym) (get freqs sym 0))
+                        (advance-step (aget cum-freqs sym) (aget freqs sym))
                         (renormalize-state bb))]
         (aset out i (byte sym))
         (aset states j state')))
     out))
 
 (defn- decode1 [bb ^long n-out]
   (let [freqs (read-frequencies1 bb)
-        cum-freqs (persistent!
-                   (reduce-kv #(assoc! %1 %2 (cumulative-frequencies %3))
-                              (transient {})
-                              freqs))
+        ^"[[I" cum-freqs (make-array int-array-type 256)
+        _ (dotimes [i 256]
+            (aset cum-freqs i (cumulative-frequencies (aget freqs i))))
+        ^"[[B" tables (make-array byte-array-type 256)
+        _ (dotimes [i 256]
+            (aset tables i (reverse-lookup-table (aget cum-freqs i))))
         quarter (quot n-out 4)
         truncated (* 4 quarter)
         states (bb/read-ints bb 4)
@@ -91,11 +101,11 @@
         (let [state (aget states j)
               f (bit-and state 0xfff)
               last-sym (aget last-syms j)
-              cfreqs (get cum-freqs last-sym)
-              sym (lookup-symbol cfreqs f)
+              ^ints cfreqs (aget cum-freqs last-sym)
+              sym (bit-and (aget ^bytes (aget tables last-sym) f) 0xff)
               state' (->> state
-                          (advance-step (nth cfreqs sym)
-                                        (get-in freqs [last-sym sym] 0))
+                          (advance-step (aget cfreqs sym)
+                                        (aget ^ints (aget freqs last-sym) sym))
                           (renormalize-state bb))]
           (aset out (+ i (* j quarter)) (byte sym))
           (aset states j state')
@@ -104,11 +114,11 @@
       (let [state (aget states 3)
             f (bit-and state 0xfff)
             last-sym (aget last-syms 3)
-            cfreq (get cum-freqs last-sym)
-            sym (lookup-symbol cfreq f)
+            ^ints cfreq (aget cum-freqs last-sym)
+            sym (bit-and (aget ^bytes (aget tables last-sym) f) 0xff)
             state' (->> state
-                        (advance-step (nth cfreq sym)
-                                      (get-in freqs [last-sym sym] 0))
+                        (advance-step (aget cfreq sym)
+                                      (aget ^ints (aget freqs last-sym) sym))
                         (renormalize-state bb))]
         (aset out (+ i truncated) (byte sym))
         (aset states 3 state')

src/cljam/io/cram/decode/record.clj

@@ -154,7 +154,7 @@
         (if (= len miss)
           "*"
           (String. (.array bb)))
-        (let [q (qs-decoder)]
+        (let [q (long (qs-decoder))]
           (.put bb (byte (qual-score q)))
           (recur (dec i) (cond-> miss (= q -1) inc)))))))
 

test/cljam/io/cram/codecs/rans4x8_test.clj

@@ -21,7 +21,8 @@
             0x63 0x174
             0x64 0x174
             0x72 0x2e8}
-           (#'rans/read-frequencies0 bb)))
+           (into {} (keep-indexed (fn [i x] (when (pos? x) [i x])))
+                 (#'rans/read-frequencies0 bb))))
     (is (zero? (.remaining bb)))))
 
 (deftest read-frequencies1-test
@@ -62,9 +63,46 @@
             0x63 {0x61 0xfff}
             0x64 {0x61 0xfff}
             0x72 {0x61 0xfff}}
-           (#'rans/read-frequencies1 bb)))
+           (into {}
+                 (keep-indexed
+                  (fn [i arr]
+                    (when-let [m (->> arr
+                                      (into {}
+                                            (keep-indexed
+                                             (fn [j x]
+                                               (when (pos? x)
+                                                 [j x]))))
+                                      not-empty)]
+                      [i m])))
+                 (#'rans/read-frequencies1 bb))))
     (is (zero? (.remaining bb)))))
 
+(deftest cumulative-frequencies-test
+  (let [freqs (-> [0x000 0x747 0x000 0x2e8 0x000 0x174 0x174 0x2e8]
+                  (concat (repeat 248 0x000))
+                  int-array)]
+    (is (= (concat [0x000 0x000 0x747 0x747 0xa2f 0xa2f 0xba3 0xd17 0xfff 0xfff]
+                   (repeat 246 0xfff))
+           (vec (#'rans/cumulative-frequencies freqs))))))
+
+(deftest reverse-lookup-table-test
+  (let [freqs (-> [0x000 0x747 0x000 0x2e8 0x000 0x174 0x174 0x2e8]
+                  (concat (repeat 248 0x000))
+                  int-array)
+        cfreqs (#'rans/cumulative-frequencies freqs)
+        table ^bytes (#'rans/reverse-lookup-table cfreqs)]
+    (is (= 1 (aget table 0x000)))
+    (is (= 1 (aget table 0x746)))
+    (is (= 3 (aget table 0x747)))
+    (is (= 3 (aget table 0xa2e)))
+    (is (= 5 (aget table 0xa2f)))
+    (is (= 5 (aget table 0xba2)))
+    (is (= 6 (aget table 0xba3)))
+    (is (= 6 (aget table 0xd16)))
+    (is (= 7 (aget table 0xd17)))
+    (is (= 7 (aget table 0xffe)))
+    (is (= 255 (bit-and (aget table 0xfff) 0xff)))))
+
 (defn- read-as-buffer ^ByteBuffer [file]
   (with-open [in (FileInputStream. (io/file file))]
     (let [ch (.getChannel in)