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)
|
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);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
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> =
|
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");
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user