add iter for reader
This commit is contained in:
parent
c28508b147
commit
a5ca0c65a7
111
src/db.rs
111
src/db.rs
@ -216,6 +216,93 @@ where
|
||||
|
||||
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)]
|
||||
@ -257,4 +344,28 @@ mod test {
|
||||
let r = fs::remove_file(tempfile);
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
16
src/main.rs
16
src/main.rs
@ -239,19 +239,15 @@ fn read_from_db2(id: u32) -> Option<Question> {
|
||||
let mut reader: db::Reader<Question> =
|
||||
db::Reader::new(NEW_DB_FILENAME, 2048).expect("new db reader");
|
||||
|
||||
let index = match id {
|
||||
let mut questions = reader.iter();
|
||||
|
||||
match id {
|
||||
0 => {
|
||||
let mut rng = rand::thread_rng();
|
||||
(1..reader.len()).into_iter().choose(&mut rng).unwrap()
|
||||
questions.choose(&mut rng)
|
||||
}
|
||||
_ => (id - 1) as usize,
|
||||
};
|
||||
|
||||
if index >= reader.len() {
|
||||
return None;
|
||||
_ => questions.nth((id - 1) as usize),
|
||||
}
|
||||
|
||||
Some(reader.get(index).expect("get question"))
|
||||
}
|
||||
fn write_db2() {
|
||||
let (tx, rx) = mpsc::channel::<Question>();
|
||||
@ -277,7 +273,7 @@ fn db_writer2_task(rx: mpsc::Receiver<Question>) {
|
||||
|
||||
writer
|
||||
.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");
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user