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