Improve performance and logging

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

Sources/macSubtitleOCR/Extensions/DataExtensions.swift

@@ -14,35 +14,9 @@ extension Data {
         self = filter { $0 != 0x00 }
     }
 
-    func value<T: BinaryInteger>(ofType _: T.Type, at offset: Int = 0, convertEndian: Bool = false) -> T? {
-        let right = offset &+ MemoryLayout<T>.size
-        guard offset >= 0, right > offset, right <= count else {
-            return nil
-        }
-        let bytes = self[offset ..< right]
-        if convertEndian {
-            return bytes.reversed().reduce(0) { T($0) << 8 + T($1) }
-        } else {
-            return bytes.reduce(0) { T($0) << 8 + T($1) }
-        }
-    }
-
-    // Extracts and removes a certain number of bytes from the beginning of the Data object.
-    //
-    // - Parameter count: The number of bytes to extract and remove.
-    // - Returns: A Data object containing the extracted bytes, or empty if there aren't enough bytes.
-    mutating func extractBytes(_ count: Int) -> Data {
-        guard count > 0, count <= self.count else {
-            return Data()
-        }
-
-        // Extract the range from the beginning
-        let extractedData = subdata(in: 0 ..< count)
-
-        // Remove the extracted bytes from the original data
-        removeSubrange(0 ..< count)
-
-        return extractedData
+    func getUInt16BE(at offset: Int = 0) -> UInt16? {
+        guard count >= offset + 2 else { return nil }
+        return withUnsafeBytes { $0.load(fromByteOffset: offset, as: UInt16.self).bigEndian }
     }
 
     /* Useful for debugging purposes

Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift

@@ -10,14 +10,13 @@ import Foundation
 import os
 
 class MKVSubtitleExtractor: MKVTrackParser {
-    func getSubtitleTrackData(trackNumber: Int, outputDirectory: URL) throws -> String? {
+    func saveSubtitleTrackData(trackNumber: Int, outputDirectory: URL) {
         let trackPath = outputDirectory.appendingPathComponent("\(trackNumber)").appendingPathExtension("sup").path
 
         if FileManager.default.createFile(atPath: trackPath, contents: tracks[trackNumber].trackData, attributes: nil) {
-            logger.debug("Created file at path: \(trackPath).")
-            return trackPath
+            logger.debug("Created file at path: \(trackPath)")
         } else {
-            fatalError("Failed to create file at path: \(trackPath).")
+            logger.error("Failed to create file at path: \(trackPath)!")
         }
     }
 }

Sources/macSubtitleOCR/MKV/MKVTrackParser.swift

@@ -153,7 +153,7 @@ class MKVTrackParser: MKVFileHandler {
                 let raw = fileHandle.readData(ofLength: Int(blockSize - (fileHandle.offsetInFile - blockStartOffset)))
                 var offset = 0
                 while (offset + 3) <= raw.count {
-                    let segmentSize = min(Int(raw.value(ofType: UInt16.self, at: offset + 1)! + 3), raw.count - offset)
+                    let segmentSize = min(Int(raw.getUInt16BE(at: offset + 1)! + 3), raw.count - offset)
                     logger.debug("Segment size \(segmentSize) at \(offset) type 0x\(String(format: "%02x", raw[offset]))")
 
                     blockData.append(pgsHeader)

Sources/macSubtitleOCR/Subtitles/PGS/PGS.swift

@@ -14,7 +14,7 @@ struct PGS {
 
     private(set) var subtitles = [Subtitle]()
     private let logger = Logger(subsystem: "github.ecdye.macSubtitleOCR", category: "PGS")
-    private var data: Data
+    private var data: Data?
     private let pgsHeaderLength = 13
 
     // MARK: - Lifecycle
@@ -23,107 +23,113 @@ struct PGS {
         let fileHandle = try FileHandle(forReadingFrom: url)
         defer { fileHandle.closeFile() }
         data = try fileHandle.readToEnd() ?? Data()
-        guard data.count > pgsHeaderLength else {
-            fatalError("Failed to read file data from: \(url.path)")
+        guard data!.count > pgsHeaderLength else {
+            fatalError("Failed to read valid subtitle data from: \(url.path)")
         }
         fileHandle.closeFile()
-        try data.withUnsafeBytes { (pointer: UnsafeRawBufferPointer) in
-            try parseData(pointer)
+        try data!.withUnsafeBytes { (buffer: UnsafeRawBufferPointer) in
+            try parseData(buffer)
         }
-
-        // try parseData()
     }
 
-    init(_ pointer: UnsafeRawBufferPointer) throws {
-        self.data = Data()
-        try parseData(pointer)
+    init(_ buffer: UnsafeRawBufferPointer) throws {
+        try parseData(buffer)
     }
 
     // MARK: - Methods
 
-    private mutating func parseData(_ pointer: UnsafeRawBufferPointer) throws {
+    private mutating func parseData(_ buffer: UnsafeRawBufferPointer) throws {
         var offset = 0
-        while offset + pgsHeaderLength < pointer.count {
-            guard let subtitle = try parseNextSubtitle(pointer, &offset)
+        while offset + pgsHeaderLength < buffer.count {
+            logger.debug("Parsing subtitle at offset: \(offset)")
+            guard let subtitle = try parseNextSubtitle(buffer, &offset)
             else {
-                if offset + pgsHeaderLength > pointer.count { break }
+                if offset + pgsHeaderLength > buffer.count { break }
                 continue
             }
 
             // Find the next timestamp to use as our end timestamp
-            subtitle.endTimestamp = parseTimestamp(pointer.loadUnaligned(fromByteOffset: offset + 2, as: UInt32.self).bigEndian)
+            subtitle.endTimestamp = getSegmentTimestamp(from: buffer, offset: offset)
 
             subtitles.append(subtitle)
         }
     }
 
-    private func parseTimestamp(_ timestamp: UInt32) -> TimeInterval {
-        return TimeInterval(timestamp) / 90000 // 90 kHz clock
-    }
-
-    private mutating func parseNextSubtitle(_ pointer: UnsafeRawBufferPointer, _ offset: inout Int) throws -> Subtitle? {
-        var multipleODS = false
+    private func parseNextSubtitle(_ buffer: UnsafeRawBufferPointer, _ offset: inout Int) throws -> Subtitle? {
+        var hasMultipleODS = false
         var ods: ODS?
         var pds: PDS?
 
         while true {
-            guard offset + pgsHeaderLength < pointer.count else {
+            guard offset + pgsHeaderLength < buffer.count else {
                 return nil // End of data
             }
 
-            let segmentType = pointer[offset + 10]
-            let segmentLength = Int(pointer.loadUnaligned(fromByteOffset: offset + 11, as: UInt16.self).bigEndian)
-            let startTimestamp = parseTimestamp(pointer.loadUnaligned(fromByteOffset: offset + 2, as: UInt32.self).bigEndian)
+            let segmentType = buffer[offset + 10]
+            let segmentLength = getSegmentLength(from: buffer, offset: offset)
+            let startTimestamp = getSegmentTimestamp(from: buffer, offset: offset)
             offset += pgsHeaderLength
 
-            // Check for the end of the subtitle stream (0x80 segment type and 0 length)
+            // End of stream check
             guard segmentType != 0x80, segmentLength != 0 else { return nil }
 
             // Parse the segment based on the type (0x14 for PCS, 0x15 for WDS, 0x16 for PDS, 0x17 for ODS)
             switch segmentType {
-            case 0x14: // PDS (Palette Definition Segment)
+            case 0x14:
                 do {
-                    pds = try PDS(pointer, offset, segmentLength)
+                    pds = try PDS(buffer, offset, segmentLength)
                     offset += segmentLength
                 } catch let macSubtitleOCRError.invalidPDSDataLength(length) {
-                    fatalError("Invalid Palette Data Segment length: \(length)")
+                    logger.error("Invalid PDS length: \(length), abandoning remaining segments!")
+                    offset = buffer.count
+                    return nil
                 }
-            case 0x15: // ODS (Object Definition Segment)
+            case 0x15:
                 do {
-                    if pointer[offset + 3] == 0x80 {
-                        ods = try ODS(pointer, offset, segmentLength)
+                    if buffer[offset + 3] == 0x80 {
+                        ods = try ODS(buffer, offset, segmentLength)
                         offset += segmentLength
-                        multipleODS = true
+                        hasMultipleODS = true
                         break
-                    } else if multipleODS {
-                        try ods?.appendSegment(pointer, offset, segmentLength)
-                        if pointer[offset + 3] != 0x40 { break }
+                    } else if hasMultipleODS {
+                        try ods!.appendSegment(buffer, offset, segmentLength)
                         offset += segmentLength
+                        if buffer[offset - 10] != 0x40 { break }
                     } else {
-                        ods = try ODS(pointer, offset, segmentLength)
+                        ods = try ODS(buffer, offset, segmentLength)
                         offset += segmentLength
                     }
                 } catch let macSubtitleOCRError.invalidODSDataLength(length) {
-                    fatalError("Invalid Object Data Segment length: \(length)")
+                    logger.error("Invalid ODS length: \(length), abandoning remaining segments!")
+                    offset = buffer.count
+                    return nil
                 }
-            case 0x16, 0x17: // PCS (Presentation Composition Segment), WDS (Window Definition Segment)
+            case 0x16, 0x17:
                 offset += segmentLength
-                break // PCS and WDS parsing not required for basic rendering
             default:
                 logger.warning("Unknown segment type: \(segmentType, format: .hex), skipping...")
                 offset += segmentLength
                 return nil
             }
+
             guard let pds, let ods else { continue }
             offset += pgsHeaderLength // Skip the end segment
             return Subtitle(
                 index: subtitles.count + 1,
                 startTimestamp: startTimestamp,
                 imageWidth: ods.objectWidth,
                 imageHeight: ods.objectHeight,
-                imageData: ods.imageData,
+                imageData: ods.decodeRLEData(),
                 imagePalette: pds.palette,
                 numberOfColors: 256)
         }
     }
+
+    private func getSegmentTimestamp(from pointer: UnsafeRawBufferPointer, offset: Int) -> TimeInterval {
+        TimeInterval(pointer.loadUnaligned(fromByteOffset: offset + 2, as: UInt32.self).bigEndian) / 90000 // 90 kHz clock
+    }
+
+    private func getSegmentLength(from pointer: UnsafeRawBufferPointer, offset: Int) -> Int {
+        Int(pointer.loadUnaligned(fromByteOffset: offset + 11, as: UInt16.self).bigEndian)
+    }
 }

Sources/macSubtitleOCR/Subtitles/PGS/Parsers/ODS.swift

@@ -13,17 +13,28 @@ struct ODS {
 
     private(set) var objectWidth: Int = 0
     private(set) var objectHeight: Int = 0
-    private var rawImageData: Data = .init()
-    private(set) var imageData: Data = .init()
+    private var encodedImageData: Data = .init()
 
     // MARK: - Lifecycle
 
-    init(_ data: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) throws {
-        try parseODS(data, offset, segmentLength)
+    init(_ buffer: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) throws {
+        guard segmentLength > 11 else {
+            throw macSubtitleOCRError.invalidODSDataLength(length: segmentLength)
+        }
+        try parseODS(buffer, offset, segmentLength)
+    }
+
+    mutating func appendSegment(_ buffer: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) throws {
+        guard segmentLength > 11 else {
+            throw macSubtitleOCRError.invalidODSDataLength(length: segmentLength)
+        }
+        try parseODS(buffer, offset, segmentLength)
     }
 
-    mutating func appendSegment(_ data: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) throws {
-        try parseODS(data, offset, segmentLength)
+    // Decodes the run-length encoded (RLE) image data
+    func decodeRLEData() -> Data {
+        let rleImageData = RLEData(data: encodedImageData, width: objectWidth, height: objectHeight)
+        return rleImageData.decodePGS()
     }
 
     // MARK: - Methods
@@ -39,32 +50,17 @@ struct ODS {
     //   2 bytes: Object width
     //   2 bytes: Object height
     //   Rest: Image data (run-length encoded, RLE)
-    private mutating func parseODS(_ data: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) throws {
-        let sequenceFlag = data[offset + 3]
-        if sequenceFlag != 0x40 {
-            objectWidth = Int(data.loadUnaligned(fromByteOffset: offset + 7, as: UInt16.self).bigEndian)
-            objectHeight = Int(data.loadUnaligned(fromByteOffset: offset + 9, as: UInt16.self).bigEndian)
-        }
-
-        // PGS includes the width and height as part of the image data length calculations
-        guard data.count - offset > 7 else {
-            throw macSubtitleOCRError.invalidODSDataLength(length: data.count - offset)
-        }
+    private mutating func parseODS(_ buffer: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) throws {
+        let sequenceFlag = buffer[offset + 3]
 
-        switch sequenceFlag {
-        case 0x40:
-            rawImageData.append(contentsOf: data[(offset + 4)..<(offset + segmentLength)])
-            imageData = decodeRLEData()
-        case 0x80:
-            rawImageData.append(contentsOf: data[(offset + 11)..<(offset + segmentLength)])
-        default:
-            rawImageData.append(contentsOf: data[(offset + 11)..<(offset + segmentLength)])
-            imageData = decodeRLEData()
+        // Only update object dimensions if this is the first part of the sequence
+        if sequenceFlag != 0x40 {
+            objectWidth = Int(buffer.loadUnaligned(fromByteOffset: offset + 7, as: UInt16.self).bigEndian)
+            objectHeight = Int(buffer.loadUnaligned(fromByteOffset: offset + 9, as: UInt16.self).bigEndian)
         }
-    }
 
-    private func decodeRLEData() -> Data {
-        let rleImageData = RLEData(data: rawImageData, width: objectWidth, height: objectHeight)
-        return rleImageData.decodePGS()
+        // Append image data to the encoded image data buffer
+        let dataOffset = sequenceFlag == 0x40 ? offset + 4 : offset + 11
+        encodedImageData.append(contentsOf: buffer[dataOffset ..< (offset + segmentLength)])
     }
 }

Sources/macSubtitleOCR/Subtitles/PGS/Parsers/PDS.swift

@@ -16,12 +16,12 @@ struct PDS {
 
     // MARK: - Lifecycle
 
-    init(_ data: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) throws {
-        let count = data.count - offset
+    init(_ buffer: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) throws {
+        let count = buffer.count - offset
         guard count >= 7, (segmentLength - 2) % 5 == 0 else {
             throw macSubtitleOCRError.invalidPDSDataLength(length: count)
         }
-        parsePDS(data, offset, segmentLength)
+        parsePDS(buffer, offset, segmentLength)
     }
 
     // MARK: - Methods
@@ -33,15 +33,15 @@ struct PDS {
     //   1 byte: Palette Version (unused by us)
     //   Followed by a series of palette entries:
     //       Each entry is 5 bytes: (Index, Y, Cr, Cb, Alpha)
-    private mutating func parsePDS(_ data: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) {
+    private mutating func parsePDS(_ buffer: UnsafeRawBufferPointer, _ offset: Int, _ segmentLength: Int) {
         // Start reading after the first 2 bytes (Palette ID and Version)
         var i = offset + 2
         while i + 4 <= (offset + segmentLength) {
-            let index = data[i]
-            let y = data[i + 1]
-            let cr = data[i + 2]
-            let cb = data[i + 3]
-            let alpha = data[i + 4]
+            let index = buffer[i]
+            let y = buffer[i + 1]
+            let cr = buffer[i + 2]
+            let cb = buffer[i + 3]
+            let alpha = buffer[i + 4]
 
             // Convert YCrCb to RGB
             let rgb = yCrCbToRGB(y: y, cr: cr, cb: cb)

Sources/macSubtitleOCR/Subtitles/VobSub/VobSub.swift

@@ -23,26 +23,26 @@ struct VobSub {
         subFile.closeFile()
         let idx = VobSubIDX(URL(filePath: idx))
         subData.withUnsafeBytes { (pointer: UnsafeRawBufferPointer) in
-            var pointer = pointer
-            extractSubtitleImages(subData: &pointer, idx: idx)
+            extractSubtitleImages(buffer: pointer, idx: idx)
         }
     }
 
     // MARK: - Methods
 
-    private mutating func extractSubtitleImages(subData: inout UnsafeRawBufferPointer, idx: VobSubIDX) {
+    private mutating func extractSubtitleImages(buffer: UnsafeRawBufferPointer, idx: VobSubIDX) {
         for index in idx.offsets.indices {
-            logger.debug("Index \(index), offset: \(idx.offsets[index]), timestamp: \(idx.timestamps[index])")
             let offset = idx.offsets[index]
             let timestamp = idx.timestamps[index]
+            logger.debug("Parsing subtitle \(index + 1), offset: \(offset), timestamp: \(timestamp)")
+
             let nextOffset: UInt64 = if index + 1 < idx.offsets.count {
                 idx.offsets[index + 1]
             } else {
-                UInt64(subData.count)
+                UInt64(buffer.count)
             }
             let subtitle = VobSubParser(
                 index: index + 1,
-                subData: &subData,
+                buffer: buffer,
                 timestamp: timestamp,
                 offset: offset,
                 nextOffset: nextOffset,