bonsai-eval-utils 0.1.1

use alloc::vec::Vec;
use core::fmt;
use core::ops::Range;

use bonsai_disk::Disk;
use embedded_io::ErrorType;
use generic_array::ArrayLength;
use generic_array::GenericArray;
use generic_array::typenum;
use typenum::marker_traits::Unsigned;



#[derive(Debug, Clone)]
pub struct RecordDisk<D: Disk> {
	pub before_storage: D,

	pub records: Vec<Record<D::WRITE_GRANULARITY>>,
}
impl<D: Disk> RecordDisk<D> {
	pub fn new(before_storage: D) -> Self {
		Self {
			before_storage,
			records: Vec::new(),
		}
	}
}
impl<D: Disk> ErrorType for RecordDisk<D> {
	type Error = <D as ErrorType>::Error;
}
impl<D: Disk> Disk for RecordDisk<D> {
	type WRITE_GRANULARITY = D::WRITE_GRANULARITY;

	const ERASE_BLOCK_SIZE: usize = D::ERASE_BLOCK_SIZE;

	fn block_count(&self) -> usize {
		self.before_storage.block_count()
	}

	fn read(&mut self, offset: usize, buf: &mut [u8]) -> Result<usize, Self::Error> {
		if buf.is_empty() {
			return Ok(0);
		}

		assert!(offset % Self::WRITE_GRANULARITY::USIZE == 0);
		assert!(buf.len() >= Self::WRITE_GRANULARITY::USIZE);

		// Iter in reverse chronological order, i.e. return the most recent state
		for rec in self.records.iter().rev() {
			match rec {
				Record::Write(WriteRecord {
					block,
					ibo,
					data,
				}) => {
					if block * Self::ERASE_BLOCK_SIZE + ibo == offset {
						buf[..Self::WRITE_GRANULARITY::USIZE].copy_from_slice(data.as_slice());
						return Ok(Self::WRITE_GRANULARITY::USIZE);
					}
				},
				Record::Erase(EraseRecord {
					block,
				}) => {
					if *block == offset / Self::ERASE_BLOCK_SIZE {
						// Emulate NOR Flash
						buf[..Self::WRITE_GRANULARITY::USIZE].fill(0xff);
						return Ok(Self::WRITE_GRANULARITY::USIZE);
					}
				},
			}
		}

		self.before_storage
			.read(offset, &mut buf[..Self::WRITE_GRANULARITY::USIZE])
	}

	fn write(&mut self, offset: usize, buf: &[u8]) -> Result<usize, Self::Error> {
		if buf.is_empty() {
			return Ok(0);
		}

		assert!(offset % Self::WRITE_GRANULARITY::USIZE == 0);
		assert!(buf.len() >= Self::WRITE_GRANULARITY::USIZE);

		// Emulate NOR Flash, i.e. read the current state and AND the new data
		// into it
		let mut buffer = GenericArray::default();
		let len = self.read(offset, &mut buffer)?;
		assert_eq!(len, Self::WRITE_GRANULARITY::USIZE);
		// AND together the latest state and the new data
		for (b, input) in buffer.iter_mut().zip(buf) {
			*b &= input;
		}

		let rec = WriteRecord {
			block: offset / Self::ERASE_BLOCK_SIZE,
			ibo: offset % Self::ERASE_BLOCK_SIZE,
			data: buffer,
		};

		self.records.push(Record::Write(rec));

		Ok(Self::WRITE_GRANULARITY::USIZE)
	}

	fn erase(&mut self, block: usize) -> Result<(), Self::Error> {
		self.records.push(Record::Erase(EraseRecord {
			block,
		}));

		Ok(())
	}
}


