use std::cmp::Ordering::{Equal, Greater, Less};
use std::fmt;
use std::mem::{forget, needs_drop};
use std::ops::{Bound, RangeBounds};
use std::ptr::from_ref;
use std::sync::atomic::Ordering::{Acquire, Relaxed, Release};

use saa::Lock;
use sdd::{AtomicRaw, Owned, PrivateCollector, RawPtr};

use super::Leaf;
use super::leaf::{
    Array, ArrayIter, DIMENSION, InsertResult, Iter, RemoveResult, RevIter, range_contains,
};
use super::node::Node;
use crate::exit_guard::ExitGuard;
use crate::utils::{
    LockPager, deref_unchecked, drop_in_place, fake_ref, get_owned, retire, snapshot, undefined,
    unwrap_unchecked,
};
use crate::{Comparable, Guard};

/// [`LeafNode`] is an ordered collection or [`Leaf`].
#[repr(C)]
pub struct LeafNode<K, V> {
    /// A child [`Leaf`] that has no upper key bound.
    ///
    /// It stores the maximum key in the node, and key-value pairs are first pushed to this
    /// [`Leaf`] until it splits.
    pub(super) unbounded_child: AtomicRaw<Leaf<K, V>>,
    /// Children of the [`LeafNode`].
    pub(super) children: Array<K, AtomicRaw<Leaf<K, V>>>,
    /// [`Lock`] to protect the [`LeafNode`].
    pub(super) lock: Lock,
}

/// [`Locker`] holds exclusive ownership of a [`LeafNode`].
pub(super) struct Locker<'n, K, V> {
    pub(super) node: &'n LeafNode<K, V>,
}

/// A state machine to keep track of the progress of a bulk removal operation.
#[derive(Clone, Copy, Eq, PartialEq)]
pub(super) enum RemoveRangeState {
    /// The maximum key of the node is less than the start bound of the range.
    Below,
    /// The maximum key of the node is contained in the range, but it is not clear whether the
    /// minimum key of the node is contained in the range.
    MaybeBelow,
    /// The maximum key and the minimum key of the node are contained in the range.
    FullyContained,
    /// The maximum key of the node is not contained in the range, but it is not clear whether
    /// the minimum key of the node is contained in the range.
    MaybeAbove,
}

impl<K, V> LeafNode<K, V> {
    /// Creates a new empty [`LeafNode`].
    #[inline]
    pub(super) fn new() -> LeafNode<K, V> {
        LeafNode {
            unbounded_child: AtomicRaw::new(unsafe {
                Owned::new_with_unchecked(Leaf::new).into_raw()
            }),
            children: Array::new(),
            lock: Lock::default(),
        }
    }

    /// Creates a new empty [`LeafNode`] in a frozen state.
    #[inline]
    pub(super) fn new_frozen() -> LeafNode<K, V> {
        LeafNode {
            unbounded_child: AtomicRaw::null(),
            children: Array::new_frozen(),
            lock: Lock::default(),
        }
    }

    /// Returns `true` if the [`LeafNode`] has retired.
    #[inline]
    pub(super) fn is_retired(&self) -> bool {
        self.lock.is_poisoned(Acquire)
    }
}

