PitchCorrector.h

// DSPark - Professional Audio DSP Framework
// Copyright (c) 2026 Cristian Moresi - MIT License
 
#pragma once
 
#include "../Analysis/PitchDetector.h"
#include "../Core/AudioBuffer.h"
#include "../Core/AudioSpec.h"
#include "../Core/DenormalGuard.h"
#include "../Core/DspMath.h"
#include "../Core/StateBlob.h"
#include "../Music/HarmonyConstants.h"
#include "PitchShifter.h"
 
#include <algorithm>
#include <atomic>
#include <cmath>
#include <cstdint>
#include <span>
#include <vector>
 
namespace dspark {
 
template <FloatType T>
class PitchCorrector final
{
public:
    // The published words are read on the audio thread every control interval,
    // and a word that is not lock-free would take a mutex inside the callback.
    static_assert(std::atomic<T>::is_always_lock_free,
                  "audio-thread stores must not lock");
    static_assert(std::atomic<std::uint32_t>::is_always_lock_free,
                  "audio-thread stores must not lock");
 
    // -- Lifecycle -------------------------------------------------------------
 
    void prepare(const AudioSpec& spec, int fftSize = 0)
    {
        if (!std::isfinite(spec.sampleRate) || !spec.isValid()) return;
 
        prepared_ = false;
 
        const int frame = fftSize > 0 ? sanitizeFrame(fftSize)
                                      : automaticFrame(spec.sampleRate);
 
        detector_.prepare(spec.sampleRate);
        // The corrector keeps the Standard engine: its retune dynamics and the
        // 2 * frame latency are tuned to it, and it moves the shift on every
        // control tick, where Studio's strike anchoring buys nothing.
        shifter_.setQuality(PitchShifter<T>::Quality::Standard);
        shifter_.prepare(spec, frame);
 
        sampleRate_ = spec.sampleRate;
        frameSize_ = frame;
        latency_ = shifter_.getLatency();
 
        prepared_ = true;
        reset();
    }
 
    void reset() noexcept
    {
        if (!prepared_) return;
 
        detector_.reset();
        target_ = 0.0;
        correction_ = 0.0;
        publishedCorrection_ = 0.0;
        shifter_.setSemitones(T(0));
        shifter_.reset();
        controlPhase_ = 0;
        smoothingForMs_ = -1.0;   // rebuild the glide coefficient on next block
        warmupRemaining_ = detector_.getWindowSize();
    }
 
    // -- Parameters (any thread) -------------------------------------------------
 
    void setScale(std::uint16_t scaleBitmask, int rootPitchClass) noexcept
    {
        const auto mask = static_cast<std::uint32_t>(scaleBitmask & 0x0FFFu);
        const auto root = static_cast<std::uint32_t>((rootPitchClass % 12 + 12) % 12);
        const auto absolute = static_cast<std::uint32_t>(harmony::scaleAtRoot(
            static_cast<harmony::NoteSet>(mask), static_cast<int>(root)));
        // One packed word: the audio thread must never pair a fresh mask with a
        // stale root, so the mask, the root and the pre-rotated absolute
        // pitch-class set travel together.
        scaleWord_.store(absolute | (root << 12) | (mask << 16),
                         std::memory_order_relaxed);
    }
 
    void setRetuneSpeedMs(T ms) noexcept
    {
        if (!std::isfinite(ms)) return;
        retuneSpeedMs_.store(std::clamp(ms, T(0), T(kMaxRetuneMs)),
                             std::memory_order_relaxed);
    }
 
    void setFormantPreserve(bool on) noexcept
    {
        formantPreserve_.store(on, std::memory_order_relaxed);
    }
 
    [[nodiscard]] std::uint16_t getScaleMask() const noexcept
    {
        return static_cast<std::uint16_t>(
            (scaleWord_.load(std::memory_order_relaxed) >> 16) & 0x0FFFu);
    }
 
    [[nodiscard]] int getRootPitchClass() const noexcept
    {
        return static_cast<int>(
            (scaleWord_.load(std::memory_order_relaxed) >> 12) & 0x0Fu);
    }
 