#[derive(Debug)]
pub struct ReplayDisk<'a, D: Disk> {
	pub before_storage: &'a mut D,

	records: &'a [Record<D::WRITE_GRANULARITY>],

	new_actions: Vec<Record<D::WRITE_GRANULARITY>>,

	/// The highest index which will cause a different result upon the first
	/// action taken by the user.
	///
	/// Zero means the state before the very first records.
	pub first_undiscovered: Option<usize>,

	merged_records: usize,
}
impl<'a, D: Disk> ReplayDisk<'a, D> {
	pub fn new(before_storage: &'a mut D, records: &'a [Record<D::WRITE_GRANULARITY>]) -> Self {
		Self {
			before_storage,
			records,
			new_actions: Vec::new(),
			first_undiscovered: None,
			merged_records: 0,
		}
	}

	pub fn from_sub_range(
		before_storage: &'a mut D,
		records: &'a [Record<D::WRITE_GRANULARITY>],
		range: Range<usize>,
	) -> Self {
		let mut me = Self {
			before_storage,
			records: &records,
			new_actions: Vec::new(),
			first_undiscovered: None,
			merged_records: 0,
		};

		me.cut_off_top(range.end);
		me.merge_bottom(range.start);

		me
	}

	pub fn with_limit(
		before_storage: &'a mut D,
		records: &'a [Record<D::WRITE_GRANULARITY>],
		range: Range<usize>,
		limit: Option<usize>,
	) -> Self {
		let mut me = Self::from_sub_range(before_storage, records, range);
		me.first_undiscovered = limit;

		me
	}

	pub fn current_range(&self) -> Range<usize> {
		let start = self.merged_records;
		let end = self.merged_records + self.records.len();
		start..end
	}

	pub fn before_range(&self) -> Range<usize> {
		let start = self.merged_records;
		if let Some(first_undiscovered) = self.first_undiscovered {
			start..(start + first_undiscovered)
		} else {
			start..start
		}
	}

	pub fn after_range(&self) -> Range<usize> {
		let start = self.merged_records;
		if let Some(first_undiscovered) = self.first_undiscovered {
			(start + first_undiscovered + 1)..(start + self.records.len())
		} else {
			(start + self.records.len())..(start + self.records.len())
		}
	}

	pub fn cut_off_top(&mut self, top: usize) {
		self.records = &self.records[..top];
	}

