use core::{
cmp::Reverse,
convert::Infallible,
mem,
ops::{BitOrAssign, BitXorAssign},
};
use alloc::borrow::Cow;
use crate::{MultiOps, RoaringBitmap};
use super::{container::Container, store::Store};
#[cfg(not(feature = "std"))]
use alloc::vec::Vec;
const BASE_COLLECT: usize = 10;
const MAX_COLLECT: usize = 50;
impl<I> MultiOps<RoaringBitmap> for I
where
I: IntoIterator<Item = RoaringBitmap>,
{
type Output = RoaringBitmap;
fn union(self) -> Self::Output {
try_multi_or_owned(self.into_iter().map(Ok::<_, Infallible>)).unwrap()
}
fn intersection(self) -> Self::Output {
try_multi_and_owned(self.into_iter().map(Ok::<_, Infallible>)).unwrap()
}
fn difference(self) -> Self::Output {
try_multi_sub_owned(self.into_iter().map(Ok::<_, Infallible>)).unwrap()
}
fn symmetric_difference(self) -> Self::Output {
try_multi_xor_owned(self.into_iter().map(Ok::<_, Infallible>)).unwrap()
}
}
impl<I, E> MultiOps<Result<RoaringBitmap, E>> for I
where
I: IntoIterator<Item = Result<RoaringBitmap, E>>,
{
type Output = Result<RoaringBitmap, E>;
fn union(self) -> Self::Output {
try_multi_or_owned(self)
}
fn intersection(self) -> Self::Output {
try_multi_and_owned(self)
}
fn difference(self) -> Self::Output {
try_multi_sub_owned(self)
}
fn symmetric_difference(self) -> Self::Output {
try_multi_xor_owned(self)
}
}
impl<'a, I> MultiOps<&'a RoaringBitmap> for I
where
I: IntoIterator<Item = &'a RoaringBitmap>,
{
type Output = RoaringBitmap;
fn union(self) -> Self::Output {
try_multi_or_ref(self.into_iter().map(Ok::<_, Infallible>)).unwrap()
}
fn intersection(self) -> Self::Output {
try_multi_and_ref(self.into_iter().map(Ok::<_, Infallible>)).unwrap()
}
fn difference(self) -> Self::Output {
try_multi_sub_ref(self.into_iter().map(Ok::<_, Infallible>)).unwrap()
}
fn symmetric_difference(self) -> Self::Output {
try_multi_xor_ref(self.into_iter().map(Ok::<_, Infallible>)).unwrap()
}
}
impl<'a, I, E: 'a> MultiOps<Result<&'a RoaringBitmap, E>> for I
where
I: IntoIterator<Item = Result<&'a RoaringBitmap, E>>,
{
type Output = Result<RoaringBitmap, E>;
fn union(self) -> Self::Output {
try_multi_or_ref(self)
}
fn intersection(self) -> Self::Output {
try_multi_and_ref(self)
}
fn difference(self) -> Self::Output {
try_multi_sub_ref(self)
}
fn symmetric_difference(self) -> Self::Output {
try_multi_xor_ref(self)
}
}
#[inline]
fn try_multi_and_owned<E>(
bitmaps: impl IntoIterator<Item = Result<RoaringBitmap, E>>,
) -> Result<RoaringBitmap, E> {
let mut iter = bitmaps.into_iter();
let mut start = collect_starting_elements(iter.by_ref())?;
start.sort_unstable_by_key(|bitmap| bitmap.containers.len());
let mut start = start.into_iter();
if let Some(mut lhs) = start.next() {
for rhs in start.map(Ok).chain(iter) {
if lhs.is_empty() {
return Ok(lhs);
}
lhs &= rhs?;
}
Ok(lhs)
} else {
Ok(RoaringBitmap::new())
}
}
#[inline]
fn try_multi_and_ref<'a, E>(
bitmaps: impl IntoIterator<Item = Result<&'a RoaringBitmap, E>>,
) -> Result<RoaringBitmap, E> {
let mut iter = bitmaps.into_iter();
let mut start = collect_starting_elements(iter.by_ref())?;
start.sort_unstable_by_key(|bitmap| bitmap.containers.len());
let mut start = start.into_iter();
if let Some(mut lhs) = start.next().cloned() {
for rhs in start.map(Ok).chain(iter) {
if lhs.is_empty() {
return Ok(lhs);
}
lhs &= rhs?;
}
Ok(lhs)
} else {
Ok(RoaringBitmap::new())
}
}
#[inline]
fn try_multi_sub_owned<E>(
bitmaps: impl IntoIterator<Item = Result<RoaringBitmap, E>>,
) -> Result<RoaringBitmap, E> {
let mut iter = bitmaps.into_iter();
match iter.next().transpose()? {
Some(mut lhs) => {
for rhs in iter {
if lhs.is_empty() {
return Ok(lhs);
}
lhs -= rhs?;
}
Ok(lhs)
}
None => Ok(RoaringBitmap::default()),
}
}
#[inline]
fn try_multi_sub_ref<'a, E>(
bitmaps: impl IntoIterator<Item = Result<&'a RoaringBitmap, E>>,
) -> Result<RoaringBitmap, E> {
let mut iter = bitmaps.into_iter();
match iter.next().transpose()?.cloned() {
Some(mut lhs) => {
for rhs in iter {
if lhs.is_empty() {
return Ok(lhs);
}
lhs -= rhs?;
}
Ok(lhs)
}
None => Ok(RoaringBitmap::default()),
}
}
#[inline]
fn try_multi_or_owned<E>(
bitmaps: impl IntoIterator<Item = Result<RoaringBitmap, E>>,
) -> Result<RoaringBitmap, E> {
let mut iter = bitmaps.into_iter();
let mut start = collect_starting_elements(iter.by_ref())?;
start.sort_unstable_by_key(|bitmap| Reverse(bitmap.containers.len()));
let start_size = start.len();
let mut start = start.into_iter();
let mut containers = if let Some(c) = start.next() {
if c.is_empty() {
start.by_ref().nth(start_size);
}
c.containers
} else {
return Ok(RoaringBitmap::new());
};
for bitmap in start.map(Ok).chain(iter) {
merge_container_owned(&mut containers, bitmap?.containers, BitOrAssign::bitor_assign);
}
containers.retain_mut(|container| {
if !container.is_empty() {
container.ensure_correct_store();
true
} else {
false
}
});
Ok(RoaringBitmap { containers })
}
#[inline]
fn try_multi_xor_owned<E>(
bitmaps: impl IntoIterator<Item = Result<RoaringBitmap, E>>,
) -> Result<RoaringBitmap, E> {
let mut iter = bitmaps.into_iter();
let mut containers = match iter.next().transpose()? {
None => Vec::new(),
Some(v) => v.containers,
};
for bitmap in iter {
merge_container_owned(&mut containers, bitmap?.containers, BitXorAssign::bitxor_assign);
}
containers.retain_mut(|container| {
if !container.is_empty() {
container.ensure_correct_store();
true
} else {
false
}
});
Ok(RoaringBitmap { containers })
}
fn merge_container_owned(
lhs: &mut Vec<Container>,
rhs: Vec<Container>,
op: impl Fn(&mut Store, Store),
) {
for mut rhs in rhs {
match lhs.binary_search_by_key(&rhs.key, |c| c.key) {
Err(loc) => lhs.insert(loc, rhs),
Ok(loc) => {
let lhs = &mut lhs[loc];
match (&lhs.store, &rhs.store) {
(Store::Array(..), Store::Array(..)) => lhs.store = lhs.store.to_bitmap(),
(Store::Array(..), Store::Bitmap(..)) => mem::swap(lhs, &mut rhs),
_ => (),
};
op(&mut lhs.store, rhs.store);
}
}
}
}
#[inline]
fn try_multi_or_ref<'a, E: 'a>(
bitmaps: impl IntoIterator<Item = Result<&'a RoaringBitmap, E>>,
) -> Result<RoaringBitmap, E> {
let mut iter = bitmaps.into_iter();
let mut start = collect_starting_elements(iter.by_ref())?;
let start_size = start.len();
start.sort_unstable_by_key(|bitmap| Reverse(bitmap.containers.len()));
let mut start = start.into_iter();
let mut containers = match start.next() {
Some(c) => {
let c: Vec<Cow<Container>> = c.containers.iter().map(Cow::Borrowed).collect();
if c.is_empty() {
start.by_ref().nth(start_size);
}
c
}
None => {
return Ok(RoaringBitmap::new());
}
};
for bitmap in start.map(Ok).chain(iter) {
merge_container_ref(&mut containers, &bitmap?.containers, |a, b| *a |= b);
}
let containers: Vec<_> = containers
.into_iter()
.filter(|container| container.len() > 0)
.map(|c| {
let mut container = c.into_owned();
container.ensure_correct_store();
container
})
.collect();
Ok(RoaringBitmap { containers })
}
#[inline]
fn try_multi_xor_ref<'a, E: 'a>(
bitmaps: impl IntoIterator<Item = Result<&'a RoaringBitmap, E>>,
) -> Result<RoaringBitmap, E> {
let mut iter = bitmaps.into_iter();
let mut containers: Vec<Cow<Container>> = match iter.next().transpose()? {
None => Vec::new(),
Some(v) => v.containers.iter().map(Cow::Borrowed).collect(),
};
for bitmap in iter {
merge_container_ref(&mut containers, &bitmap?.containers, |a, b| *a ^= b);
}
let containers: Vec<_> = containers
.into_iter()
.filter(|container| container.len() > 0)
.map(|c| {
let mut container = c.into_owned();
container.ensure_correct_store();
container
})
.collect();
Ok(RoaringBitmap { containers })
}
fn merge_container_ref<'a>(
containers: &mut Vec<Cow<'a, Container>>,
rhs: &'a [Container],
op: impl Fn(&mut Store, &Store),
) {
for rhs in rhs {
match containers.binary_search_by_key(&rhs.key, |c| c.key) {
Err(loc) => {
containers.insert(loc, Cow::Borrowed(rhs))
}
Ok(loc) => {
let lhs = &mut containers[loc];
match (&lhs.store, &rhs.store) {
(Store::Array(..), Store::Array(..)) => {
let mut store = lhs.store.to_bitmap();
op(&mut store, &rhs.store);
*lhs = Cow::Owned(Container { key: lhs.key, store });
}
(Store::Array(..), Store::Bitmap(..)) => {
let mut store = rhs.store.clone();
op(&mut store, &lhs.store);
*lhs = Cow::Owned(Container { key: lhs.key, store });
}
(Store::Bitmap(..), _) => {
op(&mut lhs.to_mut().store, &rhs.store);
}
};
}
}
}
}
#[inline]
fn collect_starting_elements<I, El, Er>(iter: I) -> Result<Vec<El>, Er>
where
I: IntoIterator<Item = Result<El, Er>>,
{
let iter = iter.into_iter();
let mut to_collect = iter.size_hint().1.unwrap_or(BASE_COLLECT);
if to_collect > MAX_COLLECT {
to_collect = BASE_COLLECT;
}
let mut ret = Vec::with_capacity(to_collect);
for el in iter.take(to_collect) {
ret.push(el?);
}
Ok(ret)
}