bass-foundry/plugins/rivet/app/src/lib.rs

//! Shared application state and bounded control messages. The audio callback never locks a session.
use crossbeam_queue::ArrayQueue;
use rivet_core::{
    Audition, Error, EventBatch, MessageKind, MidiMessage, Performance, PerformanceClip,
    PerformancePlayer, Result, Session,
};
use std::{
    fs::File,
    io::{Read, Write},
    path::{Path, PathBuf},
    sync::{
        atomic::{AtomicBool, AtomicU32, AtomicU64, Ordering},
        Arc, Mutex, RwLock,
    },
};
pub mod dialog;

#[derive(Clone, Copy)]
pub struct PlayRequest {
    pub epoch: u64,
    pub clip: PerformanceClip,
    pub quantized: bool,
}

pub struct Shared {
    pub session: RwLock<Session>,
    pub epoch: AtomicU64,
    pub revision: AtomicU64,
    pub tempo: AtomicU64,
    pub audition_tempo: AtomicU64,
    pub follow_host: AtomicBool,
    pub midi: AtomicBool,
    pub looping: AtomicBool,
    pub queued: AtomicBool,
    pub playing: AtomicBool,
    pub playhead: AtomicU64,
    /// Bar length in ticks; queued auditions switch on the next bar line.
    pub bar: AtomicU32,
    pub commands: ArrayQueue<PlayRequest>,
    pub busy: AtomicBool,
    results: Mutex<Vec<Completion>>,
}
pub enum Output {
    Imported(rivet_midi::MidiFile),
    Exported(PathBuf),
    Saved(PathBuf),
    Loaded(Box<Session>),
    Canceled,
}
pub struct Completion {
    pub revision: u64,
    pub result: Result<Output>,
}
impl Default for Shared {
    fn default() -> Self {
        Self {
            session: RwLock::new(Session::default()),
            epoch: AtomicU64::new(0),
            revision: AtomicU64::new(0),
            tempo: AtomicU64::new(140f64.to_bits()),
            audition_tempo: AtomicU64::new(f64::NAN.to_bits()),
            follow_host: AtomicBool::new(false),
            midi: AtomicBool::new(false),
            looping: AtomicBool::new(true),
            queued: AtomicBool::new(false),
            playing: AtomicBool::new(false),
            playhead: AtomicU64::new(0),
            bar: AtomicU32::new(rivet_core::BAR),
            commands: ArrayQueue::new(2),
            busy: AtomicBool::new(false),
            results: Mutex::new(Vec::new()),
        }
    }
}
impl Shared {
    pub fn snapshot(&self) -> (u64, Session) {
        let state = self.session.read().unwrap_or_else(|p| p.into_inner());
        (self.revision.load(Ordering::Acquire), state.clone())
    }
    pub fn stop(&self) {
        self.epoch.fetch_add(1, Ordering::AcqRel);
        self.playing.store(false, Ordering::Release);
        self.queued.store(false, Ordering::Release);
    }
    fn publish(&self, s: &Session) {
        self.tempo.store(s.tempo.to_bits(), Ordering::Release);
        self.follow_host.store(s.follow_host, Ordering::Release);
        self.midi
            .store(s.audition == Audition::Midi, Ordering::Release);
        self.looping.store(s.looping, Ordering::Release);
        self.bar.store(s.settings.meter.bar(), Ordering::Release);
        self.revision.fetch_add(1, Ordering::AcqRel);
    }
    pub fn replace(&self, mut session: Session) -> Result<()> {
        session.migrate()?;
        session.validate()?;
        self.stop();
        let mut current = self.session.write().unwrap_or_else(|p| p.into_inner());
        *current = session;
        self.publish(&current);
        Ok(())
    }
    pub fn edit(&self, revision: u64, mut session: Session) -> Result<()> {
        session.migrate()?;
        session.validate()?;
        let mut current = self.session.write().unwrap_or_else(|p| p.into_inner());
        if self.revision.load(Ordering::Acquire) != revision {
            return Err(Error(
                "Session changed while this operation was running. Try again.".into(),
            ));
        }
        if current.settings.meter != session.settings.meter
            || current.performance != session.performance
            || current.audition != session.audition
        {
            self.stop();
        }
        *current = session;
        self.publish(&current);
        Ok(())
    }
    fn compile_clip(&self, phrase: &rivet_core::Phrase) -> Result<PerformanceClip> {
        let state = self.session.read().unwrap_or_else(|p| p.into_inner());
        let observed = f64::from_bits(self.audition_tempo.load(Ordering::Acquire));
        let tempo = if state.follow_host && observed.is_finite() {
            observed
        } else {
            state.tempo
        };
        let settings = state.performance.clone();
        let compensate = state.audition == Audition::Midi;
        let meter = state.settings.meter;
        drop(state);
        Performance::compile(phrase, &settings, tempo, compensate, meter)?.clip()
    }
    pub fn preview(&self, phrase: &rivet_core::Phrase) -> Result<()> {
        let epoch = self.epoch.load(Ordering::Acquire);
        let clip = self.compile_clip(phrase)?;
        let next = epoch.wrapping_add(1);
        self.epoch
            .compare_exchange(epoch, next, Ordering::AcqRel, Ordering::Acquire)
            .map_err(|_| Error("Session changed while preparing audition. Try again.".into()))?;
        self.playing.store(false, Ordering::Release);
        self.queued.store(false, Ordering::Release);
        self.commands.force_push(PlayRequest {
            epoch: next,
            clip,
            quantized: false,
        });
        Ok(())
    }
    pub fn audition(&self, phrase: &rivet_core::Phrase) -> Result<()> {
        if !self.playing.load(Ordering::Acquire) {
            return self.preview(phrase);
        }
        let epoch = self.epoch.load(Ordering::Acquire);
        let clip = self.compile_clip(phrase)?;
        if self.epoch.load(Ordering::Acquire) != epoch {
            return Err(Error(
                "Session changed while preparing audition. Try again.".into(),
            ));
        }
        self.queued.store(true, Ordering::Release);
        self.commands.force_push(PlayRequest {
            epoch,
            clip,
            quantized: true,
        });
        Ok(())
    }
    /// All jobs run off the audio/UI thread; one outstanding job per instance bounds resource use.
    pub fn work(
        self: &Arc<Self>,
        revision: u64,
        job: impl FnOnce() -> Result<Output> + Send + 'static,
    ) -> Result<()> {
        self.busy
            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
            .map_err(|_| Error("Wait for the current operation to finish.".into()))?;
        let shared = self.clone();
        let result = std::thread::Builder::new()
            .name("motif-worker".into())
            .spawn(move || {
                let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(job))
                    .unwrap_or_else(|_| Err(Error("The operation failed unexpectedly.".into())));
                shared
                    .results
                    .lock()
                    .unwrap_or_else(|p| p.into_inner())
                    .push(Completion { revision, result });
                // Busy is cleared when the GUI consumes the result, so completions cannot accumulate.
            });
        if let Err(e) = result {
            self.busy.store(false, Ordering::Release);
            return Err(Error(e.to_string()));
        }
        Ok(())
    }
    pub fn take_completion(&self) -> Option<Completion> {
        let result = self.results.lock().unwrap_or_else(|p| p.into_inner()).pop();
        if result.is_some() {
            self.busy.store(false, Ordering::Release);
        }
        result
    }
}

