Merge pull request #66 from garro95/extract-if

Implement Extract if

Author: garro95

src/double_priority_queue/iterators.rs

@@ -40,6 +40,89 @@ use std::iter::*;
 
 use crate::DoublePriorityQueue;
 
+use super::Index;
+
+/// An `Iterator` in arbitrary order which uses a `predicate` to determine if
+/// an element should be removed from the `DoublePriorityQueue`.
+///
+/// It can be obtained calling the [`extract_if`](DoublePriorityQueue::extract_if) method.
+///
+/// The `predicate` has mutable access to the `(item, priority)` pairs.
+///
+/// It can update the priorities of the elements in the queue and the items
+/// in a way that does not change the result of any of `hash` or `eq`.
+///
+/// When the iterator goes out of scope, the heap is rebuilt to restore the
+/// structural properties.
+#[cfg(feature = "std")]
+pub struct ExtractIf<'a, I: 'a, P: 'a, F, H: 'a = RandomState>
+where
+    P: Ord,
+{
+    pq: &'a mut DoublePriorityQueue<I, P, H>,
+    predicate: F,
+    idx: Index,
+}
+
+#[cfg(not(feature = "std"))]
+pub struct ExtractIf<'a, I: 'a, P: 'a, F, H: 'a>
+where
+    P: Ord,
+{
+    pq: &'a mut DoublePriorityQueue<I, P, H>,
+    predicate: F,
+    idx: Index,
+}
+
+impl<'a, I: 'a, P: 'a, F, H: 'a> ExtractIf<'a, I, P, F, H>
+where
+    P: Ord,
+{
+    pub(crate) fn new(pq: &'a mut DoublePriorityQueue<I, P, H>, predicate: F) -> Self {
+        ExtractIf {
+            pq,
+            predicate,
+            idx: Index(0),
+        }
+    }
+}
+
+impl<'a, I: 'a, P: 'a, F, H: 'a> Iterator for ExtractIf<'a, I, P, F, H>
+where
+    P: Ord,
+    F: FnMut(&mut I, &mut P) -> bool,
+    H: BuildHasher,
+{
+    type Item = (I, P);
+    fn next(&mut self) -> Option<Self::Item> {
+        use indexmap::map::MutableKeys;
+
+        loop {
+            let r: Option<bool> = self
+                .pq
+                .store
+                .map
+                .get_index_mut2(self.idx.0)
+                .map(|(i, p)| (self.predicate)(i, p));
+
+            match r {
+                Some(true) => return self.pq.store.swap_remove_index(self.idx),
+                Some(false) => self.idx.0 += 1,
+                None => return None,
+            }
+        }
+    }
+}
+
+impl<'a, I: 'a, P: 'a, F, H: 'a> Drop for ExtractIf<'a, I, P, F, H>
+where
+    P: Ord,
+{
+    fn drop(&mut self) {
+        self.pq.heap_build();
+    }
+}
+
 /// A mutable iterator over the couples `(item, priority)` of the `DoublePriorityQueue`
 /// in arbitrary order.
 ///

src/double_priority_queue/mod.rs

@@ -454,6 +454,29 @@ where
         self.heap_build();
     }
 
+    /// Returns an `Iterator` removing from the queue the `(item, priority)`
+    /// pairs for which the `predicate` returns `true`, in arbitraty order.
+    ///
+    /// The `predicate` receives mutable references to both the item and
+    /// the priority.
+    ///
+    /// It's a logical error to change the item in a way
+    /// that changes the result of `Hash` or `Eq`.
+    ///
+    /// The `predicate` can change the priority. If it returns `true`, the
+    /// extracted pair will have the updated priority, otherwise, the
+    /// heap structural property will be restored once the iterator is `Drop`ped.
+    ///
+    /// # Example
+    /// ```
+    /// ```
+    pub fn extract_if<F>(&mut self, predicate: F) -> ExtractIf<I, P, F, H>
+    where
+        F: FnMut(&mut I, &mut P) -> bool,
+    {
+        ExtractIf::new(self, predicate)
+    }
+
     /// Removes the item with the lowest priority from
     /// the priority queue if the predicate returns `true`.
     ///

src/priority_queue/iterators.rs

@@ -40,6 +40,89 @@ use std::iter::*;
 
 use crate::PriorityQueue;
 
