fix: add unique index per status so bounds work correctly

contract/src/msgs/data_requests/test_helpers.rs

@@ -250,9 +250,9 @@ impl TestInfo {
 
     #[track_caller]
     pub fn get_data_result(&self, dr_id: &str) -> DataResult {
-        dbg!(self.query(query::QueryMsg::GetDataResult {
+        self.query(query::QueryMsg::GetDataResult {
             dr_id: dr_id.to_string(),
-        }))
+        })
         .unwrap()
     }
 

contract/src/msgs/data_requests/tests.rs

@@ -667,16 +667,37 @@ fn get_data_requests_by_status_with_more_drs_in_pool() {
         exit_code: 0,
     };
 
-    for i in 1..25 {
+    for i in 0..25 {
         let dr = test_helpers::calculate_dr_id_and_args(i, 1);
         let dr_id = test_info.post_data_request(&alice, dr, vec![], vec![], 1).unwrap();
 
         if i < 15 {
-            test_info.commit_result(&alice, &dr_id, alice_reveal.try_hash().unwrap()).unwrap();
+            test_info
+                .commit_result(&alice, &dr_id, alice_reveal.try_hash().unwrap())
+                .unwrap();
         }
 
+        if i < 3 {
+            test_info.reveal_result(&alice, &dr_id, alice_reveal.clone()).unwrap();
+        }
     }
 
-    let comitting = test_info.get_data_requests_by_status(DataRequestStatus::Committing, 0, 10);
-    assert_eq!(comitting.len(), 10);
-}
+    assert_eq!(
+        10,
+        test_info
+            .get_data_requests_by_status(DataRequestStatus::Committing, 0, 10)
+            .len()
+    );
+    assert_eq!(
+        12,
+        test_info
+            .get_data_requests_by_status(DataRequestStatus::Revealing, 0, 15)
+            .len()
+    );
+    assert_eq!(
+        3,
+        test_info
+            .get_data_requests_by_status(DataRequestStatus::Tallying, 0, 15)
+            .len()
+    );
+}

contract/src/msgs/data_requests/types.rs