impl<K, V> LeafNode<K, V>
where
    K: Clone + Ord,
{
    /// Searches for an entry containing the specified key.
    #[inline]
    pub(super) fn search_entry<'g, Q>(&'g self, key: &Q, guard: &'g Guard) -> Option<(&'g K, &'g V)>
    where
        Q: Comparable<K> + ?Sized,
    {
        loop {
            let (child, metadata) = self.children.min_greater_equal(key);
            if let Some(child) = child {
                if let Some(child) = deref_unchecked(child.load(Acquire, guard)) {
                    if self.children.validate(metadata) {
                        return child.search_entry(key);
                    }
                }

                // The child leaf must have been just removed.
                //
                // The `LeafNode` metadata is updated before a leaf is removed. This implies that
                // the new `metadata` will be different from the current `metadata`.
            } else if let Some(unbounded) =
                deref_unchecked(self.unbounded_child.load(Acquire, guard))
            {
                if self.children.validate(metadata) {
                    return unbounded.search_entry(key);
                }
            } else {
                return None;
            }
        }
    }

    /// Searches for the value associated with the specified key.
    #[inline]
    pub(super) fn search_value<'g, Q>(&'g self, key: &Q, guard: &'g Guard) -> Option<&'g V>
    where
        Q: Comparable<K> + ?Sized,
    {
        loop {
            let (child, metadata) = self.children.min_greater_equal(key);
            if let Some(child) = child {
                if let Some(child) = deref_unchecked(child.load(Acquire, guard)) {
                    if self.children.validate(metadata) {
                        return child.search_val(key);
                    }
                }
            } else if let Some(unbounded) =
                deref_unchecked(self.unbounded_child.load(Acquire, guard))
            {
                if self.children.validate(metadata) {
                    return unbounded.search_val(key);
                }
            } else {
                return None;
            }
        }
    }

    /// Reads an entry using the supplied closure.
    #[inline]
    pub fn read_entry<Q, R, F: FnOnce(&K, &V) -> R, P: LockPager>(
        &self,
        key: &Q,
        reader: F,
        pager: &mut P,
        guard: &Guard,
    ) -> Result<Option<R>, F>
    where
        Q: Comparable<K> + ?Sized,
    {
        loop {
            match pager.try_acquire::<true>(&self.lock) {
                Ok(true) => break,
                Ok(false) => return Err(reader),
                Err(()) => {
                    if !pager.try_wait::<true>(&self.lock) {
                        return Err(reader);
                    }
                }
            }
        }
        let _locker = ExitGuard::new((), |()| {
            self.lock.release_share();
        });

        let (child, _) = self.children.min_greater_equal(key);
        if let Some(child) = child.and_then(|c| deref_unchecked(c.load(Acquire, guard))) {
            if let Some((k, v)) = child.search_entry(key) {
                return Ok(Some(reader(k, v)));
            }
        } else if let Some(unbounded) = deref_unchecked(self.unbounded_child.load(Acquire, guard)) {
            if let Some((k, v)) = unbounded.search_entry(key) {
                return Ok(Some(reader(k, v)));
            }
        }
        Ok(None)
    }

    /// Returns an [`Iter`] pointing to the left-most leaf in the entire tree.
    #[inline]
    pub(super) fn min<'g>(&self, guard: &'g Guard) -> Option<Iter<'g, K, V>> {
        let mut min_leaf = None;
        for i in ArrayIter::new(&self.children) {
            if let Some(child) = deref_unchecked(self.children.val(i).load(Acquire, guard)) {
                min_leaf.replace(child);
                break;
            }
        }
        if min_leaf.is_none() {
            if let Some(unbounded) = deref_unchecked(self.unbounded_child.load(Acquire, guard)) {
                min_leaf.replace(unbounded);
            }
        }

        let Some(min_leaf) = min_leaf else {
            // `unbounded_child` being null means that the leaf was retired of empty.
            return None;
        };

        let mut rev_iter = RevIter::new(min_leaf);
        while rev_iter.jump(guard).is_some() {}
        rev_iter.rewind();
        Some(rev_iter.rev())
    }

    /// Returns a [`RevIter`] pointing to the right-most leaf in the entire tree.
    #[inline]
    pub(super) fn max<'g>(&self, guard: &'g Guard) -> Option<RevIter<'g, K, V>> {
        if let Some(unbounded) = deref_unchecked(self.unbounded_child.load(Acquire, guard)) {
            let mut iter = Iter::new(unbounded);
            while iter.jump(guard).is_some() {}
            iter.rewind();
            return Some(iter.rev());
        }
        // `unbounded_child` being null means that the leaf was retired of empty.
        None
    }

    /// Returns a [`Iter`] pointing to an entry that is close enough to the specified key.
    #[inline]
    pub(super) fn approximate<'g, Q, const LE: bool>(
        &self,
        key: &Q,
        guard: &'g Guard,
    ) -> Option<Iter<'g, K, V>>
    where
        Q: Comparable<K> + ?Sized,
    {
        let leaf = loop {
            let (child, metadata) = self.children.min_greater_equal(key);
            if let Some(child) = child {
                if self.children.validate(metadata) {
                    if let Some(child) = deref_unchecked(child.load(Acquire, guard)) {
                        break child;
                    }
                }
                // It is not a hot loop - see `Array::validate`.
                continue;
            }
            if let Some(unbounded) = deref_unchecked(self.unbounded_child.load(Acquire, guard)) {
                if self.children.validate(metadata) {
                    break unbounded;
                }
                continue;
            }
            // `unbounded_child` being null means that the leaf was retired of empty.
            return None;
        };

        // Tries to find "any" leaf that contains a reachable entry.
        let origin = Iter::new(leaf);
        let mut iter = origin.clone();
        if iter.next().is_none() && iter.jump(guard).is_none() {
            let mut rev_iter = origin.rev();
            if rev_iter.jump(guard).is_some() {
                iter = rev_iter.rev();
            } else {
                return None;
            }
        }
        iter.rewind();

        if LE {
            while let Some((k, _)) = iter.next() {
                if let Equal | Greater = key.compare(k) {
                    return Some(iter);
                }
                // Go to the prev leaf node that shall contain smaller keys.
                let mut rev_iter = iter.rev();
                rev_iter.jump(guard)?;
                iter = rev_iter.rev();
                // Rewind the iterator to point to the smallest key in the leaf.
                iter.rewind();
            }
        } else {
            let mut rev_iter = iter.rev();
            while let Some((k, _)) = rev_iter.next() {
                if let Less | Equal = key.compare(k) {
                    return Some(rev_iter.rev());
                }
                // Go to the next leaf node that shall contain larger keys.
                iter = rev_iter.rev();
                iter.jump(guard)?;
                rev_iter = iter.rev();
                // Rewind the iterator to point to the largest key in the leaf.
                rev_iter.rewind();
            }
        }

        // Reached the end of the linked list.
        None
    }

    /// Inserts a key-value pair.
    ///
    /// # Errors
    ///
    /// Returns an error if a retry is required.
    #[inline]
    pub(super) fn insert<P: LockPager, const UPSERT: bool>(
        &self,
        mut key: K,
        mut val: V,
        private_collector: &PrivateCollector,
        pager: &mut P,
        guard: &Guard,
    ) -> Result<InsertResult<K, V>, (K, V)> {
        loop {
            let (child, metadata) = self.children.min_greater_equal(&key);
            if let Some(child) = child {
                let child_ptr = child.load(Acquire, guard);
                if let Some(child_ref) = deref_unchecked(child_ptr) {
                    if self.children.validate(metadata) {
                        let result = child_ref.insert::<UPSERT>(key, val);
                        match result {
                            InsertResult::Duplicate(..) if UPSERT => undefined(),
                            InsertResult::Success | InsertResult::Duplicate(..) => {
                                return Ok(result);
                            }
                            InsertResult::Full(k, v) => {
                                match self.split_insert::<P, UPSERT>(
                                    child_ptr,
                                    child,
                                    k,
                                    v,
                                    pager,
                                    private_collector,
                                    guard,
                                ) {
                                    Ok(result) => return Ok(result),
                                    Err((k, v, false)) => return Err((k, v)),
                                    Err((k, v, true)) => (key, val) = (k, v),
                                }
                            }
                        }
                    }
                }
                // It is not a hot loop - see `Array::validate`.
                continue;
            }

            let unbounded_ptr = self.unbounded_child.load(Acquire, guard);
            if let Some(unbounded) = deref_unchecked(unbounded_ptr) {
                if self.children.validate(metadata) {
                    let result = unbounded.insert::<UPSERT>(key, val);
                    match result {
                        InsertResult::Duplicate(..) if UPSERT => undefined(),
                        InsertResult::Success | InsertResult::Duplicate(..) => return Ok(result),
                        InsertResult::Full(k, v) => {
                            match self.split_insert::<P, UPSERT>(
                                unbounded_ptr,
                                &self.unbounded_child,
                                k,
                                v,
                                pager,
                                private_collector,
                                guard,
                            ) {
                                Ok(result) => return Ok(result),
                                Err((k, v, false)) => return Err((k, v)),
                                Err((k, v, true)) => (key, val) = (k, v),
                            }
                        }
                    }
                }
                continue;
            }

            return Ok(InsertResult::Full(key, val));
        }
    }

    /// Removes an entry associated with the given key.
    ///
    /// # Errors
    ///
    /// Returns an error if a retry is required.
    #[inline]
    pub(super) fn remove_if<Q, F: FnMut(&V) -> bool, P: LockPager>(
        &self,
        key: &Q,
        condition: &mut F,
        private_collector: &PrivateCollector,
        pager: &mut P,
        guard: &Guard,
    ) -> Result<RemoveResult, ()>
    where
        Q: Comparable<K> + ?Sized,
    {
        loop {
            let (child, metadata) = self.children.min_greater_equal(key);
            if let Some(child) = child {
                if let Some(child) = deref_unchecked(child.load(Acquire, guard)) {
                    if self.children.validate(metadata) {
                        let result = child.remove_if(key, condition);
                        if result == RemoveResult::Frozen {
                            // Its entries may be being relocated.
                            if pager.try_wait::<false>(&self.lock)
                                && !self.children.validate(metadata)
                            {
                                continue;
                            }
                            return Err(());
                        } else if result == RemoveResult::Retired {
                            return Ok(self.post_remove(private_collector, guard));
                        }
                        return Ok(result);
                    }
                }
                // It is not a hot loop - see `Array::validate`.
                continue;
            }
            if let Some(unbounded) = deref_unchecked(self.unbounded_child.load(Acquire, guard)) {
                if self.children.validate(metadata) {
                    let result = unbounded.remove_if(key, condition);
                    if result == RemoveResult::Frozen {
                        if pager.try_wait::<false>(&self.lock) && !self.children.validate(metadata)
                        {
                            continue;
                        }
                        return Err(());
                    } else if result == RemoveResult::Retired {
                        return Ok(self.post_remove(private_collector, guard));
                    }
                    return Ok(result);
                }
                continue;
            }
            return Ok(RemoveResult::Fail);
        }
    }

    /// Removes a range of entries.
    ///
    /// Returns the number of remaining children.
    #[allow(clippy::too_many_arguments)]
    #[inline]
    pub(super) fn remove_range<'g, Q, R: RangeBounds<Q>, P: LockPager>(
        &self,
        range: &R,
        start_unbounded: bool,
        valid_lower_max_leaf: Option<&'g Leaf<K, V>>,
        valid_upper_min_node: Option<&'g Node<K, V>>,
        private_collector: &PrivateCollector,
        pager: &mut P,
        guard: &'g Guard,
    ) -> Result<usize, ()>
    where
        Q: Comparable<K> + ?Sized,
    {
        debug_assert!(valid_lower_max_leaf.is_none() || start_unbounded);
        debug_assert!(valid_lower_max_leaf.is_none() || valid_upper_min_node.is_none());

        let _locker = if pager.try_acquire::<false>(&self.lock)? {
            Locker { node: self }
        } else {
            // The leaf node was retired: retry.
            return Err(());
        };

        let mut current_state = RemoveRangeState::Below;
        let mut num_leaves = 1;
        let mut min_max_leaf = None;

        for i in ArrayIter::new(&self.children) {
            current_state = current_state.next(self.children.key(i), range, start_unbounded);
            let child = self.children.val(i);
            match current_state {
                RemoveRangeState::Below | RemoveRangeState::MaybeBelow => {
                    if let Some(leaf) = deref_unchecked(child.load(Acquire, guard)) {
                        leaf.remove_range(range);
                    }
                    num_leaves += 1;
                    if min_max_leaf.is_none() {
                        min_max_leaf.replace(child);
                    }
                }
                RemoveRangeState::FullyContained => {
                    // There can be another thread inserting keys into the leaf, and this may render
                    // those operations completely ineffective.
                    self.children.remove_unchecked(self.children.metadata(), i);
                    if let Some(leaf) = get_owned(child.load(Acquire, guard)) {
                        child.store(RawPtr::null(), Release);
                        leaf.unlink();
                        if cfg!(feature = "loom") || needs_drop::<(K, V)>() {
                            unsafe {
                                private_collector.collect_owned(leaf, guard);
                            }
                        }
                    }
                }
                RemoveRangeState::MaybeAbove => {
                    if let Some(leaf) = deref_unchecked(child.load(Acquire, guard)) {
                        leaf.remove_range(range);
                    }
                    num_leaves += 1;
                    if min_max_leaf.is_none() {
                        min_max_leaf.replace(child);
                    }
                    break;
                }
            }
        }

        if let Some(unbounded) = deref_unchecked(self.unbounded_child.load(Acquire, guard)) {
            unbounded.remove_range(range);
        }

        if let Some(valid_lower_max_leaf) = valid_lower_max_leaf {
            // Splices the max min leaf with the min max leaf.
            let min_max = min_max_leaf
                .unwrap_or(&self.unbounded_child)
                .load(Acquire, guard);
            Leaf::<K, V>::splice_link(Some(valid_lower_max_leaf), deref_unchecked(min_max), guard);
        } else if let Some(valid_upper_min_node) = valid_upper_min_node {
            // Connect the unbounded child with the minimum valid leaf in the node.
            valid_upper_min_node.remove_range(
                range,
                true,
                deref_unchecked(self.unbounded_child.load(Acquire, guard)),
                None,
                private_collector,
                pager,
                guard,
            )?;
        } else if start_unbounded {
            // `min_max_leaf` becomes the first leaf in the entire tree.
            let min_max = min_max_leaf
                .unwrap_or(&self.unbounded_child)
                .load(Acquire, guard);
            Leaf::<K, V>::splice_link(None, deref_unchecked(min_max), guard);
        }

        Ok(num_leaves)
    }

    /// Splits a full leaf and inserts the new key-value pair into a new leaf.
    #[allow(clippy::too_many_arguments, clippy::too_many_lines)]
    fn split_insert<P: LockPager, const UPSERT: bool>(
        &self,
        target_ptr: RawPtr<Leaf<K, V>>,
        target_ref: &AtomicRaw<Leaf<K, V>>,
        key: K,
        val: V,
        pager: &mut P,
        private_collector: &PrivateCollector,
        guard: &Guard,
    ) -> Result<InsertResult<K, V>, (K, V, bool)> {
        let Some(_locker) = Locker::try_lock(self) else {
            if self.is_retired() {
                // Let the parent node clean up this node.
                return Ok(InsertResult::Full(key, val));
            } else if pager.try_wait::<false>(&self.lock) {
                // Do not wait-and-acquire the lock as it is most likely that the leaf has already
                // been split by the time it acquires the lock.
                return Err((key, val, true));
            }
            return Err((key, val, false));
        };
        if target_ptr != target_ref.load(Relaxed, guard) {
            return Err((key, val, true));
        }

        let (mut low_key_leaf, mut high_key_leaf) = (None, None);
        let mut boundary_pos = 0;
        let mut downgrade = false;
        let target = unwrap_unchecked(deref_unchecked(target_ptr));
        let exit_guard = ExitGuard::new((), |()| target.unfreeze());
        if !target.distribute::<_, _, false>(
            #[inline]
            |boundary, len, ordered| {
                // E.g., `boundary == 2, len == 2`, then `i` can be as large as `1`: `high_key_leaf`
                // is not needed.
                if (boundary as usize) < len && self.children.is_full() {
                    // No space for new leaves.
                    false
                } else if ordered
                    && self.unbounded_child.load(Relaxed, guard) == target_ptr
                    && target.key(DIMENSION.ideal_max_key_pos()) < &key
                {
                    // Heuristics to avoid a leaf split.
                    debug_assert!(!self.children.is_full());
                    downgrade = true;
                    false
                } else {
                    true
                }
            },
            #[inline]
            |k, v, boundary, i| {
                // `v` is moved, not cloned; those new leaves do not own them until unfrozen.
                if i < boundary {
                    let low_key_leaf = low_key_leaf.get_or_insert_with(|| unsafe {
                        Owned::new_with_unchecked(Leaf::new_frozen)
                    });
                    boundary_pos = low_key_leaf.insert_unchecked(snapshot(k), snapshot(v), i);
                } else {
                    let high_key_leaf = high_key_leaf.get_or_insert_with(|| unsafe {
                        Owned::new_with_unchecked(Leaf::new_frozen)
                    });
                    high_key_leaf.insert_unchecked(snapshot(k), snapshot(v), i - boundary);
                }
            },
        ) {
            if downgrade {
                // Heuristics based on the assumption that the input data is ordered if the current
                // input is larger than the maximum key of the full leaf and entries in the full
                // leaf are perfectly ordered.
                let max_key = target.key(DIMENSION.ideal_max_key_pos()).clone();
                let new_leaf_ptr =
                    unsafe { Owned::new_with_unchecked(Leaf::new_frozen).into_raw() };
                let new_leaf = unwrap_unchecked(deref_unchecked(new_leaf_ptr));
                new_leaf.insert_unchecked(key, val, 0);
                new_leaf.unfreeze();
                new_leaf.prev.store(from_ref(target).cast_mut(), Relaxed);
                new_leaf.lock.lock_sync();

                target.update_link(
                    #[inline]
                    |_, next_prev, _, next_ptr| {
                        new_leaf.next.store(next_ptr.cast_mut(), Relaxed);
                        target.next.store(from_ref(new_leaf).cast_mut(), Release);
                        if let Some(next_prev) = next_prev {
                            next_prev.store(from_ref(new_leaf).cast_mut(), Release);
                        }
                    },
                );
                let result = self
                    .children
                    .insert::<false>(max_key, AtomicRaw::new(target_ptr));
                debug_assert!(matches!(result, InsertResult::Success));

                self.unbounded_child.store(new_leaf_ptr, Release);
                let released = new_leaf.lock.release_lock();
                debug_assert!(released);
                return Ok(InsertResult::Success);
            }
            return Ok(InsertResult::Full(key, val));
        }

        // Data race with iterators if the following code is executed without new leaves locked.
        // - T1 and T2 both observe, L1 -> L2.
        // - T2 splits L1 into L1_1 and L1_2: L1_1 <-> L1_2 (not reachable via tree) <-> L2.
        // - T1 splits L2 into L2_1 and L2_2: L1_1 <-> L1_2 <-> L2_1 <-> L2_2.
        // - T1 inserts entries into L2_1.
        // - T1 range queries get L1, instead of L1_2.
        // - T1 iterates over entries from L1 and L2, and cannot see entries in L2_1.
        //
        // The locking prevents T1 from splitting L2 until L1_2 becomes reachable via tree.
        let low_key_leaf =
            low_key_leaf.unwrap_or_else(|| unsafe { Owned::new_with_unchecked(Leaf::new_frozen) });
        low_key_leaf.lock.lock_sync();
        let low_key_leaf_lock = &low_key_leaf.get_guarded_ref(guard).lock;

        let result = if let Some(high_key_leaf) = high_key_leaf {
            let low_key_max = low_key_leaf.key(boundary_pos).clone();

            // Unfreeze the leaves; those leaves now take ownership of the copied values.
            low_key_leaf.unfreeze();
            high_key_leaf.unfreeze();
            let result = if low_key_max < key {
                high_key_leaf.insert::<UPSERT>(key, val)
            } else {
                low_key_leaf.insert::<UPSERT>(key, val)
            };
            debug_assert!(!matches!(result, InsertResult::Full(..)));

            low_key_leaf
                .next
                .store(high_key_leaf.as_ptr().cast_mut(), Relaxed);
            high_key_leaf
                .prev
                .store(low_key_leaf.as_ptr().cast_mut(), Relaxed);
            let high_key_leaf_lock = &high_key_leaf.get_guarded_ref(guard).lock;
            high_key_leaf_lock.lock_sync();

            target.update_link(
                #[inline]
                |prev_next, next_prev, prev_ptr, next_ptr| {
                    low_key_leaf.prev.store(prev_ptr.cast_mut(), Relaxed);
                    high_key_leaf.next.store(next_ptr.cast_mut(), Relaxed);
                    if let Some(prev_next) = prev_next {
                        prev_next.store(low_key_leaf.as_ptr().cast_mut(), Release);
                    }
                    if let Some(next_prev) = next_prev {
                        next_prev.store(high_key_leaf.as_ptr().cast_mut(), Release);
                    }
                },
            );

            // Take the max key-value stored in the low key leaf as the leaf key.
            let inserted = self
                .children
                .insert::<false>(low_key_max, AtomicRaw::new(low_key_leaf.into_raw()));
            debug_assert!(matches!(inserted, InsertResult::Success));
            let released = low_key_leaf_lock.release_lock();
            debug_assert!(released);

            debug_assert!(target_ref.load(Relaxed, guard) == target_ptr);
            target_ref.store(high_key_leaf.into_raw(), Release);
            let released = high_key_leaf_lock.release_lock();
            debug_assert!(released);

            result
        } else {
            // Unfreeze the leaf; it now takes ownership of the copied values.
            low_key_leaf.unfreeze();
            let result = low_key_leaf.insert::<UPSERT>(key, val);
            debug_assert!(!matches!(result, InsertResult::Full(..)));

            target.update_link(
                #[inline]
                |prev_next, next_prev, prev_ptr, next_ptr| {
                    low_key_leaf.prev.store(prev_ptr.cast_mut(), Relaxed);
                    low_key_leaf.next.store(next_ptr.cast_mut(), Relaxed);
                    if let Some(prev_next) = prev_next {
                        prev_next.store(low_key_leaf.as_ptr().cast_mut(), Release);
                    }
                    if let Some(next_prev) = next_prev {
                        next_prev.store(low_key_leaf.as_ptr().cast_mut(), Release);
                    }
                },
            );

            debug_assert!(target_ref.load(Relaxed, guard) == target_ptr);
            target_ref.store(low_key_leaf.into_raw(), Release);
            let released = low_key_leaf_lock.release_lock();
            debug_assert!(released);

            result
        };

        // The removed leaf stays frozen: ownership of the copied values is transferred.
        exit_guard.forget();
        retire::<(K, V), _>(target_ptr, private_collector, guard);

        Ok(result)
    }

    /// Tries to delete retired leaves after a successful removal of an entry.
    fn post_remove(&self, private_collector: &PrivateCollector, guard: &Guard) -> RemoveResult {
        let Some(lock) = Locker::try_lock(self) else {
            if self.is_retired() {
                return RemoveResult::Retired;
            }
            return RemoveResult::Success;
        };

        let mut fully_empty = true;
        for i in ArrayIter::new(&self.children) {
            let leaf = self.children.val(i);
            let leaf_ptr = leaf.load(Acquire, guard);
            let leaf_ref = unwrap_unchecked(deref_unchecked(leaf_ptr));
            if leaf_ref.is_retired() {
                leaf_ref.unlink();

                // As soon as the leaf is removed from the leaf node, the next leaf can store keys
                // that are smaller than those that were previously stored in the removed leaf node.
                //
                // Iterators cope with this by checking the prev/next pointers; right after
                // `unlink`, the prev/next leaves will not point to this leaf anymore.
                let result = self.children.remove_unchecked(self.children.metadata(), i);
                debug_assert_ne!(result, RemoveResult::Fail);
                leaf.store(RawPtr::null(), Release);
                retire::<(K, V), _>(leaf_ptr, private_collector, guard);
            } else {
                fully_empty = false;
            }
        }

        // The unbounded leaf becomes unreachable after all the other leaves are gone.
        if fully_empty {
            let unbounded_ptr = self.unbounded_child.load(Acquire, guard);
            if let Some(unbounded) = deref_unchecked(unbounded_ptr) {
                if unbounded.is_retired() {
                    unbounded.unlink();
                    self.unbounded_child.store(RawPtr::null(), Release);
                    retire::<(K, V), _>(unbounded_ptr, private_collector, guard);
                } else {
                    fully_empty = false;
                }
            }
        }

        if fully_empty {
            lock.unlock_retire();
            RemoveResult::Retired
        } else {
            RemoveResult::Success
        }
    }
}

