From cbe43c1bc60ed7494febef514fd6f4ddd4976481 Mon Sep 17 00:00:00 2001
From: Adam Rocska <adam.rocska@adams.solutions>
Date: Thu, 30 Apr 2020 18:11:56 +0200
Subject: [PATCH] ControlByte payloadSize determined for values between 29 and
 284. On to the next challenge.

Signed-off-by: Adam Rocska <adam.rocska@adams.solutions>
---
 Sources/MaxMindDBReader/ControlByte.swift     | 15 ++++--
 .../ControlByteTest.swift                     | 49 +++++++++++++++++--
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/Sources/MaxMindDBReader/ControlByte.swift b/Sources/MaxMindDBReader/ControlByte.swift
index e72c574..0623ed6 100644
--- a/Sources/MaxMindDBReader/ControlByte.swift
+++ b/Sources/MaxMindDBReader/ControlByte.swift
@@ -10,18 +10,27 @@ struct ControlByte {
 
     let firstByte                 = bytes.first!
     let typeDefinitionOnFirstByte = firstByte &>> 5
+    let isExtendedType            = typeDefinitionOnFirstByte == 0b0000_0000
 
-    guard let type = typeDefinitionOnFirstByte != 0b0000_0000
-                     ? DataType(rawValue: typeDefinitionOnFirstByte)
-                     : DataType(rawValue: bytes[bytes.index(after: bytes.startIndex)] + 7)
+    guard let type = isExtendedType
+                     ? DataType(rawValue: bytes[bytes.index(after: bytes.startIndex)] + 7)
+                     : DataType(rawValue: typeDefinitionOnFirstByte)
       else {
       return nil
     }
 
+    let sliceFrom = isExtendedType
+                    ? bytes.index(bytes.startIndex, offsetBy: 2, limitedBy: bytes.endIndex) ?? bytes.startIndex
+                    : bytes.index(bytes.startIndex, offsetBy: 1, limitedBy: bytes.endIndex) ?? bytes.startIndex
+
+    let bytesAfterTypeSpecifyingBytes = bytes[sliceFrom...]
+
     let payloadSizeDefinition = firstByte & 0b0001_1111
     switch payloadSizeDefinition {
       case _ where payloadSizeDefinition < 29:
         payloadSize = UInt32(payloadSizeDefinition)
+      case 29:
+        payloadSize = 29 + UInt32(bytesAfterTypeSpecifyingBytes.first ?? 0)
       default:
         return nil
     }
diff --git a/Tests/MaxMindDBReaderTests/ControlByteTest.swift b/Tests/MaxMindDBReaderTests/ControlByteTest.swift
index 48608f3..f71afc0 100644
--- a/Tests/MaxMindDBReaderTests/ControlByteTest.swift
+++ b/Tests/MaxMindDBReaderTests/ControlByteTest.swift
@@ -55,21 +55,62 @@ class ControlByteTest: XCTestCase {
     XCTAssertNil(ControlByte(bytes: Data([0b0000_1111, 0b0000_1111, 0b0000_0000])))
   }
 
+  fileprivate typealias PayloadSizeTestDefinition = (expectedPayloadSize: UInt32, bytes: Data)
+
   func testInit_payloadSizeDefinition_lessThan29() {
-    let testInput: [(expectedPayloadSize: UInt32, byte: UInt8)] = ControlByteTest
+    let nonExtendedRawValues: [PayloadSizeTestDefinition] = ControlByteTest
       .nonExtendedRawValues
       .reduce([]) { byteSequence, typeDefinition in
       byteSequence + (0..<29).map({
-        (expectedPayloadSize: UInt32($0), byte: $0 | (typeDefinition << 5))
+        (expectedPayloadSize: UInt32($0), bytes: Data([$0 | (typeDefinition << 5)]))
+      })
+    }
+    let extendedRawValues: [PayloadSizeTestDefinition] = ControlByteTest
+      .extendedRawValues
+      .reduce([]) { byteSequence, typeDefinition in
+      byteSequence + (0..<29).map({
+        (expectedPayloadSize: UInt32($0), bytes: Data([$0, typeDefinition - 7]))
       })
     }
 
-    for (expectedPayloadSize, byte) in testInput {
+    for (expectedPayloadSize, bytes) in (nonExtendedRawValues + extendedRawValues) {
       XCTAssertEqual(
         expectedPayloadSize,
-        ControlByte(bytes: Data([byte]))?.payloadSize
+        ControlByte(bytes: bytes)?.payloadSize
       )
     }
 
   }
+
+  func testInit_payloadSizeDefinition_exactly29() {
+    let nonExtendedRawValues: [PayloadSizeTestDefinition] = ControlByteTest
+      .nonExtendedRawValues
+      .reduce([]) { byteSequence, typeDefinition in
+      byteSequence + (29..<285).map({
+        (
+          expectedPayloadSize: UInt32($0),
+          bytes: Data([UInt8(29) | (typeDefinition << 5), UInt8($0 - 29)])
+        )
+      })
+    }
+    let extendedRawValues: [PayloadSizeTestDefinition] = ControlByteTest
+      .extendedRawValues
+      .reduce([]) { byteSequence, typeDefinition in
+      byteSequence + (29..<285).map({
+        (
+          expectedPayloadSize: UInt32($0),
+          bytes: Data([UInt8(29), typeDefinition - 7, UInt8($0 - 29)])
+        )
+      })
+    }
+
+    for (expectedPayloadSize, bytes) in (nonExtendedRawValues + extendedRawValues) {
+      XCTAssertEqual(
+        expectedPayloadSize,
+        ControlByte(bytes: bytes)?.payloadSize,
+        "Expected a payload size of \(expectedPayloadSize), but instead got \(ControlByte(bytes: bytes)?.payloadSize)"
+      )
+    }
+  }
+
 }