apple_quant_algorithmic/binned_method/

fixed.rs

use std::{array, collections::BTreeMap};

use smallvec::SmallVec;

use crate::binned_method::{BinnedDataError, FixedSectorDef};

/// Fixed collection of data binned by an index.
pub struct FixedBinnedData<
	T: Clone,
	const SECTOR_DATA_COUNT: usize,
> {
	sectors: BTreeMap<u64, [Option<T>; SECTOR_DATA_COUNT]>,
	last_sector_def: Option<FixedSectorDef>,
}

impl<
	T: Clone,
	const SECTOR_DATA_COUNT: usize,
> FixedBinnedData<T, SECTOR_DATA_COUNT> {
	pub fn last_sector_def(
		&self,
	) -> &Option<FixedSectorDef> {
		&self.last_sector_def
	}

	pub fn remove_last(
		&mut self,
	) -> Result<Option<T>, BinnedDataError> {
		let Some(
			last_sector_def,
		) = self.last_sector_def else {
			return Ok(None);
		};

		let sector = self.get_sector_mut(last_sector_def.sector_idx)?;

		#[cfg(debug_assertions)]
		if last_sector_def.data_idx >= SECTOR_DATA_COUNT {
			return Err(BinnedDataError::SectorDataOutOfRange {
				data_index: last_sector_def.data_idx,
				sector_data_count: SECTOR_DATA_COUNT,
			});
		}

		let Some(
			data,
		) = sector[last_sector_def.data_idx].take() else {
			return Err(BinnedDataError::InternalError);
		};

		if last_sector_def.data_idx == 0 {
			if last_sector_def.sector_idx == 0 {
				self.last_sector_def = None;
				self.sectors.remove(&last_sector_def.sector_idx);

				return Ok(Some(data));
			}

			if self.get_sector(last_sector_def.sector_idx - 1).is_err() {
				self.last_sector_def = None;
				self.sectors.remove(&last_sector_def.sector_idx);

				return Ok(Some(data));
			}

			self.last_sector_def = Some(FixedSectorDef {
				sector_idx: last_sector_def.sector_idx - 1,
				data_idx: SECTOR_DATA_COUNT - 1,
			});

			self.sectors.remove(&last_sector_def.sector_idx);
			return Ok(Some(data));
		}

		self.last_sector_def = Some(FixedSectorDef {
			sector_idx: last_sector_def.sector_idx,
			data_idx: last_sector_def.data_idx - 1,
		});

		Ok(Some(data))
	}

	pub fn append(
		&mut self,
		first_sector_idx: u64,
		data: impl IntoIterator<Item = T>,
	) -> Result<(), ()> {
		let data = data.into_iter();

		let mut appending_sector_def = self.last_sector_def.unwrap_or(FixedSectorDef {
			sector_idx: first_sector_idx / SECTOR_DATA_COUNT as u64,
			data_idx: (first_sector_idx -
				((first_sector_idx / SECTOR_DATA_COUNT as u64) *
				SECTOR_DATA_COUNT as u64)) as usize,
		});

		let mut sector = match self.get_sector_mut(appending_sector_def.sector_idx) {
			Ok(
				sector,
			) => sector,
			Err(BinnedDataError::InvalidSector {
				sector_idx: _,
			}) => {
				self.sectors.insert(
					appending_sector_def.sector_idx,
					array::from_fn(|_| None),
				);

				let Ok(
					sector,
				) = self.get_sector_mut(
					appending_sector_def.sector_idx,
				) else {
					return Err(());
				};

				sector
			},
			Err(_) => unimplemented!(),
		};

		for next_data in data {
			appending_sector_def.data_idx += 1;

			if appending_sector_def.data_idx == SECTOR_DATA_COUNT {
				appending_sector_def.data_idx = 0;
				appending_sector_def.sector_idx += 1;

				self.sectors.insert(
					appending_sector_def.sector_idx,
					array::from_fn(|_| None),
				);

				let Ok(
					new_sector,
				) = self.get_sector_mut(
					appending_sector_def.sector_idx,
				) else {
					return Err(());
				};

				sector = new_sector;
			}

			sector[appending_sector_def.data_idx] = Some(next_data);
		}

		Ok(())
	}

	pub fn get(
		&mut self,
		index: usize,
	) -> Result<&Option<T>, BinnedDataError> {
		let sector_id = index / SECTOR_DATA_COUNT;
		let data_idx = index - (sector_id * SECTOR_DATA_COUNT);
		let sector = self.get_sector_mut(sector_id as u64)?;

		Ok(&sector[data_idx])
	}

	pub fn iter(
		&mut self,
		min_idx: usize,
		max_idx: usize,
	) -> Result<impl Iterator<Item = &Option<T>>, BinnedDataError> {
		if max_idx < min_idx {
			return Err(BinnedDataError::InvalidMinMax {
				min: min_idx as u64,
				max: max_idx as u64,
			});
		}

		let Some(
			last_sector_def,
		) = &self.last_sector_def else {
			return Err(BinnedDataError::InternalError);
		};

		let iter_len = max_idx - min_idx + 1;
		let mut data = SmallVec::<[&Option<T>; 1024]>::with_capacity(iter_len);
		let mut last_idx = min_idx;

		loop {
			let sector_id = last_idx / SECTOR_DATA_COUNT;
			let data_min_idx = last_idx - (sector_id * SECTOR_DATA_COUNT);
			let mut data_max_idx = (sector_id + 1) * SECTOR_DATA_COUNT - 1;
			let sector_id = sector_id as u64;

			let Ok(
				sector,
			) = self.get_sector(sector_id) else {
				break;
			};

			if last_sector_def.sector_idx == sector_id as u64 {
				data_max_idx = last_sector_def.data_idx;
			}

			data.extend(sector.iter().take(data_max_idx + 1).skip(data_min_idx));
			last_idx = data_max_idx + 1;
		}

		Ok(data.into_iter())
	}

	fn get_sector(
		&self,
		sector_idx: u64,
	) -> Result<&[Option<T>; SECTOR_DATA_COUNT], BinnedDataError> {
		if let Some(
			sector,
		) = self.sectors.get(&sector_idx) {
			return Ok(sector);
		}

		Err(BinnedDataError::InvalidSector { sector_idx })
	}

	fn get_sector_mut(
		&mut self,
		sector_idx: u64,
	) -> Result<&mut [Option<T>; SECTOR_DATA_COUNT], BinnedDataError> {
		if let Some(
			sector,
		) = self.sectors.get_mut(&sector_idx) {
			return Ok(sector);
		}

		Err(BinnedDataError::InvalidSector { sector_idx })
	}
}

impl<
	T: Clone,
	const SECTOR_DATA_COUNT: usize,
> Default for FixedBinnedData<T, SECTOR_DATA_COUNT> {
	fn default() -> Self {
		Self {
			sectors: BTreeMap::default(),
			last_sector_def: None,
		}
	}
}