impl<K, V> fmt::Debug for LeafNode<K, V>
where
    K: Clone + fmt::Debug + Ord,
    V: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let guard = Guard::new();
        f.write_str("LeafNode { ")?;
        let mut state = (false, false);
        self.children.for_each_all(
            |frozen, retired| {
                state.0 = frozen;
                state.1 = retired;
                true
            },
            |i, rank, entry, removed| {
                if let Some((k, l)) = entry {
                    if let Some(l) = deref_unchecked(l.load(Acquire, &guard)) {
                        write!(f, "{i}: ({k:?}, {rank}, removed: {removed}, {l:?}), ")?;
                    } else {
                        write!(f, "{i}: ({k:?}, {rank}, removed: {removed}, null), ")?;
                    }
                }
                Ok(())
            },
        )?;
        if let Some(unbounded) = deref_unchecked(self.unbounded_child.load(Acquire, &guard)) {
            write!(f, "unbounded: {unbounded:?}, ")?;
        } else {
            write!(f, "unbounded: null, ")?;
        }
        write!(f, "frozen: {}, ", state.0)?;
        write!(f, "retired: {}", state.1)?;
        f.write_str(" }")
    }
}

impl<K, V> Drop for LeafNode<K, V> {
    #[inline]
    fn drop(&mut self) {
        // Need to manually drop each leaf node.
        self.children.for_each_pos(
            #[inline]
            |frozen| {
                if frozen {
                    // This internal node has no ownership of any nodes.
                    return false;
                }
                drop_in_place(self.unbounded_child.load(Acquire, fake_ref(self)));
                true
            },
            #[inline]
            |pos| {
                drop_in_place(self.children.val(pos).load(Acquire, fake_ref(self)));
            },
        );
    }
}

