From a5ca0c65a79f807aa71e3a4960077528c4060bf3 Mon Sep 17 00:00:00 2001
From: Dmitry <b4tm4n@mail.ru>
Date: Mon, 10 Oct 2022 00:08:57 +0300
Subject: [PATCH] add iter for reader

---
 src/db.rs   | 111 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/main.rs |  16 +++-----
 2 files changed, 117 insertions(+), 10 deletions(-)

diff --git a/src/db.rs b/src/db.rs
index bc5b0b2..f87ce8c 100644
--- a/src/db.rs
+++ b/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);
+    }
 }
diff --git a/src/main.rs b/src/main.rs
index 1bab6d5..4608f02 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -238,20 +238,16 @@ fn main() {
 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 mut questions = reader.iter();
 
-    let index = match id {
+    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");