Simplify functions to remove complexity

Signed-off-by: Ethan Dye <[email protected]>

Sources/macSubtitleOCR/MKV/EBML/EBMLParser.swift

@@ -13,46 +13,33 @@ private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category:
 
 // 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
-    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
-    }
+    // 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)

Sources/macSubtitleOCR/MKV/MKVFileHandler.swift

@@ -10,11 +10,15 @@ 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
@@ -28,6 +32,8 @@ class MKVFileHandler {
         fileHandle.closeFile()
     }
 
+    // MARK: - Functions
+
     func locateSegment() -> UInt64? {
         if let (segmentSize, _) = findElement(withID: EBML.segmentID, avoidCluster: true) as? (UInt64, UInt32) {
             return segmentSize
@@ -52,9 +58,7 @@ class MKVFileHandler {
 
             // 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)))
+                timestampScale = Double(readFixedLengthNumber(fileHandle: fileHandle, length: Int(elementSize)))
                 // swiftformat:disable:next redundantSelf
                 logger.debug("Found timestamp scale: \(self.timestampScale)")
                 return (nil, nil)
@@ -76,7 +80,6 @@ class MKVFileHandler {
                 fileHandle.seek(toFileOffset: fileHandle.offsetInFile + elementSize)
             }
         }
-
         return (nil, nil)
     }
 

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

Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift

@@ -10,16 +10,19 @@ import Foundation
 import os
 
 class MKVSubtitleExtractor: MKVTrackParser {
-    func getSubtitleTrackData(trackNumber: Int, outPath: String) throws -> String? {
-        let tmpSup = URL(fileURLWithPath: outPath).deletingPathExtension().appendingPathExtension("sup").lastPathComponent
-        let manager = FileManager.default
-        let tmpFilePath = (manager.temporaryDirectory.path + "/\(trackNumber)" + tmpSup)
+    private var stderr = StandardErrorOutputStream()
 
-        if manager.createFile(atPath: tmpFilePath, contents: tracks[trackNumber].trackData, attributes: nil) {
+    func getSubtitleTrackData(trackNumber: Int) throws -> String? {
+        let tmpFilePath = FileManager.default.temporaryDirectory
+            .appendingPathComponent("\(trackNumber)")
+            .appendingPathExtension("sup")
+            .path
+
+        if FileManager.default.createFile(atPath: tmpFilePath, contents: tracks[trackNumber].trackData, attributes: nil) {
             logger.debug("Created file at path: \(tmpFilePath).")
             return tmpFilePath
         } else {
-            logger.debug("Failed to create file at path: \(tmpFilePath).")
+            print("Failed to create file at path: \(tmpFilePath).", to: &stderr)
             throw PGSError.fileReadError
         }
     }

Sources/macSubtitleOCR/MKV/MKVTrackParser.swift

@@ -10,9 +10,13 @@ import Foundation
 import os
 
 class MKVTrackParser: MKVFileHandler {
+    // MARK: - Properties
+
     var tracks: [MKVTrack] = []
     private var stderr = StandardErrorOutputStream()
 
+    // MARK: - Functions
+
     func parseTracks(codec: String) throws {
         guard let _ = findElement(withID: EBML.segmentID) as? (UInt64, UInt32) else {
             print("Error: Segment element not found", to: &stderr)
@@ -48,39 +52,6 @@ class MKVTrackParser: MKVFileHandler {
         }
     }
 
-    private func parseTrackEntry(codec: String) -> Int? {
-        var trackNumber: Int?
-        var trackType: UInt8?
-        var codecId: String?
-
-        while let (elementID, elementSize, _) = tryParseElement() {
-            switch elementID {
-            case EBML.trackNumberID:
-                trackNumber = Int((readBytes(from: fileHandle, length: 1)?.first)!)
-                logger.debug("Found track number: \(trackNumber!)")
-            case EBML.trackTypeID: // Unused by us, left for debugging
-                trackType = readBytes(from: fileHandle, length: 1)?.first
-                logger.debug("Found track type: \(trackType!)")
-            case EBML.codecID:
-                var data = readBytes(from: fileHandle, length: Int(elementSize))
-                data?.removeNullBytes()
-                codecId = data.flatMap { String(data: $0, encoding: .ascii) }
-                logger.debug("Found codec ID: \(codecId!)")
-            default:
-                fileHandle.seek(toFileOffset: fileHandle.offsetInFile + elementSize)
-            }
-            if trackNumber != nil, trackType != nil, codecId != nil { break }
-        }
-
-        if let trackNumber, let codecId {
-            if codecId == codec {
-                return trackNumber
-            }
-        }
-        return nil
-    }
-
-    // Implement track extraction logic (e.g., `extractTrackData`) here
     func extractTrackData(trackNumber: [Int]) -> [Data]? {
         fileHandle.seek(toFileOffset: 0)
 
@@ -110,6 +81,40 @@ class MKVTrackParser: MKVFileHandler {
         return trackData.isEmpty ? nil : trackData
     }
 
+    // MARK: - Methods
+
+    private func parseTrackEntry(codec: String) -> Int? {
+        var trackNumber: Int?
+        var trackType: UInt8?
+        var codecId: String?
+
+        while let (elementID, elementSize, _) = tryParseElement() {
+            switch elementID {
+            case EBML.trackNumberID:
+                trackNumber = Int((fileHandle.readData(ofLength: 1).first)!)
+                logger.debug("Found track number: \(trackNumber!)")
+            case EBML.trackTypeID: // Unused by us, left for debugging
+                trackType = fileHandle.readData(ofLength: 1).first
+                logger.debug("Found track type: \(trackType!)")
+            case EBML.codecID:
+                var data = fileHandle.readData(ofLength: Int(elementSize))
+                data.removeNullBytes()
+                codecId = String(data: data, encoding: .ascii)
+                logger.debug("Found codec ID: \(codecId ?? "nil")")
+            default:
+                fileHandle.seek(toFileOffset: fileHandle.offsetInFile + elementSize)
+            }
+            if trackNumber != nil, trackType != nil, codecId != nil { break }
+        }
+
+        if let trackNumber, let codecId {
+            if codecId == codec {
+                return trackNumber
+            }
+        }
+        return nil
+    }
+
     private func extractClusterTimestamp() -> Int64? {
         if let (timestampElementSize, _) = findElement(withID: EBML.timestamp) as? (UInt64, UInt32) {
             return readFixedLengthNumber(fileHandle: fileHandle, length: Int(timestampElementSize))
@@ -154,14 +159,14 @@ class MKVTrackParser: MKVFileHandler {
 
                 trackData[trackNumber.firstIndex { $0 == Int(blockTrackNumber) }!].append(blockData)
             } else {
-                // Skip this block if it's for a different track
+                // Skip this block because it's for a different track
                 fileHandle.seek(toFileOffset: blockStartOffset + blockSize)
             }
         }
     }
 
     // Function to read the track number, timestamp, and lacing type (if any) from a Block or SimpleBlock header
-    func readTrackNumber(from fileHandle: FileHandle) -> (UInt64?, Int64) {
+    private func readTrackNumber(from fileHandle: FileHandle) -> (UInt64?, Int64) {
         let trackNumber = readVINT(from: fileHandle, unmodified: true)
         let timestamp = readFixedLengthNumber(fileHandle: fileHandle, length: 2)
         let suffix = fileHandle.readData(ofLength: 1).first ?? 0

Sources/macSubtitleOCR/macSubtitleOCR.swift

@@ -66,7 +66,7 @@ struct macSubtitleOCR: ParsableCommand {
             for track in mkvStream.tracks {
                 subIndex = 1 // reset counter for each track
                 logger.debug("Found subtitle track: \(track.trackNumber), Codec: \(track.codecId)")
-                intermediateFiles[track.trackNumber] = try mkvStream.getSubtitleTrackData(trackNumber: track.trackNumber, outPath: input)!
+                intermediateFiles[track.trackNumber] = try mkvStream.getSubtitleTrackData(trackNumber: track.trackNumber)!
 
                 // Open the PGS data stream
                 let PGS = try PGS(URL(fileURLWithPath: intermediateFiles[track.trackNumber]!))