impl<'n, K, V> Locker<'n, K, V> {
    /// Acquires exclusive lock on the [`LeafNode`].
    #[inline]
    pub(super) fn try_lock(leaf_node: &'n LeafNode<K, V>) -> Option<Locker<'n, K, V>> {
        if leaf_node.lock.try_lock().is_ok() {
            Some(Locker { node: leaf_node })
        } else {
            None
        }
    }

    /// Retires the leaf node by poisoning the lock.
    #[inline]
    pub(super) fn unlock_retire(self) {
        self.node.lock.poison_lock();
        forget(self);
    }
}

impl<K, V> Drop for Locker<'_, K, V> {
    #[inline]
    fn drop(&mut self) {
        self.node.lock.release_lock();
    }
}

impl RemoveRangeState {
    /// Returns the next state.
    pub(super) fn next<K, Q, R: RangeBounds<Q>>(
        self,
        key: &K,
        range: &R,
        start_unbounded: bool,
    ) -> Self
    where
        Q: Comparable<K> + ?Sized,
    {
        if range_contains(range, key) {
            match self {
                RemoveRangeState::Below => {
                    if start_unbounded {
                        RemoveRangeState::FullyContained
                    } else {
                        RemoveRangeState::MaybeBelow
                    }
                }
                RemoveRangeState::MaybeBelow | RemoveRangeState::FullyContained => {
                    RemoveRangeState::FullyContained
                }
                RemoveRangeState::MaybeAbove => unreachable!(),
            }
        } else {
            match self {
                RemoveRangeState::Below => match range.start_bound() {
                    Bound::Included(k) => match k.compare(key) {
                        Less | Equal => RemoveRangeState::MaybeAbove,
                        Greater => RemoveRangeState::Below,
                    },
                    Bound::Excluded(k) => match k.compare(key) {
                        Less => RemoveRangeState::MaybeAbove,
                        Greater | Equal => RemoveRangeState::Below,
                    },
                    Bound::Unbounded => RemoveRangeState::MaybeAbove,
                },
                RemoveRangeState::MaybeBelow | RemoveRangeState::FullyContained => {
                    RemoveRangeState::MaybeAbove
                }
                RemoveRangeState::MaybeAbove => unreachable!(),
            }
        }
    }
}

