PitchShifter.h

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Core/AudioBuffer.h"
#include "../Core/AudioSpec.h"
#include "../Core/DenormalGuard.h"
#include "../Core/DspMath.h"
#include "../Core/Interpolation.h"
#include "../Core/StateBlob.h"
#include "detail/PhaseVocoderEngine.h"
#include "detail/StudioVocoder.h"
 
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class PitchShifter
{
public:
    // -- Lifecycle -------------------------------------------------------------
 
    void prepare(const AudioSpec& spec, int fftSize = 0)
    {
        if (!spec.isValid()) return;
        if (fftSize != 0 && ((fftSize & (fftSize - 1)) != 0
                             || fftSize < 256 || fftSize > (1 << 20)))
            return;
 
        prepared_.store(false, std::memory_order_relaxed);
 
        numChannels_ = std::max(1, spec.numChannels);
        // 0 selects each engine's default frame: 2048 for Standard/High (the
        // published rendering), about 43 ms for Studio (2048 at 44.1 and
        // 48 kHz, 4096 at 88.2 and 96 kHz).
        const int fftStudio = fftSize != 0 ? fftSize : studioDefaultFrame(spec.sampleRate);
        if (fftSize == 0) fftSize = 2048;
 
        // The reader trails the write head by readOffset_; the window/OLA
        // chain adds another fftSize - synthHop, so the measured wet latency
        // is readOffset_ + fftSize - synthHop = 2 * fftSize (exact at unity
        // ratio). The dry path must delay by the SAME value or partial mixes
        // comb-filter.
        const int synthHop = fftSize / 4;
        readOffset_ = fftSize + synthHop;
        latency_    = readOffset_ + fftSize - synthHop;
        // Studio: the resample reader runs at the pitch ratio behind the
        // completed stream. At ratio a an input sample reaches the completed
        // stream after half a frame, a frame of analysis, the onset
        // lookahead and one analysis hop, and the reader needs its sinc
        // reach beyond that; the worst case is the lowest ratio, 0.5:
        // N/(2a) + N/2 + lookahead + taps/a = 1.5 N + lookahead + 32. A
        // strike anchor may move an analysis frame ahead of the nominal
        // timeline by up to a quarter frame at that ratio, which the last
        // term covers. One fixed latency for every ratio, so the dry path
        // and host compensation hold at any pitch.
        studio_.prepare(spec.sampleRate, numChannels_, fftStudio, true);
        studio_.setAnchorLeadLimit(static_cast<double>(fftStudio / 4));
        latencyStudio_ = fftStudio + fftStudio / 2 + studio_.lookahead() + fftStudio / 4 + 64;
        latencyLegacy_ = latency_;
        drySize_ = 1;
        while (drySize_ < std::max(latencyLegacy_, latencyStudio_) + 1) drySize_ <<= 1;
        dryMask_ = drySize_ - 1;
 
        // The engine owns the analysis rings, spectral state and OLA ring;
        // resample-back compensation stages (anti-alias taper, formant
        // pre-warp target) are enabled because this owner resamples.
        //
        // The three remaining capabilities are refused, and the last two are
        // written out rather than defaulted because a future default that
        // flips would change what this effect sounds like. The locked
        // analysis hop breaks the resample-back reader's Ra = Rs/ratio
        // assumption and buys nothing here. The spectral-flux onset detector
        // fires earlier and more often than the frame-energy test, and every
        // firing resets phase, so adopting it smears energy ahead of a strike
        // - audible pre-echo on percussive material - and re-renders every
        // release already in a user's hands. A mix finished last month is not
        // restored by a release note. How large the smear is depends on the
        // material and on the estimator it is measured with, so no figure is
        // quoted here; that it is audible on percussive material is the
        // durable part, and it is why this effect keeps the detector it
        // shipped with.
        engine_.prepare(spec.sampleRate, numChannels_, fftSize, true, false, false, false);
        mixMaxStep_ = static_cast<T>(1.0 / std::max(1.0, spec.sampleRate * 0.02));
        accumMask_ = engine_.olaMask();
 
        dryRing_.assign(static_cast<size_t>(numChannels_), {});
        for (int ch = 0; ch < numChannels_; ++ch)
            dryRing_[static_cast<size_t>(ch)].assign(static_cast<size_t>(drySize_), T(0));
 
        publishEngineParams();
        prepared_.store(true, std::memory_order_relaxed);
        reset();
    }
 
    void reset() noexcept
    {
        if (!prepared_.load(std::memory_order_relaxed)) return;
        studioActive_ = quality_.load(std::memory_order_relaxed) == Quality::Studio;
        latency_ = studioActive_ ? latencyStudio_ : latencyLegacy_;
        engine_.reset();
        studio_.reset();
        for (auto& r : dryRing_) std::fill(r.begin(), r.end(), T(0));
        studioOut_ = 0;
        {
            // The reader starts where output 0 shows input -latency, and
            // advances at the pitch ratio from there.
            const double start = studio_.streamPositionOf(-static_cast<double>(latencyStudio_));
            studioReadInt_ = static_cast<int64_t>(std::floor(start));
            studioReadFrac_ = start - static_cast<double>(studioReadInt_);
        }
 
        dryPos_ = 0;
        readPosInt_ = engine_.writeHead() - readOffset_;
        readPosFrac_ = 0.0;
        currentMix_ = mix_.load(std::memory_order_relaxed);
        highReader_ = quality_.load(std::memory_order_relaxed) != Quality::Standard;
        readerFadeLeft_ = 0;   // start settled on the selected reader
    }
 
    // -- Parameters (thread-safe) -----------------------------------------------
 
    enum class Quality
    {
        Standard,   
        High,       
        Studio      
    };
 
    void setQuality(Quality quality) noexcept
    {
        quality = static_cast<Quality>(std::clamp(static_cast<int>(quality), 0,
                                                  static_cast<int>(Quality::Studio)));
        quality_.store(quality, std::memory_order_relaxed);
    }
 
    [[nodiscard]] Quality getQuality() const noexcept
    {
        return quality_.load(std::memory_order_relaxed);
    }
 
    void setSemitones(T st) noexcept
    {
        if (!std::isfinite(st)) return;
        semitones_.store(std::clamp(st, T(-12), T(12)), std::memory_order_relaxed);
        publishEngineParams();
    }
 
    void setPitchRatio(T ratio) noexcept
    {
        if (!std::isfinite(ratio)) return;
        ratio = std::clamp(ratio, T(0.5), T(2));
        semitones_.store(static_cast<T>(12.0 * std::log2(static_cast<double>(ratio))),
                         std::memory_order_relaxed);
        publishEngineParams();
    }
 
    void setMix(T mix) noexcept
    {
        if (!std::isfinite(mix)) return;
        mix_.store(std::clamp(mix, T(0), T(1)), std::memory_order_relaxed);
    }
 
    void setTransientPreserve(bool enabled) noexcept
    {
        transientPreserve_.store(enabled, std::memory_order_relaxed);
        publishEngineParams();
    }
 
    void setFormantPreserve(bool enabled) noexcept
    {
        formantPreserve_.store(enabled, std::memory_order_relaxed);
        publishEngineParams();
    }
 
    [[nodiscard]] T getSemitones() const noexcept
    {
        return semitones_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] T getMix() const noexcept { return mix_.load(std::memory_order_relaxed); }
 
    [[nodiscard]] bool getTransientPreserve() const noexcept
    {
        return transientPreserve_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] bool getFormantPreserve() const noexcept
    {
        return formantPreserve_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] int getLatency() const noexcept { return latency_; }
 
    [[nodiscard]] std::vector<uint8_t> getState() const
    {
        StateWriter w(stateId("PSHF"), 1);
        w.write("semitones", static_cast<float>(semitones_.load(std::memory_order_relaxed)));
        w.write("mix", static_cast<float>(mix_.load(std::memory_order_relaxed)));
        w.write("transient", transientPreserve_.load(std::memory_order_relaxed));
        w.write("formant", formantPreserve_.load(std::memory_order_relaxed));
        w.write("quality", static_cast<int32_t>(quality_.load(std::memory_order_relaxed)));
        return w.blob();
    }
 
    bool setState(const uint8_t* data, size_t size)
    {
        StateReader r(data, size);
        if (!r.isValid() || r.processorId() != stateId("PSHF")) return false;
        setSemitones(static_cast<T>(r.read("semitones", 0.0f)));
        setMix(static_cast<T>(r.read("mix", 1.0f)));
        setTransientPreserve(r.read("transient", true));
        setFormantPreserve(r.read("formant", false));
        setQuality(static_cast<Quality>(r.read("quality", 0)));   // clamped inside
        return true;
    }
 
    // -- Processing --------------------------------------------------------------
 
    void processBlock(AudioBufferView<T> buffer) noexcept
    {
        if (!prepared_.load(std::memory_order_relaxed)) return;
        DenormalGuard guard;
 
        const int nCh = std::min(buffer.getNumChannels(), numChannels_);
        const int nS  = buffer.getNumSamples();
        // Crossing between Studio and the Standard/High engine restarts the stream:
        // the two run different latencies, so no crossfade can join them.
        if ((quality_.load(std::memory_order_relaxed) == Quality::Studio) != studioActive_)
            reset();
        if (studioActive_)
        {
            processStudio(buffer, nCh, nS);
            return;
        }
 
        // Rate-limited mix ramp (moveTowards, exact landing; settled it
        // reduces to the constant, bit-identically). A per-block ramp landed
        // in 0.7 ms with 32-sample blocks.
        const T mixTarget = mix_.load(std::memory_order_relaxed);
        const T mixStart  = currentMix_;
 
        // Reader selection: a change crossfades from the old reader to the
        // new one over kReaderFade samples (the readers differ at HF).
        const bool highQuality = quality_.load(std::memory_order_relaxed) == Quality::High;
        if (highQuality != highReader_)
        {
            highReader_ = highQuality;
            readerFadeLeft_ = kReaderFade;
        }
 
        int i = 0;
        while (i < nS)
        {
            const int chunk = std::min(nS - i, engine_.samplesToNextHop());
 
            // 1. Push input into the analysis ring and the dry-compensation ring.
            for (int ch = 0; ch < nCh; ++ch)
            {
                const T* in = buffer.getChannel(ch) + i;
                engine_.pushInput(ch, in, chunk);
                auto& dry = dryRing_[static_cast<size_t>(ch)];
                int dp = dryPos_;
                for (int k = 0; k < chunk; ++k)
                {
                    dry[static_cast<size_t>(dp)] = in[k];
                    dp = (dp + 1) & dryMask_;
                }
            }
 
            // 2. Produce output: fractional read of the synthesis stream + mix.
            const double ratio = engine_.activeRatio();
            int64_t rpEnd = readPosInt_;
            double  rfEnd = readPosFrac_;
            for (int ch = 0; ch < nCh; ++ch)
            {
                T* out = buffer.getChannel(ch) + i;
                const T* acc = engine_.olaData(ch);
                const auto& dry = dryRing_[static_cast<size_t>(ch)];
 
                int64_t rp = readPosInt_;
                double  rf = readPosFrac_;
                int     dp = dryPos_;
                int     fadeLeft = readerFadeLeft_;
 
                for (int k = 0; k < chunk; ++k)
                {
                    T wet = readWet(highReader_, acc, rp, rf);
                    if (fadeLeft > 0)
                    {
                        const T w = static_cast<T>(fadeLeft) * (T(1) / T(kReaderFade));
                        const T old = readWet(!highReader_, acc, rp, rf);
                        wet += (old - wet) * w;
                        --fadeLeft;
                    }
                    const int dryIdx = (dp - latency_) & dryMask_;
                    const T drySample = dry[static_cast<size_t>(dryIdx)];
                    const T mixVal = moveTowards(mixStart, mixTarget, mixMaxStep_ * static_cast<T>(i + k + 1));
                    // Two-product blend: exact at both ends (mix 1 emits the
                    // wet stream bit-exactly, mix 0 the delayed dry).
                    out[k] = drySample * (T(1) - mixVal) + wet * mixVal;
 
                    rf += ratio;
                    const auto adv = static_cast<int64_t>(rf);
                    rp += adv;
                    rf -= static_cast<double>(adv);
                    dp = (dp + 1) & dryMask_;
                }
                if (ch == 0) { rpEnd = rp; rfEnd = rf; }   // recurrence result
            }
 
            // Commit shared positions once per chunk, using the SAME
            // per-sample recurrence result the output loops computed (taken
            // from channel 0, whose loop ran it already): a single
            // frac + ratio * chunk product rounds differently for different
            // chunk sizes, and chunk boundaries follow the host block size,
            // so committing the product form made the output depend on how
            // the host chopped the stream (~1 ulp per flip, but a bit-exact
            // contract is a bit-exact contract).
            {
                if (nCh == 0)   // channel-less call: advance the stream anyway
                {
                    for (int k = 0; k < chunk; ++k)
                    {
                        rfEnd += ratio;
                        const auto adv = static_cast<int64_t>(rfEnd);
                        rpEnd += adv;
                        rfEnd -= static_cast<double>(adv);
                    }
                }
                readPosInt_  = rpEnd;
                readPosFrac_ = rfEnd;
                dryPos_ = (dryPos_ + chunk) & dryMask_;
                readerFadeLeft_ = std::max(0, readerFadeLeft_ - chunk);
            }
 
            // 3. Advance the engine (runs an STFT hop at the analysis boundary).
            engine_.commitInput(chunk, nCh);
 
            i += chunk;
        }
 
        currentMix_ = moveTowards(mixStart, mixTarget, mixMaxStep_ * static_cast<T>(nS));
    }
 
private:
    void processStudio(AudioBufferView<T> buffer, int nCh, int nS) noexcept
    {
        const T mixTarget = mix_.load(std::memory_order_relaxed);
        const T mixStart  = currentMix_;
        const int64_t mask = studio_.olaMask();
        int i = 0;
        while (i < nS)
        {
            const int need = studio_.samplesToNextHop();
            if (need == 0)
            {
                studio_.commitInput(0, nCh);
                continue;
            }
            const int chunk = std::min(nS - i, need);
            for (int ch = 0; ch < nCh; ++ch)
            {
                const T* in = buffer.getChannel(ch) + i;
                studio_.pushInput(ch, in, chunk);
                auto& dry = dryRing_[static_cast<size_t>(ch)];
                int dp = dryPos_;
                for (int k = 0; k < chunk; ++k)
                {
                    dry[static_cast<size_t>(dp)] = in[k];
                    dp = (dp + 1) & dryMask_;
                }
            }
 
            const double ratio = studio_.activeRatio();
            int64_t rpEnd = studioReadInt_;
            double rfEnd = studioReadFrac_;
            for (int ch = 0; ch < nCh; ++ch)
            {
                T* out = buffer.getChannel(ch) + i;
                const T* acc = studio_.olaData(ch);
                const auto& dry = dryRing_[static_cast<size_t>(ch)];
                int64_t rp = studioReadInt_;
                double rf = studioReadFrac_;
                int dp = dryPos_;
                for (int k = 0; k < chunk; ++k)
                {
                    const T wet = reader_.readRing(acc, mask, rp, rf);
                    const T drySample = dry[static_cast<size_t>((dp - latency_) & dryMask_)];
                    const T mixVal = moveTowards(mixStart, mixTarget, mixMaxStep_ * static_cast<T>(i + k + 1));
                    out[k] = drySample * (T(1) - mixVal) + wet * mixVal;
                    rf += ratio;
                    const auto adv = static_cast<int64_t>(rf);
                    rp += adv;
                    rf -= static_cast<double>(adv);
                    dp = (dp + 1) & dryMask_;
                }
                if (ch == 0) { rpEnd = rp; rfEnd = rf; }
            }
            if (nCh == 0)
            {
                for (int k = 0; k < chunk; ++k)
                {
                    rfEnd += ratio;
                    const auto adv = static_cast<int64_t>(rfEnd);
                    rpEnd += adv;
                    rfEnd -= static_cast<double>(adv);
                }
            }
            studioReadInt_ = rpEnd;
            studioReadFrac_ = rfEnd;
            dryPos_ = (dryPos_ + chunk) & dryMask_;
            studioOut_ += chunk;
 
            // Steer the frame after the next one: where the reader will stand
            // on the input timeline when it reaches that frame's centre.
            const double c = studio_.nextStreamCentre() + static_cast<double>(studio_.synthHop());
            const double here = static_cast<double>(studioReadInt_) + studioReadFrac_;
            studio_.steerTimeline(static_cast<double>(studioOut_) + (c - here) / ratio
                                  - static_cast<double>(latencyStudio_));
            studio_.commitInput(chunk, nCh);
            i += chunk;
        }
        currentMix_ = moveTowards(mixStart, mixTarget, mixMaxStep_ * static_cast<T>(nS));
    }
 
    [[nodiscard]] static int studioDefaultFrame(double sampleRate) noexcept
    {
        int n = 256;
        while (n < (1 << 16) && static_cast<double>(n) * 1.5 < 0.0427 * sampleRate) n <<= 1;
        return n;
    }
 
    [[nodiscard]] T readWet(bool high, const T* acc, int64_t ip, double frac) const noexcept
    {
        return high ? reader_.readRing(acc, accumMask_, ip, frac)
                    : readCatmullRom(acc, ip, frac);
    }
 
    [[nodiscard]] T readCatmullRom(const T* acc, int64_t ip, double frac) const noexcept
    {
        const int64_t m = accumMask_;
        const T x0 = acc[static_cast<size_t>((ip - 1) & m)];
        const T x1 = acc[static_cast<size_t>(ip & m)];
        const T x2 = acc[static_cast<size_t>((ip + 1) & m)];
        const T x3 = acc[static_cast<size_t>((ip + 2) & m)];
        const T f  = static_cast<T>(frac);
        return x1 + T(0.5) * f * (x2 - x0
                 + f * (T(2) * x0 - T(5) * x1 + T(4) * x2 - x3
                 + f * (T(3) * (x1 - x2) + x3 - x0)));
    }
 
    void publishEngineParams() noexcept
    {
        typename detail::PhaseVocoderEngine<T>::Params p;
        p.targetSemitones = static_cast<double>(semitones_.load(std::memory_order_relaxed));
        p.transientPreserve = transientPreserve_.load(std::memory_order_relaxed);
        p.formantPreserve = formantPreserve_.load(std::memory_order_relaxed);
        engine_.publishParams(p);
        typename detail::StudioVocoder<T>::Params q;
        q.targetSemitones = p.targetSemitones;
        q.transientPreserve = p.transientPreserve;
        q.formantPreserve = p.formantPreserve;
        studio_.publishParams(q);
    }
 
    // -- Members -----------------------------------------------------------------
    int numChannels_ = 0;
    std::atomic<bool> prepared_ { false };
 
    int readOffset_ = 2560;   
    int latency_ = 4096;      
    int drySize_ = 8192;
    int dryMask_ = 8191;
    int64_t accumMask_ = 8191;   
 
    detail::PhaseVocoderEngine<T> engine_;   
    detail::StudioVocoder<T> studio_;        
    bool studioActive_ = true;               
    int latencyStudio_ = 5184;               
    int latencyLegacy_ = 4096;               
    int64_t studioReadInt_ = 0;              
    double studioReadFrac_ = 0.0;            
    int64_t studioOut_ = 0;                  
    SincInterpolator<T> reader_;             
    static constexpr int kReaderFade = 64;   
    bool highReader_ = false;                
    int readerFadeLeft_ = 0;                 
 
    std::vector<std::vector<T>> dryRing_;    
 
    int dryPos_ = 0;
    int64_t readPosInt_ = 0;
    double readPosFrac_ = 0.0;
    T currentMix_ = T(1);     
    T mixMaxStep_ = T(1.0 / 960.0); 
 
    std::atomic<T> semitones_ { T(0) };
    std::atomic<T> mix_ { T(1) };
    std::atomic<bool> transientPreserve_ { true };
    std::atomic<bool> formantPreserve_ { false };
    std::atomic<Quality> quality_ { Quality::Studio };
};
 
} // namespace dspark