SpectralProcessor.h

Go to the documentation of this file.

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "FFT.h"
#include "WindowFunctions.h"
#include "AudioSpec.h"
#include "AudioBuffer.h"
#include "DenormalGuard.h"
 
#include <algorithm>
#include <cassert>
#include <cmath>
#include <cstddef>
#include <memory>
#include <vector>
 
namespace dspark {
 
template <typename T>
class SpectralProcessor
{
public:
    // -- Lifecycle -----------------------------------------------------------
 
    void prepare(const AudioSpec& spec, int fftSize = 2048, int hopSize = 0)
    {
        if (!(spec.sampleRate > 0.0) || spec.numChannels < 1) return;
 
        // Disarm first: if an allocation below throws, the processor is a
        // safe no-op instead of running half-mutated state (basic guarantee).
        prepared_ = false;
        spec_ = spec;
 
        assert((fftSize & (fftSize - 1)) == 0 && fftSize >= kMinFftSize
               && "FFT size must be a power of 2");
        // Release-safe: clamp, then round UP to a power of two (never down,
        // so a requested resolution is always met). The clamp bounds the
        // loop; kMaxFftSize is itself a power of two, so no overflow.
        fftSize = std::clamp(fftSize, kMinFftSize, kMaxFftSize);
        int pow2 = kMinFftSize;
        while (pow2 < fftSize) pow2 <<= 1;
        fftSize_ = pow2;
        mask_ = fftSize_ - 1;
 
        hopSize_ = (hopSize > 0) ? hopSize : fftSize_ / 2;
        hopSize_ = std::clamp(hopSize_, 1, fftSize_ / 2);
        if (fftSize_ % hopSize_ != 0) hopSize_ = fftSize_ / 2;
        numBins_ = fftSize_ / 2 + 1;
 
        fft_ = std::make_unique<FFTReal<T>>(fftSize_);
 
        // WOLA window: periodic sqrt-Hann (analysis and synthesis).
        window_.resize(static_cast<std::size_t>(fftSize_));
        WindowFunctions<T>::hann(window_.data(), fftSize_, true);
        for (auto& w : window_) w = std::sqrt(w);
 
        computeWolaNorm();
 
        const int nCh = spec.numChannels;
        inputRing_.resize(static_cast<std::size_t>(nCh));
        outputAccum_.resize(static_cast<std::size_t>(nCh));
        inputPos_.assign(static_cast<std::size_t>(nCh), 0);
        outputReadPos_.assign(static_cast<std::size_t>(nCh), 0);
 
        // The overlap-add accumulator holds twice the FFT size (also pow2).
        accumMask_ = (fftSize_ * 2) - 1;
 
        for (int ch = 0; ch < nCh; ++ch)
        {
            inputRing_[static_cast<std::size_t>(ch)]
                .assign(static_cast<std::size_t>(fftSize_), T(0));
            outputAccum_[static_cast<std::size_t>(ch)]
                .assign(static_cast<std::size_t>(fftSize_ * 2), T(0));
        }
 
        // Plain vectors are fine here: the FFT kernels are unaligned-safe
        // by design, so no aligned allocator is required.
        fftIn_.resize(static_cast<std::size_t>(fftSize_));
        fftOut_.resize(static_cast<std::size_t>(fftSize_ + 2));
        fftResult_.resize(static_cast<std::size_t>(fftSize_));
 
        hopCounter_ = 0;
        prepared_ = true;
    }
 
    template <typename Func>
    void processBlock(AudioBufferView<T> buffer, Func&& processFunc) noexcept
    {
        if (!prepared_) return;
        DenormalGuard guard;
 
        const int nCh = std::min(buffer.getNumChannels(),
                                 static_cast<int>(inputRing_.size()));
        const int nS  = buffer.getNumSamples();
 
        // Chunked processing between hop boundaries. Order per chunk:
        // push input, drain output, then fire the hop. Draining BEFORE the
        // hop is what anchors every frame to its absolute stream position
        // (the accumulator write position equals sample index m*hop at the
        // m-th hop no matter how the caller chops the stream); it is also
        // what makes the pipeline in-place safe (input is read into the
        // ring before the drain overwrites it).
        int i = 0;
        while (i < nS)
        {
            const int chunk = std::min(nS - i, hopSize_ - hopCounter_);
 
            // 1. Push the input chunk into the ring buffers (two contiguous
            //    spans instead of a per-sample masked loop).
            for (int ch = 0; ch < nCh; ++ch)
            {
                const T* data = buffer.getChannel(ch) + i;
                auto& ring = inputRing_[static_cast<std::size_t>(ch)];
                const int wp = inputPos_[static_cast<std::size_t>(ch)];
                const int first = std::min(chunk, fftSize_ - wp);
                std::copy_n(data, first, ring.data() + wp);
                std::copy_n(data + first, chunk - first, ring.data());
                inputPos_[static_cast<std::size_t>(ch)] = (wp + chunk) & mask_;
            }
 
            // 2. Drain the overlap-add accumulator for the chunk (read and
            //    clear in one pass, two contiguous spans).
            for (int ch = 0; ch < nCh; ++ch)
            {
                T* data = buffer.getChannel(ch) + i;
                T* acc  = outputAccum_[static_cast<std::size_t>(ch)].data();
                const int rp = outputReadPos_[static_cast<std::size_t>(ch)];
                const int first = std::min(chunk, fftSize_ * 2 - rp);
                for (int k = 0; k < first; ++k)
                {
                    data[k] = acc[rp + k];
                    acc[rp + k] = T(0);
                }
                for (int k = first; k < chunk; ++k)
                {
                    data[k] = acc[k - first];
                    acc[k - first] = T(0);
                }
                outputReadPos_[static_cast<std::size_t>(ch)] = (rp + chunk) & accumMask_;
            }
 
            // 3. Fire the STFT hop at the boundary.
            hopCounter_ += chunk;
            if (hopCounter_ >= hopSize_)
            {
                hopCounter_ = 0;
                for (int ch = 0; ch < nCh; ++ch)
                    processHop(ch, processFunc);
            }
 
            i += chunk;
        }
    }
 