#[cfg(not(feature = "loom"))]
#[cfg(test)]
mod test {
    use super::*;
    use std::sync::Arc;
    use std::sync::atomic::AtomicBool;
    use tokio::sync::Barrier;

    #[test]
    fn basic() {
        let private_collector = PrivateCollector::new();
        let guard = Guard::new();
        let leaf_node: LeafNode<String, String> = LeafNode::new();
        assert!(matches!(
            leaf_node.insert::<(), false>(
                "MY GOODNESS!".to_owned(),
                "OH MY GOD!!".to_owned(),
                &private_collector,
                &mut (),
                &guard
            ),
            Ok(InsertResult::Success)
        ));
        assert!(matches!(
            leaf_node.insert::<(), false>(
                "GOOD DAY".to_owned(),
                "OH MY GOD!!".to_owned(),
                &private_collector,
                &mut (),
                &guard
            ),
            Ok(InsertResult::Success)
        ));
        assert_eq!(
            leaf_node.search_entry("MY GOODNESS!", &guard).unwrap().1,
            "OH MY GOD!!"
        );
        assert_eq!(
            leaf_node.search_entry("GOOD DAY", &guard).unwrap().1,
            "OH MY GOD!!"
        );
        assert!(matches!(
            leaf_node.remove_if(
                "GOOD DAY",
                &mut |v| v == "OH MY",
                &private_collector,
                &mut (),
                &guard
            ),
            Ok(RemoveResult::Fail)
        ));
        assert!(matches!(
            leaf_node.remove_if(
                "GOOD DAY",
                &mut |v| v == "OH MY GOD!!",
                &private_collector,
                &mut (),
                &guard
            ),
            Ok(RemoveResult::Success)
        ));
        assert!(matches!(
            leaf_node.remove_if(
                "GOOD",
                &mut |v| v == "OH MY",
                &private_collector,
                &mut (),
                &guard
            ),
            Ok(RemoveResult::Fail)
        ));
        assert!(matches!(
            leaf_node.remove_if(
                "MY GOODNESS!",
                &mut |_| true,
                &private_collector,
                &mut (),
                &guard
            ),
            Ok(RemoveResult::Retired)
        ));
        assert!(matches!(
            leaf_node.insert::<(), false>(
                "HI".to_owned(),
                "HO".to_owned(),
                &private_collector,
                &mut (),
                &guard
            ),
            Ok(InsertResult::Full(..))
        ));
    }

