audio: add scheduling, and optimize atomic operations

5 files changed, 216 insertions, 143 deletions

diff --git a/audio/Cargo.toml b/audio/Cargo.toml
index 58f18a2..4bd92b2 100644
--- a/audio/Cargo.toml
+++ b/audio/Cargo.toml
@@ -8,6 +8,8 @@ publish.workspace = true
 rust-version.workspace = true
 
 [dependencies]
+atomic.workspace = true
+bytemuck.workspace = true
 raylib.workspace = true
 
 [lints]
diff --git a/audio/src/channel.rs b/audio/src/channel.rs
index 64970f2..48d416b 100644
--- a/audio/src/channel.rs
+++ b/audio/src/channel.rs
@@ -1,20 +1,35 @@
 #![expect(clippy::cast_possible_truncation)]
 
-use std::sync::{Arc, Mutex};
-
+use atomic::Atomic;
+use bytemuck::NoUninit;
 use raylib::{
 	audio::{AudioStream, RaylibAudio},
 	prelude::audio_stream_callback::set_audio_stream_callback,
 };
+use std::sync::{Arc, atomic::Ordering};
 
 const SAMPLE_RATE: u32 = 44100;
 const SAMPLE_SIZE: u32 = i16::BITS;
 
-#[derive(Debug, Clone, Copy)]
+fn map_pitch(pitch: u8) -> f32 {
+	const HALF_STEP: f32 = 1.059_463_1;
+	HALF_STEP.powf(pitch as f32 - 128.0)
+}
+
+const fn map_volume(volume: u8) -> f32 {
+	if volume > 100 {
+		1.0
+	} else {
+		(volume as f32) / 100.0
+	}
+}
+
+#[repr(u8)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, NoUninit)]
 pub enum ChannelKind {
 	Pulse { duty_cycle: u8 },
-	Triangle,
-	Noise { mode: bool, lsr: i16 },
+	Triangle { _pad: u8 },
+	Noise { mode: bool },
 }
 impl ChannelKind {
 	pub const fn pulse() -> Self {
@@ -22,75 +37,98 @@ impl ChannelKind {
 	}
 
 	pub const fn triangle() -> Self {
-		Self::Triangle
+		Self::Triangle { _pad: 0 }
 	}
 
 	pub const fn noise() -> Self {
-		Self::Noise {
-			mode: false,
-			lsr: 1,
-		}
+		Self::Noise { mode: false }
 	}
 
-	fn sample(&mut self, phase: f32) -> i16 {
+	fn sample(self, phase: f32, lsr_ptr: &mut i16) -> i16 {
 		match self {
 			Self::Pulse { duty_cycle } => {
-				let duty = *duty_cycle as f32 / 100.0;
+				let duty = duty_cycle as f32 / 100.0;
 				if phase < duty { i16::MAX } else { i16::MIN }
 			}
-			Self::Triangle => {
+			Self::Triangle { .. } => {
 				let steps = 32;
 				let tphase = if phase < 0.5 { phase } else { 1.0 - phase };
 				let step = (tphase * steps as f32).floor() as u32;
 				let value = ((step as f32 / (steps - 1) as f32) * 4.0) - 1.0;
 				(value * i16::MAX as f32) as i16
 			}
-			Self::Noise { mode, lsr } => {
-				let feedback = if *mode {
-					(*lsr & 1) ^ ((*lsr >> 1) & 1)
+			Self::Noise { mode } => {
+				let mut lsr = *lsr_ptr;
+				let feedback = if mode {
+					(lsr & 1) ^ ((lsr >> 1) & 1)
 				} else {
-					(*lsr & 1) ^ ((*lsr >> 6) & 1)
+					(lsr & 1) ^ ((lsr >> 6) & 1)
 				};
-				*lsr >>= 1;
-				*lsr |= feedback << 14;
-				*lsr
+				lsr >>= 1;
+				lsr |= feedback << 14;
+				*lsr_ptr = lsr;
+				lsr
 			}
 		}
 	}
 }
 
-#[derive(Debug)]
+#[repr(C)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, NoUninit)]
 pub struct Channel {
-	pub volume: f32,
-	pub pitch: f32,
+	pub volume: u8,
+	pub pitch: u8,
 	pub kind: ChannelKind,
-	phase: f32,
 }
 impl Channel {
 	const fn new(kind: ChannelKind) -> Self {
 		Self {
-			volume: 0.0,
-			pitch: 1.0,
+			volume: 0,
+			pitch: 128,
 			kind,
-			phase: 0.0,
 		}
 	}
 
-	fn sample(&mut self, buffer: &mut [i16]) {
-		let freq = self.pitch * 440.0;
+	fn sample(self, buffer: &mut [i16], phase_ptr: &mut f32, lsr_ptr: &mut i16) {
+		let freq = map_pitch(self.pitch) * 440.0;
+		let volume = map_volume(self.volume);
 		let step = freq / SAMPLE_RATE as f32;
+		let mut phase = *phase_ptr;
 		for sample in buffer {
-			self.phase += step;
-			if self.phase >= 1.0 {
-				self.phase -= 1.0;
+			phase += step;
+			if phase >= 1.0 {
+				phase -= 1.0;
 			}
-			let real_note = self.kind.sample(self.phase);
-			let note = (real_note as f32 * self.volume) as i16;
+			let real_note = self.kind.sample(phase, lsr_ptr);
+			let note = (real_note as f32 * volume) as i16;
 			*sample += note / 4;
 		}
+		*phase_ptr = phase;
 	}
 }
 
+#[derive(Debug, Clone, Copy)]
+struct RuntimeState {
+	pulse_a_phase: f32,
+	pulse_b_phase: f32,
+	triangle_phase: f32,
+	noise_phase: f32,
+	noise_lsfr: i16,
+}
+impl RuntimeState {
+	const fn new() -> Self {
+		Self {
+			pulse_a_phase: 0.0,
+			pulse_b_phase: 0.0,
+			triangle_phase: 0.0,
+			noise_phase: 0.0,
+			noise_lsfr: 1,
+		}
+	}
+}
+
+#[repr(C)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, NoUninit)]
 pub struct Channels {
 	pub pulse_a: Channel,
 	pub pulse_b: Channel,
@@ -98,7 +136,7 @@ pub struct Channels {
 	pub noise: Channel,
 }
 impl Channels {
-	const fn new() -> Self {
+	pub const fn new() -> Self {
 		let pulse_a = Channel::new(ChannelKind::pulse());
 		let pulse_b = Channel::new(ChannelKind::pulse());
 		let triangle = Channel::new(ChannelKind::triangle());
@@ -113,23 +151,27 @@ impl Channels {
 }
 
 pub struct Device<'s> {
-	pub channels: Arc<Mutex<Channels>>,
+	pub channels: Arc<Atomic<Channels>>,
 	#[expect(dead_code)]
 	stream: AudioStream<'s>,
 }
 
-#[expect(clippy::unwrap_used)]
 impl<'s> Device<'s> {
 	pub(crate) fn load(handle: &'s RaylibAudio) -> crate::Result<Self> {
-		let channels = Arc::new(Mutex::new(Channels::new()));
+		let channels = Arc::new(Atomic::new(Channels::new()));
 		let stream = handle.new_audio_stream(SAMPLE_RATE, SAMPLE_SIZE, 1);
 		let cb_data = Arc::clone(&channels);
+		let mut state = RuntimeState::new();
 		set_audio_stream_callback(&stream, move |buffer| {
-			let mut state = cb_data.lock().unwrap();
-			state.pulse_a.sample(buffer);
-			state.pulse_b.sample(buffer);
-			state.triangle.sample(buffer);
-			state.noise.sample(buffer);
+			let chan = cb_data.load(Ordering::Relaxed);
+			chan.pulse_a
+				.sample(buffer, &mut state.pulse_a_phase, &mut state.noise_lsfr);
+			chan.pulse_b
+				.sample(buffer, &mut state.pulse_b_phase, &mut state.noise_lsfr);
+			chan.triangle
+				.sample(buffer, &mut state.triangle_phase, &mut state.noise_lsfr);
+			chan.noise
+				.sample(buffer, &mut state.noise_phase, &mut state.noise_lsfr);
 		})?;
 		stream.set_volume(1.0);
 		stream.set_pitch(1.0);
diff --git a/audio/src/data.rs b/audio/src/data.rs
index 9ed8019..f130c2d 100644
--- a/audio/src/data.rs
+++ b/audio/src/data.rs
@@ -1,43 +1,48 @@
-use crate::program::Program;
+use crate::program::Track;
 use std::fs;
 
 macro_rules! load_asm {
-	($path:tt) => {{
+	($looping:tt, $path:tt) => {{
 		let res = if cfg!(any(feature = "static", target_arch = "wasm32")) {
 			let src = include_str!(concat!("../../", $path));
-			Program::parse(src, true)
+			Track::parse(src, $looping)
 		} else {
 			let src = fs::read_to_string($path)?;
-			Program::parse(&src, true)
+			Track::parse(&src, $looping)
 		};
 		res.map_err(|mut err| {
 			err.file = Some(String::from($path));
 			err
 		})?
 	}};
-    ($first:tt, $($arg:tt),*) => {
-		load_asm!($first)$(.merge(load_asm!($arg)))*
+    ($looping:tt, $first:tt, $($arg:tt),*) => {
+		load_asm!($looping, $first)$(.merge(load_asm!($looping, $arg)))*
 	};
 }
 
 pub struct Data {
-	pub explore: Program,
-	pub megalovania: Program,
+	pub explore: Track,
+	pub megalovania: Track,
+	pub test: Track,
 }
 impl Data {
 	pub fn load() -> crate::Result<Self> {
 		let explore = load_asm!(
+			true,
 			"assets/asm/explore_melody.asm",
 			"assets/asm/explore_harmony1.asm",
 			"assets/asm/explore_harmony2.asm"
 		);
 		let megalovania = load_asm!(
+			true,
 			"assets/asm/megalovania_melody.asm",
 			"assets/asm/megalovania_base.asm"
 		);
+		let test = load_asm!(false, "assets/asm/test.asm");
 		Ok(Self {
 			explore,
 			megalovania,
+			test,
 		})
 	}
 }
diff --git a/audio/src/lib.rs b/audio/src/lib.rs
index 02800b9..2f0c4ed 100644
--- a/audio/src/lib.rs
+++ b/audio/src/lib.rs
@@ -1,11 +1,20 @@
 //! The `audio` crate stores all audio assets that need to be loaded during runtime
 
 use raylib::audio::RaylibAudio;
-use std::time::{Duration, Instant};
+use std::{
+	collections::BinaryHeap,
+	sync::atomic::Ordering,
+	time::{Duration, Instant},
+};
 
 use channel::Device;
 use data::Data;
 
+use crate::{
+	channel::Channels,
+	program::{Program, Track},
+};
+
 mod channel;
 mod data;
 mod parse;
@@ -26,6 +35,7 @@ const TIME_SLICE: Duration = Duration::from_millis(1000 / AUDIO_FPS as u64);
 pub struct Audio {
 	device: Device<'static>,
 	last: Instant,
+	queue: BinaryHeap<Program>,
 	pub data: Data,
 }
 impl Audio {
@@ -42,17 +52,43 @@ impl Audio {
 		let handle = Box::leak(Box::new(RaylibAudio::init_audio_device()?));
 		let device = Device::load(handle)?;
 		let last = Instant::now();
+		let queue = BinaryHeap::new();
 		let data = Data::load()?;
 
-		Ok(Self { device, last, data })
+		Ok(Self {
+			device,
+			last,
+			queue,
+			data,
+		})
+	}
+
+	pub fn schedule(&mut self, track: Track, priority: u32) {
+		let program = Program::new(track, priority);
+		self.queue.push(program);
+	}
+
+	pub fn clear(&mut self) {
+		self.queue.clear();
 	}
 
-	#[expect(clippy::unwrap_used)]
 	pub fn update(&mut self) {
 		if self.last.elapsed() >= TIME_SLICE {
-			let mut channels = self.device.channels.lock().unwrap();
-			self.data.explore.exec(&mut channels);
-			self.data.megalovania.exec(&mut channels);
+			let res_opt = self
+				.queue
+				.peek_mut()
+				.map(|mut program| (program.exec(), program.is_finished()));
+			let channels = match res_opt {
+				Some((res, finished)) => {
+					if finished {
+						self.queue.pop();
+					}
+					res
+				}
+				// make the output quiet!
+				None => Channels::new(),
+			};
+			self.device.channels.store(channels, Ordering::Relaxed);
 			self.last = Instant::now();
 		}
 	}
diff --git a/audio/src/program.rs b/audio/src/program.rs
index e5d31bb..c6f048d 100644
--- a/audio/src/program.rs
+++ b/audio/src/program.rs
@@ -1,4 +1,4 @@
-use std::fmt;
+use std::{fmt, rc::Rc};
 
 use crate::{
 	channel::{ChannelKind, Channels},
@@ -50,48 +50,19 @@ impl fmt::Display for Instruction {
 }
 use Instruction as I;
 
-fn map_pitch(pitch: u8) -> f32 {
-	const HALF_STEP: f32 = 1.059_463_1;
-	HALF_STEP.powf(pitch as f32 - 128.0)
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Track {
+	pub ins: Rc<Vec<Instruction>>,
+	pub looping: bool,
 }
-
-const fn map_volume(volume: u8) -> f32 {
-	if volume > 100 {
-		1.0
-	} else {
-		(volume as f32) / 100.0
-	}
-}
-
-#[derive(Debug, Clone)]
-pub struct Program {
-	ins: Vec<Instruction>,
-	pc: usize,
-	pause_cnt: usize,
-	pause_len: usize,
-	looping: bool,
-	playing: bool,
-}
-impl Program {
-	pub const fn new(ins: Vec<Instruction>, looping: bool) -> Self {
-		Self {
-			ins,
-			pc: 0,
-			pause_cnt: 0,
-			pause_len: 4,
-			looping,
-			playing: false,
-		}
-	}
-
+impl Track {
 	pub fn parse(src: &str, looping: bool) -> parse::Result<Self> {
-		let ins = parse::parse(src)?;
-		Ok(Self::new(ins, looping))
+		let ins = Rc::new(parse::parse(src)?);
+		Ok(Self { ins, looping })
 	}
 
 	pub fn merge(self, other: Self) -> Self {
 		let mut res = vec![];
-
 		let mut l = 0;
 		let mut r = 0;
 		let mut l_pause = 0;
@@ -154,10 +125,46 @@ impl Program {
 			}
 		}
 
-		Self::new(res, self.looping)
+		Self {
+			ins: Rc::new(res),
+			looping: self.looping,
+		}
+	}
+}
+impl fmt::Display for Track {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		for (idx, ins) in self.ins.iter().enumerate() {
+			if idx != 0 {
+				writeln!(f)?;
+			}
+			write!(f, "{ins}")?;
+		}
+		Ok(())
 	}
+}
 
-	fn exec_ins(&mut self, channels: &mut Channels, ins: Instruction) {
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Program {
+	track: Track,
+	priority: u32,
+	pc: usize,
+	pause_cnt: usize,
+	pause_len: usize,
+	channels: Channels,
+}
+impl Program {
+	pub const fn new(track: Track, priority: u32) -> Self {
+		Self {
+			track,
+			priority,
+			pc: 0,
+			pause_cnt: 0,
+			pause_len: 4,
+			channels: Channels::new(),
+		}
+	}
+
+	const fn exec_ins(&mut self, ins: Instruction) {
 		use Instruction as I;
 		match ins {
 			I::Pause(cnt) => {
@@ -172,93 +179,74 @@ impl Program {
 			I::SetVolume(chan_spec, volume) => {
 				// set the volume (amplitude) on a given channel
 				use ChanSpec as C;
-				let v = map_volume(volume);
 				match chan_spec {
-					C::PulseA => channels.pulse_a.volume = v,
-					C::PulseB => channels.pulse_b.volume = v,
-					C::Triangle => channels.triangle.volume = v,
-					C::Noise => channels.noise.volume = v,
+					C::PulseA => self.channels.pulse_a.volume = volume,
+					C::PulseB => self.channels.pulse_b.volume = volume,
+					C::Triangle => self.channels.triangle.volume = volume,
+					C::Noise => self.channels.noise.volume = volume,
 				}
 			}
 			I::SetPitch(chan_spec, pitch) => {
 				// set the pitch (freq) on a given channel
 				use ChanSpec as C;
-				let p = map_pitch(pitch);
 				match chan_spec {
-					C::PulseA => channels.pulse_a.pitch = p,
-					C::PulseB => channels.pulse_b.pitch = p,
-					C::Triangle => channels.triangle.pitch = p,
-					C::Noise => channels.noise.pitch = p,
+					C::PulseA => self.channels.pulse_a.pitch = pitch,
+					C::PulseB => self.channels.pulse_b.pitch = pitch,
+					C::Triangle => self.channels.triangle.pitch = pitch,
+					C::Noise => self.channels.noise.pitch = pitch,
 				}
 			}
 			I::SetPulseDutyA(dc) => {
-				if let ChannelKind::Pulse { duty_cycle } = &mut channels.pulse_a.kind {
+				if let ChannelKind::Pulse { duty_cycle } = &mut self.channels.pulse_a.kind {
 					*duty_cycle = dc;
 				}
 			}
 			I::SetPulseDutyB(dc) => {
-				if let ChannelKind::Pulse { duty_cycle } = &mut channels.pulse_b.kind {
+				if let ChannelKind::Pulse { duty_cycle } = &mut self.channels.pulse_b.kind {
 					*duty_cycle = dc;
 				}
 			}
 			I::SetNoiseMode(m) => {
-				if let ChannelKind::Noise { mode, .. } = &mut channels.noise.kind {
+				if let ChannelKind::Noise { mode, .. } = &mut self.channels.noise.kind {
 					*mode = m;
 				}
 			}
 		}
 	}
 
-	pub fn exec(&mut self, channels: &mut Channels) {
-		if !self.playing {
-			return;
-		}
+	pub fn exec(&mut self) -> Channels {
 		if self.pause_cnt > 0 {
 			self.pause_cnt -= 1;
 		}
 		loop {
-			if self.pause_cnt > 0 {
-				break;
-			}
 			if self.is_finished() {
-				if self.looping {
+				if self.track.looping {
 					self.pc = 0;
 				} else {
-					self.playing = false;
 					break;
 				}
 			}
-			let ins = self.ins[self.pc];
-			self.exec_ins(channels, ins);
+			if self.pause_cnt > 0 {
+				break;
+			}
+			let ins = self.track.ins[self.pc];
+			self.exec_ins(ins);
 			self.pc += 1;
 		}
+		self.channels
 	}
 
-	pub const fn is_finished(&self) -> bool {
-		self.pc >= self.ins.len()
-	}
-
-	pub const fn is_playing(&self) -> bool {
-		self.playing
+	pub fn is_finished(&self) -> bool {
+		self.pc >= self.track.ins.len()
 	}
-
-	pub const fn play(&mut self) {
-		self.playing = true;
-		self.pc = 0;
-	}
-
-	pub const fn set_looping(&mut self, looping: bool) {
-		self.looping = looping;
+}
+impl PartialOrd for Program {
+	fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+		Some(self.cmp(other))
 	}
 }
-impl fmt::Display for Program {
-	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-		for (idx, ins) in self.ins.iter().enumerate() {
-			if idx != 0 {
-				writeln!(f)?;
-			}
-			write!(f, "{ins}")?;
-		}
-		Ok(())
+impl Ord for Program {
+	fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+		other.priority.cmp(&self.priority)
 	}
 }