stemify/ios/Stemify/Engine/StemDecoder.swift

import AudioToolbox
import Foundation
import os

/// os_unfair_lock in a stable heap slot. Short critical sections only; the render thread takes it.
final class UnfairLock: @unchecked Sendable {
    private let p: UnsafeMutablePointer<os_unfair_lock>
    init() {
        p = .allocate(capacity: 1)
        p.initialize(to: os_unfair_lock())
    }
    deinit { p.deallocate() }
    @inline(__always) func lock() { os_unfair_lock_lock(p) }
    @inline(__always) func unlock() { os_unfair_lock_unlock(p) }
    @inline(__always) func with<T>(_ f: () throws -> T) rethrows -> T {
        os_unfair_lock_lock(p); defer { os_unfair_lock_unlock(p) }
        return try f()
    }
}

/// Single-producer single-consumer ring of interleaved stereo float frames.
final class FloatRing: @unchecked Sendable {
    private let buf: UnsafeMutablePointer<Float>
    private let cap: Int
    private var readPos: Int64 = 0
    private var writePos: Int64 = 0
    private let lock = UnfairLock()
    private let notFull = DispatchSemaphore(value: 0)
    private var eosFlag = false

    init(capacityFrames: Int) {
        cap = capacityFrames * 2
        buf = .allocate(capacity: cap)
        buf.initialize(repeating: 0, count: cap)
    }
    deinit { buf.deallocate() }

    var eos: Bool {
        get { lock.with { eosFlag } }
        set { lock.with { eosFlag = newValue } }
    }

    var availableFrames: Int { lock.with { Int((writePos - readPos) / 2) } }
    var freeFrames: Int { lock.with { Int((Int64(cap) - (writePos - readPos)) / 2) } }

    /// Blocks while full; returns false if [abort] became true first. [len] counts floats.
    func write(_ src: UnsafePointer<Float>, _ len: Int, abort: () -> Bool) -> Bool {
        var done = 0
        while done < len {
            lock.lock()
            let free = Int(Int64(cap) - (writePos - readPos))
            if free <= 0 {
                lock.unlock()
                if abort() { return false }
                _ = notFull.wait(timeout: .now() + .milliseconds(20))
                continue
            }
            lock.unlock()
            if abort() { return false }
            lock.lock()
            let n = min(Int(Int64(cap) - (writePos - readPos)), len - done)
            let w = Int(writePos % Int64(cap))
            let first = min(n, cap - w)
            (buf + w).update(from: src + done, count: first)
            if n > first { buf.update(from: src + done + first, count: n - first) }
            writePos += Int64(n)
            done += n
            lock.unlock()
        }
        return true
    }

    /// Copies up to [frames] frames into dst; returns frames copied. Never blocks on data.
    func read(_ dst: UnsafeMutablePointer<Float>, _ frames: Int) -> Int {
        lock.lock()
        let n = Int(min(Int64(frames * 2), writePos - readPos))
        let r = Int(readPos % Int64(cap))
        let first = min(n, cap - r)
        dst.update(from: buf + r, count: first)
        if n > first { (dst + first).update(from: buf, count: n - first) }
        readPos += Int64(n)
        lock.unlock()
        if n > 0 { notFull.signal() }
        return n / 2
    }

    func clear() {
        lock.with {
            readPos = writePos
            eosFlag = false
        }
        notFull.signal()
    }

    func wake() { notFull.signal() }
}

/// The decoder's view of a cached resource. Retries brief network trouble before giving up.
final class ResourceDataSource: @unchecked Sendable {
    private let res: CachedResource
    private let onClose: () -> Void
    private let lock = NSLock()
    private var failure: Error?
    private var isClosed = false

    init(_ res: CachedResource, onClose: @escaping () -> Void = {}) {
        self.res = res
        self.onClose = onClose
    }

    var failed: Error? { lock.withLock { failure } }
    var closed: Bool { lock.withLock { isClosed } }

    /// Bytes read, 0 at end of file.
    func readAt(_ position: Int64, _ buffer: UnsafeMutableRawPointer, _ size: Int) throws -> Int {
        var attempt = 0
        while true {
            do {
                let n = try res.readAt(position, buffer, size)
                lock.withLock { failure = nil }
                if n > 0 && attempt == 0 { res.prefetch() }
                return max(n, 0)
            } catch {
                attempt += 1
                if closed || attempt > 8 {
                    lock.withLock { failure = error }
                    throw error
                }
                Thread.sleep(forTimeInterval: 0.25 * Double(attempt))
            }
        }
    }

    func size() -> Int64 {
        do { return try res.size() } catch {
            lock.withLock { failure = error }
            return -1
        }
    }

    func close() {
        let first = lock.withLock { () -> Bool in
            if isClosed { return false }
            isClosed = true
            return true
        }
        if first { onClose() }
    }
}