    #[test]
    fn bulk() {
        let private_collector = PrivateCollector::new();
        let guard = Guard::new();
        let leaf_node: LeafNode<usize, usize> = LeafNode::new();
        for k in 0..1024 {
            let mut result =
                leaf_node.insert::<(), false>(k, k, &private_collector, &mut (), &guard);
            if result.is_err() {
                result = leaf_node.insert::<(), false>(k, k, &private_collector, &mut (), &guard);
            }
            match result.unwrap() {
                InsertResult::Success => {
                    assert_eq!(leaf_node.search_entry(&k, &guard), Some((&k, &k)));
                    assert_eq!(
                        leaf_node.search_entry(&(k / 2), &guard),
                        Some((&(k / 2), &(k / 2)))
                    );
                }
                InsertResult::Full(_, _) => {
                    for r in 0..(k - 1) {
                        assert_eq!(leaf_node.search_entry(&r, &guard), Some((&r, &r)));
                        assert!(
                            leaf_node
                                .remove_if(&r, &mut |_| true, &private_collector, &mut (), &guard)
                                .is_ok()
                        );
                        assert_eq!(leaf_node.search_entry(&r, &guard), None);
                    }
                    assert_eq!(
                        leaf_node.search_entry(&(k - 1), &guard),
                        Some((&(k - 1), &(k - 1)))
                    );
                    assert_eq!(
                        leaf_node.remove_if(
                            &(k - 1),
                            &mut |_| true,
                            &private_collector,
                            &mut (),
                            &guard
                        ),
                        Ok(RemoveResult::Retired)
                    );
                    assert_eq!(leaf_node.search_entry(&(k - 1), &guard), None);
                    break;
                }
                InsertResult::Duplicate(..) => unreachable!(),
            }
        }
    }

