Add bucket queue and PriorityQueueByKey

factorio-graph/src/priority_queue/bucket_queue.rs

@@ -0,0 +1,147 @@
+use std::collections::{HashMap, VecDeque};
+
+use super::{PriorityQueue, PriorityQueueByKey};
+
+pub struct BucketQueue<Item> {
+    min: i64,
+    data: VecDeque<Vec<(usize, Item)>>,
+    map: HashMap<usize, (i64, usize)>,
+    nextfree: usize,
+}
+
+impl<Item> std::fmt::Debug for BucketQueue<Item> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        struct DebugWrapper2<'a, I>(&'a Vec<(usize, I)>);
+        impl<'a, I> std::fmt::Debug for DebugWrapper2<'a, I> {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                f.debug_list().entries(self.0.iter().map(|v| v.0)).finish()
+            }
+        }
+        struct DebugWrapper<'a, I>(&'a VecDeque<Vec<(usize, I)>>);
+        impl<'a, I> std::fmt::Debug for DebugWrapper<'a, I> {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                f.debug_list()
+                    .entries(self.0.iter().map(|v| DebugWrapper2(v)))
+                    .finish()
+            }
+        }
+        f.debug_struct("BucketQueue")
+            .field("min", &self.min)
+            .field("data", &DebugWrapper::<Item>(&self.data))
+            .field("map", &self.map)
+            .field("nextfree", &self.nextfree)
+            .finish()
+    }
+}
+
+pub struct BucketQueueHandle(usize);
+
+impl<Item> PriorityQueueByKey<Item> for BucketQueue<Item>
+where
+    for<'a> &'a Item: Into<i64>,
+{
+    type Handle = BucketQueueHandle;
+
+    fn new() -> Self {
+        Self {
+            min: 0,
+            data: VecDeque::new(),
+            map: HashMap::new(),
+            nextfree: 0,
+        }
+    }
+
+    fn insert(&mut self, item: Item) -> Self::Handle {
+        self.nextfree += 1;
+        let handle = BucketQueueHandle(self.nextfree - 1);
+
+        self.insert_with_handle(item, &handle);
+
+        handle
+    }
+
+    fn pop_min(&mut self) -> Option<Item> {
+        if let Some(bucket) = self.data.get_mut(0) {
+            if let Some((i, item)) = bucket.pop() {
+                self.map.remove(&i);
+                Some(item)
+            } else {
+                self.data.pop_front();
+                self.min += 1;
+                self.pop_min()
+            }
+        } else {
+            None
+        }
+    }
+
+    fn decrease_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item)) {
+        self.update_key(handle, f);
+    }
+
+    fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item)) {
+        self.update_key(handle, f);
+    }
+}
+
+impl<Item> BucketQueue<Item>
+where
+    for<'a> &'a Item: Into<i64>,
+{
+    fn get_index(&self, i: i64) -> i64 {
+        i - self.min
+    }
+
+    fn insert_with_handle(&mut self, item: Item, handle: &BucketQueueHandle) {
+        let value = (&item).into();
+        if self.data.is_empty() {
+            self.min = value;
+            self.map.insert(handle.0, (value, 0));
+            self.data.push_back(vec![(handle.0, item)]);
+        } else {
+            let index = self.get_index(value);
+
+            if index < 0 {
+                self.min = value;
+                for _ in 0..index.abs() {
+                    self.data.push_front(vec![]);
+                }
+
+                self.map.insert(handle.0, (value, 0));
+                self.data.get_mut(0).unwrap().push((handle.0, item));
+            } else {
+                while index as usize >= self.data.len() {
+                    self.data.push_back(Vec::new());
+                }
+
+                self.map.insert(
+                    handle.0,
+                    (value, self.data.get(index as usize).unwrap().len()),
+                );
+                self.data
+                    .get_mut(index as usize)
+                    .unwrap()
+                    .push((handle.0, item));
+            }
+        }
+    }
+
+    fn remove(&mut self, handle: &BucketQueueHandle) -> Item {
+        let (i, o) = self.map.remove(&handle.0).unwrap();
+        let index = self.get_index(i);
+
+        let bucket = self.data.get_mut(index as usize).unwrap();
+        let item = bucket.swap_remove(o).1;
+
+        if let Some((swapped, _)) = bucket.get(o) {
+            self.map.insert(*swapped, (i, o));
+        }
+
+        item
+    }
+    fn update_key(&mut self, handle: &BucketQueueHandle, f: impl Fn(&mut Item)) {
+        let mut item = self.remove(handle);
+        f(&mut item);
+        self.insert_with_handle(item, handle);
+    }
+}

factorio-graph/src/priority_queue/mod.rs

@@ -1,4 +1,5 @@
 pub mod binary_heap;
+pub mod bucket_queue;
 pub mod fibonacci_heap;
 
 pub trait PriorityQueue<Item>
@@ -17,6 +18,68 @@ where
     fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item));
 }
 