    [[nodiscard]] T getRetuneSpeedMs() const noexcept
    {
        return retuneSpeedMs_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] bool getFormantPreserve() const noexcept
    {
        return formantPreserve_.load(std::memory_order_relaxed);
    }
 
    [[nodiscard]] std::vector<uint8_t> getState() const
    {
        StateWriter w(stateId("PCOR"), 1);
        w.write("scaleMask", static_cast<int32_t>(getScaleMask()));
        w.write("root", static_cast<int32_t>(getRootPitchClass()));
        w.write("retuneMs", static_cast<float>(getRetuneSpeedMs()));
        w.write("formant", getFormantPreserve());
        return w.blob();
    }
 
    bool setState(const uint8_t* data, size_t size)
    {
        StateReader r(data, size);
        if (!r.isValid() || r.processorId() != stateId("PCOR")) return false;
        setScale(static_cast<std::uint16_t>(r.read("scaleMask", int32_t(0x0FFF)) & 0x0FFF),
                 r.read("root", int32_t(0)));
        setRetuneSpeedMs(static_cast<T>(r.read("retuneMs", 0.0f)));
        setFormantPreserve(r.read("formant", false));
        return true;
    }
 
    [[nodiscard]] int getLatency() const noexcept
    {
        return prepared_ ? latency_ : 0;
    }
 
    [[nodiscard]] int getFrameSize() const noexcept
    {
        return prepared_ ? frameSize_ : 0;
    }
 
    // -- Processing --------------------------------------------------------------
 
    void processBlock(AudioBufferView<T> buffer) noexcept
    {
        if (!prepared_) return;
        const int numSamples = buffer.getNumSamples();
        if (numSamples <= 0) return;
 
        DenormalGuard guard;
 
        // One load per block for each independent word. The scale word is
        // internally consistent by packing; the retune speed and the formant
        // flag stand alone.
        const std::uint32_t absoluteMask =
            scaleWord_.load(std::memory_order_relaxed) & 0x0FFFu;
 
        const double retuneMs = static_cast<double>(
            retuneSpeedMs_.load(std::memory_order_relaxed));
        if (retuneMs != smoothingForMs_) rebuildSmoothing(retuneMs);
 
        const bool formant = formantPreserve_.load(std::memory_order_relaxed);
        if (formant != publishedFormant_)
        {
            // Forwarded by the stream owner, which is the shifter's only
            // control writer here (see the file documentation).
            shifter_.setFormantPreserve(formant);
            publishedFormant_ = formant;
        }
 
        const T* const detectionInput =
            buffer.getNumChannels() > 0 ? buffer.getChannel(0) : nullptr;
 
        int offset = 0;
        while (offset < numSamples)
        {
            // The grid is absolute (controlPhase_ persists across blocks): the
            // correction advances once, at the START of each whole interval,
            // and every sample of that interval is then rendered with it in
            // force. The shifter's engine adopts a new value at its own
            // analysis hops, which do not land on this grid once the ratio
            // leaves unity, so publishing at the start of the interval that
            // contains a hop - rather than at whichever chunk boundary the
            // host block size happens to create - is what makes the rendered
            // output identical however the stream is chopped into blocks.
            if (controlPhase_ == 0) advanceCorrection(absoluteMask);
 
            const int chunk =
                std::min(numSamples - offset, kControlInterval - controlPhase_);
 
            if (detectionInput != nullptr)
                detector_.pushSamples(std::span<const T>(
                    detectionInput + offset, static_cast<std::size_t>(chunk)));
 
            if (warmupRemaining_ > 0)
                warmupRemaining_ = std::max(0, warmupRemaining_ - chunk);
 
            shifter_.processBlock(buffer.getSubView(offset, chunk));
 
            controlPhase_ += chunk;
            if (controlPhase_ >= kControlInterval) controlPhase_ = 0;
            offset += chunk;
        }
    }
 
private:
    static constexpr double kAutoSpanRef = 2048.0;
    static constexpr double kAutoSpanRate = 48000.0;
    static constexpr int kAutoMinFrame = 512;
    static constexpr int kAutoMaxFrame = 32768;
 