    #[allow(clippy::too_many_lines)]
    #[cfg_attr(miri, ignore)]
    #[tokio::test(flavor = "multi_thread", worker_threads = 16)]
    async fn parallel() {
        let num_tasks = 8;
        let workload_size = 64;
        let private_collector = Arc::new(PrivateCollector::new());
        let barrier = Arc::new(Barrier::new(num_tasks));
        for _ in 0..16 {
            let leaf_node = Arc::new(LeafNode::new());
            assert!(
                leaf_node
                    .insert::<(), false>(
                        usize::MAX,
                        usize::MAX,
                        &private_collector,
                        &mut (),
                        &Guard::new()
                    )
                    .is_ok()
            );
            let mut task_handles = Vec::with_capacity(num_tasks);
            for task_id in 0..num_tasks {
                let private_collector_clone = private_collector.clone();
                let barrier_clone = barrier.clone();
                let leaf_node_clone = leaf_node.clone();
                task_handles.push(tokio::task::spawn(async move {
                    barrier_clone.wait().await;
                    let guard = Guard::new();
                    let mut max_key = None;
                    let range = (task_id * workload_size)..((task_id + 1) * workload_size);
                    for id in range.clone() {
                        loop {
                            if let Ok(r) = leaf_node_clone.insert::<(), false>(
                                id,
                                id,
                                &private_collector_clone,
                                &mut (),
                                &guard,
                            ) {
                                match r {
                                    InsertResult::Success => {
                                        if let Ok(InsertResult::Success) = leaf_node_clone
                                            .insert::<(), false>(
                                                id,
                                                id,
                                                &private_collector_clone,
                                                &mut (),
                                                &guard,
                                            )
                                        {
                                            unreachable!()
                                        }
                                        break;
                                    }
                                    InsertResult::Full(..) => {
                                        max_key.replace(id);
                                        break;
                                    }
                                    InsertResult::Duplicate(..) => unreachable!(),
                                }
                            }
                        }
                        if max_key.is_some() {
                            break;
                        }
                    }
                    for id in range.clone() {
                        if max_key == Some(id) {
                            break;
                        }
                        assert_eq!(leaf_node_clone.search_entry(&id, &guard), Some((&id, &id)));
                    }
                    for id in range {
                        if max_key == Some(id) {
                            break;
                        }
                        loop {
                            if let Ok(r) = leaf_node_clone.remove_if(
                                &id,
                                &mut |_| true,
                                &private_collector_clone,
                                &mut (),
                                &guard,
                            ) {
                                match r {
                                    RemoveResult::Success | RemoveResult::Fail => break,
                                    RemoveResult::Frozen | RemoveResult::Retired => unreachable!(),
                                }
                            }
                        }
                        assert!(
                            leaf_node_clone.search_entry(&id, &guard).is_none(),
                            "{}",
                            id
                        );
                        if let Ok(RemoveResult::Success) = leaf_node_clone.remove_if(
                            &id,
                            &mut |_| true,
                            &private_collector_clone,
                            &mut (),
                            &guard,
                        ) {
                            unreachable!()
                        }
                    }
                }));
            }

            for r in futures::future::join_all(task_handles).await {
                assert!(r.is_ok());
            }
            assert!(
                leaf_node
                    .remove_if(
                        &usize::MAX,
                        &mut |_| true,
                        &private_collector,
                        &mut (),
                        &Guard::new()
                    )
                    .is_ok()
            );
        }
    }

