From b1100a759450fa940f7fb03f786bc0bbbadcc07d Mon Sep 17 00:00:00 2001
From: Timofey Solomko <tsolomko@gmail.com>
Date: Wed, 8 Feb 2017 19:34:07 +0300
Subject: [PATCH] Added new function to BitToByteWriter which writes a number
 with specified amount of bits.

---
 Sources/BitToByteWriter.swift | 12 ++++++++++++
 Sources/Deflate.swift         | 33 ++++++++-------------------------
 2 files changed, 20 insertions(+), 25 deletions(-)

diff --git a/Sources/BitToByteWriter.swift b/Sources/BitToByteWriter.swift
index e8087093..c0ece31b 100644
--- a/Sources/BitToByteWriter.swift
+++ b/Sources/BitToByteWriter.swift
@@ -78,6 +78,18 @@ final class BitToByteWriter {
         }
     }
 
+    func write(number: Int, bitsCount: Int) {
+        var mask = 1
+        for _ in 0..<bitsCount {
+            if number & mask > 0 {
+                self.write(bit: 1)
+            } else {
+                self.write(bit: 0)
+            }
+            mask <<= 1
+        }
+    }
+
     func finish() {
         self.buffer.append(self.currentByte)
         self.currentByte = 0
diff --git a/Sources/Deflate.swift b/Sources/Deflate.swift
index 8c255acf..588fad82 100644
--- a/Sources/Deflate.swift
+++ b/Sources/Deflate.swift
@@ -233,22 +233,23 @@ public final class Deflate: DecompressionAlgorithm {
     public static func compress(data: Data) throws -> Data {
         let bytes = data.toArray(type: UInt8.self)
         let bldCodes = Deflate.lengthEncode(bytes)
-        let huffmanEncodedBytes = Deflate.huffmanEncode(bldCodes)
+        let huffmanEncodedBytes = Deflate.encodeHuffmanBlock(bldCodes)
         return Data(bytes: huffmanEncodedBytes)
     }
 
-    private static func huffmanEncode(_ bldCodes: [BLDCode]) -> [UInt8] {
-        /// Empty DWP object for creating Huffman trees.
-        var pointerData = DataWithPointer(data: Data(), bitOrder: .reversed)
 
+    private static func encodeHuffmanBlock(_ bldCodes: [BLDCode]) -> [UInt8] {
         let bitWriter = BitToByteWriter(bitOrder: .reversed)
+
         // Write block header.
         // Note: For now it is only static huffman blocks.
         // Note: Only one block is supported for now.
-
         bitWriter.write(bit: 1)
         bitWriter.write(bits: [1, 0])
 
+        /// Empty DWP object for creating Huffman trees.
+        var pointerData = DataWithPointer(data: Data(), bitOrder: .reversed)
+
         // Constructing Huffman trees for the case of block with preset alphabets.
         // In this case codes for literals and distances are fixed.
         // Bootstraps for trees (first element in pair is code, second is number of bits).
@@ -272,16 +273,7 @@ public final class Deflate: DecompressionAlgorithm {
 
                 // TODO: Add check for empty returned array.
                 bitWriter.write(bits: mainLiterals.code(symbol: lengthSymbol))
-
-                var mask = 1
-                for _ in 0..<extraLengthBitsCount {
-                    if lengthExtra & mask > 0 {
-                        bitWriter.write(bit: 1)
-                    } else {
-                        bitWriter.write(bit: 0)
-                    }
-                    mask <<= 1
-                }
+                bitWriter.write(number: lengthExtra, bitsCount: extraLengthBitsCount)
 
                 let distanceSymbol = ((HuffmanTree.Constants.distanceBase.index { $0 > distance }) ?? 30) - 1
                 let distanceExtra = distance - HuffmanTree.Constants.distanceBase[distanceSymbol]
@@ -290,16 +282,7 @@ public final class Deflate: DecompressionAlgorithm {
 
                 // TODO: Add check for empty returned array.
                 bitWriter.write(bits: mainDistances.code(symbol: distanceSymbol))
-
-                mask = 1
-                for _ in 0..<extraDistanceBitsCount {
-                    if distanceExtra & mask > 0 {
-                        bitWriter.write(bit: 1)
-                    } else {
-                        bitWriter.write(bit: 0)
-                    }
-                    mask <<= 1
-                }
+                bitWriter.write(number: distanceExtra, bitsCount: extraDistanceBitsCount)
             }
         }