mea 0.6.3

// Copyright 2024 tison <wander4096@gmail.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use slab::Slab;

/// A guarded linked list.
///
/// * `guard`'s `next` points to the first node (regular head).
/// * `guard`'s `prev` points to the last node (regular tail).
#[derive(Debug)]
pub(crate) struct WaitList<T> {
    // if None, the list is uninitialized and empty
    guard: Option<usize>,
    nodes: Slab<Node<T>>,
}

#[derive(Debug)]
struct Node<T> {
    prev: usize,
    next: usize,
    stat: Option<T>,
}

impl<T> WaitList<T> {
    /// Ensures the wait list is initialized, returning the guard index.
    fn ensure_init(&mut self) -> usize {
        if let Some(guard) = self.guard {
            return guard;
        }

        let first = self.nodes.vacant_entry();
        let guard = first.key();
        first.insert(Node {
            prev: guard,
            next: guard,
            stat: None,
        });
        self.guard = Some(guard);
        guard
    }

    pub(crate) const fn new() -> Self {
        Self {
            guard: None,
            nodes: Slab::new(),
        }
    }

    /// Registers a waiter to the head of the wait list.
    ///
    /// # Panic
    ///
    /// Panics if `idx` is `Some`.
    pub(crate) fn register_waiter_to_head(
        &mut self,
        idx: &mut Option<usize>,
        f: impl FnOnce() -> Option<T>,
    ) {
        assert!(idx.is_none());

        let guard = self.ensure_init();
        let stat = f();
        let prev_head = self.nodes[guard].next;
        let new_node = Node {
            prev: guard,
            next: prev_head,
            stat,
        };
        let new_key = self.nodes.insert(new_node);
        self.nodes[guard].next = new_key;
        self.nodes[prev_head].prev = new_key;
        *idx = Some(new_key);
    }

    /// Registers a waiter to the tail of the wait list.
    ///
    /// # Panic
    ///
    /// Panics if `idx` is `Some`.
    pub(crate) fn register_waiter_to_tail(
        &mut self,
        idx: &mut Option<usize>,
        f: impl FnOnce() -> Option<T>,
    ) {
        assert!(idx.is_none());

        let guard = self.ensure_init();
        let stat = f();
        let prev_tail = self.nodes[guard].prev;
        let new_node = Node {
            prev: prev_tail,
            next: guard,
            stat,
        };
        let new_key = self.nodes.insert(new_node);
        self.nodes[guard].prev = new_key;
        self.nodes[prev_tail].next = new_key;
        *idx = Some(new_key);
    }

    /// Removes a previously registered waker from the wait list, if the predicate `f` returns
    /// `true`.
    pub(crate) fn remove_waiter(
        &mut self,
        idx: usize,
        f: impl FnOnce(&mut T) -> bool,
    ) -> Option<&mut T> {
        // SAFETY: the wait list must be initialized before any waiter can be registered
        let guard = self.guard.expect("wait list must be uninitialized");

        assert_ne!(idx, guard);

        fn retrieve_stat<T>(node: &mut Node<T>) -> &mut T {
            // SAFETY: `idx` is a valid key + non-guard node always has `Some(stat)`
            node.stat.as_mut().unwrap()
        }

        if f(retrieve_stat(&mut self.nodes[idx])) {
            let prev = self.nodes[idx].prev;
            let next = self.nodes[idx].next;
            self.nodes[prev].next = next;
            self.nodes[next].prev = prev;
            self.nodes[idx].prev = idx;
            self.nodes[idx].next = idx;
            Some(retrieve_stat(&mut self.nodes[idx]))
        } else {
            None
        }
    }

    /// Removes the first waiter from the wait list, if the predicate `f` returns `true`.
    pub(crate) fn remove_first_waiter(&mut self, f: impl FnOnce(&mut T) -> bool) -> Option<&mut T> {
        let guard = self.guard?;
        let first = self.nodes[guard].next;
        if first != guard {
            self.remove_waiter(first, f)
        } else {
            None
        }
    }

    /// Returns `true` if the wait list is empty.
    pub(crate) fn is_empty(&self) -> bool {
        self.guard
            .is_none_or(|guard| self.nodes[guard].next == guard)
    }

    pub(crate) fn with_mut(&mut self, idx: usize, drop: impl FnOnce(&mut T) -> bool) {
        let node = &mut self.nodes[idx];
        if drop(node.stat.as_mut().unwrap()) {
            self.nodes.remove(idx);
        }
    }
}