struct DecodeError: LocalizedError {
    let message: String
    init(_ m: String) { message = m }
    var errorDescription: String? { message }
}

/// Decodes one stem on its own thread into [ring], as stereo float at the file's own rate.
/// FLAC, Ogg Opus and whatever else Core Audio reads all go through ExtAudioFile, whose seeks
/// are sample-accurate, so every stem of a track resumes on exactly the same frame.
final class StemDecoder: @unchecked Sendable {
    let name: String
    let ring: FloatRing
    private let source: ResourceDataSource
    private let typeHint: AudioFileTypeID
    private let prepared = NSCondition()
    private var isPrepared = false
    private let state = UnfairLock()
    private var _sampleRate = 0
    private var _durationUs: Int64 = 0
    private var _error: Error?
    private var _readyGen = -1

    var sampleRate: Int { state.with { _sampleRate } }
    var durationUs: Int64 { state.with { _durationUs } }
    var error: Error? { state.with { _error } }
    /// Generation whose data is in the ring (or that hit EOS).
    var readyGen: Int { state.with { _readyGen } }

    private let seekLock = NSCondition()
    private var pendingGen = -1
    private var pendingFrame: Int64 = 0
    private var releasedFlag = false
    private var released: Bool { seekLock.withLock { releasedFlag } }
    private var thread: Thread?

    /// [hint] is the file's extension or format ("flac", "opus", "mp3", ...), so Core Audio
    /// doesn't have to guess from the bytes.
    init(name: String, source: ResourceDataSource, ringFrames: Int, hint: String) {
        self.name = name
        self.source = source
        ring = FloatRing(capacityFrames: ringFrames)
        typeHint = Self.fileType(hint)
    }

    static func fileType(_ hint: String) -> AudioFileTypeID {
        switch hint.lowercased() {
        case "flac": return kAudioFileFLACType
        case "opus", "ogg", "oga": return AudioFileTypeID(0x4F67_6773) // 'Oggs'
        case "mp3": return kAudioFileMP3Type
        case "m4a", "mp4", "alac": return kAudioFileM4AType
        case "aac": return kAudioFileAAC_ADTSType
        case "wav": return kAudioFileWAVEType
        case "aif", "aiff": return kAudioFileAIFFType
        case "caf": return kAudioFileCAFType
        default: return 0
        }
    }

    func start(gen: Int, frame: Int64) {
        seek(gen: gen, frame: frame)
        let t = Thread { [self] in run() }
        t.name = "stem-\(name)"
        t.qualityOfService = .userInteractive
        thread = t
        t.start()
    }

    /// Blocks until the stream's format is known. False on error or timeout.
    func awaitPrepared(timeoutMs: Int) -> Bool {
        let deadline = Date().addingTimeInterval(Double(timeoutMs) / 1000)
        prepared.lock()
        while !isPrepared { if !prepared.wait(until: deadline) { break } }
        let ok = isPrepared
        prepared.unlock()
        return ok && error == nil
    }

    private func markPrepared() {
        prepared.lock()
        isPrepared = true
        prepared.broadcast()
        prepared.unlock()
    }

    func seek(gen: Int, frame: Int64) {
        seekLock.lock()
        pendingGen = gen
        pendingFrame = frame
        seekLock.broadcast()
        seekLock.unlock()
        ring.wake()
    }

    private func seekPending(_ current: Int) -> Bool { seekLock.withLock { pendingGen != current || releasedFlag } }

    func release() {
        seekLock.lock()
        releasedFlag = true
        seekLock.broadcast()
        seekLock.unlock()
        source.close()
        ring.wake()
    }

    private static let readProc: AudioFile_ReadProc = { ud, pos, req, buf, actual in
        let src = Unmanaged<ResourceDataSource>.fromOpaque(ud).takeUnretainedValue()
        do {
            actual.pointee = UInt32(try src.readAt(pos, buf, Int(req)))
            return noErr
        } catch {
            actual.pointee = 0
            return kAudioFileUnspecifiedError
        }
    }

    private static let sizeProc: AudioFile_GetSizeProc = { ud in
        max(0, Unmanaged<ResourceDataSource>.fromOpaque(ud).takeUnretainedValue().size())
    }

