Added a check for HuffmanTree.code(…) returning empty array.

This commit is contained in:
Timofey Solomko
2017-02-18 22:20:47 +03:00
parent 9b33b81353
commit c7d982c19b
+14 -10
View File
@@ -243,7 +243,7 @@ public final class Deflate: DecompressionAlgorithm {
}
let bldCodes = Deflate.lengthEncode(bytes)
let huffmanEncodedBytes = Deflate.encodeHuffmanBlock(bldCodes)
let huffmanEncodedBytes = try Deflate.encodeHuffmanBlock(bldCodes)
return Data(bytes: huffmanEncodedBytes)
}
@@ -275,7 +275,7 @@ public final class Deflate: DecompressionAlgorithm {
return out
}
private static func encodeHuffmanBlock(_ bldCodes: [BLDCode]) -> [UInt8] {
private static func encodeHuffmanBlock(_ bldCodes: [BLDCode]) throws -> [UInt8] {
let bitWriter = BitToByteWriter(bitOrder: .reversed)
// Write block header.
@@ -300,25 +300,30 @@ public final class Deflate: DecompressionAlgorithm {
for code in bldCodes {
switch code {
case .byte(let byte):
// TODO: Add check for empty returned array.
bitWriter.write(bits: mainLiterals.code(symbol: byte.toInt()))
let codeOfByte = mainLiterals.code(symbol: byte.toInt())
guard codeOfByte.count > 0
else { throw DeflateError.SymbolNotFound }
bitWriter.write(bits: codeOfByte)
case .lengthDistance(let length, let distance):
let lengthSymbol = HuffmanTree.Constants.lengthCode[length - 3]
let lengthExtra = length - HuffmanTree.Constants.lengthBase[lengthSymbol - 257]
let extraLengthBitsCount = (257 <= lengthSymbol && lengthSymbol <= 260) || lengthSymbol == 285 ?
0 : (((lengthSymbol - 257) >> 2) - 1)
// TODO: Add check for empty returned array.
bitWriter.write(bits: mainLiterals.code(symbol: lengthSymbol))
let codeOfLength = mainLiterals.code(symbol: lengthSymbol)
guard codeOfLength.count > 0
else { throw DeflateError.SymbolNotFound }
bitWriter.write(bits: codeOfLength)
bitWriter.write(number: lengthExtra, bitsCount: extraLengthBitsCount)
let distanceSymbol = ((HuffmanTree.Constants.distanceBase.index { $0 > distance }) ?? 30) - 1
let distanceExtra = distance - HuffmanTree.Constants.distanceBase[distanceSymbol]
let extraDistanceBitsCount = distanceSymbol == 0 || distanceSymbol == 1 ? 0 : ((distanceSymbol >> 1) - 1)
// TODO: Add check for empty returned array.
bitWriter.write(bits: mainDistances.code(symbol: distanceSymbol))
let codeOfDistance = mainDistances.code(symbol: distanceSymbol)
guard codeOfDistance.count > 0
else { throw DeflateError.SymbolNotFound }
bitWriter.write(bits: codeOfDistance)
bitWriter.write(number: distanceExtra, bitsCount: extraDistanceBitsCount)
}
}
@@ -344,7 +349,6 @@ public final class Deflate: DecompressionAlgorithm {
}
}
// TODO: Expand dictionary size.
private static func lengthEncode(_ rawBytes: [UInt8]) -> [BLDCode] {
precondition(rawBytes.count >= 3, "Too small array!")