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,
		}
	}
}