From 49228a049f8acd4e17c84d2de76bfce68f9e60b7 Mon Sep 17 00:00:00 2001 From: Timofey Solomko Date: Sun, 3 Sep 2017 14:16:59 +0300 Subject: [PATCH] Fix HuffmanTree init crash when zero codeLength is encountered --- Sources/DecodingHuffmanTree.swift | 2 +- Sources/EncodingHuffmanTree.swift | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Sources/DecodingHuffmanTree.swift b/Sources/DecodingHuffmanTree.swift index e1fa9fd0..712cb760 100644 --- a/Sources/DecodingHuffmanTree.swift +++ b/Sources/DecodingHuffmanTree.swift @@ -17,7 +17,7 @@ class DecodingHuffmanTree { self.bitReader = bitReader // Sort `lengths` array to calculate canonical Huffman code. - let sortedLengths = lengths.sorted { (left: HuffmanLength, right: HuffmanLength) -> Bool in + let sortedLengths = lengths.filter { $0.codeLength > 0 }.sorted { (left: HuffmanLength, right: HuffmanLength) -> Bool in if left.codeLength == right.codeLength { return left.symbol < right.symbol } else { diff --git a/Sources/EncodingHuffmanTree.swift b/Sources/EncodingHuffmanTree.swift index 8580aac7..76444f17 100644 --- a/Sources/EncodingHuffmanTree.swift +++ b/Sources/EncodingHuffmanTree.swift @@ -16,7 +16,7 @@ class EncodingHuffmanTree { self.bitWriter = bitWriter // Sort `lengths` array to calculate canonical Huffman code. - let sortedLengths = lengths.sorted { (left: HuffmanLength, right: HuffmanLength) -> Bool in + let sortedLengths = lengths.filter { $0.codeLength > 0 }.sorted { (left: HuffmanLength, right: HuffmanLength) -> Bool in if left.codeLength == right.codeLength { return left.symbol < right.symbol } else { @@ -146,7 +146,7 @@ class EncodingHuffmanTree { func bitSize(for stats: [(Int, Int)]) -> Int { var totalSize = 0 - for (symbol, count) in stats { + for (symbol, count) in stats where count > 0 { guard symbol < self.codingIndices.count else { fatalError("Symbol is not found.") } let codingIndex = self.codingIndices[symbol]