Freelist defrag

near-sdk/src/store/free_list/mod.rs

@@ -223,6 +223,92 @@ where
     pub fn drain(&mut self) -> Drain<T> {
         Drain::new(self)
     }
+
+    /// Empty slots in the front of the list is swapped with occupied slots in back of the list.
+    /// Defrag helps reduce gas cost in certain scenarios where lot of elements in front of the list are
+    /// removed without getting replaced. Please see https://github.com/near/near-sdk-rs/issues/990
+    pub fn defrag<F>(&mut self, callback: F)
+    where
+        F: FnMut(&T, u32),
+    {
+        Defrag::new(self).defrag(callback);
+        self.first_free = None;
+    }
+}
+
+/// Defrag struct has helper functions to perform defragmentation of `FreeList`. See the
+/// documentation of function [`FreeList::defrag`] for more details.
+struct Defrag<'a, T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    elements: &'a mut Vector<Slot<T>>,
+    occupied_count: u32,
+    curr_free_slot: Option<FreeListIndex>,
+    defrag_index: u32,
+}
+
+impl<'a, T> Defrag<'a, T>
+where
+    T: BorshSerialize + BorshDeserialize,
+{
+    /// Create a new struct for defragmenting `FreeList`.
+    fn new(list: &'a mut FreeList<T>) -> Self {
+        Self {
+            elements: &mut list.elements,
+            occupied_count: list.occupied_count,
+            defrag_index: list.occupied_count,
+            curr_free_slot: list.first_free,
+        }
+    }
+
+    fn defrag<F>(&mut self, mut callback: F)
+    where
+        F: FnMut(&T, u32),
+    {
+        while let Some(curr_free_index) = self.next_free_slot() {
+            if let Some((value, occupied_index)) = self.next_occupied() {
+                callback(value, curr_free_index.0);
+                //The entry at curr_free_index.0 should have `None` by now.
+                //Moving it to `occupied_index` will make that entry empty.
+                self.elements.swap(curr_free_index.0, occupied_index);
+            } else {
+                //Could not find an occupied slot to fill the free slot
+                env::panic_str(ERR_INCONSISTENT_STATE)
+            }
+        }
+
+        // After defragmenting, these should all be `Slot::Empty`.
+        self.elements.remove_tail(self.elements.len() - self.occupied_count);
+    }
+
+    fn next_free_slot(&mut self) -> Option<FreeListIndex> {
+        while let Some(curr_free_index) = self.curr_free_slot {
+            let curr_slot = self.elements.get(curr_free_index.0);
+            self.curr_free_slot = match curr_slot {
+                Some(Slot::Empty { next_free }) => *next_free,
+                Some(Slot::Occupied(_)) => {
+                    //The free list chain should not have an occupied slot
+                    env::panic_str(ERR_INCONSISTENT_STATE)
+                }
+                _ => None,
+            };
+            if curr_free_index.0 < self.occupied_count {
+                return Some(curr_free_index);
+            }
+        }
+        None
+    }
+
+    fn next_occupied(&mut self) -> Option<(&T, u32)> {
+        while self.defrag_index < self.elements.len {
+            if let Some(Slot::Occupied(value)) = self.elements.get(self.defrag_index) {
+                return Some((value, self.defrag_index));
+            }
+            self.defrag_index += 1;
+        }
+        None
+    }
 }
 
 #[cfg(not(target_arch = "wasm32"))]