    private func run() {
        var fileId: AudioFileID?
        var ext: ExtAudioFileRef?
        let srcRef = Unmanaged.passRetained(source)
        defer {
            if let ext { ExtAudioFileDispose(ext) }
            if let fileId { AudioFileClose(fileId) }
            srcRef.release()
            markPrepared()
        }
        do {
            var st = AudioFileOpenWithCallbacks(srcRef.toOpaque(), Self.readProc, nil, Self.sizeProc, nil, typeHint, &fileId)
            if st != noErr || fileId == nil {
                throw source.failed ?? DecodeError("Can't read \(name) (\(st))")
            }
            st = ExtAudioFileWrapAudioFileID(fileId!, false, &ext)
            guard st == noErr, let ext else { throw DecodeError("Can't decode \(name) (\(st))") }

            var fmt = AudioStreamBasicDescription()
            var sz = UInt32(MemoryLayout<AudioStreamBasicDescription>.size)
            st = ExtAudioFileGetProperty(ext, kExtAudioFileProperty_FileDataFormat, &sz, &fmt)
            guard st == noErr, fmt.mSampleRate > 0 else { throw DecodeError("No audio in \(name)") }
            let rate = Int(fmt.mSampleRate.rounded())
            var frames: Int64 = 0
            sz = UInt32(MemoryLayout<Int64>.size)
            ExtAudioFileGetProperty(ext, kExtAudioFileProperty_FileLengthFrames, &sz, &frames)
            // Asking for stereo makes the converter upmix mono and drop channels past the
            // first two, the same as Android's toStereo.
            var client = AudioStreamBasicDescription(
                mSampleRate: fmt.mSampleRate, mFormatID: kAudioFormatLinearPCM,
                mFormatFlags: kAudioFormatFlagIsFloat | kAudioFormatFlagIsPacked,
                mBytesPerPacket: 8, mFramesPerPacket: 1, mBytesPerFrame: 8,
                mChannelsPerFrame: 2, mBitsPerChannel: 32, mReserved: 0)
            st = ExtAudioFileSetProperty(ext, kExtAudioFileProperty_ClientDataFormat, UInt32(MemoryLayout.size(ofValue: client)), &client)
            guard st == noErr else { throw DecodeError("Can't convert \(name) (\(st))") }
            if fmt.mChannelsPerFrame > 2 {
                // Keep L/R instead of letting the converter downmix surround.
                var conv: AudioConverterRef?
                sz = UInt32(MemoryLayout<AudioConverterRef?>.size)
                if ExtAudioFileGetProperty(ext, kExtAudioFileProperty_AudioConverter, &sz, &conv) == noErr, let conv {
                    var map: [Int32] = [0, 1]
                    AudioConverterSetProperty(conv, kAudioConverterChannelMap, UInt32(MemoryLayout<Int32>.size * 2), &map)
                    var nul: Unmanaged<CFPropertyList>?
                    ExtAudioFileSetProperty(ext, kExtAudioFileProperty_ConverterConfig, UInt32(MemoryLayout<Unmanaged<CFPropertyList>?>.size), &nul)
                }
            }
            state.with {
                _sampleRate = rate
                _durationUs = frames > 0 ? frames * 1_000_000 / Int64(rate) : 0
            }
            markPrepared()

            let chunk = 4096
            let scratch = UnsafeMutablePointer<Float>.allocate(capacity: chunk * 2)
            defer { scratch.deallocate() }
            var gen = -1
            var outputDone = false
            var pos: Int64 = 0

            while !released {
                seekLock.lock()
                if pendingGen != gen {
                    gen = pendingGen
                    let target = pendingFrame
                    seekLock.unlock()
                    ring.clear()
                    st = ExtAudioFileSeek(ext, target)
                    if st != noErr { throw source.failed ?? DecodeError("Seek failed in \(name) (\(st))") }
                    pos = target
                    outputDone = false
                } else if outputDone {
                    ring.eos = true
                    state.with { _readyGen = gen }
                    _ = seekLock.wait(until: Date().addingTimeInterval(0.1))
                    seekLock.unlock()
                    continue
                } else {
                    seekLock.unlock()
                }
                if released { break }

                var n = UInt32(chunk)
                var abl = AudioBufferList(mNumberBuffers: 1, mBuffers: AudioBuffer(mNumberChannels: 2, mDataByteSize: UInt32(chunk * 8), mData: scratch))
                st = ExtAudioFileRead(ext, &n, &abl)
                if st != noErr {
                    if released { break }
                    throw source.failed ?? DecodeError("Decode failed in \(name) (\(st))")
                }
                if n == 0 {
                    // A short file is only an error if the source gave up well before the end.
                    if let f = source.failed, frames > 0, pos < frames - Int64(rate) * 2 { throw f }
                    outputDone = true
                    continue
                }
                pos += Int64(n)
                let myGen = gen
                if ring.write(scratch, Int(n) * 2, abort: { seekPending(myGen) }) {
                    state.with { _readyGen = gen }
                }
            }
        } catch {
            if !released {
                NSLog("StemDecoder %@ failed: %@", name, error.localizedDescription)
                state.with { _error = error }
            }
            ring.eos = true
            state.with { _readyGen = Int.max }
        }
    }
}