bonsai-eval-utils 0.1.1

//! In-memory disk implementation for testing and debugging.
//!
//! This module contains a simple in-memory disk implementation that can be used
//! for testing and debugging. It is not suitable for production use.
//!
//! It will emulate NOR flash by combining the data with the buffer via bitwise
//! AND. Additonally it will enforce the specified write granularity and block
//! size.


use alloc::boxed::Box;
use alloc::vec;
use core::fmt;
use core::ops;

use bonsai_disk::Disk;
use embedded_io::ErrorType;
use generic_array::ConstArrayLength;
use generic_array::IntoArrayLength;
use generic_array::typenum::Const;

use crate::IecNumber;
use crate::MetricNumber;



extern crate alloc;


#[derive(Clone, Copy, Default)]
pub struct AccessStats {
	pub reads: usize,
	pub read_bytes: usize,
	pub writes: usize,
	pub write_bytes: usize,
	pub erases: usize,
}
impl ops::Add for AccessStats {
	type Output = Self;

	fn add(self, other: Self) -> Self {
		Self {
			reads: self.reads + other.reads,
			read_bytes: self.read_bytes + other.read_bytes,
			writes: self.writes + other.writes,
			write_bytes: self.write_bytes + other.write_bytes,
			erases: self.erases + other.erases,
		}
	}
}
impl ops::Sub for AccessStats {
	type Output = Self;

	fn sub(self, other: Self) -> Self {
		Self {
			reads: self.reads.saturating_sub(other.reads),
			read_bytes: self.read_bytes.saturating_sub(other.read_bytes),
			writes: self.writes.saturating_sub(other.writes),
			write_bytes: self.write_bytes.saturating_sub(other.write_bytes),
			erases: self.erases.saturating_sub(other.erases),
		}
	}
}
impl ops::AddAssign for AccessStats {
	fn add_assign(&mut self, other: Self) {
		*self = self.clone() + other;
	}
}
impl ops::SubAssign for AccessStats {
	fn sub_assign(&mut self, other: Self) {
		*self = self.clone() - other;
	}
}
impl ops::Mul<usize> for AccessStats {
	type Output = Self;

	fn mul(self, rhs: usize) -> Self {
		Self {
			reads: self.reads * rhs,
			read_bytes: self.read_bytes * rhs,
			writes: self.writes * rhs,
			write_bytes: self.write_bytes * rhs,
			erases: self.erases * rhs,
		}
	}
}
impl ops::Div<usize> for AccessStats {
	type Output = Self;

	fn div(self, rhs: usize) -> Self {
		Self {
			reads: self.reads / rhs,
			read_bytes: self.read_bytes / rhs,
			writes: self.writes / rhs,
			write_bytes: self.write_bytes / rhs,
			erases: self.erases / rhs,
		}
	}
}
impl ops::MulAssign<usize> for AccessStats {
	fn mul_assign(&mut self, rhs: usize) {
		*self = self.clone() * rhs;
	}
}
impl ops::DivAssign<usize> for AccessStats {
	fn div_assign(&mut self, rhs: usize) {
		*self = self.clone() / rhs;
	}
}
impl fmt::Debug for AccessStats {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		if self.writes == 0 && self.write_bytes == 0 && self.erases == 0 {
			f.debug_struct("ReadStats")
				.field("reads", &self.reads)
				.field("read_bytes", &self.read_bytes)
				.finish()
		} else {
			f.debug_struct("AccessStats")
				.field("reads", &self.reads)
				.field("read_bytes", &self.read_bytes)
				.field("writes", &self.writes)
				.field("write_bytes", &self.write_bytes)
				.field("erases", &self.erases)
				.finish()
		}
	}
}
impl fmt::Display for AccessStats {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		if self.writes == 0 && self.write_bytes == 0 && self.erases == 0 {
			write!(
				f,
				"reads {}, {}B",
				MetricNumber(self.reads),
				IecNumber(self.read_bytes)
			)
		} else {
			write!(
				f,
				"reads {}, {}B, writes {}, {}B, erases {} pages",
				MetricNumber(self.reads),
				IecNumber(self.read_bytes),
				MetricNumber(self.writes),
				IecNumber(self.write_bytes),
				MetricNumber(self.erases),
			)
		}
	}
}