@@ -236,6 +322,25 @@ mod tests {
     use super::*;
     use crate::test_utils::test_env::setup_free;
 
+    #[test]
+    fn new_bucket_is_empty() {
+        let bucket: FreeList<u8> = FreeList::new(b"b");
+        assert!(bucket.is_empty());
+    }
+
+    #[test]
+    fn occupied_count_gets_updated() {
+        let mut bucket = FreeList::new(b"b");
+        let indices: Vec<_> = (0..5).map(|i| bucket.insert(i)).collect();
+
+        assert_eq!(bucket.occupied_count, 5);
+
+        bucket.remove(indices[1]);
+        bucket.remove(indices[3]);
+
+        assert_eq!(bucket.occupied_count, 3);
+    }
+
     #[test]
     fn basic_functionality() {
         let mut bucket = FreeList::new(b"b");
@@ -252,6 +357,32 @@ mod tests {
         assert_eq!(bucket.get(i3), Some(&4));
     }
 
+    #[test]
+    fn defrag() {
+        let mut bucket = FreeList::new(b"b");
+        let indices: Vec<_> = (0..8).map(|i| bucket.insert(i)).collect();
+
+        //Empty, Empty, Empty, Empty, Occupied, Empty, Occupied, Empty
+        bucket.remove(indices[1]);
+        bucket.remove(indices[3]);
+        bucket.remove(indices[0]);
+        bucket.remove(indices[5]);
+        bucket.remove(indices[2]);
+        bucket.remove(indices[7]);
+
+        //4 should move to index 0, 6 should move to index 1
+        bucket.defrag(|_, _| {});
+
+        //Check the free slots chain is complete after defrag
+        assert_eq!(bucket.occupied_count, bucket.len());
+
+        assert_eq!(*bucket.get(indices[0]).unwrap(), 4u8);
+        assert_eq!(*bucket.get(indices[1]).unwrap(), 6u8);
+        for i in indices[2..].iter() {
+            assert_eq!(bucket.get(*i), None);
+        }
+    }
+
     #[test]
     fn bucket_iterator() {
         let mut bucket = FreeList::new(b"b");

near-sdk/src/store/unordered_map/mod.rs

@@ -630,6 +630,23 @@ where
     }
 }
 
+impl<K, V, H> UnorderedMap<K, V, H>
+where
+    K: BorshSerialize + BorshDeserialize + Ord + Clone,
+    V: BorshSerialize + BorshDeserialize,
+    H: ToKey,
+{
+    pub fn defrag(&mut self) {
+        self.keys.defrag(|key, new_index| {
+            // Check if value is in map to replace first
+            let entry = self.values.get_mut_inner(key);
+            if let Some(existing) = entry.value_mut() {
+                existing.key_index = FreeListIndex(new_index);
+            }
+        });
+    }
+}
+
 #[cfg(not(target_arch = "wasm32"))]
 #[cfg(test)]
 mod tests {
@@ -756,4 +773,37 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn defrag() {
+        let mut map = UnorderedMap::new(b"b");
+
+        let all_indices = 0..=8;
+
+        for i in all_indices {
+            map.insert(i, i);
+        }
+
+        let removed = [2, 4, 6];
+        let existing = [0, 1, 3, 5, 7, 8];
+
+        for id in removed {
+            map.remove(&id);
+        }
+
+        map.defrag();
+
+        for i in removed {
+            assert_eq!(map.get(&i), None);
+        }
+        for i in existing {
+            assert_eq!(map.get(&i), Some(&i));
+        }
+
+        //Check the elements moved during defragmentation
+        assert_eq!(map.remove_entry(&7).unwrap(), (7, 7));
+        assert_eq!(map.remove_entry(&8).unwrap(), (8, 8));
+        assert_eq!(map.remove_entry(&1).unwrap(), (1, 1));
+        assert_eq!(map.remove_entry(&3).unwrap(), (3, 3));
+    }
 }

near-sdk/src/store/vec/mod.rs

@@ -349,7 +349,7 @@ where
         self.values.get_mut(index)
     }
 
-    fn swap(&mut self, a: u32, b: u32) {
+    pub(crate) fn swap(&mut self, a: u32, b: u32) {
         if a >= self.len() || b >= self.len() {
             env::panic_str(ERR_INDEX_OUT_OF_BOUNDS);
         }
@@ -408,6 +408,23 @@ where
         prev
     }
 
+    /// Removes the last n elements from a vector. Returns `true` on success.
+    /// If the vector is shorter than n, returns `false` without committing any changes.
+    pub(crate) fn remove_tail(&mut self, n: u32) -> bool {
+        let new_idx = match self.len.checked_sub(n) {
+            Some(new_idx) => new_idx,
+            None => return false,
+        };
+
+        for ix in new_idx..self.len {
+            self.values.remove(ix);
+        }
+
+        self.len = new_idx;
+
+        true
+    }
+
     /// Inserts a element at `index`, returns an evicted element.
     ///
     /// # Panics
@@ -569,6 +586,23 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_remove_tail() {
+        let mut vec = Vector::new(b"v".to_vec());
+
+        vec.push(5);
+        vec.push(2);
+        vec.push(3);
+        vec.remove_tail(2);
+
+        assert_eq!(vec.len(), 1);
+        assert_eq!(vec.get(0), Some(&5));
+        assert_eq!(vec.get(1), None);
+        assert_eq!(vec.values.get(1), None);
+        assert_eq!(vec.get(2), None);
+        assert_eq!(vec.values.get(2), None);
+    }
+
     #[test]
     pub fn test_replace() {
         let mut rng = rand_xorshift::XorShiftRng::seed_from_u64(1);