add iter for reader

This commit is contained in:
Dmitry Belyaev 2022-10-10 00:08:57 +03:00
parent c28508b147
commit a5ca0c65a7
2 changed files with 117 additions and 10 deletions

111
src/db.rs
View File

@ -216,6 +216,93 @@ where
Ok(item.0) Ok(item.0)
} }
pub fn iter(&mut self) -> ReaderIter<'_, T> {
ReaderIter::new(self)
}
}
pub struct ReaderIter<'a, T>
where
T: bincode::Decode,
{
reader: &'a mut Reader<T>,
index: Option<usize>,
}
impl<'a, T> ReaderIter<'a, T>
where
T: bincode::Decode,
{
fn new(reader: &'a mut Reader<T>) -> Self {
ReaderIter {
reader,
index: None,
}
}
}
impl<'a, T> Iterator for ReaderIter<'a, T>
where
T: bincode::Decode,
{
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
if self.index.is_none() && self.reader.len() != 0 {
self.index = Some(0);
}
match self.index {
Some(i) if i < self.reader.len() => self.nth(i),
_ => None,
}
}
fn nth(&mut self, n: usize) -> Option<Self::Item> {
if self.reader.len() <= n {
return None;
}
self.index = Some(n + 1);
let item = self.reader.get(n);
match item {
Ok(item) => Some(item),
Err(_) => None,
}
}
fn size_hint(&self) -> (usize, Option<usize>) {
let len = self.reader.len();
if self.index.is_none() {
return (len, Some(len));
}
let index = self.index.unwrap();
let rem = if len > index + 1 {
len - (index + 1)
} else {
0
};
(rem, Some(rem))
}
fn count(self) -> usize
where
Self: Sized,
{
self.reader.len()
}
}
impl<'a, T> ExactSizeIterator for ReaderIter<'a, T>
where
T: bincode::Decode,
{
fn len(&self) -> usize {
self.reader.len()
}
} }
#[cfg(test)] #[cfg(test)]
@ -257,4 +344,28 @@ mod test {
let r = fs::remove_file(tempfile); let r = fs::remove_file(tempfile);
drop(r); drop(r);
} }
#[test]
fn test_write_read_iter() {
let mut rng = rand::thread_rng();
let tempfile = std::env::temp_dir().with_file_name(format!("test-{}.tmp", rng.next_u32()));
let opts = WriterOpts { compress_lvl: 1, data_buf_size: 10 * 1024 * 1024, out_buf_size: 10 * 1024 * 1024, current_buf_size: 4096 };
let mut writer: Writer<TestData> = Writer::new(&tempfile, opts).expect("new writer");
let items_iter = gen_data(10);
let items: Vec<TestData> = items_iter.collect();
writer.load(&mut items.clone().into_iter()).expect("load");
writer.finish().expect("finish write");
let mut reader: Reader<TestData> = Reader::new(&tempfile, 2048).expect("new reader");
assert_eq!(items.len(), reader.len());
items.into_iter().zip(reader.iter()).for_each(|pair| {
assert_eq!(pair.0, pair.1);
});
let r = fs::remove_file(tempfile);
drop(r);
}
} }

View File

@ -239,19 +239,15 @@ fn read_from_db2(id: u32) -> Option<Question> {
let mut reader: db::Reader<Question> = let mut reader: db::Reader<Question> =
db::Reader::new(NEW_DB_FILENAME, 2048).expect("new db reader"); db::Reader::new(NEW_DB_FILENAME, 2048).expect("new db reader");
let index = match id { let mut questions = reader.iter();
match id {
0 => { 0 => {
let mut rng = rand::thread_rng(); let mut rng = rand::thread_rng();
(1..reader.len()).into_iter().choose(&mut rng).unwrap() questions.choose(&mut rng)
} }
_ => (id - 1) as usize, _ => questions.nth((id - 1) as usize),
};
if index >= reader.len() {
return None;
} }
Some(reader.get(index).expect("get question"))
} }
fn write_db2() { fn write_db2() {
let (tx, rx) = mpsc::channel::<Question>(); let (tx, rx) = mpsc::channel::<Question>();
@ -277,7 +273,7 @@ fn db_writer2_task(rx: mpsc::Receiver<Question>) {
writer writer
.load(&mut iter) .load(&mut iter)
.unwrap_or_else(|e| panic!("db writer push, num={num}, {e:#?}")); .unwrap_or_else(|e| panic!("db writer load, {e:#?}"));
writer.finish().expect("db writer finish"); writer.finish().expect("db writer finish");