/// A simple in-memory disk.
///
/// Used for testing and debugging.
#[derive(Debug, Clone)]
pub struct RamDisk<const BLOCK_SIZE: usize, const WRITE_SIZE: usize> {
	pub data: Box<[u8]>,

	pub stats: AccessStats,
}
impl<const BLOCK_SIZE: usize, const WRITE_SIZE: usize> RamDisk<BLOCK_SIZE, WRITE_SIZE> {
	pub fn new(blocks: usize) -> Self {
		Self {
			data: vec![0; blocks * BLOCK_SIZE].into_boxed_slice(),
			stats: AccessStats::default(),
		}
	}

	pub fn new_erased(blocks: usize) -> Self {
		Self {
			data: vec![0xFF; blocks * BLOCK_SIZE].into_boxed_slice(),
			stats: AccessStats::default(),
		}
	}

	pub fn clear_stats(&mut self) {
		self.stats = AccessStats::default();
	}
}


impl<const BLOCK_SIZE: usize, const WRITE_SIZE: usize> AsMut<Self>
	for RamDisk<BLOCK_SIZE, WRITE_SIZE>
{
	fn as_mut(&mut self) -> &mut Self {
		self
	}
}

impl<const BLOCK_SIZE: usize, const WRITE_SIZE: usize> ErrorType
	for RamDisk<BLOCK_SIZE, WRITE_SIZE>
{
	type Error = core::convert::Infallible;
}
impl<const BLOCK_SIZE: usize, const WRITE_SIZE: usize> Disk for RamDisk<BLOCK_SIZE, WRITE_SIZE>
where
	Const<WRITE_SIZE>: IntoArrayLength,
	// Const<WRITE_SIZE>: ToUInt,
	// typenum::U<WRITE_SIZE>: ArrayLength,
{
	type WRITE_GRANULARITY = ConstArrayLength<WRITE_SIZE>;

	const ERASE_BLOCK_SIZE: usize = BLOCK_SIZE;

	#[track_caller]
	fn read(&mut self, offset: usize, buf: &mut [u8]) -> Result<usize, Self::Error> {
		assert!(BLOCK_SIZE % WRITE_SIZE == 0);
		assert!(BLOCK_SIZE >= WRITE_SIZE);

		assert!(
			offset % WRITE_SIZE == 0,
			"Unaligned read: {offset} % {WRITE_SIZE}"
		);
		assert!(
			offset + buf.len() <= self.data.len(),
			"Read out of bounds: {offset} + {} > {}",
			buf.len(),
			self.data.len()
		);

		// An empty buffer is a valid read.
		if buf.len() == 0 {
			return Ok(0);
		}

		assert!(
			buf.len() >= WRITE_SIZE,
			"Buffer too small: {} < {WRITE_SIZE}",
			buf.len()
		);

		// Round buf len down to the next multiple of WRITE_SIZE
		let buf_len = buf.len() / WRITE_SIZE * WRITE_SIZE;

		// Read at most one block.
		let block_end = ((offset / BLOCK_SIZE) + 1) * BLOCK_SIZE;
		let read_end = usize::min(block_end, offset + buf_len);
		let read_len = read_end - offset;

		buf[..read_len].copy_from_slice(&self.data[offset..read_end]);
		self.stats.reads += 1;
		self.stats.read_bytes += read_len;

		Ok(read_len)
	}

	#[track_caller]
	fn write(&mut self, offset: usize, buf: &[u8]) -> Result<usize, Self::Error> {
		assert!(BLOCK_SIZE % WRITE_SIZE == 0);
		assert!(BLOCK_SIZE >= WRITE_SIZE);

		assert!(
			offset % WRITE_SIZE == 0,
			"Unaligned write: {offset} % {WRITE_SIZE}"
		);
		assert!(
			offset + buf.len() <= self.data.len(),
			"Write out of bounds: {offset} + {} > {}",
			buf.len(),
			self.data.len()
		);

		// An empty buffer is a valid read.
		if buf.len() == 0 {
			return Ok(0);
		}

		assert!(
			buf.len() >= WRITE_SIZE,
			"Buffer too small: {} < {WRITE_SIZE}",
			buf.len()
		);

		// Round buf len down to the next multiple of WRITE_SIZE
		let buf_len = buf.len() / WRITE_SIZE * WRITE_SIZE;

		// Write at most one block.
		let block_end = ((offset / BLOCK_SIZE) + 1) * BLOCK_SIZE;
		let write_end = usize::min(block_end, offset + buf_len);
		let write_len = write_end - offset;

		// Emulate NOR flash by combining the data with the buffer via bitwise AND
		for i in 0..write_len {
			self.data[offset + i] &= buf[i];
		}

		self.stats.writes += 1;
		self.stats.write_bytes += write_len;

		Ok(write_len)
	}

	#[track_caller]
	fn erase(&mut self, block: usize) -> Result<(), Self::Error> {
		assert!(BLOCK_SIZE % WRITE_SIZE == 0);
		assert!(BLOCK_SIZE >= WRITE_SIZE);

		let offset = block * BLOCK_SIZE;
		let len = BLOCK_SIZE;

		// Emulate NOR flash by setting the block to all ones
		for i in (offset)..(offset + len) {
			self.data[i] = 0xFF;
		}

		self.stats.erases += 1;

		Ok(())
	}

	#[track_caller]
	fn block_count(&self) -> usize {
		assert!(BLOCK_SIZE % WRITE_SIZE == 0);
		assert!(BLOCK_SIZE >= WRITE_SIZE);

		debug_assert!(self.data.len() % BLOCK_SIZE == 0);

		self.data.len() / BLOCK_SIZE
	}
}