+use super::Index;
+
+/// An `Iterator` in arbitrary order which uses a `predicate` to determine if
+/// an element should be removed from the `PriorityQueue`.
+///
+/// It can be obtained calling the [`extract_if`](PriorityQueue::extract_if) method.
+///
+/// The `predicate` has mutable access to the `(item, priority)` pairs.
+///
+/// It can update the priorities of the elements in the queue and the items
+/// in a way that does not change the result of any of `hash` or `eq`.
+///
+/// When the iterator goes out of scope, the heap is rebuilt to restore the
+/// structural properties.
+#[cfg(feature = "std")]
+pub struct ExtractIf<'a, I: 'a, P: 'a, F, H: 'a = RandomState>
+where
+    P: Ord,
+{
+    pq: &'a mut PriorityQueue<I, P, H>,
+    predicate: F,
+    idx: Index,
+}
+
+#[cfg(not(feature = "std"))]
+pub struct ExtractIf<'a, I: 'a, P: 'a, F, H: 'a>
+where
+    P: Ord,
+{
+    pq: &'a mut PriorityQueue<I, P, H>,
+    predicate: F,
+    idx: Index,
+}
+
+impl<'a, I: 'a, P: 'a, F, H: 'a> ExtractIf<'a, I, P, F, H>
+where
+    P: Ord,
+{
+    pub(crate) fn new(pq: &'a mut PriorityQueue<I, P, H>, predicate: F) -> Self {
+        ExtractIf {
+            pq,
+            predicate,
+            idx: Index(0),
+        }
+    }
+}
+
+impl<'a, I: 'a, P: 'a, F, H: 'a> Iterator for ExtractIf<'a, I, P, F, H>
+where
+    P: Ord,
+    F: FnMut(&mut I, &mut P) -> bool,
+    H: BuildHasher,
+{
+    type Item = (I, P);
+    fn next(&mut self) -> Option<Self::Item> {
+        use indexmap::map::MutableKeys;
+
+        loop {
+            let r: Option<bool> = self
+                .pq
+                .store
+                .map
+                .get_index_mut2(self.idx.0)
+                .map(|(i, p)| (self.predicate)(i, p));
+
+            match r {
+                Some(true) => return self.pq.store.swap_remove_index(self.idx),
+                Some(false) => self.idx.0 += 1,
+                None => return None,
+            }
+        }
+    }
+}
+
+impl<'a, I: 'a, P: 'a, F, H: 'a> Drop for ExtractIf<'a, I, P, F, H>
+where
+    P: Ord,
+{
+    fn drop(&mut self) {
+        self.pq.heap_build();
+    }
+}
+
 /// A mutable iterator over the couples `(item, priority)` of the `PriorityQueue`
 /// in arbitrary order.
 ///

src/priority_queue/mod.rs

@@ -364,6 +364,25 @@ where
         self.heap_build();
     }
 
+    /// Returns an `Iterator` removing from the queue the `(item, priority)`
+    /// pairs for which the `predicate` returns `true`, in arbitraty order.
+    ///
+    /// The `predicate` receives mutable references to both the item and
+    /// the priority.
+    ///
+    /// It's a logical error to change the item in a way
+    /// that changes the result of `Hash` or `Eq`.
+    ///
+    /// The `predicate` can change the priority. If it returns `true`, the
+    /// extracted pair will have the updated priority, otherwise, the
+    /// heap structural property will be restored once the iterator is `Drop`ped.
+    pub fn extract_if<F>(&mut self, predicate: F) -> ExtractIf<I, P, F, H>
+    where
+        F: FnMut(&mut I, &mut P) -> bool,
+    {
+        ExtractIf::new(self, predicate)
+    }
+
     /// Removes the item with the greatest priority from
     /// the priority queue if the `predicate` returns `true`.
     ///

src/store.rs

@@ -267,6 +267,41 @@ impl<I, P, H> Store<I, P, H> {
         self.heap.swap(a.0, b.0);
     }
 