    // -- Queries -------------------------------------------------------------
 
    [[nodiscard]] int getLatency() const noexcept { return fftSize_; }
    [[nodiscard]] int getFFTSize() const noexcept { return fftSize_; }
    [[nodiscard]] int getNumBins() const noexcept { return numBins_; }
    [[nodiscard]] int getHopSize() const noexcept { return hopSize_; }
 
    void reset() noexcept
    {
        for (auto& ring : inputRing_) std::fill(ring.begin(), ring.end(), T(0));
        for (auto& acc : outputAccum_) std::fill(acc.begin(), acc.end(), T(0));
        std::fill(inputPos_.begin(), inputPos_.end(), 0);
        std::fill(outputReadPos_.begin(), outputReadPos_.end(), 0);
        hopCounter_ = 0;
    }
 
private:
    template <typename Func>
    void processHop(int ch, Func& processFunc) noexcept
    {
        // 1. Copy from ring buffer with the analysis window, as two
        //    contiguous multiply spans (auto-vectorizable).
        const T* ring = inputRing_[static_cast<std::size_t>(ch)].data();
        const T* win  = window_.data();
        const int readPos = inputPos_[static_cast<std::size_t>(ch)];
        const int head = fftSize_ - readPos;
        for (int k = 0; k < head; ++k)
            fftIn_[static_cast<std::size_t>(k)] = ring[readPos + k] * win[k];
        for (int k = head; k < fftSize_; ++k)
            fftIn_[static_cast<std::size_t>(k)] = ring[k - head] * win[k];
 
        // 2. Forward FFT
        fft_->forward(fftIn_.data(), fftOut_.data());
 
        // 3. User callback (inlined directly here)
        processFunc(fftOut_.data(), numBins_);
 
        // 4. Inverse FFT (FFTReal::inverse already applies the 1/N
        //    normalization, so the analysis->synthesis round-trip is unity
        //    before windowing).
        fft_->inverse(fftOut_.data(), fftResult_.data());
 
        // 5. Synthesis window and overlap-add, two contiguous spans.
        T* acc = outputAccum_[static_cast<std::size_t>(ch)].data();
        const T* res = fftResult_.data();
        const int writePos = outputReadPos_[static_cast<std::size_t>(ch)];
        const int headOut = std::min(fftSize_, fftSize_ * 2 - writePos);
        for (int k = 0; k < headOut; ++k)
            acc[writePos + k] += res[k] * win[k] * wolaNorm_;
        for (int k = headOut; k < fftSize_; ++k)
            acc[k - headOut] += res[k] * win[k] * wolaNorm_;
    }
 
    void computeWolaNorm() noexcept
    {
        const int numOverlaps = fftSize_ / hopSize_;
        T maxSum = T(0);
        for (int pos = 0; pos < hopSize_; ++pos)
        {
            T sumSq = T(0);
            for (int hop = 0; hop < numOverlaps; ++hop)
            {
                const int idx = pos + hop * hopSize_;
                if (idx < fftSize_)
                {
                    const T w = window_[static_cast<std::size_t>(idx)];
                    sumSq += w * w;
                }
            }
            if (sumSq > maxSum) maxSum = sumSq;
        }
        wolaNorm_ = (maxSum > T(1e-10)) ? (T(1) / maxSum) : T(1);
    }
 
    // -- Members -------------------------------------------------------------
    static constexpr int kMinFftSize = 4;       
    static constexpr int kMaxFftSize = 1 << 20; 
 
    AudioSpec spec_ {};
    bool prepared_ = false;
 
    int fftSize_ = 2048;
    int hopSize_ = 1024;
    int numBins_ = 1025;
    int mask_ = 2047;      // Bitwise mask for fftSize_ bounds
    int accumMask_ = 4095; // Bitwise mask for fftSize_*2 bounds
 
    std::unique_ptr<FFTReal<T>> fft_;
    std::vector<T> window_;
    T wolaNorm_ = T(1);
 
    std::vector<std::vector<T>> inputRing_;
    std::vector<std::vector<T>> outputAccum_;
    std::vector<int> inputPos_;
    std::vector<int> outputReadPos_;
    int hopCounter_ = 0;
 
    // Shared work buffers for sequential channel processing
    std::vector<T> fftIn_, fftOut_, fftResult_;
};
 
} // namespace dspark