Add full functionality when reading from MKV files

Sources/macSubtitleOCR/MKV/EBML/EBMLParser.swift

@@ -9,57 +9,40 @@
 import Foundation
 import os
 
-private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "main")
+private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "ebml")
 
 // Helper function to read variable-length integers (VINT) from MKV (up to 8 bytes)
 func readVINT(from fileHandle: FileHandle, unmodified: Bool = false) -> UInt64 {
-    guard let firstByte = fileHandle.readData(ofLength: 1).first else {
-        return 0
-    }
+    guard let firstByte = fileHandle.readData(ofLength: 1).first else { return 0 }
+
     var length: UInt8 = 1
     var mask: UInt8 = 0x80
 
     // Find how many bytes are needed for the VINT (variable integer)
-    while (firstByte & UInt8(mask)) == 0 {
+    while (firstByte & mask) == 0 {
         length += 1
         mask >>= 1
     }
 
-    // Extract the value
-    logger.debug("Length: \(length), Mask: 0x\(String(format: "%08X", mask))")
-    if mask - 1 == 0x0F {
-        mask = 0xFF // Hacky workaround that I still don't understand why is needed
-    } else if length == 1, !unmodified {
-        mask = firstByte
-    } else {
-        mask = mask - 1
-    }
-    logger.debug("Byte before: 0x\(String(format: "%08X", firstByte))")
-    logger.debug("Byte after: 0x\(String(format: "%08X", firstByte & mask))")
+    // Adjust mask based on length and unmodified flag
+    mask = (mask == 0x10) ? 0xFF : (length == 1 && !unmodified) ? firstByte : mask - 1
 
     var value = UInt64(firstByte & mask)
 
     if length > 1 {
         let data = fileHandle.readData(ofLength: Int(length - 1))
         for byte in data {
-            value <<= 8
-            value |= UInt64(byte)
+            value = (value << 8) | UInt64(byte)
         }
     }
     logger.debug("VINT: 0x\(String(format: "%08X", value))")
 
     return value
 }
 
-// Helper function to read a specified number of bytes
-func readBytes(from fileHandle: FileHandle, length: Int) -> Data? {
-    fileHandle.readData(ofLength: length)
-}
-
 // Helper function to read an EBML element's ID and size
 func readEBMLElement(from fileHandle: FileHandle, unmodified: Bool = false) -> (elementID: UInt32, elementSize: UInt64) {
     let elementID = readVINT(from: fileHandle, unmodified: unmodified)
     let elementSize = readVINT(from: fileHandle, unmodified: true)
-    logger.debug("elementID: 0x\(String(format: "%08X", elementID)), elementSize: \(elementSize)")
     return (UInt32(elementID), elementSize)
 }

Sources/macSubtitleOCR/MKV/MKVError.swift

@@ -0,0 +1,12 @@
+//
+// MKVError.swift
+// macSubtitleOCR
+//
+// Created by Ethan Dye on 9/19/24.
+// Copyright © 2024 Ethan Dye. All rights reserved.
+//
+
+enum MKVError: Error {
+    case segmentElementNotFound
+    case tracksElementNotFound
+}

Sources/macSubtitleOCR/MKV/MKVFileHandler.swift

@@ -0,0 +1,91 @@
+//
+// MKVFileHandler.swift
+// macSubtitleOCR
+//
+// Created by Ethan Dye on 9/20/24.
+// Copyright © 2024 Ethan Dye. All rights reserved.
+//
+
+import Foundation
+import os
+
+class MKVFileHandler {
+    // MARK: - Properties
+
+    var fileHandle: FileHandle
+    var eof: UInt64
+    var timestampScale: Double = 1000000.0 // Default value if not specified in a given MKV file
+    var logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "mkv")
+
+    // MARK: - Lifecycle
+
+    init(filePath: String) throws {
+        guard FileManager.default.fileExists(atPath: filePath) else {
+            throw macSubtitleOCRError.fileReadError
+        }
+        fileHandle = try FileHandle(forReadingFrom: URL(fileURLWithPath: filePath))
+        eof = fileHandle.seekToEndOfFile()
+        fileHandle.seek(toFileOffset: 0)
+    }
+
+    deinit {
+        fileHandle.closeFile()
+    }
+
+    // MARK: - Functions
+
+    func locateSegment() -> UInt64? {
+        if let (segmentSize, _) = findElement(withID: EBML.segmentID, avoidCluster: true) as? (UInt64, UInt32) {
+            return segmentSize
+        }
+        return nil
+    }
+
+    func locateCluster() -> UInt64? {
+        if let (clusterSize, _) = findElement(withID: EBML.cluster, avoidCluster: false) as? (UInt64, UInt32) {
+            return clusterSize
+        }
+        return nil
+    }
+
+    // Find EBML element by ID, avoiding Cluster header
+    func findElement(withID targetID: UInt32, _ tgtID2: UInt32? = nil, avoidCluster: Bool = true) -> (UInt64?, UInt32?) {
+        while let (elementID, elementSize, elementOffset) = tryParseElement() {
+            // Ensure we stop if we have reached or passed the EOF
+            if fileHandle.offsetInFile >= eof {
+                return (nil, nil)
+            }
+
+            // If, by chance, we find a TimestampScale element, update it from the default
+            if elementID == EBML.timestampScale {
+                timestampScale = Double(readFixedLengthNumber(fileHandle: fileHandle, length: Int(elementSize)))
+                // swiftformat:disable:next redundantSelf
+                logger.debug("Found timestamp scale: \(self.timestampScale)")
+                return (nil, nil)
+            }
+
+            // If a Cluster header is encountered, seek back to the start of the Cluster
+            if elementID == EBML.cluster && avoidCluster {
+                logger.debug("Encountered Cluster: seeking back to before the cluster header")
+                fileHandle.seek(toFileOffset: elementOffset)
+                return (nil, nil)
+            }
+
+            // If the element matches the target ID (or secondary ID), return its size
+            if elementID == targetID || (tgtID2 != nil && elementID == tgtID2!) {
+                return (elementSize, elementID)
+            } else {
+                // Skip over the element's data by seeking to its end
+                logger.debug("Found: \(elementID), but not \(targetID), skipping element")
+                fileHandle.seek(toFileOffset: fileHandle.offsetInFile + elementSize)
+            }
+        }
+        return (nil, nil)
+    }
+
+    func tryParseElement() -> (elementID: UInt32, elementSize: UInt64, oldOffset: UInt64)? {
+        let oldOffset = fileHandle.offsetInFile
+        let (elementID, elementSize) = readEBMLElement(from: fileHandle)
+        return (elementID, elementSize, oldOffset: oldOffset)
+    }
+}

