use std::{
collections::BTreeMap,
ops::{BitAnd, BitOr, Bound::*, RangeBounds},
};
use iter_set::{intersection, union};
pub struct QueriableDataStore<T> {
items: Vec<T>,
}
impl<T> QueriableDataStore<T> {
pub fn filter(&self, filter: DataFilter) -> impl Iterator<Item = &T> {
filter.indices.into_iter().map(move |v| &self.items[v])
}
pub fn get_index<F, U>(&self, index_provider: F) -> SortedIndex<U>
where
F: Fn(&T) -> U,
U: Ord,
{
SortedIndex::new(self, index_provider)
}
pub fn items(&self) -> impl Iterator<Item = &T> {
self.items.iter()
}
}
impl<T> From<Vec<T>> for QueriableDataStore<T> {
fn from(items: Vec<T>) -> Self {
Self { items }
}
}
#[derive(Clone, Eq, PartialEq)]
pub struct SortedIndex<T> {
pairs: BTreeMap<T, Vec<usize>>,
}
impl<T> SortedIndex<T>
where
T: Ord,
{
pub fn new<F, U>(data_store: &QueriableDataStore<U>, index_provider: F) -> Self
where
F: Fn(&U) -> T,
{
let mut pairs: BTreeMap<T, Vec<usize>> = BTreeMap::new();
for (index, item) in data_store.items().enumerate() {
let key = index_provider(item);
pairs.entry(key).or_insert_with(|| vec![]).push(index);
}
Self { pairs }
}
pub fn filter_range<R>(&self, range: R) -> DataFilter
where
R: RangeBounds<T>,
{
let filtered = self
.pairs
.range(range)
.into_iter()
.flat_map(|(_, indices)| indices.iter().cloned());
DataFilter::from_unsorted(filtered)
}
pub fn filter_between(&self, lower_inclusive: T, upper_inclusive: T) -> DataFilter {
self.filter_range((Included(lower_inclusive), Included(upper_inclusive)))
}
pub fn filter_eq(&self, value: T) -> DataFilter {
if let Some(keys) = self.pairs.get(&value) {
DataFilter::from_unsorted(keys.iter().cloned())
} else {
DataFilter::default()
}
}
pub fn filter_gt(&self, lower_limit: T) -> DataFilter {
self.filter_range((Excluded(lower_limit), Unbounded))
}
pub fn filter_gte(&self, lower_limit: T) -> DataFilter {
self.filter_range((Included(lower_limit), Unbounded))
}
pub fn filter_lt(&self, upper_limit: T) -> DataFilter {
self.filter_range((Unbounded, Excluded(upper_limit)))
}
pub fn filter_lte(&self, upper_limit: T) -> DataFilter {
self.filter_range((Unbounded, Included(upper_limit)))
}
pub fn first(&self) -> DataFilter {
DataFilter::from_unsorted(
self.pairs
.iter()
.flat_map(|v| v.1.iter())
.cloned()
.nth(0)
.into_iter(),
)
}
pub fn first_n(&self, n: usize) -> DataFilter {
DataFilter::from_unsorted(self.pairs.iter().flat_map(|v| v.1.iter()).cloned().take(n))
}
pub fn last(&self) -> DataFilter {
DataFilter::from_unsorted(
self.pairs
.iter()
.rev()
.flat_map(|v| v.1.iter())
.cloned()
.nth(0)
.into_iter(),
)
}
pub fn last_n(&self, n: usize) -> DataFilter {
DataFilter::from_unsorted(
self.pairs
.iter()
.rev()
.flat_map(|v| v.1.iter())
.cloned()
.take(n),
)
}
}
#[derive(Default)]
pub struct DataFilter {
indices: Vec<usize>,
}
impl DataFilter {
fn from_unsorted<T>(unsorted_indices: T) -> Self
where
T: Iterator<Item = usize>,
{
let mut indices: Vec<usize> = unsorted_indices.collect();
indices.sort();
Self { indices }
}
}
impl BitAnd for DataFilter {
type Output = DataFilter;
fn bitand(self, other: DataFilter) -> Self::Output {
Self {
indices: intersection(self.indices, other.indices).collect(),
}
}
}
impl BitOr for DataFilter {
type Output = DataFilter;
fn bitor(self, other: DataFilter) -> Self::Output {
Self {
indices: union(self.indices, other.indices).collect(),
}
}
}