+    /// Remove and return the element at index `idx`
+    /// and swap it with the last element keeping a consistent
+    /// state.
+    ///
+    /// Computes in **O(1)** time (average)
+    pub fn swap_remove_index(&mut self, idx: Index) -> Option<(I, P)> {
+        // swap_remove the position from the qp
+        let position = self.qp.swap_remove(idx.0);
+        self.size -= 1;
+
+        if idx.0 < self.size {
+            // SAFETY: head validity checked on the previous line.
+            // All positions point to valid heap items because we already
+            // updated the qp.
+            unsafe {
+                *self.heap.get_unchecked_mut(self.qp.get_unchecked(idx.0).0) = idx;
+            }
+        }
+        self.heap.swap_remove(position.0);
+        // Fix indexes and swap remove the old heap head from the qp
+        if position.0 < self.size {
+            // SAFETY: position validity checked on the previous line.
+            // Indexes still point to valid qp items because we didn't
+            // remove anything from qp yet
+            unsafe {
+                *self
+                    .qp
+                    .get_unchecked_mut(self.heap.get_unchecked(position.0).0) = position;
+            }
+        }
+
+        // swap remove from the map and return to the client
+        self.map.swap_remove_index(idx.0)
+    }
+
     /// Remove and return the element in position `position`
     /// and swap it with the last element keeping a consistent
     /// state.

tests/double_priority_queue.rs

@@ -786,12 +786,11 @@ mod doublepq_tests {
         assert_eq!(pq.pop_max(), Some(("b", 20)));
 
         /*
-        As expected, this does not compile
+        // As expected, this does not compile
         let iter_mut = pq.iter_mut();
-        iter_mut.for_each(|(_, p)| {*p += 2});
 
-        assert_eq!(pq.pop_max(), Some(("f", 9)));
-        */
+        assert_eq!(pq.pop_max(), Some(("a", 21)));
+        iter_mut.for_each(|(_, p)| {*p += 2}); */
     }
 
     #[test]
@@ -819,7 +818,6 @@ mod doublepq_tests {
         pq.push(Animal::new("bird".to_string(), true, false), 7);
         pq.push(Animal::new("fish".to_string(), false, true), 4);
         pq.push(Animal::new("cow".to_string(), false, false), 3);
-
         pq.retain(|i, _| i.can_swim);
 
         assert_eq!(
@@ -873,6 +871,71 @@ mod doublepq_tests {
         );
     }
 
+    #[test]
+    fn extract_if() {
+        #[derive(Hash, PartialEq, Eq, Debug)]
+        struct Animal {
+            name: String,
+            can_fly: bool,
+            can_swim: bool,
+        }
+
+        impl Animal {
+            pub fn new(name: String, can_fly: bool, can_swim: bool) -> Self {
+                Animal {
+                    name,
+                    can_fly,
+                    can_swim,
+                }
+            }
+        }
+
+        let mut pq = DoublePriorityQueue::new();
+        pq.push(Animal::new("dog".to_string(), false, true), 1);
+        pq.push(Animal::new("cat".to_string(), false, false), 2);
+        pq.push(Animal::new("bird".to_string(), true, false), 7);
+        pq.push(Animal::new("fish".to_string(), false, true), 4);
+        pq.push(Animal::new("cow".to_string(), false, false), 3);
+
+        let swimming_animals: Vec<(Animal, i32)> = pq
+            .extract_if(|i, p| {
+                if i.can_fly {
+                    *p -= 18;
+                    return false;
+                }
+
+                i.can_swim
+            })
+            .collect();
+
+        assert_eq!(
+            swimming_animals,
+            [
+                (Animal::new("dog".to_string(), false, true), 1),
+                (Animal::new("fish".to_string(), false, true), 4)
+            ]
+        );
+        assert_eq!(
+            pq.pop_max(),
+            Some((Animal::new("cow".to_string(), false, false), 3))
+        );
+        assert_eq!(
+            pq.pop_max(),
+            Some((Animal::new("cat".to_string(), false, false), 2))
+        );
+        assert_eq!(
+            pq.pop_max(),
+            Some((Animal::new("bird".to_string(), true, false), -11))
+        );
+
+        /*
+        // As expected, this does not compile
+        let extract_if = pq.extract_if(|i, p| { i.can_fly });
+
+        assert_eq!(pq.pop_max(), None);
+        extract_if.for_each(|(_, p)| println!("{:?}", p)); */
+    }
+
     #[test]
     fn into_sorted_iter() {
         let mut pq = DoublePriorityQueue::new();