unnieversal/crates/robotuna/src/lib.rs

#![allow(incomplete_features)]
#![feature(generic_associated_types)]

use baseplug::{MidiReceiver, Plugin, ProcessContext};
use pvoc::{FreqBin, PhaseVocoder};
use serde::{Deserialize, Serialize};

use utils::logs::*;
use utils::pitch::*;

const DET_LEN: usize = 2 << 8;

baseplug::model! {
    #[derive(Debug, Serialize, Deserialize)]
    struct RoboTunaModel {
        #[model(min = 0.0, max = 1.0)]
        #[parameter(name = "manual/snap")]
        manual: f32,
        #[model(min = 0.1, max = 2.1)]
        #[parameter(name = "frequency gain")]
        freq_gain: f32,
    }
}

impl Default for RoboTunaModel {
    fn default() -> Self {
        Self {
            manual: 1.0,
            freq_gain: 1.0,
        }
    }
}

struct RoboTuna {
    /// Current midi note
    note: Option<u8>,

    /// Current pitches
    pitch_l: Option<f32>,
    pitch_r: Option<f32>,

    detector_thread: pitch_detection::PitchDetectorThread<DET_LEN>,

    pvoc: PhaseVocoder,
}

impl Plugin for RoboTuna {
    const NAME: &'static str = "robotuna";
    const PRODUCT: &'static str = "robotuna";
    const VENDOR: &'static str = "unnieversal";

    const INPUT_CHANNELS: usize = 2;
    const OUTPUT_CHANNELS: usize = 2;

    type Model = RoboTunaModel;

    #[inline]
    fn new(sample_rate: f32, _model: &RoboTunaModel) -> Self {
        setup_logging("robotuna.log");

        let detector_thread = pitch_detection::PitchDetectorThread::<DET_LEN>::new();

        log::info!("finished init");

        Self {
            note: None,
            pitch_l: None,
            pitch_r: None,

            detector_thread,

            pvoc: PhaseVocoder::new(2, sample_rate as f64, 128, 4),
        }
    }

    #[inline]
    fn process(&mut self, model: &RoboTunaModelProcess, ctx: &mut ProcessContext<Self>) {
        let input = &ctx.inputs[0].buffers;
        let output = &mut ctx.outputs[0].buffers;

        for i in 0..ctx.nframes {
            // pass input to pitch detectors
            self.detector_thread
                .write(input[0][i], input[1][i], ctx.sample_rate as u32);

            // Try to get a pitch from short detector thread
            if let Some((pitch_l, pitch_r)) = self.detector_thread.try_get_pitch() {
                // Update current pitch
                // We use `or`, so we keep the old value if the current one is None
                self.pitch_l = pitch_l.or(self.pitch_l);
                self.pitch_r = pitch_r.or(self.pitch_r);
            }
        }

        let shift = self.shift(model.freq_gain[0], model.manual[0] < 0.5);
        self.pvoc.process(
            input,
            output,
            |channels: usize, bins: usize, input: &[Vec<FreqBin>], output: &mut [Vec<FreqBin>]| {
                for i in 0..channels {
                    for j in 0..bins / 2 {
                        let index = ((j as f64) * shift[i]) as usize;
                        if index < bins / 2 {
                            output[i][index].freq = input[i][j].freq * shift[i];
                            output[i][index].amp += input[i][j].amp;
                        }
                    }
                }
            },
        );
    }
}
impl RoboTuna {
    fn pitch(&self) -> (f32, f32) {
        let l = self.pitch_l.unwrap_or(220.0);
        let r = self.pitch_r.unwrap_or(220.0);

        (l, r)
    }

    fn shift(&self, freq_gain: f32, manual: bool) -> [f64; 2] {
        let (current_pitch_l, current_pitch_r) = self.pitch();

        if manual {
            // If we're on manual, get the expected frequency from the midi note
            if let Some(expected) = self.note.map(midi_note_to_pitch) {
                let l = expected / current_pitch_l;
                let r = expected / current_pitch_r;
                [(freq_gain * l) as f64, (freq_gain * r) as f64]
            } else {
                // If there's no note, we just do frequency gain
                [freq_gain as f64, freq_gain as f64]
            }
        } else {
            // If we're on snap, get the closest note
            let expected_l = closest_note_freq(current_pitch_l);
            let expected_r = closest_note_freq(current_pitch_r);

            let l = expected_l / current_pitch_l;
            let r = expected_r / current_pitch_r;
            [(freq_gain * l) as f64, (freq_gain * r) as f64]
        }
    }
}
impl MidiReceiver for RoboTuna {
    fn midi_input(&mut self, _model: &RoboTunaModelProcess, data: [u8; 3]) {
        match data[0] {
            // note on
            0x90 => {
                self.note = Some(data[1]);
            }
            // note off
            0x80 => {
                // only set note to None if it's the same one we currently have
                if let Some(n) = self.note {
                    if n == data[1] {
                        self.note = None;
                    }
                }
            }

            _ => (),
        }
    }
}

baseplug::vst2!(RoboTuna, b"tuna");