Making iterator tests parameterized so that they can test various Ite…

…rator that scans the entire DB (MystenLabs#15285)

## Description 

Making the iterator test parameterized so that tests can test both
`Rocks::Iter` and `Rocks::SafeIter`.

## Test Plan 

How did you test the new or updated feature?

---
If your changes are not user-facing and not a breaking change, you can
skip the following section. Otherwise, please indicate what changed, and
then add to the Release Notes section as highlighted during the release
process.

### Type of Change (Check all that apply)

- [ ] protocol change
- [ ] user-visible impact
- [ ] breaking change for a client SDKs
- [ ] breaking change for FNs (FN binary must upgrade)
- [ ] breaking change for validators or node operators (must upgrade
binaries)
- [ ] breaking change for on-chain data layout
- [ ] necessitate either a data wipe or data migration

### Release notes

crates/typed-store/src/rocks/tests.rs

@@ -1,6 +1,8 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 use super::*;
+use crate::rocks::iter::{Iter, RevIter};
+use crate::rocks::safe_iter::{SafeIter, SafeRevIter};
 use crate::rocks::util::{is_ref_count_value, reference_count_merge_operator};
 use crate::{reopen, retry_transaction, retry_transaction_forever};
 use rstest::rstest;
@@ -12,6 +14,82 @@ fn temp_dir() -> std::path::PathBuf {
         .into_path()
 }
 
+// A wrapper that holds different type of iterators for testing purpose. We use it to get same
+// typed key value paris from the database in parameterized tests, while varying different types
+// of underlying Iterator.
+enum TestIteratorWrapper<'a, K, V> {
+    Iter(Iter<'a, K, V>),
+    RevIter(RevIter<'a, K, V>),
+    SafeIter(SafeIter<'a, K, V>),
+    SafeRevIter(SafeRevIter<'a, K, V>),
+}
+
+// Implement Iterator for TestIteratorWrapper that returns the same type result for different types of Iterator.
+// For non-safe Iterator, it returns the key value pair. For SafeIterator, it consumes the result (assuming no error),
+// and return they key value pairs.
+impl<'a, K: DeserializeOwned, V: DeserializeOwned> Iterator for TestIteratorWrapper<'a, K, V> {
+    type Item = (K, V);
+    fn next(&mut self) -> Option<Self::Item> {
+        match self {
+            TestIteratorWrapper::Iter(iter) => iter.next(),
+            TestIteratorWrapper::RevIter(iter) => iter.next(),
+            TestIteratorWrapper::SafeIter(iter) => iter.next().map(|result| result.unwrap()),
+            TestIteratorWrapper::SafeRevIter(iter) => iter.next().map(|result| result.unwrap()),
+        }
+    }
+}
+
+// Creates an Iterator based on `use_safe_iter` on `db`.
+fn get_iter<K, V>(db: &DBMap<K, V>, use_safe_iter: bool) -> TestIteratorWrapper<'_, K, V>
+where
+    K: Serialize + DeserializeOwned,
+    V: Serialize + DeserializeOwned,
+{
+    match use_safe_iter {
+        true => TestIteratorWrapper::SafeIter(db.safe_iter()),
+        false => TestIteratorWrapper::Iter(db.unbounded_iter()),
+    }
+}
+
+impl<'a, K: Serialize, V> TestIteratorWrapper<'a, K, V> {
+    pub fn skip_to(self, key: &K) -> Result<Self, TypedStoreError> {
+        match self {
+            TestIteratorWrapper::Iter(iter) => Ok(TestIteratorWrapper::Iter(iter.skip_to(key)?)),
+            TestIteratorWrapper::SafeIter(iter) => {
+                Ok(TestIteratorWrapper::SafeIter(iter.skip_to(key)?))
+            }
+            _ => panic!("Not supported!"),
+        }
+    }
+    pub fn skip_prior_to(self, key: &K) -> Result<Self, TypedStoreError> {
+        match self {
+            TestIteratorWrapper::Iter(iter) => {
+                Ok(TestIteratorWrapper::Iter(iter.skip_prior_to(key)?))
+            }
+            TestIteratorWrapper::SafeIter(iter) => {
+                Ok(TestIteratorWrapper::SafeIter(iter.skip_prior_to(key)?))
+            }
+            _ => panic!("Not supported!"),
+        }
+    }
+    pub fn skip_to_last(self) -> Self {
+        match self {
+            TestIteratorWrapper::Iter(iter) => TestIteratorWrapper::Iter(iter.skip_to_last()),
+            TestIteratorWrapper::SafeIter(iter) => {
+                TestIteratorWrapper::SafeIter(iter.skip_to_last())
+            }
+            _ => panic!("Not supported!"),
+        }
+    }
+    pub fn reverse(self) -> Self {
+        match self {
+            TestIteratorWrapper::Iter(iter) => TestIteratorWrapper::RevIter(iter.reverse()),
+            TestIteratorWrapper::SafeIter(iter) => TestIteratorWrapper::SafeRevIter(iter.reverse()),
+            _ => panic!("Not supported!"),
+        }
+    }
+}
+
 #[rstest]
 #[tokio::test]
 async fn test_open(#[values(true, false)] is_transactional: bool) {
@@ -187,7 +265,10 @@ async fn test_chunked_multi_get(#[values(true, false)] is_transactional: bool) {
 
 #[rstest]
 #[tokio::test]
-async fn test_skip(#[values(true, false)] is_transactional: bool) {
+async fn test_skip(
+    #[values(true, false)] is_transactional: bool,
+    #[values(true, false)] use_safe_iter: bool,
+) {
     let db = open_map(temp_dir(), None, is_transactional);
 
     db.insert(&123, &"123".to_string())
@@ -198,10 +279,13 @@ async fn test_skip(#[values(true, false)] is_transactional: bool) {
         .expect("Failed to insert");
 
     // Skip all smaller
-    let key_vals: Vec<_> = db.safe_iter().skip_to(&456).expect("Seek failed").collect();
+    let key_vals: Vec<_> = get_iter(&db, use_safe_iter)
+        .skip_to(&456)
+        .expect("Seek failed")
+        .collect();
     assert_eq!(key_vals.len(), 2);
-    assert_eq!(key_vals[0], Ok((456, "456".to_string())));
-    assert_eq!(key_vals[1], Ok((789, "789".to_string())));
+    assert_eq!(key_vals[0], (456, "456".to_string()));
+    assert_eq!(key_vals[1], (789, "789".to_string()));
 
     // Skip all smaller: same for the keys iterator
     let keys: Vec<_> = db.keys().skip_to(&456).expect("Seek failed").collect();
@@ -211,31 +295,40 @@ async fn test_skip(#[values(true, false)] is_transactional: bool) {
 
     // Skip to the end
     assert_eq!(
-        db.safe_iter().skip_to(&999).expect("Seek failed").count(),
+        get_iter(&db, use_safe_iter)
+            .skip_to(&999)
+            .expect("Seek failed")
+            .count(),
         0
     );
     // same for the keys
     assert_eq!(db.keys().skip_to(&999).expect("Seek failed").count(), 0);
 
     // Skip to last
     assert_eq!(
-        db.safe_iter().skip_to_last().next(),
-        Some(Ok((789, "789".to_string())))
+        get_iter(&db, use_safe_iter).skip_to_last().next(),
+        Some((789, "789".to_string()))
     );
     // same for the keys
     assert_eq!(db.keys().skip_to_last().next(), Some(Ok(789)));
 
     // Skip to successor of first value
     assert_eq!(
-        db.safe_iter().skip_to(&000).expect("Skip failed").count(),
+        get_iter(&db, use_safe_iter)
+            .skip_to(&000)
+            .expect("Skip failed")
+            .count(),
         3
     );
     assert_eq!(db.keys().skip_to(&000).expect("Skip failed").count(), 3);
 }
 
 #[rstest]
 #[tokio::test]
-async fn test_skip_to_previous_simple(#[values(true, false)] is_transactional: bool) {
+async fn test_skip_to_previous_simple(
+    #[values(true, false)] is_transactional: bool,
+    #[values(true, false)] use_safe_iter: bool,
+) {
     let db = open_map(temp_dir(), None, is_transactional);
 
     db.insert(&123, &"123".to_string())
@@ -246,13 +339,12 @@ async fn test_skip_to_previous_simple(#[values(true, false)] is_transactional: b
         .expect("Failed to insert");
 
     // Skip to the one before the end
-    let key_vals: Vec<_> = db
-        .safe_iter()
+    let key_vals: Vec<_> = get_iter(&db, use_safe_iter)
         .skip_prior_to(&999)
         .expect("Seek failed")
         .collect();
     assert_eq!(key_vals.len(), 1);
-    assert_eq!(key_vals[0], Ok((789, "789".to_string())));
+    assert_eq!(key_vals[0], (789, "789".to_string()));
     // Same for the keys iterator
     let keys: Vec<_> = db
         .keys()
@@ -265,7 +357,7 @@ async fn test_skip_to_previous_simple(#[values(true, false)] is_transactional: b
     // Skip to prior of first value
     // Note: returns an empty iterator!
     assert_eq!(
-        db.safe_iter()
+        get_iter(&db, use_safe_iter)
             .skip_prior_to(&000)
             .expect("Seek failed")
             .count(),
@@ -280,7 +372,10 @@ async fn test_skip_to_previous_simple(#[values(true, false)] is_transactional: b
 
 #[rstest]
 #[tokio::test]
-async fn test_iter_skip_to_previous_gap(#[values(true, false)] is_transactional: bool) {
+async fn test_iter_skip_to_previous_gap(
+    #[values(true, false)] is_transactional: bool,
+    #[values(true, false)] use_safe_iter: bool,
+) {
     let db = open_map(temp_dir(), None, is_transactional);
 
     for i in 1..100 {
@@ -290,12 +385,12 @@ async fn test_iter_skip_to_previous_gap(#[values(true, false)] is_transactional:
     }
 
     // Skip prior to will return an iterator starting with an "unexpected" key if the sought one is not in the table
-    let db_iter = db.safe_iter().skip_prior_to(&50).unwrap();
+    let db_iter = get_iter(&db, use_safe_iter).skip_prior_to(&50).unwrap();
 
     assert_eq!(
         (49..50)
             .chain(51..100)
-            .map(|i| Ok((i, i.to_string())))
+            .map(|i| (i, i.to_string()))
             .collect::<Vec<_>>(),
         db_iter.collect::<Vec<_>>()
     );
@@ -323,34 +418,44 @@ async fn test_remove(#[values(true, false)] is_transactional: bool) {
 
 #[rstest]
 #[tokio::test]
-async fn test_iter(#[values(true, false)] is_transactional: bool) {
+async fn test_iter(
+    #[values(true, false)] is_transactional: bool,
+    #[values(true, false)] use_safe_iter: bool,
+) {
     let db = open_map(temp_dir(), None, is_transactional);
     db.insert(&123456789, &"123456789".to_string())
         .expect("Failed to insert");
+    db.insert(&987654321, &"987654321".to_string())
+        .expect("Failed to insert");
+
+    let mut iter = get_iter(&db, use_safe_iter);
 
-    let mut iter = db.safe_iter();
-    assert_eq!(Some(Ok((123456789, "123456789".to_string()))), iter.next());
+    assert_eq!(Some((123456789, "123456789".to_string())), iter.next());
+    assert_eq!(Some((987654321, "987654321".to_string())), iter.next());
     assert_eq!(None, iter.next());
 }
 
 #[rstest]
 #[tokio::test]
-async fn test_iter_reverse(#[values(true, false)] is_transactional: bool) {
+async fn test_iter_reverse(
+    #[values(true, false)] is_transactional: bool,
+    #[values(true, false)] use_safe_iter: bool,
+) {
     let db = open_map(temp_dir(), None, is_transactional);
 
     db.insert(&1, &"1".to_string()).expect("Failed to insert");
     db.insert(&2, &"2".to_string()).expect("Failed to insert");
     db.insert(&3, &"3".to_string()).expect("Failed to insert");
 
-    let mut iter = db.safe_iter().skip_to_last().reverse();
-    assert_eq!(Some(Ok((3, "3".to_string()))), iter.next());
-    assert_eq!(Some(Ok((2, "2".to_string()))), iter.next());
-    assert_eq!(Some(Ok((1, "1".to_string()))), iter.next());
+    let mut iter = get_iter(&db, use_safe_iter).skip_to_last().reverse();
+    assert_eq!(Some((3, "3".to_string())), iter.next());
+    assert_eq!(Some((2, "2".to_string())), iter.next());
+    assert_eq!(Some((1, "1".to_string())), iter.next());
     assert_eq!(None, iter.next());
 
-    let mut iter = db.safe_iter().skip_to(&2).unwrap().reverse();
-    assert_eq!(Some(Ok((2, "2".to_string()))), iter.next());
-    assert_eq!(Some(Ok((1, "1".to_string()))), iter.next());
+    let mut iter = get_iter(&db, use_safe_iter).skip_to(&2).unwrap().reverse();
+    assert_eq!(Some((2, "2".to_string())), iter.next());
+    assert_eq!(Some((1, "1".to_string())), iter.next());
     assert_eq!(None, iter.next());
 }