apple_quant_algorithmic/binned_method/

flex.rs

use std::collections::BTreeMap;

use smallvec::SmallVec;

use crate::binned_method::{BinnedDataError, FlexSectorDef, SectorDataIdx};

/// Flexible collection of data binned by an index.
pub struct FlexBinnedData<T: Ord> {
	sectors: BTreeMap<u64, Vec<T>>,
	last_sector_def: Option<FlexSectorDef>,
}

impl<T: Ord> FlexBinnedData<T> {
	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)?;

		let Some(
			data,
		) = sector.pop() else {
			return Err(BinnedDataError::InternalError);
		};

		if sector.is_empty() {
			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(FlexSectorDef {
				sector_idx: last_sector_def.sector_idx - 1,
			});

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

		Ok(Some(data))
	}

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

		let appending_sector_def = self.last_sector_def.unwrap_or(FlexSectorDef {
			sector_idx,
		});

		let sector = match self.get_sector_mut(appending_sector_def.sector_idx) {
			Ok(
				sector,
			) => sector,
			Err(BinnedDataError::InvalidSector {
				sector_idx: _,
			}) => {
				self.sectors.insert(sector_idx, Vec::with_capacity(256));

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

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

		sector.extend(data);
		sector.sort();

		Ok(())
	}

	pub fn get_sector(
		&self,
		sector_idx: u64,
	) -> Result<&Vec<T>, 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 Vec<T>, BinnedDataError> {
		if let Some(
			sector,
		) = self.sectors.get_mut(&sector_idx) {
			return Ok(sector);
		}

		Err(BinnedDataError::InvalidSector { sector_idx })
	}

	pub fn iter(
		&self,
		start_sector_idx: u64,
		last_sector_idx: u64,
		f_start_data: impl Fn(&T) -> bool,
		f_last_data: impl Fn(&T) -> bool,
	) -> Result<impl Iterator<Item = &T>, ()> {
		let mut start_sector_ref_data_idx = None;

		for sector_idx in start_sector_idx..=last_sector_idx {
			let Ok(
				start_sector,
			) = self.get_sector(sector_idx) else {
				continue;
			};

			let Some(
				start_data_idx,
			) = start_sector.iter().position(&f_start_data) else {
				continue;
			};

			start_sector_ref_data_idx = Some(SectorDataIdx {
				sector: start_sector,
				sector_idx: sector_idx,
				data_idx: start_data_idx,
			});

			break;
		}

		let mut last_sector_ref_data_idx = None;

		for sector_idx in (start_sector_idx..=last_sector_idx).rev() {
			let Ok(
				last_sector,
			) = self.get_sector(sector_idx) else {
				continue;
			};

			let Some(
				last_data_idx,
			) = last_sector.iter().rposition(&f_last_data) else {
				continue;
			};

			last_sector_ref_data_idx = Some(SectorDataIdx {
				sector: last_sector,
				sector_idx: sector_idx,
				data_idx: last_data_idx,
			});

			break;
		}

		let mut sector_slices: SmallVec<[&[T]; 100]> = SmallVec::default();

		if (
			(start_sector_ref_data_idx.is_some() &&
			last_sector_ref_data_idx.is_none()) ||
			(start_sector_ref_data_idx.is_none() && last_sector_ref_data_idx.is_some())
		) {
			return Err(());
		}

		let Some(
			start_sector_ref_data_idx,
		) = start_sector_ref_data_idx else {
			return Ok(sector_slices.into_iter().flatten());
		};

		let Some(
			last_sector_ref_data_idx,
		) = last_sector_ref_data_idx else {
			return Ok(sector_slices.into_iter().flatten());
		};

		if last_sector_ref_data_idx.sector_idx < start_sector_ref_data_idx.sector_idx {
			return Err(());
		}

		if start_sector_ref_data_idx.sector_idx == last_sector_ref_data_idx.sector_idx {
			let sector = start_sector_ref_data_idx.sector;
			let start_data_idx = start_sector_ref_data_idx.data_idx;
			let last_data_idx = last_sector_ref_data_idx.data_idx;
			let sector_slice = &sector[start_data_idx..=last_data_idx];

			sector_slices.push(sector_slice);

			return Ok(sector_slices.into_iter().flatten());
		}

		let start_sector_slice = &start_sector_ref_data_idx.sector[start_sector_ref_data_idx.data_idx..];

		sector_slices.push(start_sector_slice);

		// In between start and last sectors.
		for sector_idx in (start_sector_ref_data_idx.sector_idx +
			1)..last_sector_ref_data_idx.sector_idx
		{
			let Ok(
				sector,
			) = self.get_sector(sector_idx) else {
				continue;
			};

			sector_slices.push(sector);
		}

		let last_sector_slice = &last_sector_ref_data_idx.sector[..last_sector_ref_data_idx.data_idx];

		sector_slices.push(last_sector_slice);
		Ok(sector_slices.into_iter().flatten())
	}
}

impl<T: Ord> Default for FlexBinnedData<T> {
	fn default() -> Self {
		Self {
			sectors: BTreeMap::default(),
			last_sector_def: None,
		}
	}
}