use std::mem::ManuallyDrop;
use std::num::NonZeroUsize;
pub trait Groupable {
fn is_same_group(&self, other: &Self) -> bool;
}
#[derive(Debug, Copy, Clone, Eq, PartialEq)]
struct Index(NonZeroUsize);
impl Index {
fn new(v: usize) -> Self {
Self(NonZeroUsize::new(v + 1).unwrap())
}
const fn get(&self) -> usize {
self.0.get() - 1
}
}
#[derive(Debug)]
pub struct Queue<T>
where
T: Groupable,
{
head: Option<Index>,
free: Option<Index>,
tail: Option<Index>,
tail_leader: Option<Index>,
arr: Vec<Node<T>>,
}
#[derive(Debug)]
enum Node<T> {
Leader {
next_group: Option<Index>,
next: Option<Index>,
prev: Option<Index>,
val: ManuallyDrop<T>,
},
Normal {
next: Option<Index>,
prev: Option<Index>,
val: ManuallyDrop<T>,
},
}
impl<T> Node<T> {
fn next_group(&self) -> Option<Index> {
match self {
Node::Leader { next_group, .. } => *next_group,
Node::Normal { .. } => panic!("Method must be called on a leader"),
}
}
fn next_group_mut(&mut self) -> &mut Option<Index> {
match self {
Node::Leader { next_group, .. } => next_group,
Node::Normal { .. } => panic!("Method must be called on a leader"),
}
}
const fn next(&self) -> Option<Index> {
match self {
Node::Leader { next, .. } => *next,
Node::Normal { next, .. } => *next,
}
}
fn next_mut(&mut self) -> &mut Option<Index> {
match self {
Node::Leader { next, .. } | Node::Normal { next, .. } => next,
}
}
const fn prev(&self) -> Option<Index> {
match self {
Node::Leader { prev, .. } => *prev,
Node::Normal { prev, .. } => *prev,
}
}
fn prev_mut(&mut self) -> &mut Option<Index> {
match self {
Node::Leader { prev, .. } | Node::Normal { prev, .. } => prev,
}
}
const fn val(&self) -> &ManuallyDrop<T> {
match self {
Node::Leader { val, .. } => val,
Node::Normal { val, .. } => val,
}
}
fn val_mut(&mut self) -> &mut ManuallyDrop<T> {
match self {
Node::Leader { val, .. } => val,
Node::Normal { val, .. } => val,
}
}
}
impl<T: Groupable> Queue<T> {
pub fn new() -> Self {
Self {
head: None,
free: None,
tail: None,
tail_leader: None,
arr: vec![],
}
}
pub fn is_empty(&self) -> bool {
self.head.is_none()
}
pub fn push(&mut self, val: T) {
let mut cur = self.head;
let mut leader_idx = None;
while let Some(idx) = cur {
assert!(matches!(self.arr[idx.get()], Node::Leader {..}));
if val.is_same_group(&self.arr[idx.get()].val()) {
leader_idx = Some(idx);
break;
}
cur = self.arr[idx.get()].next_group();
}
let val = ManuallyDrop::new(val);
if let Some(idx) = leader_idx {
if let Some(next_gp) = self.arr[idx.get()].next_group() {
let node = Node::Normal {
next: Some(next_gp),
prev: self.arr[next_gp.get()].prev(),
val,
};
if let Some(free_idx) = self.free {
self.free = self.arr[free_idx.get()].next();
*self.arr[node.prev().unwrap().get()].next_mut() = Some(free_idx);
*self.arr[node.next().unwrap().get()].prev_mut() = Some(free_idx);
self.arr[free_idx.get()] = node;
} else {
let new_idx = Index::new(self.arr.len());
*self.arr[node.prev().unwrap().get()].next_mut() = Some(new_idx);
*self.arr[node.next().unwrap().get()].prev_mut() = Some(new_idx);
self.arr.push(node);
}
} else {
let node = Node::Normal {
next: None,
prev: self.tail,
val,
};
if let Some(free_idx) = self.free {
self.free = self.arr[free_idx.get()].next();
*self.arr[node.prev().unwrap().get()].next_mut() = Some(free_idx);
self.arr[free_idx.get()] = node;
self.tail = Some(free_idx);
} else {
let new_idx = Index::new(self.arr.len());
*self.arr[node.prev().unwrap().get()].next_mut() = Some(new_idx);
self.arr.push(node);
self.tail = Some(new_idx);
}
}
} else {
let node = Node::Leader {
next_group: None,
next: None,
prev: self.tail,
val,
};
if let Some(free_idx) = self.free {
self.free = self.arr[free_idx.get()].next();
if let Some(tail_ldr) = self.tail_leader {
*self.arr[tail_ldr.get()].next_group_mut() = Some(free_idx);
self.tail_leader = Some(free_idx);
*self.arr[self.tail.unwrap().get()].next_mut() = Some(free_idx);
self.tail = Some(free_idx);
self.arr[free_idx.get()] = node;
} else {
self.head = Some(free_idx);
self.tail = Some(free_idx);
self.tail_leader = Some(free_idx);
self.arr[free_idx.get()] = node;
}
} else {
let new_idx = Index::new(self.arr.len());
if let Some(tail_ldr) = self.tail_leader {
*self.arr[tail_ldr.get()].next_group_mut() = Some(new_idx);
self.tail_leader = Some(new_idx);
*self.arr[self.tail.unwrap().get()].next_mut() = Some(new_idx);
self.tail = Some(new_idx);
self.arr.push(node);
} else {
self.head = Some(new_idx);
self.tail = Some(new_idx);
self.tail_leader = Some(new_idx);
self.arr.push(node);
}
}
}
}
pub fn pop(&mut self) -> Option<T> {
if let Some(head_idx) = self.head {
if let Some(next_gp) = self.arr[head_idx.get()].next_group() {
let next_idx = self.arr[head_idx.get()].next().unwrap();
if next_idx == next_gp {
*self.arr[next_gp.get()].prev_mut() = None;
self.head = Some(next_gp);
} else {
let new_head = Node::Leader {
next_group: Some(next_gp),
next: self.arr[next_idx.get()].next(),
prev: None,
val: unsafe { std::ptr::read(self.arr[next_idx.get()].val()) },
};
self.arr[next_idx.get()] = new_head;
self.head = Some(next_idx);
}
} else {
if let Some(next_idx) = self.arr[head_idx.get()].next() {
let new_head = Node::Leader {
next_group: None,
next: self.arr[next_idx.get()].next(),
prev: None,
val: unsafe { std::ptr::read(self.arr[next_idx.get()].val()) },
};
self.arr[next_idx.get()] = new_head;
self.head = Some(next_idx);
} else {
self.head = None;
self.tail = None;
self.tail_leader = None;
}
}
*self.arr[head_idx.get()].next_mut() = self.free.take();
self.free = Some(head_idx);
unsafe { Some(ManuallyDrop::take(self.arr[head_idx.get()].val_mut())) }
} else {
None
}
}
pub fn pop_group(&mut self) -> Vec<T> {
let mut ret_vals = vec![];
if let Some(head_idx) = self.head {
if let Some(nxt_grp) = self.arr[head_idx.get()].next_group() {
self.head = Some(nxt_grp);
*self.arr[nxt_grp.get()].prev_mut() = None;
} else {
self.head = None;
self.tail = None;
self.tail_leader = None;
}
let mut cur = Some(head_idx);
while let Some(idx) = cur {
if cur == self.head {
break;
} else {
cur = self.arr[idx.get()].next();
}
let val = unsafe { ManuallyDrop::take(self.arr[idx.get()].val_mut()) };
ret_vals.push(val);
*self.arr[idx.get()].next_mut() = self.free.take();
self.free = Some(idx);
}
}
ret_vals
}
}
impl<T: Groupable> Drop for Queue<T> {
fn drop(&mut self) {
let mut cur = self.head;
while let Some(idx) = cur {
cur = self.arr[idx.get()].next();
unsafe { ManuallyDrop::drop(self.arr[idx.get()].val_mut()) }
}
unsafe { self.arr.set_len(0) }
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::cell::Cell;
#[derive(Debug)]
enum MyFoo<'a> {
Foo1(u32, &'a Cell<i32>),
Foo2(u32, &'a Cell<i32>),
Foo3(u32, &'a Cell<i32>),
}
impl Groupable for MyFoo<'_> {
fn is_same_group(&self, other: &Self) -> bool {
match (self, other) {
(MyFoo::Foo1(..), MyFoo::Foo1(..))
| (MyFoo::Foo2(..), MyFoo::Foo2(..))
| (MyFoo::Foo3(..), MyFoo::Foo3(..)) => true,
_ => false,
}
}
}
impl Drop for MyFoo<'_> {
fn drop(&mut self) {
match self {
MyFoo::Foo1(x, c) => {
c.set(c.get() - 1);
println!("Dropping Foo1 : {}\tref_count: {}", x, c.get())
}
MyFoo::Foo2(x, c) => {
c.set(c.get() - 1);
println!("Dropping Foo2 : {}\tref_count: {}", x, c.get())
}
MyFoo::Foo3(x, c) => {
c.set(c.get() - 1);
println!("Dropping Foo3 : {}\tref_count: {}", x, c.get())
}
}
}
}
impl PartialEq for MyFoo<'_> {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(MyFoo::Foo1(x, _), MyFoo::Foo1(y, _)) => x == y,
(MyFoo::Foo2(x, _), MyFoo::Foo2(y, _)) => x == y,
(MyFoo::Foo3(x, _), MyFoo::Foo3(y, _)) => x == y,
_ => false,
}
}
}
impl<'a> MyFoo<'a> {
fn new(kind: u8, val: u32, counter: &'a Cell<i32>) -> Self {
counter.set(counter.get() + 1);
match kind {
1 => Self::Foo1(val, counter),
2 => Self::Foo2(val, counter),
3 => Self::Foo3(val, counter),
_ => panic!("Unknown variant"),
}
}
}
#[test]
fn grouping() {
let ref_counter = Cell::new(0);
let mut q = Queue::<MyFoo>::new();
q.push(MyFoo::new(1, 1, &ref_counter));
q.push(MyFoo::new(2, 2, &ref_counter));
q.push(MyFoo::new(1, 3, &ref_counter));
assert_eq!(
q.pop_group(),
[
MyFoo::new(1, 1, &ref_counter),
MyFoo::new(1, 3, &ref_counter)
]
);
assert_eq!(q.pop_group(), [MyFoo::new(2, 2, &ref_counter),]);
assert_eq!(q.pop_group(), []);
drop(q);
assert_eq!(ref_counter.get(), 0);
}
}