Musubi/Packages/MusubiKit/Sources/AudioMatrix/OutputEngine.swift

import AudioToolbox
import CoreAudio
import os

/// Plays the routes of one output device through an AUHAL unit.
final class OutputEngine: @unchecked Sendable {
    let device: AudioDevice
    let sampleRate: Double
    let channels: Int

    private var unit: AudioUnit?
    private let routes = OSAllocatedUnfairLock<[RouteRenderer]>(initialState: [])
    private var renderRoutes: [RouteRenderer] = [] // render thread only
    /// Peak per output channel, written on the render thread, read and reset by the UI.
    let channelPeaks: UnsafeMutablePointer<Float>

    init(device: AudioDevice, bufferFrames: UInt32 = 256) throws {
        self.device = device
        self.sampleRate = device.sampleRate
        self.channels = device.outputChannels
        channelPeaks = .allocate(capacity: max(device.outputChannels, 1))
        channelPeaks.initialize(repeating: 0, count: max(device.outputChannels, 1))

        var description = AudioComponentDescription(
            componentType: kAudioUnitType_Output, componentSubType: kAudioUnitSubType_HALOutput,
            componentManufacturer: kAudioUnitManufacturer_Apple, componentFlags: 0, componentFlagsMask: 0)
        guard let component = AudioComponentFindNext(nil, &description) else { throw AudioError.noHAL }
        var newUnit: AudioUnit?
        try check(AudioComponentInstanceNew(component, &newUnit), "instance")
        guard let newUnit else { throw AudioError.noHAL }
        unit = newUnit

        var deviceID = device.deviceID
        try check(AudioUnitSetProperty(newUnit, kAudioOutputUnitProperty_CurrentDevice, kAudioUnitScope_Global, 0,
                                       &deviceID, UInt32(MemoryLayout<AudioDeviceID>.size)), "device")

        var format = AudioStreamBasicDescription.floatNonInterleaved(sampleRate: sampleRate, channels: channels)
        try check(AudioUnitSetProperty(newUnit, kAudioUnitProperty_StreamFormat, kAudioUnitScope_Input, 0,
                                       &format, UInt32(MemoryLayout<AudioStreamBasicDescription>.size)), "format")

        var frames = bufferFrames
        AudioUnitSetProperty(newUnit, kAudioDevicePropertyBufferFrameSize, kAudioUnitScope_Global, 0,
                             &frames, UInt32(MemoryLayout<UInt32>.size))

        var callback = AURenderCallbackStruct(inputProc: outputRender, inputProcRefCon: Unmanaged.passUnretained(self).toOpaque())
        try check(AudioUnitSetProperty(newUnit, kAudioUnitProperty_SetRenderCallback, kAudioUnitScope_Input, 0,
                                       &callback, UInt32(MemoryLayout<AURenderCallbackStruct>.size)), "callback")
        try check(AudioUnitInitialize(newUnit), "initialize")
        try check(AudioOutputUnitStart(newUnit), "start")
    }

    deinit {
        stop()
        channelPeaks.deallocate()
    }

    func stop() {
        guard let unit else { return }
        AudioOutputUnitStop(unit)
        AudioUnitUninitialize(unit)
        AudioComponentInstanceDispose(unit)
        self.unit = nil
    }

    func setRoutes(_ newRoutes: [RouteRenderer]) {
        routes.withLock { $0 = newRoutes }
    }

    fileprivate func render(frames: Int, buffers: UnsafeMutableAudioBufferListPointer) {
        if let current = routes.withLockIfAvailable({ $0 }) { renderRoutes = current }

        for buffer in buffers {
            if let data = buffer.mData { memset(data, 0, Int(buffer.mDataByteSize)) }
        }
        for route in renderRoutes where route.outputChannel < buffers.count {
            guard let data = buffers[route.outputChannel].mData else { continue }
            route.render(into: data.assumingMemoryBound(to: Float.self), frames: frames, outputRate: sampleRate)
        }
        for (index, buffer) in buffers.enumerated() where index < channels {
            guard let data = buffer.mData?.assumingMemoryBound(to: Float.self) else { continue }
            var peak: Float = 0
            for i in 0..<frames { peak = max(peak, abs(data[i])) }
            channelPeaks[index] = max(channelPeaks[index], peak)
        }
    }
}

private func outputRender(
    refCon: UnsafeMutableRawPointer,
    flags: UnsafeMutablePointer<AudioUnitRenderActionFlags>,
    timestamp: UnsafePointer<AudioTimeStamp>,
    bus: UInt32,
    frames: UInt32,
    data: UnsafeMutablePointer<AudioBufferList>?
) -> OSStatus {
    guard let data else { return noErr }
    let engine = Unmanaged<OutputEngine>.fromOpaque(refCon).takeUnretainedValue()
    engine.render(frames: Int(frames), buffers: UnsafeMutableAudioBufferListPointer(data))
    return noErr
}

enum AudioError: Error {
    case noHAL
    case coreAudio(String, OSStatus)
}

func check(_ status: OSStatus, _ step: String) throws {
    if status != noErr { throw AudioError.coreAudio(step, status) }
}

extension AudioStreamBasicDescription {
    static func floatNonInterleaved(sampleRate: Double, channels: Int) -> AudioStreamBasicDescription {
        AudioStreamBasicDescription(
            mSampleRate: sampleRate, mFormatID: kAudioFormatLinearPCM,
            mFormatFlags: kAudioFormatFlagIsFloat | kAudioFormatFlagIsPacked | kAudioFormatFlagIsNonInterleaved,
            mBytesPerPacket: 4, mFramesPerPacket: 1, mBytesPerFrame: 4,
            mChannelsPerFrame: UInt32(channels), mBitsPerChannel: 32, mReserved: 0)
    }
}