@@ -53,6 +53,9 @@ impl StatusValue {
 
 pub struct DataRequestsMap<'a> {
     pub len:            Item<'a, u32>,
+    pub committing_len: Item<'a, u32>,
+    pub revealing_len:  Item<'a, u32>,
+    pub tallying_len:   Item<'a, u32>,
     pub reqs:           Map<'a, &'a Hash, StatusValue>,
     pub index_to_key:   Map<'a, u32, Hash>,
     pub key_to_index:   Map<'a, &'a Hash, u32>,
@@ -64,6 +67,9 @@ macro_rules! enumerable_status_map {
     ($namespace:literal) => {
         DataRequestsMap {
             len:            Item::new(concat!($namespace, "_len")),
+            committing_len: Item::new(concat!($namespace, "_committing_len")),
+            revealing_len:  Item::new(concat!($namespace, "_revealing_len")),
+            tallying_len:   Item::new(concat!($namespace, "_tallying_len")),
             reqs:           Map::new(concat!($namespace, "_reqs")),
             index_to_key:   Map::new(concat!($namespace, "_index_to_key")),
             key_to_index:   Map::new(concat!($namespace, "_key_to_index")),
@@ -72,54 +78,58 @@ macro_rules! enumerable_status_map {
     };
 }
 
-impl<'a> DataRequestsMap<'_> {
+impl DataRequestsMap<'_> {
     pub fn initialize(&self, store: &mut dyn Storage) -> StdResult<()> {
         self.len.save(store, &0)?;
+        self.committing_len.save(store, &0)?;
+        self.revealing_len.save(store, &0)?;
+        self.tallying_len.save(store, &0)?;
         Ok(())
     }
 
+    pub fn get_status_len_item(&self, status: &DataRequestStatus) -> &Item<u32> {
+        match status {
+            DataRequestStatus::Committing => &self.committing_len,
+            DataRequestStatus::Revealing => &self.revealing_len,
+            DataRequestStatus::Tallying => &self.tallying_len,
+        }
+    }
+
     pub fn has(&self, store: &dyn Storage, key: &Hash) -> bool {
         self.key_to_index.has(store, key)
     }
 
     pub fn update(
-        &'a self,
+        &self,
         store: &mut dyn Storage,
         key: &Hash,
         dr: DataRequest,
         status: Option<DataRequestStatus>,
     ) -> StdResult<()> {
-        let Some(index) = self.key_to_index.may_load(store, key)? else {
+        let Some(old) = self.reqs.may_load(store, key)? else {
             return Err(StdError::generic_err("Key does not exist"));
         };
 
-        // remove the old status key
-        let old = self.reqs.load(store, key)?;
-        let old_status_key = StatusIndexKey::new(index, Some(old.status.clone()));
-        let status = if let Some(status) = status {
-            self.status_to_keys.remove(store, &old_status_key);
-            self.status_to_keys
-                .save(store, &StatusIndexKey::new(index, Some(status.clone())), key)?;
-            status
-        } else {
-            old.status
-        };
-        self.reqs.save(store, key, &StatusValue::with_status(dr, status))
+        let status = status.unwrap_or(old.status.clone());
+
+        self.swap_remove(store, key)?;
+        self.insert_with_status(store, key, dr, status)?;
+        Ok(())
     }
 
-    pub fn len(&'a self, store: &dyn Storage) -> Result<u32, StdError> {
+    pub fn len(&self, store: &dyn Storage) -> Result<u32, StdError> {
         self.len.load(store)
     }
 
-    pub fn may_get_by_key(&'a self, store: &dyn Storage, key: &Hash) -> StdResult<Option<DataRequest>> {
+    pub fn may_get_by_key(&self, store: &dyn Storage, key: &Hash) -> StdResult<Option<DataRequest>> {
         self.reqs.may_load(store, key).map(|opt| opt.map(|req| req.req))
     }
 
-    pub fn get_by_key(&'a self, store: &dyn Storage, key: &Hash) -> StdResult<DataRequest> {
+    pub fn get_by_key(&self, store: &dyn Storage, key: &Hash) -> StdResult<DataRequest> {
         self.reqs.load(store, key).map(|req| req.req)
     }
 
-    pub fn may_get_by_index(&'a self, store: &dyn Storage, index: u32) -> StdResult<Option<DataRequest>> {
+    pub fn may_get_by_index(&self, store: &dyn Storage, index: u32) -> StdResult<Option<DataRequest>> {
         if let Some(key) = self.index_to_key.may_load(store, index)? {
             self.may_get_by_key(store, &key)
         } else {
@@ -133,12 +143,49 @@ impl<'a> DataRequestsMap<'_> {
         }
 
         let index = self.len.load(store)?;
+        let status_len_item = self.get_status_len_item(&DataRequestStatus::Committing);
+        let status_index = status_len_item.load(store)?;
+
+        // save the request
         self.reqs.save(store, key, &StatusValue::new(req))?;
         self.index_to_key.save(store, index, key)?;
         self.key_to_index.save(store, key, &index)?;
         self.status_to_keys
-            .save(store, &StatusIndexKey::new(index, None), key)?;
+            .save(store, &StatusIndexKey::new(status_index, None), key)?;
+
+        // increment the indexes
         self.len.save(store, &(index + 1))?;
+        status_len_item.save(store, &(status_index + 1))?;
+        Ok(())
+    }
+
+    // only to be internally used by update
+    fn insert_with_status(
+        &self,
+        store: &mut dyn Storage,
+        key: &Hash,
+        req: DataRequest,
+        status: DataRequestStatus,
+    ) -> StdResult<()> {
+        if self.key_to_index.may_load(store, key)?.is_some() {
+            return Err(StdError::generic_err("Key already exists"));
+        }
+
+        let index = self.len.load(store)?;
+        let status_len_item = self.get_status_len_item(&status);
+        let status_index = status_len_item.load(store)?;
+
+        // Correctly save the new status key
+        self.reqs
+            .save(store, key, &StatusValue::with_status(req, status.clone()))?;
+        self.index_to_key.save(store, index, key)?;
+        self.key_to_index.save(store, key, &index)?;
+        self.status_to_keys
+            .save(store, &StatusIndexKey::new(status_index, Some(status)), key)?;
+
+        // Increment the indexes
+        self.len.save(store, &(index + 1))?;
+        status_len_item.save(store, &(status_index + 1))?;
         Ok(())
     }
 
@@ -153,14 +200,17 @@ impl<'a> DataRequestsMap<'_> {
             .ok_or_else(|| StdError::generic_err("Key does not exist"))?;
         let req_status = self.reqs.load(store, key)?.status;
 
-        // Get the current lengt
+        // Get the current length
         let total_items = self.len.load(store)?;
 
         // Shouldn't be reachable
         if total_items == 0 {
             unreachable!("No items in the map, so key should not exist");
         }
 
+        let status_len_item = self.get_status_len_item(&req_status);
+        let req_status_index = status_len_item.load(store)? - 1;
+
         // Handle case when removing the last or only item
         // means we can just remove the key and return
         if total_items == 1 || req_index == total_items - 1 {
@@ -171,17 +221,20 @@ impl<'a> DataRequestsMap<'_> {
                 store,
                 &StatusIndexKey {
                     status: req_status,
-                    index:  req_index,
+                    index:  req_status_index,
                 },
             );
             self.len.save(store, &(total_items - 1))?;
+            status_len_item.save(store, &req_status_index)?;
             return Ok(());
         }
 
         // Swap the last item into the position of the removed item
         let last_index = total_items - 1;
         let last_key = self.index_to_key.load(store, last_index)?;
         let last_status = self.reqs.load(store, &last_key)?.status;
+        let last_status_len_item = self.get_status_len_item(&last_status);
+        let last_status_index = last_status_len_item.load(store)? - 1;
 
         // Update mapping for the swapped item
         self.index_to_key.save(store, req_index, &last_key)?;
@@ -196,15 +249,15 @@ impl<'a> DataRequestsMap<'_> {
                 store,
                 &StatusIndexKey {
                     status: last_status.clone(),
-                    index:  last_index,
+                    index:  last_status_index,
                 },
             );
             // Remove the entry for the element that is being removed from the status_to_keys map.
             self.status_to_keys.remove(
                 store,
                 &StatusIndexKey {
                     status: req_status,
-                    index:  req_index,
+                    index:  req_status_index,
                 },
             );
             // Add a new entry for the element that was previously last in the status_to_keys map at its new index.
@@ -225,6 +278,8 @@ impl<'a> DataRequestsMap<'_> {
 
         // Update length
         self.len.save(store, &last_index)?;
+        status_len_item.save(store, &req_status_index)?;
+
         Ok(())
     }
 

contract/src/msgs/data_requests/types_tests.rs

@@ -23,6 +23,14 @@ impl TestInfo<'_> {
         assert_eq!(self.map.len(&self.store).unwrap(), expected);
     }
 
+    #[track_caller]
+    fn assert_status_len(&self, expected: u32, status: &DataRequestStatus) {
+        assert_eq!(
+            self.map.get_status_len_item(status).load(&self.store).unwrap(),
+            expected
+        );
+    }
+
     #[track_caller]
     fn assert_index_to_key(&self, index: u32, key: Option<Hash>) {
         assert_eq!(self.map.index_to_key.may_load(&self.store, index).unwrap(), key);
@@ -94,6 +102,7 @@ fn enum_map_insert() {
     let (key, val) = create_test_dr(1);
     test_info.insert(&key, val);
     test_info.assert_len(1);
+    test_info.assert_status_len(1, &DataRequestStatus::Committing);
     test_info.assert_index_to_key(0, Some(key));
     test_info.assert_key_to_index(&key, Some(0));
     test_info.assert_status_index_to_key(StatusIndexKey::new(0, None), Some(key));
@@ -133,13 +142,15 @@ fn enum_map_update() {
 
     test_info.insert(&key1, dr1.clone());
     test_info.assert_len(1);
+    test_info.assert_status_len(1, &DataRequestStatus::Committing);
     assert_eq!(test_info.get_by_index(0), Some(dr1));
-    test_info.assert_status_index_to_key(dbg!(StatusIndexKey::new(0, None)), Some(key1));
+    test_info.assert_status_index_to_key(StatusIndexKey::new(0, None), Some(key1));
 
     test_info.update(&key1, dr2.clone(), None);
     test_info.assert_len(1);
+    test_info.assert_status_len(1, &DataRequestStatus::Committing);
     assert_eq!(test_info.get_by_index(0), Some(dr2));
-    test_info.assert_status_index_to_key(dbg!(StatusIndexKey::new(0, None)), Some(key1));
+    test_info.assert_status_index_to_key(StatusIndexKey::new(0, None), Some(key1));
 }
 
 #[test]
@@ -163,6 +174,7 @@ fn enum_map_remove_first() {
 
     test_info.swap_remove(&key1);
     test_info.assert_len(2);
+    test_info.assert_status_len(2, &DataRequestStatus::Committing);
     test_info.assert_index_to_key(0, Some(key3));
     test_info.assert_key_to_index(&key3, Some(0));
     test_info.assert_status_index_to_key(StatusIndexKey::new(0, None), Some(key3));
@@ -188,9 +200,11 @@ fn enum_map_remove_last() {
     test_info.insert(&key2, req2.clone()); // 1
     test_info.insert(&key3, req3.clone()); // 2
     test_info.assert_len(3);
+    test_info.assert_status_len(3, &DataRequestStatus::Committing);
 
     test_info.swap_remove(&key3);
     test_info.assert_len(2);
+    test_info.assert_status_len(2, &DataRequestStatus::Committing);
     test_info.assert_index_to_key(2, None);
     test_info.assert_key_to_index(&key3, None);
     test_info.assert_status_index_to_key(StatusIndexKey::new(2, None), None);
@@ -226,9 +240,11 @@ fn enum_map_remove() {
     test_info.insert(&key3, req3.clone()); // 2
     test_info.insert(&key4, req4.clone()); // 3
     test_info.assert_len(4);
+    test_info.assert_status_len(4, &DataRequestStatus::Committing);
 
     test_info.swap_remove(&key2);
     test_info.assert_len(3);
+    test_info.assert_status_len(3, &DataRequestStatus::Committing);
     test_info.assert_index_to_key(1, Some(key4));
     test_info.assert_key_to_index(&key4, Some(1));
 
@@ -323,6 +339,7 @@ fn remove_only_item() {
     test_info.insert(&key, req.clone());
     test_info.swap_remove(&key);
     test_info.assert_len(0);
+    test_info.assert_status_len(0, &DataRequestStatus::Committing);
     test_info.assert_index_to_key(0, None);
     test_info.assert_key_to_index(&key, None);
     test_info.assert_status_index_to_key(StatusIndexKey::new(0, None), None);