diff --git a/Sources/macSubtitleOCR/MKV/EBML/EBMLParser.swift b/Sources/macSubtitleOCR/MKV/EBML/EBMLParser.swift index 7ab0a1e..a57923b 100644 --- a/Sources/macSubtitleOCR/MKV/EBML/EBMLParser.swift +++ b/Sources/macSubtitleOCR/MKV/EBML/EBMLParser.swift @@ -13,34 +13,26 @@ 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))") @@ -48,11 +40,6 @@ func readVINT(from fileHandle: FileHandle, unmodified: Bool = false) -> UInt64 { 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) diff --git a/Sources/macSubtitleOCR/MKV/MKVFileHandler.swift b/Sources/macSubtitleOCR/MKV/MKVFileHandler.swift index 25168fe..d6fc838 100644 --- a/Sources/macSubtitleOCR/MKV/MKVFileHandler.swift +++ b/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) } diff --git a/Sources/macSubtitleOCR/MKV/MKVHelpers.swift b/Sources/macSubtitleOCR/MKV/MKVHelpers.swift index 0a7dbad..c1bef9c 100644 --- a/Sources/macSubtitleOCR/MKV/MKVHelpers.swift +++ b/Sources/macSubtitleOCR/MKV/MKVHelpers.swift @@ -15,18 +15,14 @@ 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 @@ -34,13 +30,7 @@ func readFixedLengthNumber(fileHandle: FileHandle, length: Int, signed: Bool = f // 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 diff --git a/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift b/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift index 0c62e6f..b672138 100644 --- a/Sources/macSubtitleOCR/MKV/MKVSubtitleExtractor.swift +++ b/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 } } diff --git a/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift b/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift index 1645dc3..1e7cdb7 100644 --- a/Sources/macSubtitleOCR/MKV/MKVTrackParser.swift +++ b/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 diff --git a/Sources/macSubtitleOCR/macSubtitleOCR.swift b/Sources/macSubtitleOCR/macSubtitleOCR.swift index 7194dce..9fa07ab 100644 --- a/Sources/macSubtitleOCR/macSubtitleOCR.swift +++ b/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]!))