unnieversal/crates/utils/src/pitch/pitch_detection.rs

use pitch_detection::detector::yin::YINDetector;
use pitch_detection::detector::PitchDetector;

use crate::buffers::Buffers;
use ringbuf::{Consumer, Producer, RingBuffer};

pub fn generate_pitch_detector(size: usize) -> impl PitchDetector<f32> {
    let padding = size / 2;

    YINDetector::new(size, padding)
}

/// Returns an option with (Frequency, Clarity)
pub fn pitch_detect(
    detector: &mut dyn PitchDetector<f32>,
    signal: &[f32],
    sample_rate: u32,
) -> Option<(f32, f32)> {
    const POWER_THRESHOLD: f32 = 0.15;
    const CLARITY_THRESHOLD: f32 = 0.5;

    let pitch = detector.get_pitch(
        &signal,
        sample_rate as usize,
        POWER_THRESHOLD,
        CLARITY_THRESHOLD,
    );

    pitch.map(|a| (a.frequency, a.clarity))
}

pub struct DetectionInput<const LEN: usize> {
    pub buffers: Buffers<LEN>,
    pub sample_rate: u32,
}

pub struct DetectionOutput<const LEN: usize> {
    pub buffers: Buffers<LEN>,
    pub pitch_l: Option<f32>,
    pub pitch_r: Option<f32>,
}

pub fn detect<const LEN: usize>(
    mut inputs: Consumer<DetectionInput<LEN>>,
    mut outputs: Producer<DetectionOutput<LEN>>,
) {
    let mut detector_l = generate_pitch_detector(LEN);
    let mut detector_r = generate_pitch_detector(LEN);

    loop {
        if let Some(DetectionInput::<LEN> {
            buffers,
            sample_rate,
        }) = inputs.pop()
        {
            let pitch_l = pitch_detect(&mut detector_l, &buffers.l, sample_rate).map(|a| a.0);
            let pitch_r = pitch_detect(&mut detector_r, &buffers.r, sample_rate).map(|a| a.0);

            let _ = outputs.push(DetectionOutput {
                buffers,
                pitch_l,
                pitch_r,
            });
        }
    }
}

pub struct PitchDetectorThread<const LEN: usize> {
    /// Current recording buffer
    /// Input goes here
    recording_buffer: Buffers<LEN>,

    /// Ringbuf producer so we can send audio chunks to the processing thread
    recordings: Producer<DetectionInput<LEN>>,
    /// Ringbuf consumer so we can receive processed buffers from the processing threads
    processed: Consumer<DetectionOutput<LEN>>,

    /// Contains some empty buffers so we can reuse them instead of doing allocations
    /// Buffers here are not actually empty, since we don't spend any time clearing them
    /// But since they will be overwritten, this isn't an issue
    empty_buffers: Vec<Buffers<LEN>>,
}

impl<const LEN: usize> PitchDetectorThread<LEN> {
    pub fn new() -> Self {
        let (recordings, recording_rx) = RingBuffer::<DetectionInput<LEN>>::new(30).split();
        let (processed_tx, processed) = RingBuffer::<DetectionOutput<LEN>>::new(30).split();

        // Spawn analysis thread
        std::thread::spawn(move || {
            detect(recording_rx, processed_tx);
        });

        // keep some empty buffer around so we can swap them
        let mut empty_buffers = Vec::with_capacity(80);
        empty_buffers.append(&mut vec![Buffers::new(); 30]);

        Self {
            recordings,
            processed,
            empty_buffers,
            recording_buffer: Buffers::new(),
        }
    }

    pub fn write(&mut self, l: f32, r: f32, sample_rate: u32) {
        let full = self.recording_buffer.write_advance(l, r);
        // If we fill the buffer, switch it with an empty one
        if full {
            // we have to loop here, cause when the daw renders audio it tries to do it faster than
            // real time. if we don't loop and wait, the processing thread gets stuck with all of the buffers,
            // and we run out of empty ones to switch to
            // the loop-wait ensures that we don't panic when there isn't an empty buffer
            loop {
                // get the empty buffer from unused buffer list
                if let Some(mut buf) = self.empty_buffers.pop() {
                    buf.reset();
                    // swap it with recording buffer
                    std::mem::swap(&mut buf, &mut self.recording_buffer);
                    buf.reset();

                    // pass it to the processor thread
                    let _ = self.recordings.push(DetectionInput::<LEN> {
                        buffers: buf,
                        sample_rate,
                    });
                    break;
                }
                std::thread::sleep(std::time::Duration::from_micros(10));
            }
        }
    }

    pub fn try_get_pitch(&mut self) -> Option<(Option<f32>, Option<f32>)> {
        let DetectionOutput::<LEN> {
            buffers,
            pitch_l,
            pitch_r,
        } = self.processed.pop()?;
        self.empty_buffers.push(buffers);

        Some((pitch_l, pitch_r))
    }
}