Files
stemify/ios/Stemify/Engine/StemEngine.swift
T
cdnutterandClaude Opus 5.5 2ddfe051d8 iOS 1.0.1: mini player, Now Playing spacing, play spinner fix
- Mini player floats as a solid card and every screen leaves room for it,
  so it no longer covers the bottom of lists
- More space between title and scrubber; panel switches animate
- Play button no longer spins forever after changing tracks: the engine
  now always reports buffering on load
- Springier button presses and play/pause symbol transitions

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-20 08:50:42 -07:00

674 lines
23 KiB
Swift

import AVFoundation
import Foundation
@inline(__always) func monoNanos() -> Int64 { Int64(DispatchTime.now().uptimeNanoseconds) }
/// One track's stems, decoding in sync. Built off the main thread, then handed to the engine.
final class TrackSession: @unchecked Sendable {
let token: Int64
let names: [String]
let decoders: [StemDecoder]
var sampleRate: Int { decoders.first?.sampleRate ?? 0 }
var durationFrames: Int64 { (decoders.map(\.durationUs).max() ?? 0) * Int64(sampleRate) / 1_000_000 }
fileprivate var gen = 0
fileprivate var framePos: Int64 = 0
fileprivate var target: [Float]
fileprivate var current: [Float]
fileprivate var mixVersion = -1
init(token: Int64, names: [String], decoders: [StemDecoder]) {
self.token = token
self.names = names
self.decoders = decoders
target = Array(repeating: 1, count: names.count)
current = Array(repeating: 1, count: names.count)
}
func error() -> Error? { decoders.lazy.compactMap(\.error).first }
fileprivate func ready(_ need: Int) -> Bool {
decoders.allSatisfy { d in
d.error != nil || (d.readyGen == gen || d.readyGen == Int.max) && (d.ring.eos || d.ring.availableFrames >= need)
}
}
fileprivate func finished() -> Bool {
decoders.allSatisfy { d in
d.error != nil || (d.readyGen == gen || d.readyGen == Int.max) && d.ring.eos && d.ring.availableFrames == 0
}
}
/// Re-targets decoders; only valid before the session is handed to the engine.
func startAt(_ frame: Int64) {
gen += 1
framePos = frame
decoders.forEach { $0.seek(gen: gen, frame: frame) }
}
func release() { decoders.forEach { $0.release() } }
}
/// Per-stem RMS of short blocks as they're mixed, stamped with their track frame, so meters can
/// show what's coming out of the speaker now rather than what was just mixed.
final class MeterHistory: @unchecked Sendable {
static let block = 256
static let maxStems = 8
/// ~5 s at 48 kHz: well past the output buffer.
private static let cap = 1024
private static let staleFrames: Int64 = 8192
private var tokens = [Int64](repeating: 0, count: cap)
private var starts = [Int64](repeating: 0, count: cap)
private var values = [Float](repeating: 0, count: cap * maxStems)
private var head = 0
private var count = 0
private let lock = UnfairLock()
func clear() { lock.with { count = 0 } }
func add(_ token: Int64, _ start: Int64, _ rms: UnsafePointer<Float>, _ stems: Int) {
lock.with {
tokens[head] = token
starts[head] = start
let base = head * Self.maxStems
for k in 0..<Self.maxStems { values[base + k] = k < stems ? rms[k] : 0 }
head = (head + 1) % Self.cap
count = min(count + 1, Self.cap)
}
}
/// Fills [out] with the levels of [token]'s block playing at [frame]; false if there's none.
func read(_ token: Int64, _ frame: Int64, _ out: inout [Float]) -> Bool {
lock.with {
guard count > 0 else { return false }
for k in 1...count {
let i = (head - k + Self.cap) % Self.cap
if tokens[i] != token || starts[i] > frame { continue }
if frame - starts[i] > Self.staleFrames { return false }
for j in 0..<min(out.count, Self.maxStems) { out[j] = values[i * Self.maxStems + j] }
return true
}
return false
}
}
}
/// The far end of the output ring: the render callback takes mixed frames from it and counts
/// what it took, so the engine knows what has reached the speaker.
private final class OutputTap: @unchecked Sendable {
let ring: FloatRing
let rate: Int
private let lock = UnfairLock()
private var consumed: Int64 = 0
private var lastStart: Int64 = 0
private var lastNanos: Int64 = 0
private var runningFlag = false
private var latencyNanos: Int64 = 0
private let scratch: UnsafeMutablePointer<Float>
private static let scratchFrames = 4096
init(rate: Int, frames: Int) {
self.rate = rate
ring = FloatRing(capacityFrames: frames)
scratch = .allocate(capacity: Self.scratchFrames * 2)
}
deinit { scratch.deallocate() }
var running: Bool {
get { lock.with { runningFlag } }
set { lock.with { runningFlag = newValue } }
}
func setLatency(_ seconds: Double) { lock.with { latencyNanos = Int64(seconds * 1e9) } }
func flush() {
lock.with {
ring.clear()
consumed = 0
lastStart = 0
lastNanos = 0
}
}
/// Frames heard so far (since the last flush), never more than [written].
func presented(_ now: Int64, _ written: Int64) -> Int64 {
lock.with {
if lastNanos == 0 { return 0 }
let elapsed = runningFlag ? max(0, now - lastNanos) : 0
let est = lastStart + elapsed * Int64(rate) / 1_000_000_000 - latencyNanos * Int64(rate) / 1_000_000_000
return min(max(0, min(est, consumed)), written)
}
}
func render(_ silence: UnsafeMutablePointer<ObjCBool>, _ ts: UnsafePointer<AudioTimeStamp>, _ frames: AVAudioFrameCount, _ abl: UnsafeMutablePointer<AudioBufferList>) -> OSStatus {
let list = UnsafeMutableAudioBufferListPointer(abl)
let n = Int(frames)
guard list.count >= 2, let l = list[0].mData?.assumingMemoryBound(to: Float.self),
let r = list[1].mData?.assumingMemoryBound(to: Float.self) else { return noErr }
lock.lock()
if !runningFlag {
lock.unlock()
l.update(repeating: 0, count: n)
r.update(repeating: 0, count: n)
silence.pointee = true
return noErr
}
let host = ts.pointee.mFlags.contains(.hostTimeValid) ? Int64(AVAudioTime.seconds(forHostTime: ts.pointee.mHostTime) * 1e9) : monoNanos()
lastStart = consumed
lastNanos = host
var done = 0
while done < n {
let want = min(n - done, Self.scratchFrames)
let got = ring.read(scratch, want)
for i in 0..<got {
l[done + i] = scratch[i * 2]
r[done + i] = scratch[i * 2 + 1]
}
consumed += Int64(got)
done += got
if got < want { break }
}
lock.unlock()
if done < n {
(l + done).update(repeating: 0, count: n - done)
(r + done).update(repeating: 0, count: n - done)
}
return noErr
}
}
protocol StemEngineListener: AnyObject {
func onEngineState()
/// [token]'s audio is now what's coming out of the speaker.
func onTransition(_ token: Int64)
func onEnded(_ token: Int64)
func onError(_ token: Int64, _ message: String)
}
/// Mixes a [TrackSession]'s stems into one output: per-stem gain (ramped), EQ with anti-clip
/// preamp, master gain, then a peak limiter. All output calls happen on the engine thread; the
/// public methods just queue commands.
final class StemEngine: @unchecked Sendable {
static let chunk = 1024
/// Per stem, decoded ahead of the mixer.
static let ringSeconds = 3
/// Frames each stem must have before (re)starting output.
static let prefill = 8192
static let bufferMs = 120
/// Where playback is; position = frame + (now - nanos) while [running].
struct Snapshot {
var token: Int64
var frame: Int64
var rate: Int
var nanos: Int64
var running: Bool
func positionMs(_ now: Int64 = monoNanos()) -> Int64 {
if rate <= 0 { return 0 }
let f = running ? frame + (now - nanos) * Int64(rate) / 1_000_000_000 : frame
return f * 1000 / Int64(rate)
}
}
private struct Segment { let writtenStart: Int64; let token: Int64; let frameStart: Int64 }
private weak var listener: StemEngineListener?
private let cmdLock = NSCondition()
private var commands: [() -> Void] = []
private var released = false
private var session: TrackSession?
private var next: TrackSession?
private var wantPlay = false
private let av = AVAudioEngine()
private var node: AVAudioSourceNode?
private var tap: OutputTap?
private var trackRate = 0
private var written: Int64 = 0
private var segments: [Segment] = []
private var drainUntil: Int64 = -1
private var drainDeadline: Int64 = 0
private var drainToken: Int64 = -1
private var waitingSince: Int64 = 0
private var errorReported: Int64 = -1
private let eq = Equalizer()
private let limiter = Limiter()
private var masterCur: Float = 1
private let out = UnsafeMutablePointer<Float>.allocate(capacity: chunk * 2)
private let tmp = UnsafeMutablePointer<Float>.allocate(capacity: chunk * 2)
private let meters = MeterHistory()
private var meterSums = [Float](repeating: 0, count: chunk / MeterHistory.block * MeterHistory.maxStems)
private let shared = UnfairLock()
private var _eqParams: EqParams?
private var _masterGain: Float = 1
private var _mix: [String: Float] = [:]
private var _mixVersion = 0
private var _snapshot = Snapshot(token: -1, frame: 0, rate: 44100, nanos: 0, running: false)
private var _playing = false
private var _buffering = false
var eqParams: EqParams? {
get { shared.with { _eqParams } }
set { shared.with { _eqParams = newValue } }
}
var masterGain: Float {
get { shared.with { _masterGain } }
set { shared.with { _masterGain = newValue } }
}
private(set) var snapshot: Snapshot {
get { shared.with { _snapshot } }
set { shared.with { _snapshot = newValue } }
}
var playing: Bool { shared.with { _playing } }
var buffering: Bool { shared.with { _buffering } }
init(listener: StemEngineListener) {
self.listener = listener
out.initialize(repeating: 0, count: Self.chunk * 2)
tmp.initialize(repeating: 0, count: Self.chunk * 2)
NotificationCenter.default.addObserver(forName: .AVAudioEngineConfigurationChange, object: av, queue: nil) { [weak self] _ in
// The route or hardware format changed and the engine stopped; start it again.
self?.post { self?.restartOutput() }
}
#if os(iOS)
NotificationCenter.default.addObserver(forName: AVAudioSession.routeChangeNotification, object: nil, queue: nil) { [weak self] _ in
self?.post { self?.tap?.setLatency(Self.outputLatency()) }
}
#endif
let t = Thread { [weak self] in self?.loop() }
t.name = "stem-engine"
t.qualityOfService = .userInteractive
t.start()
}
// ---- commands (any thread) --------------------------------------------------------
private func post(_ cmd: @escaping () -> Void) {
cmdLock.lock()
commands.append(cmd)
cmdLock.broadcast()
cmdLock.unlock()
}
/// Replaces whatever is playing with [s], which must already be prepared.
func load(_ s: TrackSession, play: Bool) {
post { [self] in
dropAll()
session = s
wantPlay = play
resetOutput(s)
// Starts out buffering so the step that begins playback always reports a change,
// even when the last track was already playing.
setFlags(false, play)
publish()
}
}
/// The track to continue into gaplessly when the current one ends (or nil).
func queueNext(_ s: TrackSession?) {
post { [self] in
if next !== s { next?.release() }
next = s
}
}
func setPlaying(_ play: Bool) {
post { [self] in
wantPlay = play
if !play { pauseOutput() }
publish()
}
}
func seekTo(_ ms: Int64) {
post { [self] in
guard let s = session else { return }
let frame = min(max(ms * Int64(s.sampleRate) / 1000, 0), max(0, s.durationFrames - 1))
s.gen += 1
s.framePos = frame
s.decoders.forEach { $0.seek(gen: s.gen, frame: frame) }
resetOutput(s)
publish()
}
}
func stop() {
post { [self] in
dropAll()
wantPlay = false
publish()
}
}
func setMix(_ gains: [String: Float]) {
shared.with {
_mix = gains
_mixVersion += 1
}
}
func release() {
cmdLock.lock()
released = true
cmdLock.broadcast()
cmdLock.unlock()
}
/// The meters' levels for what's audible right now, into [out] (one per stem). False when
/// nothing is playing, so the meters can fall.
func levelsNow(_ out: inout [Float]) -> Bool {
let snap = snapshot
if !snap.running || snap.rate <= 0 { return false }
let frame = snap.frame + (monoNanos() - snap.nanos) * Int64(snap.rate) / 1_000_000_000
return meters.read(snap.token, frame, &out)
}
// ---- engine thread ----------------------------------------------------------------
private var isReleased: Bool { cmdLock.withLock { released } }
private func dropAll() {
session?.release()
next?.release()
session = nil
next = nil
drainUntil = -1
pauseOutput()
tap?.flush()
written = 0
segments.removeAll()
}
private func resetOutput(_ s: TrackSession) {
pauseOutput()
tap?.flush()
written = 0
segments = [Segment(writtenStart: 0, token: s.token, frameStart: s.framePos)]
drainUntil = -1
waitingSince = 0
meters.clear()
errorReported = -1
}
private func loop() {
while !isReleased {
while true {
cmdLock.lock()
let cmd = commands.isEmpty ? nil : commands.removeFirst()
cmdLock.unlock()
guard let cmd else { break }
cmd()
}
do {
try step()
} catch {
let tok = session?.token ?? -1
dropAll()
wantPlay = false
publish()
let msg = error.localizedDescription
DispatchQueue.main.async { [weak self] in self?.listener?.onError(tok, msg) }
}
}
dropAll()
av.stop()
}
private func idle(_ ms: Int) {
cmdLock.lock()
if commands.isEmpty && !released { _ = cmdLock.wait(until: Date().addingTimeInterval(Double(ms) / 1000)) }
cmdLock.unlock()
}
private func step() throws {
guard let s = session, wantPlay else {
setFlags(false, false)
publish()
idle(100)
return
}
if let e = s.error() {
if errorReported != s.token {
errorReported = s.token
wantPlay = false
pauseOutput()
setFlags(false, false)
let tok = s.token, msg = e.localizedDescription
DispatchQueue.main.async { [weak self] in self?.listener?.onError(tok, msg.isEmpty ? "Playback failed" : msg) }
}
idle(50)
return
}
if drainUntil >= 0 {
drain(s)
return
}
let need = written == 0 || buffering ? Self.prefill : Self.chunk
if !s.ready(need) && !s.finished() {
let now = monoNanos()
if waitingSince == 0 { waitingSince = now }
if !buffering && (written == 0 || now - waitingSince > 60_000_000) {
tap?.running = false
setFlags(false, true)
}
publish()
idle(5)
return
}
waitingSince = 0
try ensureOutput(s.sampleRate)
try startOutput()
setFlags(true, false)
guard let tap else { return }
if tap.ring.freeFrames < Self.chunk {
publish()
idle(5)
return
}
let n = mixChunk(s)
if n == 0 {
if s.finished() { endOfTrack(s) } else { idle(2) }
return
}
eq.configure(eqParams, sampleRate: s.sampleRate)
eq.process(out, n)
let m = masterGain
rampGain(out, n, masterCur, m)
masterCur = m
limiter.process(out, n)
_ = tap.ring.write(out, n * 2, abort: { false })
written += Int64(n)
publish()
}
private func endOfTrack(_ s: TrackSession) {
if let nx = next, nx.sampleRate == s.sampleRate, nx.error() == nil {
// Gapless: keep writing into the same output.
segments.append(Segment(writtenStart: written, token: nx.token, frameStart: nx.framePos))
session = nx
next = nil
s.release()
return
}
// Let the output play out, then hand over.
drainUntil = written
drainToken = s.token
drainDeadline = monoNanos() + written * 1_000_000_000 / Int64(max(1, trackRate)) + 1_000_000_000
}
private func drain(_ s: TrackSession) {
let (presented, _) = self.presented()
publish()
if presented < drainUntil - 64 && monoNanos() < drainDeadline {
idle(10)
return
}
drainUntil = -1
if let nx = next, nx.error() == nil {
session = nx
next = nil
s.release()
resetOutput(nx)
let tok = nx.token
DispatchQueue.main.async { [weak self] in self?.listener?.onTransition(tok) }
} else {
let tok = drainToken
dropAll()
wantPlay = false
setFlags(false, false)
publish()
DispatchQueue.main.async { [weak self] in self?.listener?.onEnded(tok) }
}
}
private func mixChunk(_ s: TrackSession) -> Int {
if s.finished() { return 0 }
var n = Self.chunk
for d in s.decoders {
if d.error != nil || d.ring.eos && d.ring.availableFrames == 0 { continue }
n = min(n, d.ring.availableFrames)
}
if n <= 0 { return 0 }
let (gains, version) = shared.with { (_mix, _mixVersion) }
if s.mixVersion != version {
for i in s.names.indices { s.target[i] = gains[s.names[i]] ?? 1 }
s.mixVersion = version
}
out.update(repeating: 0, count: n * 2)
let stride = MeterHistory.maxStems
let block = MeterHistory.block
let blocks = (n + block - 1) / block
if meterSums.count < blocks * stride { meterSums = [Float](repeating: 0, count: blocks * stride) }
for k in meterSums.indices { meterSums[k] = 0 }
for (i, d) in s.decoders.enumerated() {
let got = d.ring.read(tmp, n)
if got < n { (tmp + got * 2).update(repeating: 0, count: (n - got) * 2) }
let g0 = s.current[i]
let g1 = s.target[i]
s.current[i] = g1
if g0 == 0 && g1 == 0 { continue }
let step = (g1 - g0) / Float(n)
var g = g0
for b in 0..<blocks {
let end = min(n, (b + 1) * block)
var sum: Float = 0
for f in (b * block)..<end {
let l = tmp[f * 2] * g
let r = tmp[f * 2 + 1] * g
out[f * 2] += l
out[f * 2 + 1] += r
sum += l * l + r * r
g += step
}
if i < stride { meterSums[b * stride + i] = (sum / Float(2 * (end - b * block))).squareRoot() }
}
}
meterSums.withUnsafeBufferPointer { p in
for b in 0..<blocks { meters.add(s.token, s.framePos + Int64(b * block), p.baseAddress! + b * stride, s.names.count) }
}
s.framePos += Int64(n)
return n
}
private func rampGain(_ buf: UnsafeMutablePointer<Float>, _ n: Int, _ from: Float, _ to: Float) {
if from == 1 && to == 1 { return }
let step = (to - from) / Float(n)
var g = from
for f in 0..<n {
buf[f * 2] *= g
buf[f * 2 + 1] *= g
g += step
}
}
private static func outputLatency() -> Double {
#if os(iOS)
let s = AVAudioSession.sharedInstance()
return s.outputLatency + s.ioBufferDuration
#else
return 0.01
#endif
}
private func ensureOutput(_ rate: Int) throws {
if node != nil && trackRate == rate { return }
av.stop()
if let node { av.detach(node) }
#if os(iOS)
try? AVAudioSession.sharedInstance().setPreferredSampleRate(Double(rate))
#endif
let t = OutputTap(rate: rate, frames: max(4096, rate * Self.bufferMs / 1000))
t.setLatency(Self.outputLatency())
guard let fmt = AVAudioFormat(standardFormatWithSampleRate: Double(rate), channels: 2) else {
throw DecodeError("Unsupported sample rate \(rate)")
}
let n = AVAudioSourceNode(format: fmt) { silence, ts, frames, abl in t.render(silence, ts, frames, abl) }
av.attach(n)
av.connect(n, to: av.mainMixerNode, format: fmt)
av.prepare()
node = n
tap = t
trackRate = rate
limiter.setSampleRate(rate)
}
private func startOutput() throws {
guard let tap else { return }
if !av.isRunning {
#if os(iOS)
try AVAudioSession.sharedInstance().setActive(true)
#endif
try av.start()
tap.setLatency(Self.outputLatency())
}
tap.running = true
}
private func pauseOutput() {
tap?.running = false
if av.isRunning { av.pause() }
}
private func restartOutput() {
tap?.setLatency(Self.outputLatency())
// A paused engine also gets stopped by a route change; step() starts it again when it plays.
if wantPlay, tap?.running == true, !av.isRunning { try? av.start() }
}
/// (frames presented since the last flush, at monoNanos).
private func presented() -> (Int64, Int64) {
let now = monoNanos()
guard let tap else { return (0, now) }
return (tap.presented(now, written), now)
}
private func publish() {
let (p, nanos) = presented()
while segments.count > 1 && segments[1].writtenStart <= p {
segments.removeFirst()
let tok = segments[0].token
DispatchQueue.main.async { [weak self] in self?.listener?.onTransition(tok) }
}
let s = session
let rate = s.map { $0.decoders.isEmpty ? trackRate : $0.sampleRate } ?? trackRate
if let seg = segments.first {
let running = playing && drainUntil < 0 || drainUntil >= 0 && p < drainUntil
snapshot = Snapshot(token: seg.token, frame: seg.frameStart + max(0, p - seg.writtenStart), rate: rate, nanos: nanos, running: running)
} else {
snapshot = Snapshot(token: s?.token ?? -1, frame: s?.framePos ?? 0, rate: rate, nanos: nanos, running: false)
}
}
private func setFlags(_ playing: Bool, _ buffering: Bool) {
let changed = shared.with { () -> Bool in
if playing == _playing && buffering == _buffering { return false }
_playing = playing
_buffering = buffering
return true
}
if changed { DispatchQueue.main.async { [weak self] in self?.listener?.onEngineState() } }
}
}