pub fn read_limited(path: &Path, limit: usize) -> Result<Vec<u8>> {
    let file =
        File::open(path).map_err(|e| Error(format!("Cannot open {}: {e}", path.display())))?;
    let mut bytes = Vec::new();
    file.take(limit as u64 + 1)
        .read_to_end(&mut bytes)
        .map_err(|e| Error(e.to_string()))?;
    if bytes.len() > limit {
        return Err(Error(format!(
            "File exceeds the {} MiB limit.",
            limit / 1024 / 1024
        )));
    }
    Ok(bytes)
}
/// Never silently overwrite a user's existing MIDI/session file.
pub fn write_new(path: &Path, bytes: &[u8]) -> Result<()> {
    let mut file = File::options()
        .write(true)
        .create_new(true)
        .open(path)
        .map_err(|e| {
            Error(format!(
                "Cannot create {}: {e}. Choose a new filename if it already exists.",
                path.display()
            ))
        })?;
    file.write_all(bytes)
        .and_then(|_| file.sync_all())
        .map_err(|e| Error(e.to_string()))
}

/// Fixed-size scheduling and MIDI state; no session locks or heap work on this path.
#[derive(Default)]
pub struct AudioEngine {
    pub player: PerformancePlayer,
    epoch: u64,
    pending: Option<(f64, PerformanceClip)>,
    active: [Option<u8>; 16],
    controls: [Option<(u8, u8)>; 16],
    clear: bool,
    was_midi: bool,
}
pub struct AudioFrame {
    pub sample: f32,
    pub events: EventBatch,
    /// Channel-one summaries retained for simple monophonic callers.
    pub off: Option<u8>,
    pub on: Option<(u8, u8)>,
}
impl AudioEngine {
    pub fn stop(&mut self) {
        self.player.stop();
        self.pending = None;
        self.clear = true;
    }
    pub fn begin_block(&mut self, shared: &Shared) {
        let epoch = shared.epoch.load(Ordering::Acquire);
        if self.epoch != epoch {
            self.stop();
            self.epoch = epoch;
        }
        self.player
            .set_looping(shared.looping.load(Ordering::Acquire));
        for _ in 0..2 {
            let Some(request) = shared.commands.pop() else {
                break;
            };
            if request.epoch != epoch {
                continue;
            }
            if request.quantized && self.player.is_playing() {
                let now = self.player.absolute_tick();
                let lead = request.clip.lead_in as f64;
                let bar = shared.bar.load(Ordering::Acquire).max(1) as f64;
                let boundary = ((now + lead) / bar).floor().mul_add(bar, bar);
                self.pending = Some((boundary - lead, request.clip));
            } else {
                self.player.play(request.clip);
                self.pending = None;
                self.clear = true;
                shared.queued.store(false, Ordering::Release);
            }
        }
    }
    fn release(&mut self, batch: &mut EventBatch) {
        for channel in 0..16 {
            if let Some(key) = self.active[channel].take() {
                batch.push(MidiMessage {
                    kind: MessageKind::Off,
                    channel: channel as u8,
                    key,
                    value: 0,
                    rest: 0,
                });
            }
            if let Some((key, rest)) = self.controls[channel].take() {
                batch.push(MidiMessage {
                    kind: MessageKind::Control,
                    channel: channel as u8,
                    key,
                    value: rest,
                    rest,
                });
            }
        }
    }
    fn forward(&mut self, message: MidiMessage, batch: &mut EventBatch) {
        let channel = message.channel as usize;
        match message.kind {
            MessageKind::On => self.active[channel] = Some(message.key),
            MessageKind::Off => {
                if self.active[channel] == Some(message.key) {
                    self.active[channel] = None;
                }
            }
            MessageKind::Control => self.controls[channel] = Some((message.key, message.rest)),
        }
        batch.push(message);
    }
    pub fn sample(&mut self, bpm: f64) -> f32 {
        self.frame(bpm, false).sample
    }
    pub fn frame(&mut self, bpm: f64, midi: bool) -> AudioFrame {
        if self.pending.as_ref().is_some_and(|(boundary, _)| {
            self.player.absolute_tick() >= *boundary || !self.player.is_playing()
        }) {
            let (_, clip) = self.pending.take().unwrap();
            self.player.play(clip);
            self.clear = true;
        }
        let mut events = EventBatch::default();
        if self.clear || (!midi && self.was_midi) {
            self.release(&mut events);
            self.clear = false;
        }
        let (sample, raw) = self.player.next_frame(bpm);
        if midi {
            if !self.was_midi {
                let current = self.player.current_messages();
                for message in current.events[..current.len].iter().flatten().copied() {
                    self.forward(message, &mut events);
                }
            } else {
                for message in raw.events[..raw.len].iter().flatten().copied() {
                    self.forward(message, &mut events);
                }
            }
            if !self.player.is_playing() {
                self.release(&mut events);
            }
        }
        self.was_midi = midi;
        let off = events
            .events
            .iter()
            .flatten()
            .find(|m| m.channel == 0 && m.kind == MessageKind::Off)
            .map(|m| m.key);
        let on = events
            .events
            .iter()
            .flatten()
            .find(|m| m.channel == 0 && m.kind == MessageKind::On)
            .map(|m| (m.key, m.value));
        AudioFrame {
            sample,
            events,
            off,
            on,
        }
    }
    pub fn end_block(&self, shared: &Shared) {
        shared.queued.store(
            self.pending.is_some() || !shared.commands.is_empty(),
            Ordering::Release,
        );
        shared
            .playing
            .store(self.player.is_playing(), Ordering::Release);
        shared
            .playhead
            .store(self.player.tick().to_bits(), Ordering::Relaxed);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn stop_cancels_pending_and_active_preview() {
        let shared = Shared::default();
        let motif = rivet_core::demo_motif();
        let mut audio = AudioEngine::default();
        shared.preview(&motif).unwrap();
        shared.stop();
        audio.begin_block(&shared);
        assert!(!audio.player.is_playing());
        shared.preview(&motif).unwrap();
        audio.begin_block(&shared);
        assert!(audio.player.is_playing());
        shared.replace(Session::default()).unwrap();
        audio.begin_block(&shared);
        assert!(!audio.player.is_playing());
    }
    #[test]
    fn stale_edits_do_not_replace_new_session() {
        let s = Shared::default();
        let (revision, old) = s.snapshot();
        let mut new = old.clone();
        new.tempo = 160.;
        s.replace(new).unwrap();
        assert!(s.edit(revision, old).is_err());
        assert_eq!(s.snapshot().1.tempo, 160.);
    }
    #[test]
    fn bounded_preview_queue_keeps_latest_request() {
        let s = Shared::default();
        for _ in 0..8 {
            s.preview(&rivet_core::demo_motif()).unwrap();
        }
        assert_eq!(s.commands.len(), 2);
        let mut a = AudioEngine::default();
        a.begin_block(&s);
        assert!(a.player.is_playing());
    }
    #[test]
    fn replacing_session_publishes_audio_controls() {
        let s = Shared::default();
        s.replace(Session {
            tempo: 172.,
            follow_host: true,
            ..Session::default()
        })
        .unwrap();
        assert_eq!(f64::from_bits(s.tempo.load(Ordering::Acquire)), 172.);
        assert!(s.follow_host.load(Ordering::Acquire));
    }

    #[test]
    fn worker_completion_keeps_the_revision_of_its_input() {
        let shared = Arc::new(Shared::default());
        let revision = shared.revision.load(Ordering::Acquire);
        let (release, wait) = std::sync::mpsc::channel();
        shared
            .work(revision, move || {
                wait.recv().unwrap();
                Ok(Output::Canceled)
            })
            .unwrap();
        shared.replace(Session::default()).unwrap();
        release.send(()).unwrap();
        let completion = (0..1000)
            .find_map(|_| {
                let completion = shared.take_completion();
                if completion.is_none() {
                    std::thread::sleep(std::time::Duration::from_millis(1));
                }
                completion
            })
            .expect("worker should complete");
        assert_eq!(completion.revision, revision);
        assert_ne!(completion.revision, shared.revision.load(Ordering::Acquire));
        assert!(!shared.busy.load(Ordering::Acquire));
    }
}