use crate::common::codec_reader::CodecReader;
use crate::fixed::FixedVec;
use crate::variable::traits::Storable;
use dsi_bitstream::{
codes::params::DefaultReadParams,
dispatch::{Codes, CodesRead, StaticCodeRead},
impls::{BufBitReader, MemWordReader},
prelude::{BitRead, BitSeek, Endianness},
};
use std::marker::PhantomData;
pub(crate) type SeqVecBitReader<'a, E> =
BufBitReader<E, MemWordReader<u64, &'a [u64], true>, DefaultReadParams>;
pub struct SeqIter<'a, T: Storable, E: Endianness>
where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
reader: SeqVecBitReader<'a, E>,
code_reader: CodecReader<'a, E>,
end_bit: u64,
known_len: Option<usize>,
_marker: PhantomData<T>,
}
impl<'a, T: Storable, E: Endianness> SeqIter<'a, T, E>
where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
#[inline]
pub(crate) fn new(data: &'a [u64], start_bit: u64, end_bit: u64, encoding: Codes) -> Self {
let mut reader = SeqVecBitReader::<E>::new(MemWordReader::new_inf(data));
let _ = reader.set_bit_pos(start_bit);
let code_reader = CodecReader::new(encoding);
Self {
reader,
code_reader,
end_bit,
known_len: None,
_marker: PhantomData,
}
}
#[inline]
pub(crate) fn new_with_len(
data: &'a [u64],
start_bit: u64,
end_bit: u64,
encoding: Codes,
len: Option<usize>,
) -> Self {
let mut iter = Self::new(data, start_bit, end_bit, encoding);
iter.known_len = len;
iter
}
#[inline]
pub fn end_bit(&self) -> u64 {
self.end_bit
}
}
impl<'a, T: Storable, E: Endianness> Iterator for SeqIter<'a, T, E>
where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
type Item = T;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
if let Some(ref mut remaining) = self.known_len {
if *remaining == 0 {
return None;
}
*remaining -= 1;
let word = self.code_reader.read(&mut self.reader).unwrap();
return Some(T::from_word(word));
}
if self.reader.bit_pos().unwrap() >= self.end_bit {
return None;
}
let word = self.code_reader.read(&mut self.reader).unwrap();
Some(T::from_word(word))
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
if let Some(len) = self.known_len {
return (len, Some(len));
}
(0, None)
}
}
impl<'a, T: Storable, E: Endianness> std::iter::FusedIterator for SeqIter<'a, T, E> where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>
{
}
impl<'a, T: Storable, E: Endianness> ExactSizeIterator for SeqIter<'a, T, E>
where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
#[inline]
fn len(&self) -> usize {
self.known_len.expect(
"SeqIter::len() called on iterator with unknown length. \
This can only happen if the iterator was not created with a pre-stored \
length. Use count() instead.",
)
}
}
pub struct SeqVecIter<'a, T: Storable, E: Endianness, B: AsRef<[u64]>>
where
SeqVecBitReader<'a, E>: CodesRead<E>,
{
data: &'a [u64],
bit_offsets: &'a FixedVec<u64, u64, E, B>,
seq_lengths: Option<&'a FixedVec<u64, u64, E, Vec<u64>>>,
encoding: Codes,
front: usize,
back: usize,
_marker: PhantomData<(T, E)>,
}
impl<'a, T: Storable, E: Endianness, B: AsRef<[u64]>> SeqVecIter<'a, T, E, B>
where
SeqVecBitReader<'a, E>: CodesRead<E>,
{
#[inline]
pub(crate) fn new(
data: &'a [u64],
bit_offsets: &'a FixedVec<u64, u64, E, B>,
seq_lengths: Option<&'a FixedVec<u64, u64, E, Vec<u64>>>,
encoding: Codes,
num_sequences: usize,
) -> Self {
Self {
data,
bit_offsets,
seq_lengths,
encoding,
front: 0,
back: num_sequences,
_marker: PhantomData,
}
}
#[inline]
pub fn remaining(&self) -> usize {
self.back.saturating_sub(self.front)
}
}
impl<'a, T: Storable, E: Endianness, B: AsRef<[u64]>> Iterator for SeqVecIter<'a, T, E, B>
where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
type Item = SeqIter<'a, T, E>;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
if self.front >= self.back {
return None;
}
let start_bit = unsafe { self.bit_offsets.get_unchecked(self.front) };
let end_bit = unsafe { self.bit_offsets.get_unchecked(self.front + 1) };
let len = self
.seq_lengths
.map(|lengths| unsafe { lengths.get_unchecked(self.front) as usize });
self.front += 1;
Some(SeqIter::new_with_len(
self.data,
start_bit,
end_bit,
self.encoding,
len,
))
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let remaining = self.remaining();
(remaining, Some(remaining))
}
#[inline]
fn count(self) -> usize {
self.remaining()
}
#[inline]
fn nth(&mut self, n: usize) -> Option<Self::Item> {
if n >= self.remaining() {
self.front = self.back;
return None;
}
self.front += n;
self.next()
}
}
impl<'a, T: Storable, E: Endianness, B: AsRef<[u64]>> DoubleEndedIterator
for SeqVecIter<'a, T, E, B>
where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
#[inline]
fn next_back(&mut self) -> Option<Self::Item> {
if self.front >= self.back {
return None;
}
self.back -= 1;
let start_bit = unsafe { self.bit_offsets.get_unchecked(self.back) };
let end_bit = unsafe { self.bit_offsets.get_unchecked(self.back + 1) };
let len = self
.seq_lengths
.map(|lengths| unsafe { lengths.get_unchecked(self.back) as usize });
Some(SeqIter::new_with_len(
self.data,
start_bit,
end_bit,
self.encoding,
len,
))
}
#[inline]
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
if n >= self.remaining() {
self.back = self.front;
return None;
}
self.back -= n;
self.next_back()
}
}
impl<'a, T: Storable, E: Endianness, B: AsRef<[u64]>> ExactSizeIterator for SeqVecIter<'a, T, E, B>
where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
#[inline]
fn len(&self) -> usize {
self.remaining()
}
}
impl<'a, T: Storable, E: Endianness, B: AsRef<[u64]>> std::iter::FusedIterator
for SeqVecIter<'a, T, E, B>
where
for<'b> SeqVecBitReader<'b, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
}
pub struct SeqVecIntoIter<T, E>
where
T: Storable + 'static,
E: Endianness + 'static,
for<'a> SeqVecBitReader<'a, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
data: &'static [u64],
encoding: Codes,
bit_offsets_data: &'static [u64],
bit_offsets_num_bits: u64,
bit_offsets_len: usize,
seq_lengths: Option<FixedVec<u64, u64, E, Vec<u64>>>,
_data_owner: Vec<u64>,
_bit_offsets_owner: Vec<u64>,
num_sequences: usize,
current_index: usize,
last_end_bit: u64,
_markers: PhantomData<(T, E)>,
}
impl<T, E> SeqVecIntoIter<T, E>
where
T: Storable + 'static,
E: Endianness + 'static,
for<'a> SeqVecBitReader<'a, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
pub(crate) fn new(vec: super::SeqVec<T, E, Vec<u64>>) -> Self {
let encoding = vec.encoding;
let num_sequences = vec.num_sequences();
let seq_lengths = vec.seq_lengths;
let _data_owner = vec.data;
let _bit_offsets_owner = vec.bit_offsets.as_limbs().to_vec();
let bit_offsets_num_bits = vec.bit_offsets.bit_width() as u64;
let bit_offsets_len = vec.bit_offsets.len();
let data_ref: &'static [u64] = unsafe { std::mem::transmute(_data_owner.as_slice()) };
let bit_offsets_ref: &'static [u64] =
unsafe { std::mem::transmute(_bit_offsets_owner.as_slice()) };
Self {
data: data_ref,
encoding,
bit_offsets_data: bit_offsets_ref,
bit_offsets_num_bits,
bit_offsets_len,
seq_lengths,
_data_owner,
_bit_offsets_owner,
num_sequences,
current_index: 0,
last_end_bit: 0,
_markers: PhantomData,
}
}
#[inline]
fn get_offset(&self, index: usize) -> u64 {
if index >= self.bit_offsets_len {
return 0; }
let offset_bits = self.bit_offsets_num_bits;
let start_bit = (index as u64) * offset_bits;
let mut reader = SeqVecBitReader::<E>::new(MemWordReader::new_inf(self.bit_offsets_data));
let _ = reader.set_bit_pos(start_bit);
reader.read_bits(offset_bits as usize).unwrap_or(0)
}
}
impl<T, E> Iterator for SeqVecIntoIter<T, E>
where
T: Storable + 'static,
E: Endianness + 'static,
for<'a> SeqVecBitReader<'a, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
type Item = SeqIter<'static, T, E>;
#[inline]
fn next(&mut self) -> Option<Self::Item> {
if self.current_index >= self.num_sequences {
return None;
}
let start_bit = if self.current_index == 0 {
self.get_offset(0)
} else {
self.last_end_bit
};
let end_bit = self.get_offset(self.current_index + 1);
self.last_end_bit = end_bit;
let len = self
.seq_lengths
.as_ref()
.map(|lengths| unsafe { lengths.get_unchecked(self.current_index) as usize });
self.current_index += 1;
Some(SeqIter::new_with_len(
self.data,
start_bit,
end_bit,
self.encoding,
len,
))
}
#[inline]
fn size_hint(&self) -> (usize, Option<usize>) {
let remaining = self.num_sequences.saturating_sub(self.current_index);
(remaining, Some(remaining))
}
}
impl<T, E> ExactSizeIterator for SeqVecIntoIter<T, E>
where
T: Storable + 'static,
E: Endianness + 'static,
for<'a> SeqVecBitReader<'a, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
#[inline]
fn len(&self) -> usize {
self.num_sequences.saturating_sub(self.current_index)
}
}
impl<T, E> DoubleEndedIterator for SeqVecIntoIter<T, E>
where
T: Storable + 'static,
E: Endianness + 'static,
for<'a> SeqVecBitReader<'a, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
#[inline]
fn next_back(&mut self) -> Option<Self::Item> {
if self.current_index >= self.num_sequences {
return None;
}
self.num_sequences -= 1;
let start_bit = self.get_offset(self.num_sequences);
let end_bit = self.get_offset(self.num_sequences + 1);
let len = self
.seq_lengths
.as_ref()
.map(|lengths| unsafe { lengths.get_unchecked(self.num_sequences) as usize });
Some(SeqIter::new_with_len(
self.data,
start_bit,
end_bit,
self.encoding,
len,
))
}
}
impl<T, E> std::iter::FusedIterator for SeqVecIntoIter<T, E>
where
T: Storable + 'static,
E: Endianness + 'static,
for<'a> SeqVecBitReader<'a, E>: BitRead<E, Error = core::convert::Infallible>
+ CodesRead<E>
+ BitSeek<Error = core::convert::Infallible>,
{
}