#[cfg(test)]
mod test {
	use super::*;

	const BLOCK_SIZE: usize = 512;
	const WRITE_SIZE: usize = 8;

	#[test]
	fn test_ramdisk_simple_write() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// Write data to the disk
		let write_data = [1, 2, 3, 4, 5, 6, 7, 8];
		let write_result = disk.write(2 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_data.len());

		// Check the content of the array in the ram disk directly
		for i in 0..write_data.len() {
			assert_eq!(disk.data[2 * WRITE_SIZE + i], write_data[i]);
		}

		// Assert write stats
		assert_eq!(disk.stats.writes, 1);
		assert_eq!(disk.stats.write_bytes, write_data.len());
		assert_eq!(disk.stats.reads, 0);
		assert_eq!(disk.stats.read_bytes, 0);
		assert_eq!(disk.stats.erases, 0);
	}

	#[test]
	fn test_ramdisk_partial_write() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// Write data to the disk with a 12 byte buffer
		let write_data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12];
		let write_result = disk.write(2 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), 8); // Only the first 8 bytes should be written

		// Check the content of the array in the ram disk directly
		for i in 0..8 {
			assert_eq!(disk.data[2 * WRITE_SIZE + i], write_data[i]);
		}

		// Ensure the remaining bytes are unchanged (still erased)
		for i in 8..12 {
			assert_eq!(disk.data[2 * WRITE_SIZE + i], 0xFF);
		}

		// Assert write stats
		assert_eq!(disk.stats.writes, 1);
		assert_eq!(disk.stats.write_bytes, 8);
		assert_eq!(disk.stats.reads, 0);
		assert_eq!(disk.stats.read_bytes, 0);
		assert_eq!(disk.stats.erases, 0);
	}

	#[test]
	fn test_ramdisk_simple_read() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// Directly set some bytes in the disk array
		let direct_data = [10, 20, 30, 40, 50, 60, 70, 80];
		for i in 0..direct_data.len() {
			disk.data[2 * WRITE_SIZE + i] = direct_data[i];
		}

		// Read data from the disk
		let mut read_data = [0u8; 8];
		let read_result = disk.read(2 * WRITE_SIZE, &mut read_data);
		assert!(read_result.is_ok());
		assert_eq!(read_result.unwrap(), direct_data.len());
		assert_eq!(read_data, direct_data);

		// Assert read stats
		assert_eq!(disk.stats.reads, 1);
		assert_eq!(disk.stats.read_bytes, direct_data.len());
		assert_eq!(disk.stats.writes, 0);
		assert_eq!(disk.stats.write_bytes, 0);
		assert_eq!(disk.stats.erases, 0);
	}

	#[test]
	fn test_ramdisk_partial_read() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// Directly set some bytes in the disk array
		let direct_data = [10, 20, 30, 40, 50, 60, 70, 80];
		for i in 0..direct_data.len() {
			disk.data[2 * WRITE_SIZE + i] = direct_data[i];
		}

		// Read data from the disk with a 12 byte buffer
		let mut read_data = [0u8; 12];
		let read_result = disk.read(2 * WRITE_SIZE, &mut read_data);
		assert!(read_result.is_ok());
		assert_eq!(read_result.unwrap(), 8); // Only the first 8 bytes should be read

		// Check the content of the read data
		for i in 0..8 {
			assert_eq!(read_data[i], direct_data[i]);
		}

		// Ensure the remaining bytes are unchanged (still zero)
		for i in 8..12 {
			assert_eq!(read_data[i], 0);
		}

		// Assert read stats
		assert_eq!(disk.stats.reads, 1);
		assert_eq!(disk.stats.read_bytes, 8);
		assert_eq!(disk.stats.writes, 0);
		assert_eq!(disk.stats.write_bytes, 0);
		assert_eq!(disk.stats.erases, 0);
	}

	#[test]
	fn test_ramdisk_read_write() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// Write data to the disk
		let write_data = [1, 2, 3, 4, 5, 6, 7, 8];
		let write_result = disk.write(3 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_data.len());

		// Read data from the disk
		let mut read_data = [0u8; 8];
		let read_result = disk.read(3 * WRITE_SIZE, &mut read_data);
		assert!(read_result.is_ok());
		assert_eq!(read_result.unwrap(), write_data.len());
		assert_eq!(read_data, write_data);

		// Assert read/write stats
		assert_eq!(disk.stats.reads, 1);
		assert_eq!(disk.stats.read_bytes, write_data.len());
		assert_eq!(disk.stats.writes, 1);
		assert_eq!(disk.stats.write_bytes, write_data.len());
		assert_eq!(disk.stats.erases, 0);
	}

	#[test]
	fn test_ramdisk_single_write_two_reads() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// Write data to the disk
		let write_data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
		let write_result = disk.write(2 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_data.len());

		// First read from the disk
		let mut read_data1 = [0u8; 8];
		let read_result1 = disk.read(2 * WRITE_SIZE, &mut read_data1);
		assert!(read_result1.is_ok());
		assert_eq!(read_result1.unwrap(), 8);
		assert_eq!(read_data1, [1, 2, 3, 4, 5, 6, 7, 8]);

		// Second read from the disk
		let mut read_data2 = [0u8; 8];
		let read_result2 = disk.read(2 * WRITE_SIZE + 8, &mut read_data2);
		assert!(read_result2.is_ok());
		assert_eq!(read_result2.unwrap(), 8);
		assert_eq!(read_data2, [9, 10, 11, 12, 13, 14, 15, 16]);

		// Assert read/write stats
		assert_eq!(disk.stats.reads, 2);
		assert_eq!(disk.stats.read_bytes, 16);
		assert_eq!(disk.stats.writes, 1);
		assert_eq!(disk.stats.write_bytes, 16);
		assert_eq!(disk.stats.erases, 0);
	}

	#[test]
	fn test_ramdisk_two_writes_single_read() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// First write to the disk
		let write_data1 = [1, 2, 3, 4, 5, 6, 7, 8];
		let write_result1 = disk.write(2 * WRITE_SIZE, &write_data1);
		assert!(write_result1.is_ok());
		assert_eq!(write_result1.unwrap(), write_data1.len());

		// Second write to the disk
		let write_data2 = [9, 10, 11, 12, 13, 14, 15, 16];
		let write_result2 = disk.write(3 * WRITE_SIZE, &write_data2);
		assert!(write_result2.is_ok());
		assert_eq!(write_result2.unwrap(), write_data2.len());

		// Read data from the disk
		let mut read_data = [0u8; 16];
		let read_result = disk.read(2 * WRITE_SIZE, &mut read_data);
		assert!(read_result.is_ok());
		assert_eq!(read_result.unwrap(), 16);
		assert_eq!(
			read_data,
			[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
		);

		// Assert read/write stats
		assert_eq!(disk.stats.reads, 1);
		assert_eq!(disk.stats.read_bytes, 16);
		assert_eq!(disk.stats.writes, 2);
		assert_eq!(disk.stats.write_bytes, 16);
		assert_eq!(disk.stats.erases, 0);
	}

	#[test]
	fn test_ramdisk_repeated_write() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// Write data to the disk
		let write_data = [1, 2, 3, 4, 5, 6, 7, 8];
		let write_result = disk.write(3 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_data.len());

		// Write the same data again to the same location
		let write_result = disk.write(3 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_data.len());

		// Read data from the disk
		let mut read_data = [0u8; 8];
		let read_result = disk.read(3 * WRITE_SIZE, &mut read_data);
		assert!(read_result.is_ok());
		assert_eq!(read_result.unwrap(), write_data.len());
		assert_eq!(read_data, write_data);

		// Assert read/write stats
		assert_eq!(disk.stats.reads, 1);
		assert_eq!(disk.stats.read_bytes, write_data.len());
		assert_eq!(disk.stats.writes, 2);
		assert_eq!(disk.stats.write_bytes, write_data.len() * 2);
		assert_eq!(disk.stats.erases, 0);
	}

	#[test]
	fn test_ramdisk_erase() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new_erased(4);

		// Write data to the disk
		let write_data = [1, 2, 3, 4, 5, 6, 7, 8];
		let write_result = disk.write(3 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_data.len());

		// Erase the block
		let erase_result = disk.erase(0);
		assert!(erase_result.is_ok());

		// Read data from the disk
		let mut read_data = [0u8; 8];
		let read_result = disk.read(3 * WRITE_SIZE, &mut read_data);
		assert!(read_result.is_ok());
		assert_eq!(read_result.unwrap(), write_data.len());
		// No longer what we wrote
		assert_ne!(read_data, write_data);

		// Assert read/write stats
		assert_eq!(disk.stats.reads, 1);
		assert_eq!(disk.stats.read_bytes, 8);
		assert_eq!(disk.stats.writes, 1);
		assert_eq!(disk.stats.write_bytes, 8);
		assert_eq!(disk.stats.erases, 1);
	}

	#[test]
	fn test_ramdisk_write_without_erase() {
		let mut disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new(4);

		// Just write something to make it non-erased
		let write_garbage = [0; 8];
		let write_result = disk.write(3 * WRITE_SIZE, &write_garbage);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_garbage.len());

		// Write data to the disk without erasing yields garbage
		let write_data = [1, 2, 3, 4, 5, 6, 7, 8];
		let write_result = disk.write(3 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_data.len());

		// Read data from the disk
		let mut read_data = [0u8; 8];
		let read_result = disk.read(3 * WRITE_SIZE, &mut read_data);
		assert!(read_result.is_ok());
		assert_eq!(read_result.unwrap(), write_data.len());
		// Not what we wrote
		assert_ne!(read_data, write_data);

		// Erase the block
		let erase_result = disk.erase(0);
		assert!(erase_result.is_ok());

		// Write data to the disk after erasing works
		let write_result = disk.write(3 * WRITE_SIZE, &write_data);
		assert!(write_result.is_ok());
		assert_eq!(write_result.unwrap(), write_data.len());

		// Read data from the disk
		let read_result = disk.read(3 * WRITE_SIZE, &mut read_data);
		assert!(read_result.is_ok());
		assert_eq!(read_result.unwrap(), write_data.len());
		// Exactly what we wrote
		assert_eq!(read_data, write_data);

		// Assert read/write stats
		assert_eq!(disk.stats.reads, 2);
		assert_eq!(disk.stats.read_bytes, 16);
		assert_eq!(disk.stats.writes, 3);
		assert_eq!(disk.stats.write_bytes, 24);
		assert_eq!(disk.stats.erases, 1);
	}

	#[test]
	fn test_ramdisk_block_count() {
		let disk = RamDisk::<BLOCK_SIZE, WRITE_SIZE>::new(3);

		assert_eq!(disk.block_count(), 3);
	}
}