Sources/macSubtitleOCR/MKV/MKVHelpers.swift

@@ -15,32 +15,22 @@ func getUInt16BE(buffer: Data, offset: Int) -> UInt16 {
 // Function to read a fixed length number of bytes and convert in into a (Un)signed integer
 func readFixedLengthNumber(fileHandle: FileHandle, length: Int, signed: Bool = false) -> Int64 {
     let data = fileHandle.readData(ofLength: length)
-    let pos = 0
-
     var result: Int64 = 0
-    for i in 0 ..< length {
-        result = result * 0x100 + Int64(data[pos + i])
+
+    for byte in data {
+        result = result << 8 | Int64(byte)
     }
 
-    if signed {
-        let signBitMask: UInt8 = 0x80
-        if data[pos] & signBitMask != 0 {
-            result -= Int64(1) << (8 * length) // Apply two's complement for signed numbers
-        }
+    if signed, data.first! & 0x80 != 0 {
+        result -= Int64(1) << (8 * length) // Apply two's complement for signed integers
     }
 
     return result
 }
 
 // Encode the absolute timestamp as 4 bytes in big-endian format for PGS
 func encodePTSForPGS(_ timestamp: Int64) -> [UInt8] {
-    let timestamp = UInt32(timestamp) // Convert to unsigned 32-bit value
-    return [
-        UInt8((timestamp >> 24) & 0xFF),
-        UInt8((timestamp >> 16) & 0xFF),
-        UInt8((timestamp >> 8) & 0xFF),
-        UInt8(timestamp & 0xFF),
-    ]
+    withUnsafeBytes(of: UInt32(timestamp).bigEndian) { Array($0) }
 }
 
 // Calculate the absolute timestamp with 90 kHz accuracy for PGS format