From 9da406fc53fb79f1a1dd191a14aa2b0ac1443030 Mon Sep 17 00:00:00 2001
From: Timofey Solomko <tsolomko@gmail.com>
Date: Sun, 19 Nov 2017 12:47:46 +0300
Subject: [PATCH] Small improvement to EncHuffTree.code function

---
 Sources/Common/BitWriter.swift           | 7 ++++---
 Sources/Common/EncodingHuffmanTree.swift | 9 +--------
 2 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/Sources/Common/BitWriter.swift b/Sources/Common/BitWriter.swift
index cb86148b..dfea59e0 100644
--- a/Sources/Common/BitWriter.swift
+++ b/Sources/Common/BitWriter.swift
@@ -75,11 +75,12 @@ class BitWriter {
         }
     }
 
-    func write(number: Int, bitsCount: Int) {
-        var mask = bitOrder == .straight ? 1 << (bitsCount - 1) : 1
+    func write(number: Int, bitsCount: Int, bitOrder: BitReader.BitOrder? = nil) {
+        let bitOrderToUse = (bitOrder ?? self.bitOrder)
+        var mask = bitOrderToUse == .straight ? 1 << (bitsCount - 1) : 1
         for _ in 0..<bitsCount {
             self.write(bit: number & mask > 0 ? 1 : 0)
-            switch bitOrder {
+            switch bitOrderToUse {
             case .straight:
                 mask >>= 1
             case .reversed:
diff --git a/Sources/Common/EncodingHuffmanTree.swift b/Sources/Common/EncodingHuffmanTree.swift
index c37b06e9..7323a83d 100644
--- a/Sources/Common/EncodingHuffmanTree.swift
+++ b/Sources/Common/EncodingHuffmanTree.swift
@@ -61,14 +61,7 @@ class EncodingHuffmanTree {
         guard codingIndex[0] > -1
             else { fatalError("Symbol is not found.") }
 
-        var treeCode = codingIndex[0]
-        let bits = codingIndex[1]
-
-        for _ in 0..<bits {
-            let bit = treeCode & 1
-            self.bitWriter.write(bit: bit == 0 ? 0 : 1)
-            treeCode >>= 1
-        }
+        self.bitWriter.write(number: codingIndex[0], bitsCount: codingIndex[1], bitOrder: .reversed)
     }
 
     func bitSize(for stats: [Int]) -> Int {