AutoGain.h

Go to the documentation of this file.

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Core/AudioBuffer.h"
#include "../Core/AudioSpec.h"
#include "../Core/DspMath.h"
#include "../Core/SimdOps.h"
#include "../Core/StateBlob.h"
 
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class AutoGain
{
    // Ensure the atomic type won't trigger a hidden mutex lock in the audio thread
    static_assert(std::atomic<T>::is_always_lock_free,
        "AutoGain requires a lock-free float type for thread safety in the audio path.");
 
public:
    void prepare(const AudioSpec& spec) noexcept
    {
        if (!spec.isValid()) return; // release-safe: keep previous state
 
        sampleRate_ = spec.sampleRate;
        numChannels_ = spec.numChannels;
        reset();
    }
 
    void pushReference(AudioBufferView<T> buffer) noexcept
    {
        if (std::min(buffer.getNumChannels(), numChannels_) <= 0 ||
            buffer.getNumSamples() <= 0)
            return;
        refLevelDb_ = measureRmsDb(buffer);
    }
 
    void compensate(AudioBufferView<T> buffer) noexcept
    {
        const int numCh = std::min(buffer.getNumChannels(), numChannels_);
        const int numSamples = buffer.getNumSamples();
 
        if (numSamples == 0 || numCh == 0) return;
 
        T outLevelDb = measureRmsDb(buffer);
        T targetDb = refLevelDb_ - outLevelDb;
 
        // Safety: Prevent NaN propagation if both levels are -Inf
        if (std::isnan(targetDb))
            targetDb = T(0);
 
        // Clamp to safety limits
        const T maxComp = maxCompensation_.load(std::memory_order_relaxed);
        targetDb = std::clamp(targetDb, -maxComp, maxComp);
 
        // Silence bypass (-90 dB threshold)
        if (refLevelDb_ < SILENCE_THRESH_DB && outLevelDb < SILENCE_THRESH_DB)
            targetDb = T(0);
 
        // Calculate analytical end-state of the one-pole filter for the current block size:
        // alpha = exp(-N / (Fs * tau))
        const T smoothSecs = smoothTimeSecs_.load(std::memory_order_relaxed);
        const T alpha = static_cast<T>(std::exp(-static_cast<double>(numSamples)
                                  / (sampleRate_ * static_cast<double>(smoothSecs))));
        const T endCompensationDb = targetDb + (compensationDb_ - targetDb) * alpha;
 
        // Convert dB to linear gain for interpolation
        const T startGain = decibelsToGain(compensationDb_);
        const T endGain = decibelsToGain(endCompensationDb);
        const T gainStep = (endGain - startGain) / static_cast<T>(numSamples);
 
        // Apply linearly interpolated gain.
        // This loop structure guarantees no loop-carried dependencies, enabling strict SIMD vectorization.
        for (int ch = 0; ch < numCh; ++ch)
        {
            T* data = buffer.getChannel(ch);
            for (int i = 0; i < numSamples; ++i)
            {
                data[i] *= (startGain + static_cast<T>(i) * gainStep);
            }
        }
 
        // Update internal state for the next block
        compensationDb_ = endCompensationDb;
    }
 
    void reset() noexcept
    {
        refLevelDb_ = SILENCE_THRESH_DB;
        compensationDb_ = T(0);
    }
 
    [[nodiscard]] T getCompensationDb() const noexcept { return compensationDb_; }
 
    void setMaxCompensation(T dB) noexcept
    {
        if (!std::isfinite(dB)) return;
        maxCompensation_.store(std::abs(dB), std::memory_order_relaxed);
    }
 
    void setSmoothingTime(T ms) noexcept
    {
        if (!std::isfinite(ms)) return;
        smoothTimeSecs_.store(std::max<T>(ms * T(0.001), T(0.001)),
                              std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getMaxCompensation() const noexcept
    {
        return maxCompensation_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getSmoothingTime() const noexcept
    {
        return smoothTimeSecs_.load(std::memory_order_relaxed) * T(1000);
    }
 
    [[nodiscard]] std::vector<uint8_t> getState() const
    {
        // The blob stores float (setState reads float back); the explicit
        // casts also keep this overload resolvable when T is double.
        StateWriter w(stateId("AGAN"), 1);
        w.write("maxComp", static_cast<float>(maxCompensation_.load(std::memory_order_relaxed)));
        w.write("smoothMs", static_cast<float>(getSmoothingTime()));
        return w.blob();
    }
 
    bool setState(const uint8_t* data, size_t size)
    {
        StateReader r(data, size);
        if (!r.isValid() || r.processorId() != stateId("AGAN")) return false;
        setMaxCompensation(static_cast<T>(r.read("maxComp", 12.0f)));
        setSmoothingTime(static_cast<T>(r.read("smoothMs", 100.0f)));
        return true;
    }
 
private:
    [[nodiscard]] T measureRmsDb(AudioBufferView<T> buffer) const noexcept
    {
        const int numCh = std::min(buffer.getNumChannels(), numChannels_);
        const int numSamples = buffer.getNumSamples();
        if (numCh <= 0 || numSamples <= 0) return SILENCE_THRESH_DB;
 
        T sumSq = T(0);
        const int totalSamples = numSamples * numCh;
 
        for (int ch = 0; ch < numCh; ++ch)
            sumSq += simd::sumOfSquares(buffer.getChannel(ch), numSamples);
 
        // Apply a strict epsilon floor (approx -150dB) to avoid std::sqrt(0) and gainToDecibels(0) -> -Inf
        T meanSq = std::max<T>(sumSq / static_cast<T>(totalSamples), T(1e-15));
        return gainToDecibels(std::sqrt(meanSq));
    }
 
    static constexpr T SILENCE_THRESH_DB = T(-90);   
 
    double sampleRate_ = 44100.0;                    
    int numChannels_ = 0;                            
 
    std::atomic<T> smoothTimeSecs_{ T(0.100) };      
    T refLevelDb_ = SILENCE_THRESH_DB;               
    T compensationDb_ = T(0);                        
 
    std::atomic<T> maxCompensation_{ T(12) };        
};
 
} // namespace dspark