+pub trait PriorityQueueByKey<Item>
+where
+    for<'a> &'a Item: Into<i64>,
+{
+    type Handle;
+    fn new() -> Self;
+
+    fn insert(&mut self, item: Item) -> Self::Handle;
+    fn pop_min(&mut self) -> Option<Item>;
+    fn decrease_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item));
+    fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item));
+}
+
+#[derive(Debug, Clone, Copy)]
+pub struct ByKey<Item>(Item);
+
+impl<Item> PartialEq for ByKey<Item>
+where
+    for<'a> &'a Item: Into<i64>,
+{
+    fn eq(&self, other: &Self) -> bool {
+        (&self.0).into().eq(&(&other.0).into())
+    }
+}
+
+impl<Item> PartialOrd for ByKey<Item>
+where
+    for<'a> &'a Item: Into<i64>,
+{
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        (&self.0).into().partial_cmp(&(&other.0).into())
+    }
+}
+
+impl<P, Item> PriorityQueueByKey<Item> for P
+where
+    P: PriorityQueue<ByKey<Item>>,
+    for<'a> &'a Item: Into<i64>,
+{
+    type Handle = P::Handle;
+
+    fn new() -> Self {
+        P::new()
+    }
+
+    fn insert(&mut self, item: Item) -> Self::Handle {
+        self.insert(ByKey(item))
+    }
+
+    fn pop_min(&mut self) -> Option<Item> {
+        self.pop_min().map(|r| r.0)
+    }
+
+    fn decrease_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item)) {
+        self.decrease_key(handle, |i| f(&mut i.0));
+    }
+
+    fn increase_key(&mut self, handle: &Self::Handle, f: impl Fn(&mut Item)) {
+        self.increase_key(handle, |i| f(&mut i.0));
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::PriorityQueue;

factorio-graph/src/wheighted_graph/shortest_path.rs

@@ -2,7 +2,7 @@ use std::{collections::HashMap, fmt::Debug, hash::Hash, hash::Hasher};
 
 use tracing::{field::Empty, trace_span};
 
-use crate::priority_queue::PriorityQueue;
+use crate::priority_queue::{PriorityQueue, PriorityQueueByKey};
 
 use super::WheightedGraph;
 
@@ -18,6 +18,12 @@ impl<N> QueueObject<N> {
     }
 }
 
+impl<N> Into<i64> for &QueueObject<N> {
+    fn into(self) -> i64 {
+        self.score
+    }
+}
+
 impl<N> PartialOrd for QueueObject<N> {
     fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
         self.score.partial_cmp(&other.score)
@@ -58,9 +64,8 @@ impl<N, H> MapObject<N, H> {
 
 pub fn dijkstra<G, P, E>(graph: &G, start: G::Node, end: E) -> Option<Vec<G::Node>>
 where
-    P: PriorityQueue<QueueObject<G::Node>> + Debug,
-    P::Handle: Debug,
-    G::Node: Eq + Hash + Clone + Debug,
+    P: PriorityQueueByKey<QueueObject<G::Node>>,
+    G::Node: Eq + Hash + Clone,
     G: WheightedGraph,
     E: Fn(&'_ G::Node) -> bool,
 {
@@ -159,9 +164,8 @@ impl<G: WheightedGraph, F: Fn(&G::Node) -> i64> WheightedGraph for GraphWrapper<
 
 pub fn a_star<G, P, E, F>(graph: &G, start: G::Node, end: E, dist: F) -> Option<Vec<G::Node>>
 where
-    P: PriorityQueue<QueueObject<G::Node>> + Debug,
-    P::Handle: Debug,
-    G::Node: Eq + Hash + Clone + Debug,
+    P: PriorityQueueByKey<QueueObject<G::Node>>,
+    G::Node: Eq + Hash + Clone,
     G: WheightedGraph,
     E: Fn(&'_ G::Node) -> bool,
     F: Fn(&G::Node) -> i64,

factorio-graph/src/wheighted_graph/steiner_tree.rs

@@ -1,6 +1,6 @@
 use tracing::trace_span;
 
-use crate::priority_queue::PriorityQueue;
+use crate::priority_queue::{PriorityQueue, PriorityQueueByKey};
 use std::hash::Hash;
 use std::{collections::HashSet, fmt::Debug};
 
@@ -32,9 +32,8 @@ where
 
 pub fn takaheshi_matsuyama<G, P>(graph: &G, nodes: &[G::Node]) -> Option<Vec<Vec<G::Node>>>
 where
-    P: PriorityQueue<QueueObject<Option<G::Node>>> + Debug,
-    P::Handle: Debug,
-    G::Node: Eq + Hash + Clone + Debug,
+    P: PriorityQueueByKey<QueueObject<Option<G::Node>>>,
+    G::Node: Eq + Hash + Clone,
     G: WheightedGraph,
 {
     let _complete_span = trace_span!("takaheshi_matsuyama").entered();