    #[cfg_attr(miri, ignore)]
    #[tokio::test(flavor = "multi_thread", worker_threads = 16)]
    async fn durability() {
        let num_tasks = 8_usize;
        let workload_size = 64_usize;
        let private_collector = Arc::new(PrivateCollector::new());
        for _ in 0..16 {
            for k in 0..=workload_size {
                let barrier = Arc::new(Barrier::new(num_tasks));
                let leaf_node: Arc<LeafNode<usize, usize>> = Arc::new(LeafNode::new());
                let inserted: Arc<AtomicBool> = Arc::new(AtomicBool::new(false));
                let mut task_handles = Vec::with_capacity(num_tasks);
                for _ in 0..num_tasks {
                    let private_collector_clone = private_collector.clone();
                    let barrier_clone = barrier.clone();
                    let leaf_node_clone = leaf_node.clone();
                    let inserted_clone = inserted.clone();
                    task_handles.push(tokio::spawn(async move {
                        {
                            barrier_clone.wait().await;
                            let guard = Guard::new();
                            if let Ok(InsertResult::Success) = leaf_node_clone.insert::<(), false>(
                                k,
                                k,
                                &private_collector_clone,
                                &mut (),
                                &guard,
                            ) {
                                assert!(!inserted_clone.swap(true, Relaxed));
                            }
                        }
                        {
                            barrier_clone.wait().await;
                            let guard = Guard::new();
                            for i in 0..workload_size {
                                if i != k {
                                    let result = leaf_node_clone.insert::<(), false>(
                                        i,
                                        i,
                                        &private_collector_clone,
                                        &mut (),
                                        &guard,
                                    );
                                    drop(result);
                                }
                                assert_eq!(
                                    leaf_node_clone.search_entry(&k, &guard).unwrap(),
                                    (&k, &k)
                                );
                            }
                            for i in 0..workload_size {
                                let max_iter =
                                    leaf_node_clone.approximate::<_, true>(&k, &guard).unwrap();
                                assert!(*max_iter.get().unwrap().0 <= k);
                                let mut min_iter = leaf_node_clone.min(&guard).unwrap();
                                if let Some((k_ref, v_ref)) = min_iter.next() {
                                    assert_eq!(*k_ref, *v_ref);
                                    assert!(*k_ref <= k);
                                } else {
                                    let (k_ref, v_ref) = min_iter.jump(&guard).unwrap();
                                    assert_eq!(*k_ref, *v_ref);
                                    assert!(*k_ref <= k);
                                }
                                let _result = leaf_node_clone.remove_if(
                                    &i,
                                    &mut |v| *v != k,
                                    &private_collector_clone,
                                    &mut (),
                                    &guard,
                                );
                                assert_eq!(
                                    leaf_node_clone.search_entry(&k, &guard).unwrap(),
                                    (&k, &k)
                                );
                            }
                        }
                    }));
                }
                for r in futures::future::join_all(task_handles).await {
                    assert!(r.is_ok());
                }
                assert!((*inserted).load(Relaxed));
            }
        }
    }
}