	pub fn merge_bottom(&mut self, bottom: usize) {
		let mergable = &self.records[..bottom];
		self.records = &self.records[bottom..];

		self.merged_records += mergable.len();

		// Modify the base image to incorporate the merged records
		for rec in mergable.iter() {
			match rec {
				Record::Write(WriteRecord {
					block,
					ibo,
					data,
				}) => {
					self.before_storage
						.write(block * D::ERASE_BLOCK_SIZE + ibo, data.as_slice())
						.unwrap();
				},
				Record::Erase(EraseRecord {
					block,
				}) => {
					self.before_storage.erase(*block).unwrap();
				},
			}
		}
	}
}
impl<D: Disk> ErrorType for ReplayDisk<'_, D> {
	type Error = <D as ErrorType>::Error;
}
impl<D: Disk> Disk for ReplayDisk<'_, D> {
	type WRITE_GRANULARITY = D::WRITE_GRANULARITY;

	const ERASE_BLOCK_SIZE: usize = D::ERASE_BLOCK_SIZE;

	fn block_count(&self) -> usize {
		self.before_storage.block_count()
	}

	fn read(&mut self, offset: usize, buf: &mut [u8]) -> Result<usize, Self::Error> {
		if buf.is_empty() {
			return Ok(0);
		}

		assert!(offset % Self::WRITE_GRANULARITY::USIZE == 0);
		assert!(buf.len() >= Self::WRITE_GRANULARITY::USIZE);

		// Iter in reverse chronological order, i.e. return the most recent state
		for rec in self.new_actions.iter().rev() {
			match rec {
				Record::Write(WriteRecord {
					block,
					ibo,
					data,
				}) => {
					if block * Self::ERASE_BLOCK_SIZE + ibo == offset {
						buf[..Self::WRITE_GRANULARITY::USIZE].copy_from_slice(data.as_slice());
						return Ok(Self::WRITE_GRANULARITY::USIZE);
					}
				},
				Record::Erase(EraseRecord {
					block,
				}) => {
					if *block == offset / Self::ERASE_BLOCK_SIZE {
						// Emulate NOR Flash
						buf[..Self::WRITE_GRANULARITY::USIZE].fill(0xff);
						return Ok(Self::WRITE_GRANULARITY::USIZE);
					}
				},
			}
		}

		let limit = if let Some(limit) = self.first_undiscovered {
			limit + 1
		} else {
			self.records.len()
		};

		// Iter in reverse chronological order, i.e. return the most recent state
		for (i, rec) in self.records[..limit].iter().enumerate().rev() {
			match rec {
				Record::Write(WriteRecord {
					block,
					ibo,
					data,
				}) => {
					if block * Self::ERASE_BLOCK_SIZE + ibo == offset {
						if self.first_undiscovered.is_none() {
							self.first_undiscovered = Some(i);
						}

						buf[..Self::WRITE_GRANULARITY::USIZE].copy_from_slice(data.as_slice());
						return Ok(Self::WRITE_GRANULARITY::USIZE);
					}
				},
				Record::Erase(EraseRecord {
					block,
				}) => {
					if *block == offset / Self::ERASE_BLOCK_SIZE {
						if self.first_undiscovered.is_none() {
							self.first_undiscovered = Some(i + 1);
						}

						// Emulate NOR Flash
						buf[..Self::WRITE_GRANULARITY::USIZE].fill(0xff);
						return Ok(Self::WRITE_GRANULARITY::USIZE);
					}
				},
			}
		}

		self.before_storage
			.read(offset, &mut buf[..Self::WRITE_GRANULARITY::USIZE])
	}

	fn write(&mut self, offset: usize, buf: &[u8]) -> Result<usize, Self::Error> {
		if buf.is_empty() {
			return Ok(0);
		}

		assert!(offset % Self::WRITE_GRANULARITY::USIZE == 0);
		assert!(buf.len() >= Self::WRITE_GRANULARITY::USIZE);

		// Emulate NOR Flash, i.e. read the current state and AND the new data
		// into it
		let mut buffer = GenericArray::default();
		let len = self.read(offset, &mut buffer[..Self::WRITE_GRANULARITY::USIZE])?;
		assert_eq!(len, Self::WRITE_GRANULARITY::USIZE);
		// AND together the latest state and the new data
		for (b, input) in buffer.iter_mut().zip(buf) {
			*b &= input;
		}

		let rec = WriteRecord {
			block: offset / Self::ERASE_BLOCK_SIZE,
			ibo: offset % Self::ERASE_BLOCK_SIZE,
			data: buffer,
		};

		self.new_actions.push(Record::Write(rec));

		Ok(Self::WRITE_GRANULARITY::USIZE)
	}

	fn erase(&mut self, block: usize) -> Result<(), Self::Error> {
		self.new_actions.push(Record::Erase(EraseRecord {
			block,
		}));

		Ok(())
	}
}

use derivative::Derivative;

#[derive(Derivative)]
#[derivative(Debug(bound = ""), Clone(bound = ""))]
pub enum Record<N: ArrayLength> {
	Write(WriteRecord<N>),
	Erase(EraseRecord),
}
impl<N: ArrayLength> fmt::Display for Record<N> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		match self {
			Record::Write(rec) => {
				let block = rec.block;
				let range = rec.ibo..(rec.ibo + rec.data.len());
				let data = &rec.data;

				write!(f, "Write at {block} @ {range:?}, data: {data:02x?}")
			},
			Record::Erase(rec) => write!(f, "Erase at {}", rec.block),
		}
	}
}


#[derive(Derivative)]
#[derivative(Debug(bound = ""), Clone(bound = ""))]
pub struct WriteRecord<N: ArrayLength> {
	block: usize,
	ibo: usize,
	data: GenericArray<u8, N>,
}

#[derive(Debug, Clone)]
pub struct EraseRecord {
	block: usize,
}