    [[nodiscard]] static int automaticFrame(double sampleRate) noexcept
    {
        const double target = sampleRate * (kAutoSpanRef / kAutoSpanRate);
        int frame = kAutoMinFrame;
        while (frame < kAutoMaxFrame && static_cast<double>(frame) < target)
            frame <<= 1;
        return frame;
    }
 
    [[nodiscard]] static int sanitizeFrame(int requested) noexcept
    {
        constexpr int kMax = 1 << 20;
        int frame = 256;
        while (frame < requested && frame < kMax) frame <<= 1;
        return frame;
    }
 
    static constexpr int kControlInterval = 64;
 
    static constexpr double kLandingSemitones = 0.001;
 
    static constexpr double kMaxRetuneMs = 10000.0;
 
    void rebuildSmoothing(double retuneMs) noexcept
    {
        smoothingForMs_ = retuneMs;
        if (retuneMs <= 0.0 || sampleRate_ <= 0.0)
        {
            smoothing_ = 1.0;   // hard snap
            return;
        }
        const double tauSamples = retuneMs * 0.001 * sampleRate_;
        smoothing_ = 1.0 - std::exp(-static_cast<double>(kControlInterval) / tauSamples);
    }
 
    void advanceCorrection(std::uint32_t absoluteMask) noexcept
    {
        if (absoluteMask == 0u)
        {
            target_ = 0.0;   // no scale to snap to: correction disabled
        }
        else if (warmupRemaining_ == 0)
        {
            // Two gates before a report may establish a held correction: the
            // analysis window must be full (its first, mostly-zero window
            // reports the shortest lag the search allows, at full confidence),
            // and the fundamental must be one this effect could act on - above
            // a quarter of the sample rate not even its second harmonic fits
            // below Nyquist. Unvoiced, or refused: keep the previous target.
            const double frequency = static_cast<double>(detector_.getFrequencyHz());
            if (frequency > 0.0 && frequency <= 0.25 * sampleRate_)
            {
                const double midi = 69.0 + 12.0 * std::log2(frequency / 440.0);
                if (std::isfinite(midi))
                    target_ = std::clamp(nearestInScale(midi, absoluteMask) - midi,
                                         -12.0, 12.0);
            }
        }
 
        if (smoothing_ >= 1.0)
        {
            correction_ = target_;
        }
        else
        {
            correction_ += (target_ - correction_) * smoothing_;
            if (std::abs(target_ - correction_) < kLandingSemitones)
                correction_ = target_;   // exact landing
        }
 
        if (correction_ != publishedCorrection_)
        {
            shifter_.setSemitones(static_cast<T>(correction_));
            publishedCorrection_ = correction_;
        }
    }
 
    [[nodiscard]] static double nearestInScale(double midi,
                                               std::uint32_t absoluteMask) noexcept
    {
        const int base = static_cast<int>(std::floor(midi));
        int bestNote = base;
        double bestDistance = 1.0e9;
        for (int note = base - 12; note <= base + 13; ++note)
        {
            const int pitchClass = ((note % 12) + 12) % 12;
            if ((absoluteMask & (1u << pitchClass)) == 0u) continue;
            const double distance = std::abs(static_cast<double>(note) - midi);
            if (distance < bestDistance)
            {
                bestDistance = distance;
                bestNote = note;
            }
        }
        return static_cast<double>(bestNote);
    }
 
    // -- Members -----------------------------------------------------------------
 
    PitchDetector<T> detector_;   
    PitchShifter<T> shifter_;     
 
    double sampleRate_ = 48000.0;
    int frameSize_ = 0;
    int latency_ = 0;
    bool prepared_ = false;
 
    // Stream-owner state (audio thread only).
    double target_ = 0.0;                
    double correction_ = 0.0;            
    double publishedCorrection_ = 0.0;   
    double smoothing_ = 1.0;             
    double smoothingForMs_ = -1.0;       
    bool publishedFormant_ = false;      
    int controlPhase_ = 0;               
    int warmupRemaining_ = 0;            
 
    std::atomic<std::uint32_t> scaleWord_ { 0x0FFF0FFFu };
    std::atomic<T> retuneSpeedMs_ { T(0) };
    std::atomic<bool> formantPreserve_ { false };
};
 
} // namespace dspark