libdictenstein 0.1.0

//! ART Node Types for Persistent Adaptive Radix Trie
//!
//! This module implements the four adaptive node types used in ART:
//!
//! | Node Type | Children | Search Method | Size   | When Used |
//! |-----------|----------|---------------|--------|-----------|
//! | Node4     | 1-4      | Linear scan   | ~48B   | Sparse nodes |
//! | Node16    | 5-16     | SIMD (SSE4.1) | ~160B  | Moderate fanout |
//! | Node48    | 17-48    | Index array   | ~656B  | High fanout |
//! | Node256   | 49-256   | Direct array  | ~2KB   | Dense nodes |
//!
//! # Node Layout
//!
//! All nodes share a common header followed by type-specific data:
//!
//! ```text
//! ┌───────────────────────────────────────────────────────────┐
//! │                     NodeHeader (16 bytes)                  │
//! ├───────────────────────────────────────────────────────────┤
//! │ node_type: u8        │ Node type discriminant             │
//! │ num_children: u8     │ Current number of children         │
//! │ prefix_len: u8       │ Compressed prefix length (0-12)    │
//! │ flags: u8            │ Flags (is_final, dirty, etc.)      │
//! │ _padding: [u8; 4]    │ Alignment padding                  │
//! │ version: u64         │ Version for optimistic locking     │
//! └───────────────────────────────────────────────────────────┘
//! ```
//!
//! # Path Compression
//!
//! Each node can store up to 12 bytes of compressed prefix. When traversing
//! the trie, if the prefix matches, we skip those bytes in the search key.
//! This significantly reduces tree height for common patterns.
//!
//! # SIMD Optimization
//!
//! Node16 uses SSE4.1 SIMD instructions (`_mm_cmpeq_epi8`) for parallel
//! key comparison, finding matches in ~3 cycles vs ~16 for linear scan.

pub mod node16;
pub mod node256;
pub mod node4;
pub mod node48;

// Lock-free persistent node types (requires `im` crate)
#[cfg(feature = "persistent-artrie")]
pub mod atomic_ptr;
#[cfg(feature = "persistent-artrie")]
pub mod persistent_node;

pub use node16::Node16;
pub use node256::Node256;
pub use node4::Node4;
pub use node48::Node48;

// Lock-free persistent node exports
#[cfg(feature = "persistent-artrie")]
pub use atomic_ptr::AtomicNodePtr;
#[cfg(feature = "persistent-artrie")]
pub use persistent_node::PersistentNode;

use std::sync::atomic::{AtomicU16, AtomicU64, AtomicU8, Ordering};

use super::swizzled_ptr::SwizzledPtr;

/// Maximum path compression prefix length (12 bytes)
pub const MAX_PREFIX_LEN: usize = 12;

/// Node flags
pub mod flags {
    /// Node represents a valid dictionary entry (is_final)
    pub const IS_FINAL: u8 = 0b0000_0001;
    /// Node has been modified (dirty)
    pub const IS_DIRTY: u8 = 0b0000_0010;
    /// Node is a leaf (bucket pointer)
    pub const IS_LEAF: u8 = 0b0000_0100;
    /// Node has dirty descendants (for selective persistence optimization)
    ///
    /// When set, indicates that at least one descendant node in the subtree
    /// rooted at this node has been modified. This flag propagates upward
    /// during mutations, enabling `persist_to_disk()` to skip entire clean
    /// subtrees during checkpoint.
    pub const HAS_DIRTY_DESCENDANTS: u8 = 0b0000_1000;
}

/// Common header for all ART node types
///
/// This header is shared by Node4, Node16, Node48, and Node256.
/// Total size: 16 bytes (2 cache lines would hold header + small node).
#[repr(C)]
#[derive(Debug, Clone)]
pub struct NodeHeader {
    /// Node type discriminant (4, 16, 48, or 0 for Node256)
    pub node_type: u8,
    /// Length of the compressed prefix (0-12)
    pub prefix_len: u8,
    /// Node flags (IS_FINAL, IS_DIRTY, IS_LEAF)
    pub flags: u8,
    /// Padding for alignment
    pub _padding: u8,
    /// Number of children currently stored (u16 to support Node256 with 256 children)
    pub num_children: u16,
    /// More padding for 8-byte alignment
    pub _padding2: [u8; 2],
    /// Version counter for optimistic locking
    pub version: u64,
}

impl NodeHeader {
    /// Create a new node header
    pub fn new(node_type: u8) -> Self {
        Self {
            node_type,
            prefix_len: 0,
            flags: 0,
            _padding: 0,
            num_children: 0,
            _padding2: [0; 2],
            version: 0,
        }
    }

    /// Check if this node represents a final (accepting) state
    #[inline]
    pub fn is_final(&self) -> bool {
        self.flags & flags::IS_FINAL != 0
    }

    /// Set the final flag
    #[inline]
    pub fn set_final(&mut self, final_state: bool) {
        if final_state {
            self.flags |= flags::IS_FINAL;
        } else {
            self.flags &= !flags::IS_FINAL;
        }
    }

    /// Check if this node is dirty (modified since last write-back)
    #[inline]
    pub fn is_dirty(&self) -> bool {
        self.flags & flags::IS_DIRTY != 0
    }

    /// Set the dirty flag
    #[inline]
    pub fn set_dirty(&mut self, dirty: bool) {
        if dirty {
            self.flags |= flags::IS_DIRTY;
        } else {
            self.flags &= !flags::IS_DIRTY;
        }
    }

    /// Check if this is a leaf node (points to a bucket)
    #[inline]
    pub fn is_leaf(&self) -> bool {
        self.flags & flags::IS_LEAF != 0
    }

    /// Set the leaf flag
    #[inline]
    pub fn set_leaf(&mut self, leaf: bool) {
        if leaf {
            self.flags |= flags::IS_LEAF;
        } else {
            self.flags &= !flags::IS_LEAF;
        }
    }

    /// Check if this node has dirty descendants
    ///
    /// When true, indicates that at least one descendant node in the subtree
    /// rooted at this node has been modified since the last checkpoint.
    #[inline]
    pub fn has_dirty_descendants(&self) -> bool {
        self.flags & flags::HAS_DIRTY_DESCENDANTS != 0
    }

    /// Set the has_dirty_descendants flag
    ///
    /// This flag should be set when any descendant node is modified,
    /// and propagated upward to all ancestors.
    #[inline]
    pub fn set_has_dirty_descendants(&mut self, has_dirty: bool) {
        if has_dirty {
            self.flags |= flags::HAS_DIRTY_DESCENDANTS;
        } else {
            self.flags &= !flags::HAS_DIRTY_DESCENDANTS;
        }
    }

    /// Check if this node or any descendant needs to be persisted
    ///
    /// Returns true if either:
    /// - The node itself is dirty (IS_DIRTY flag set)
    /// - The node has dirty descendants (HAS_DIRTY_DESCENDANTS flag set)
    ///
    /// This is used by `persist_to_disk()` to determine whether to traverse
    /// into a subtree or skip it entirely.
    #[inline]
    pub fn needs_persistence(&self) -> bool {
        (self.flags & (flags::IS_DIRTY | flags::HAS_DIRTY_DESCENDANTS)) != 0
    }

    /// Clear both dirty flags (IS_DIRTY and HAS_DIRTY_DESCENDANTS)
    ///
    /// This should be called after successfully persisting a node to disk
    /// to reset the dirty tracking state.
    #[inline]
    pub fn clear_dirty_flags(&mut self) {
        self.flags &= !(flags::IS_DIRTY | flags::HAS_DIRTY_DESCENDANTS);
    }

    /// Increment the version counter (for optimistic locking)
    #[inline]
    pub fn increment_version(&mut self) {
        self.version = self.version.wrapping_add(1);
    }

    /// Check if version matches expected value (for optimistic validation)
    ///
    /// Returns true if the version matches, false otherwise.
    /// Used by MVCC-Lite to validate that a node hasn't been modified.
    #[inline]
    pub fn check_version(&self, expected: u64) -> bool {
        self.version == expected
    }

    /// Get the current version number
    #[inline]
    pub fn version(&self) -> u64 {
        self.version
    }
}

impl Default for NodeHeader {
    fn default() -> Self {
        Self::new(0)
    }
}

// =============================================================================
// Atomic Node Header for Lock-Free Operations
// =============================================================================

/// Atomic header for lock-free concurrent access to ART nodes.
///
/// This provides CAS (compare-and-swap) operations for the critical fields
/// that need to be updated atomically during concurrent inserts:
///
/// - `flags`: For atomically setting IS_FINAL when claiming a node
/// - `num_children`: For atomically incrementing child count
/// - `version`: For optimistic locking validation
///
/// The non-atomic `node_type` and `prefix_len` are immutable once a node
/// is created, so they don't need atomic access.
///
/// # Memory Layout
///
/// Same as `NodeHeader` (16 bytes, C-repr) for safe transmutation.
#[repr(C)]
pub struct AtomicNodeHeader {
    /// Node type discriminant (immutable after creation)
    pub node_type: u8,
    /// Length of the compressed prefix (immutable after creation)
    pub prefix_len: u8,
    /// Node flags (atomic for CAS on IS_FINAL)
    flags: AtomicU8,
    /// Padding for alignment
    _padding: u8,
    /// Number of children (atomic for concurrent child updates)
    num_children: AtomicU16,
    /// Padding for 8-byte alignment
    _padding2: [u8; 2],
    /// Version counter for optimistic locking
    version: AtomicU64,
}

impl AtomicNodeHeader {
    /// Create a new atomic node header.
    pub fn new(node_type: u8) -> Self {
        Self {
            node_type,
            prefix_len: 0,
            flags: AtomicU8::new(0),
            _padding: 0,
            num_children: AtomicU16::new(0),
            _padding2: [0; 2],
            version: AtomicU64::new(0),
        }
    }

    /// Create an atomic header from a regular header.
    ///
    /// This is used when converting a node for lock-free operations.
    pub fn from_header(header: &NodeHeader) -> Self {
        Self {
            node_type: header.node_type,
            prefix_len: header.prefix_len,
            flags: AtomicU8::new(header.flags),
            _padding: 0,
            num_children: AtomicU16::new(header.num_children),
            _padding2: [0; 2],
            version: AtomicU64::new(header.version),
        }
    }

    /// Convert to a regular header for serialization.
    ///
    /// This performs atomic loads to get a consistent snapshot.
    pub fn to_header(&self) -> NodeHeader {
        NodeHeader {
            node_type: self.node_type,
            prefix_len: self.prefix_len,
            flags: self.flags.load(Ordering::Acquire),
            _padding: 0,
            num_children: self.num_children.load(Ordering::Acquire),
            _padding2: [0; 2],
            version: self.version.load(Ordering::Acquire),
        }
    }

    // =========================================================================
    // Atomic Flag Operations
    // =========================================================================

    /// Atomically try to set the IS_FINAL flag.
    ///
    /// Uses fetch_or to set the flag and returns true if we were the first
    /// to set it (i.e., the flag was not previously set).
    ///
    /// This is the core primitive for lock-free insert: the thread that
    /// successfully sets IS_FINAL "wins" and gets to assign the value/index.
    ///
    /// # Returns
    ///
    /// - `true` if we set the flag (winner)
    /// - `false` if the flag was already set (another thread won)
    #[inline]
    pub fn try_set_final(&self) -> bool {
        let old = self.flags.fetch_or(flags::IS_FINAL, Ordering::AcqRel);
        (old & flags::IS_FINAL) == 0
    }

    /// Check if this node is final (atomic read).
    #[inline]
    pub fn is_final(&self) -> bool {
        (self.flags.load(Ordering::Acquire) & flags::IS_FINAL) != 0
    }

    /// Atomically set the dirty flag.
    #[inline]
    pub fn set_dirty(&self) {
        self.flags.fetch_or(flags::IS_DIRTY, Ordering::Release);
    }

    /// Check if this node is dirty (atomic read).
    #[inline]
    pub fn is_dirty(&self) -> bool {
        (self.flags.load(Ordering::Acquire) & flags::IS_DIRTY) != 0
    }

    /// Atomically set the has_dirty_descendants flag.
    #[inline]
    pub fn set_has_dirty_descendants(&self) {
        self.flags
            .fetch_or(flags::HAS_DIRTY_DESCENDANTS, Ordering::Release);
    }

    /// Check if this node has dirty descendants (atomic read).
    #[inline]
    pub fn has_dirty_descendants(&self) -> bool {
        (self.flags.load(Ordering::Acquire) & flags::HAS_DIRTY_DESCENDANTS) != 0
    }

    /// Check if this node or any descendant needs persistence (atomic read).
    #[inline]
    pub fn needs_persistence(&self) -> bool {
        (self.flags.load(Ordering::Acquire) & (flags::IS_DIRTY | flags::HAS_DIRTY_DESCENDANTS)) != 0
    }

    /// Atomically clear both dirty flags.
    #[inline]
    pub fn clear_dirty_flags(&self) {
        self.flags.fetch_and(
            !(flags::IS_DIRTY | flags::HAS_DIRTY_DESCENDANTS),
            Ordering::Release,
        );
    }

    /// Load all flags atomically.
    #[inline]
    pub fn flags(&self) -> u8 {
        self.flags.load(Ordering::Acquire)
    }

    // =========================================================================
    // Atomic Child Count Operations
    // =========================================================================

    /// Atomically get the number of children.
    #[inline]
    pub fn num_children(&self) -> u16 {
        self.num_children.load(Ordering::Acquire)
    }

    /// Atomically increment the child count.
    ///
    /// Returns the previous value.
    #[inline]
    pub fn fetch_add_children(&self, count: u16) -> u16 {
        self.num_children.fetch_add(count, Ordering::AcqRel)
    }

    /// Atomically decrement the child count.
    ///
    /// Returns the previous value.
    #[inline]
    pub fn fetch_sub_children(&self, count: u16) -> u16 {
        self.num_children.fetch_sub(count, Ordering::AcqRel)
    }

    /// Compare-and-swap the child count.
    ///
    /// # Returns
    ///
    /// - `Ok(current)` if successful
    /// - `Err(actual)` if the current value didn't match expected
    #[inline]
    pub fn compare_exchange_children(&self, expected: u16, new: u16) -> Result<u16, u16> {
        self.num_children
            .compare_exchange(expected, new, Ordering::AcqRel, Ordering::Acquire)
    }

    // =========================================================================
    // Atomic Version Operations
    // =========================================================================

    /// Atomically get the current version.
    #[inline]
    pub fn version(&self) -> u64 {
        self.version.load(Ordering::Acquire)
    }

    /// Atomically increment the version counter.
    ///
    /// Returns the previous value.
    #[inline]
    pub fn fetch_add_version(&self) -> u64 {
        self.version.fetch_add(1, Ordering::AcqRel)
    }

    /// Check if version matches expected value (for optimistic validation).
    #[inline]
    pub fn check_version(&self, expected: u64) -> bool {
        self.version.load(Ordering::Acquire) == expected
    }

    /// Begin a write operation (increment to odd for optimistic locking).
    ///
    /// Returns the version before incrementing.
    #[inline]
    pub fn begin_write(&self) -> u64 {
        self.version.fetch_add(1, Ordering::AcqRel)
    }

    /// End a write operation (increment to even for optimistic locking).
    #[inline]
    pub fn end_write(&self) {
        self.version.fetch_add(1, Ordering::Release);
    }

    /// Check if the version is stable (even = not being modified).
    #[inline]
    pub fn is_version_stable(&self) -> bool {
        self.version.load(Ordering::Acquire) % 2 == 0
    }
}

impl std::fmt::Debug for AtomicNodeHeader {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AtomicNodeHeader")
            .field("node_type", &self.node_type)
            .field("prefix_len", &self.prefix_len)
            .field("flags", &self.flags.load(Ordering::Relaxed))
            .field("num_children", &self.num_children.load(Ordering::Relaxed))
            .field("version", &self.version.load(Ordering::Relaxed))
            .finish()
    }
}

impl Clone for AtomicNodeHeader {
    fn clone(&self) -> Self {
        Self {
            node_type: self.node_type,
            prefix_len: self.prefix_len,
            flags: AtomicU8::new(self.flags.load(Ordering::Acquire)),
            _padding: 0,
            num_children: AtomicU16::new(self.num_children.load(Ordering::Acquire)),
            _padding2: [0; 2],
            version: AtomicU64::new(self.version.load(Ordering::Acquire)),
        }
    }
}

impl Default for AtomicNodeHeader {
    fn default() -> Self {
        Self::new(0)
    }
}

/// Compressed prefix for path compression
///
/// Stores up to 12 bytes of shared prefix that can be skipped during traversal.
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct CompressedPrefix {
    /// The prefix bytes
    pub bytes: [u8; MAX_PREFIX_LEN],
}

impl CompressedPrefix {
    /// Create an empty prefix
    pub const fn empty() -> Self {
        Self {
            bytes: [0; MAX_PREFIX_LEN],
        }
    }

    /// Create a prefix from a byte slice
    ///
    /// # Panics
    /// Panics if the slice is longer than MAX_PREFIX_LEN
    pub fn from_bytes(bytes: &[u8]) -> Self {
        assert!(
            bytes.len() <= MAX_PREFIX_LEN,
            "prefix too long: {} > {}",
            bytes.len(),
            MAX_PREFIX_LEN
        );
        let mut prefix = Self::empty();
        prefix.bytes[..bytes.len()].copy_from_slice(bytes);
        prefix
    }

    /// Check if a key matches this prefix
    ///
    /// Returns the number of matching bytes (up to prefix_len)
    pub fn match_key(&self, key: &[u8], prefix_len: usize) -> usize {
        let check_len = prefix_len.min(key.len()).min(MAX_PREFIX_LEN);
        for i in 0..check_len {
            if self.bytes[i] != key[i] {
                return i;
            }
        }
        check_len
    }

    /// Get the prefix as a slice
    pub fn as_slice(&self, len: usize) -> &[u8] {
        &self.bytes[..len.min(MAX_PREFIX_LEN)]
    }
}

impl Default for CompressedPrefix {
    fn default() -> Self {
        Self::empty()
    }
}

/// Unified ART node enum
///
/// This enum wraps all four node types for type-safe dispatch.
#[derive(Debug, Clone)]
pub enum Node {
    /// Node with 1-4 children
    N4(Box<Node4>),
    /// Node with 5-16 children
    N16(Box<Node16>),
    /// Node with 17-48 children
    N48(Box<Node48>),
    /// Node with 49-256 children
    N256(Box<Node256>),
}

impl Node {
    /// Get the node header
    pub fn header(&self) -> &NodeHeader {
        match self {
            Node::N4(n) => &n.header,
            Node::N16(n) => &n.header,
            Node::N48(n) => &n.header,
            Node::N256(n) => &n.header,
        }
    }

    /// Get the node header mutably
    pub fn header_mut(&mut self) -> &mut NodeHeader {
        match self {
            Node::N4(n) => &mut n.header,
            Node::N16(n) => &mut n.header,
            Node::N48(n) => &mut n.header,
            Node::N256(n) => &mut n.header,
        }
    }

    /// Get the compressed prefix
    pub fn prefix(&self) -> &CompressedPrefix {
        match self {
            Node::N4(n) => &n.prefix,
            Node::N16(n) => &n.prefix,
            Node::N48(n) => &n.prefix,
            Node::N256(n) => &n.prefix,
        }
    }

    /// Get the compressed prefix mutably
    pub fn prefix_mut(&mut self) -> &mut CompressedPrefix {
        match self {
            Node::N4(n) => &mut n.prefix,
            Node::N16(n) => &mut n.prefix,
            Node::N48(n) => &mut n.prefix,
            Node::N256(n) => &mut n.prefix,
        }
    }

    /// Check if this node is final (accepting state)
    #[inline]
    pub fn is_final(&self) -> bool {
        self.header().is_final()
    }

    /// Look up a child by key byte
    pub fn find_child(&self, key: u8) -> Option<&SwizzledPtr> {
        match self {
            Node::N4(n) => n.find_child(key),
            Node::N16(n) => n.find_child(key),
            Node::N48(n) => n.find_child(key),
            Node::N256(n) => n.find_child(key),
        }
    }

    /// Look up a child by key byte (mutable)
    pub fn find_child_mut(&mut self, key: u8) -> Option<&mut SwizzledPtr> {
        match self {
            Node::N4(n) => n.find_child_mut(key),
            Node::N16(n) => n.find_child_mut(key),
            Node::N48(n) => n.find_child_mut(key),
            Node::N256(n) => n.find_child_mut(key),
        }
    }

    /// Get an iterator over all (key, child) pairs
    pub fn iter_children(&self) -> Box<dyn Iterator<Item = (u8, &SwizzledPtr)> + '_> {
        match self {
            Node::N4(n) => Box::new(n.iter_children()),
            Node::N16(n) => Box::new(n.iter_children()),
            Node::N48(n) => Box::new(n.iter_children()),
            Node::N256(n) => Box::new(n.iter_children()),
        }
    }

    /// Get the number of children
    #[inline]
    pub fn num_children(&self) -> usize {
        self.header().num_children as usize
    }

    /// Check if the node is full
    pub fn is_full(&self) -> bool {
        match self {
            Node::N4(n) => n.is_full(),
            Node::N16(n) => n.is_full(),
            Node::N48(n) => n.is_full(),
            Node::N256(n) => n.is_full(),
        }
    }

    /// Check if the node should grow to a larger type
    pub fn should_grow(&self) -> bool {
        self.is_full()
    }

    /// Check if the node should shrink to a smaller type
    pub fn should_shrink(&self) -> bool {
        match self {
            Node::N4(_) => false,
            Node::N16(n) => n.header.num_children <= 4,
            Node::N48(n) => n.header.num_children <= 16,
            Node::N256(n) => n.header.num_children <= 48,
        }
    }

    /// Create an empty Node4
    pub fn new_node4() -> Self {
        Node::N4(Box::new(Node4::new()))
    }

    /// Create an empty Node16
    pub fn new_node16() -> Self {
        Node::N16(Box::new(Node16::new()))
    }

    /// Create an empty Node48
    pub fn new_node48() -> Self {
        Node::N48(Box::new(Node48::new()))
    }

    /// Create an empty Node256
    pub fn new_node256() -> Self {
        Node::N256(Box::new(Node256::new()))
    }

    /// Grow this node to the next larger type.
    ///
    /// This is called when a node becomes full and needs more capacity.
    ///
    /// - Node4 → Node16
    /// - Node16 → Node48
    /// - Node48 → Node256
    /// - Node256 → None (already at max size)
    ///
    /// Returns `None` if the node is already Node256 (cannot grow further).
    pub fn grow(&self) -> Option<Node> {
        match self {
            Node::N4(n) => Some(Node::N16(Box::new(n.grow()))),
            Node::N16(n) => Some(Node::N48(Box::new(n.grow()))),
            Node::N48(n) => Some(Node::N256(Box::new(n.grow()))),
            Node::N256(_) => None, // Cannot grow further
        }
    }

    /// Shrink this node to the next smaller type.
    ///
    /// This is called when a node has too few children for its type.
    ///
    /// - Node16 → Node4 (when ≤4 children)
    /// - Node48 → Node16 (when ≤16 children)
    /// - Node256 → Node48 (when ≤48 children)
    /// - Node4 → None (already at min size)
    ///
    /// Returns `None` if the node is already Node4 (cannot shrink further).
    pub fn shrink(&self) -> Option<Node> {
        match self {
            Node::N4(_) => None, // Cannot shrink further
            Node::N16(n) => Some(Node::N4(Box::new(n.shrink()))),
            Node::N48(n) => Some(Node::N16(Box::new(n.shrink()))),
            Node::N256(n) => Some(Node::N48(Box::new(n.shrink()))),
        }
    }

    /// Add a child to this node, automatically growing if necessary.
    ///
    /// If the node is full, this method will grow the node to the next
    /// larger type and return the new node. Otherwise, it adds the child
    /// to the existing node and returns None.
    ///
    /// # Returns
    ///
    /// - `Ok(None)` - Child added successfully to existing node
    /// - `Ok(Some(new_node))` - Node was grown and child added to new node
    /// - `Err(AddChildError::KeyExists)` - Key already exists
    pub fn add_child_growing(
        &mut self,
        key: u8,
        child: SwizzledPtr,
    ) -> Result<Option<Node>, AddChildError> {
        match self {
            Node::N4(n) => {
                if n.is_full() {
                    let mut grown = n.grow();
                    grown.add_child(key, child)?;
                    Ok(Some(Node::N16(Box::new(grown))))
                } else {
                    n.add_child(key, child)?;
                    Ok(None)
                }
            }
            Node::N16(n) => {
                if n.is_full() {
                    let mut grown = n.grow();
                    grown.add_child(key, child)?;
                    Ok(Some(Node::N48(Box::new(grown))))
                } else {
                    n.add_child(key, child)?;
                    Ok(None)
                }
            }
            Node::N48(n) => {
                if n.is_full() {
                    let mut grown = n.grow();
                    grown.add_child(key, child)?;
                    Ok(Some(Node::N256(Box::new(grown))))
                } else {
                    n.add_child(key, child)?;
                    Ok(None)
                }
            }
            Node::N256(n) => {
                n.add_child(key, child)?;
                Ok(None)
            }
        }
    }

    /// Remove a child from this node, automatically shrinking if appropriate.
    ///
    /// If the node has too few children after removal, this method will
    /// shrink the node to the next smaller type and return the new node.
    /// Otherwise, it removes the child and returns None.
    ///
    /// # Returns
    ///
    /// - `Some((removed_child, None))` - Child removed, node unchanged
    /// - `Some((removed_child, Some(new_node)))` - Child removed and node shrunk
    /// - `None` - Key not found
    pub fn remove_child_shrinking(&mut self, key: u8) -> Option<(SwizzledPtr, Option<Node>)> {
        match self {
            Node::N4(n) => n.remove_child(key).map(|removed| (removed, None)),
            Node::N16(n) => {
                if let Some(removed) = n.remove_child(key) {
                    if n.header.num_children <= 4 {
                        Some((removed, Some(Node::N4(Box::new(n.shrink())))))
                    } else {
                        Some((removed, None))
                    }
                } else {
                    None
                }
            }
            Node::N48(n) => {
                if let Some(removed) = n.remove_child(key) {
                    if n.header.num_children <= 16 {
                        Some((removed, Some(Node::N16(Box::new(n.shrink())))))
                    } else {
                        Some((removed, None))
                    }
                } else {
                    None
                }
            }
            Node::N256(n) => {
                if let Some(removed) = n.remove_child(key) {
                    if n.header.num_children <= 48 {
                        Some((removed, Some(Node::N48(Box::new(n.shrink())))))
                    } else {
                        Some((removed, None))
                    }
                } else {
                    None
                }
            }
        }
    }
}

/// Trait for common ART node operations
pub trait ArtNode {
    /// Find a child by key byte
    fn find_child(&self, key: u8) -> Option<&SwizzledPtr>;

    /// Find a child mutably by key byte
    fn find_child_mut(&mut self, key: u8) -> Option<&mut SwizzledPtr>;

    /// Add a child with the given key
    fn add_child(&mut self, key: u8, child: SwizzledPtr) -> Result<(), AddChildError>;

    /// Remove a child by key
    fn remove_child(&mut self, key: u8) -> Option<SwizzledPtr>;

    /// Check if the node is full
    fn is_full(&self) -> bool;

    /// Get an iterator over (key, child) pairs
    fn iter_children(&self) -> impl Iterator<Item = (u8, &SwizzledPtr)>;
}

/// Error when adding a child fails
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AddChildError {
    /// Node is full, needs to grow to a larger type
    NodeFull,
    /// Key already exists
    KeyExists,
}

impl std::fmt::Display for AddChildError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            AddChildError::NodeFull => write!(f, "node is full"),
            AddChildError::KeyExists => write!(f, "key already exists"),
        }
    }
}

impl std::error::Error for AddChildError {}

// =============================================================================
// Sequential Sibling Storage Support
// =============================================================================

/// Child storage mode for serialization
///
/// This enum specifies how children should be encoded when serializing a node.
/// It supports two modes:
///
/// 1. **Direct**: Each child pointer is encoded individually (legacy mode or
///    when children are in different arenas)
/// 2. **Sequential**: Children are stored contiguously in the same arena, so
///    we only need to store `(first_slot, count)` instead of N pointers
///
/// Sequential mode provides significant space savings:
/// - Node4 with 3 children: 24 bytes → 3 bytes (88% reduction)
/// - Node16 with 10 children: 80 bytes → 3 bytes (96% reduction)
/// - Node256 with 100 children: 800 bytes → 4 bytes (99.5% reduction)
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChildStorage {
    /// Individual pointers for each child
    ///
    /// Used when:
    /// - Children are in different arenas (cross-arena references)
    /// - Backward compatibility with older format
    /// - Node modifications during runtime
    Direct,

    /// Sequential sibling storage
    ///
    /// When children are allocated contiguously, we only store:
    /// - `first_slot`: The slot ID of the first child
    /// - `arena_id`: The arena containing all children
    /// - `count`: Number of children (from header.num_children)
    ///
    /// Children can be reconstructed as: first_slot, first_slot+1, ..., first_slot+(count-1)
    Sequential {
        /// Arena containing all children
        arena_id: u32,
        /// Slot ID of the first child
        first_slot: u32,
    },
}

impl ChildStorage {
    /// Check if this is direct (individual pointer) storage
    #[inline]
    pub fn is_direct(&self) -> bool {
        matches!(self, ChildStorage::Direct)
    }

    /// Check if this is sequential (contiguous) storage
    #[inline]
    pub fn is_sequential(&self) -> bool {
        matches!(self, ChildStorage::Sequential { .. })
    }

    /// Get the first slot for sequential storage
    ///
    /// Returns `None` if this is direct storage.
    pub fn first_slot(&self) -> Option<u32> {
        match self {
            ChildStorage::Sequential { first_slot, .. } => Some(*first_slot),
            ChildStorage::Direct => None,
        }
    }

    /// Get the arena ID for sequential storage
    ///
    /// Returns `None` if this is direct storage.
    pub fn arena_id(&self) -> Option<u32> {
        match self {
            ChildStorage::Sequential { arena_id, .. } => Some(*arena_id),
            ChildStorage::Direct => None,
        }
    }

    /// Create a sequential storage reference
    pub fn sequential(arena_id: u32, first_slot: u32) -> Self {
        ChildStorage::Sequential {
            arena_id,
            first_slot,
        }
    }
}

impl Default for ChildStorage {
    fn default() -> Self {
        ChildStorage::Direct
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_node_header_flags() {
        let mut header = NodeHeader::new(4);
        assert!(!header.is_final());
        assert!(!header.is_dirty());
        assert!(!header.is_leaf());

        header.set_final(true);
        assert!(header.is_final());

        header.set_dirty(true);
        assert!(header.is_dirty());

        header.set_leaf(true);
        assert!(header.is_leaf());

        header.set_final(false);
        assert!(!header.is_final());
        assert!(header.is_dirty()); // Other flags unchanged
    }

    #[test]
    fn test_node_header_dirty_descendants_flag() {
        let mut header = NodeHeader::new(4);

        // Initially no dirty flags
        assert!(!header.has_dirty_descendants());
        assert!(!header.is_dirty());
        assert!(!header.needs_persistence());

        // Set has_dirty_descendants
        header.set_has_dirty_descendants(true);
        assert!(header.has_dirty_descendants());
        assert!(!header.is_dirty());
        assert!(header.needs_persistence()); // Should need persistence

        // Set dirty as well
        header.set_dirty(true);
        assert!(header.has_dirty_descendants());
        assert!(header.is_dirty());
        assert!(header.needs_persistence());

        // Clear dirty but keep has_dirty_descendants
        header.set_dirty(false);
        assert!(header.has_dirty_descendants());
        assert!(!header.is_dirty());
        assert!(header.needs_persistence()); // Still needs persistence due to descendants

        // Clear has_dirty_descendants
        header.set_has_dirty_descendants(false);
        assert!(!header.has_dirty_descendants());
        assert!(!header.needs_persistence());
    }

    #[test]
    fn test_node_header_clear_dirty_flags() {
        let mut header = NodeHeader::new(4);

        // Set both dirty flags
        header.set_dirty(true);
        header.set_has_dirty_descendants(true);
        assert!(header.is_dirty());
        assert!(header.has_dirty_descendants());
        assert!(header.needs_persistence());

        // Clear both at once
        header.clear_dirty_flags();
        assert!(!header.is_dirty());
        assert!(!header.has_dirty_descendants());
        assert!(!header.needs_persistence());

        // Other flags should be unchanged
        header.set_final(true);
        header.set_leaf(true);
        header.set_dirty(true);
        header.set_has_dirty_descendants(true);
        header.clear_dirty_flags();
        assert!(header.is_final()); // Preserved
        assert!(header.is_leaf()); // Preserved
        assert!(!header.is_dirty());
        assert!(!header.has_dirty_descendants());
    }

    #[test]
    fn test_node_header_version() {
        let mut header = NodeHeader::new(4);
        assert_eq!(header.version, 0);

        header.increment_version();
        assert_eq!(header.version, 1);

        header.increment_version();
        assert_eq!(header.version, 2);
    }

    #[test]
    fn test_compressed_prefix() {
        let prefix = CompressedPrefix::from_bytes(b"hello");
        assert_eq!(prefix.as_slice(5), b"hello");
        assert_eq!(prefix.as_slice(3), b"hel");
    }

    #[test]
    fn test_prefix_matching() {
        let prefix = CompressedPrefix::from_bytes(b"hello");

        // Full match
        assert_eq!(prefix.match_key(b"hello world", 5), 5);

        // Partial match
        assert_eq!(prefix.match_key(b"help", 5), 3);

        // No match
        assert_eq!(prefix.match_key(b"world", 5), 0);

        // Key shorter than prefix
        assert_eq!(prefix.match_key(b"hel", 5), 3);
    }

    #[test]
    #[should_panic(expected = "prefix too long")]
    fn test_prefix_too_long() {
        CompressedPrefix::from_bytes(b"this is way too long for the prefix");
    }

    // =========================================================================
    // ChildStorage Tests
    // =========================================================================

    #[test]
    fn test_child_storage_direct() {
        let storage = ChildStorage::Direct;
        assert!(storage.is_direct());
        assert!(!storage.is_sequential());
        assert_eq!(storage.first_slot(), None);
        assert_eq!(storage.arena_id(), None);
    }

    #[test]
    fn test_child_storage_sequential() {
        let storage = ChildStorage::sequential(5, 100);
        assert!(!storage.is_direct());
        assert!(storage.is_sequential());
        assert_eq!(storage.arena_id(), Some(5));
        assert_eq!(storage.first_slot(), Some(100));
    }

    #[test]
    fn test_child_storage_default() {
        let storage = ChildStorage::default();
        assert!(storage.is_direct());
    }

    #[test]
    fn test_child_storage_equality() {
        let s1 = ChildStorage::sequential(1, 10);
        let s2 = ChildStorage::sequential(1, 10);
        let s3 = ChildStorage::sequential(1, 20);
        let s4 = ChildStorage::Direct;

        assert_eq!(s1, s2);
        assert_ne!(s1, s3);
        assert_ne!(s1, s4);
        assert_eq!(s4, ChildStorage::Direct);
    }

    // =========================================================================
    // AtomicNodeHeader Tests
    // =========================================================================

    #[test]
    fn test_atomic_header_new() {
        let header = AtomicNodeHeader::new(4);
        assert_eq!(header.node_type, 4);
        assert_eq!(header.prefix_len, 0);
        assert_eq!(header.flags(), 0);
        assert_eq!(header.num_children(), 0);
        assert_eq!(header.version(), 0);
    }

    #[test]
    fn test_atomic_header_from_regular() {
        let mut regular = NodeHeader::new(16);
        regular.set_final(true);
        regular.set_dirty(true);
        regular.num_children = 5;
        regular.version = 42;

        let atomic = AtomicNodeHeader::from_header(&regular);
        assert_eq!(atomic.node_type, 16);
        assert!(atomic.is_final());
        assert!(atomic.is_dirty());
        assert_eq!(atomic.num_children(), 5);
        assert_eq!(atomic.version(), 42);
    }

    #[test]
    fn test_atomic_header_to_regular() {
        let atomic = AtomicNodeHeader::new(48);
        atomic.try_set_final();
        atomic.fetch_add_children(3);
        atomic.fetch_add_version();

        let regular = atomic.to_header();
        assert_eq!(regular.node_type, 48);
        assert!(regular.is_final());
        assert_eq!(regular.num_children, 3);
        assert_eq!(regular.version, 1);
    }

    #[test]
    fn test_atomic_header_try_set_final() {
        let header = AtomicNodeHeader::new(4);

        // First call should succeed (winner)
        assert!(header.try_set_final());
        assert!(header.is_final());

        // Second call should fail (already set)
        assert!(!header.try_set_final());
        assert!(header.is_final());
    }

    #[test]
    fn test_atomic_header_try_set_final_concurrent() {
        use std::sync::Arc;
        use std::thread;

        let header = Arc::new(AtomicNodeHeader::new(4));
        let num_threads = 10;

        let handles: Vec<_> = (0..num_threads)
            .map(|_| {
                let h = Arc::clone(&header);
                thread::spawn(move || h.try_set_final())
            })
            .collect();

        let results: Vec<bool> = handles
            .into_iter()
            .map(|h| h.join().expect("thread should complete"))
            .collect();

        // Exactly one thread should have won
        let winners = results.iter().filter(|&&x| x).count();
        assert_eq!(winners, 1, "exactly one thread should win try_set_final");

        // The flag should be set
        assert!(header.is_final());
    }

    #[test]
    fn test_atomic_header_dirty_flags() {
        let header = AtomicNodeHeader::new(4);

        assert!(!header.is_dirty());
        assert!(!header.has_dirty_descendants());
        assert!(!header.needs_persistence());

        header.set_dirty();
        assert!(header.is_dirty());
        assert!(header.needs_persistence());

        header.set_has_dirty_descendants();
        assert!(header.has_dirty_descendants());

        header.clear_dirty_flags();
        assert!(!header.is_dirty());
        assert!(!header.has_dirty_descendants());
        assert!(!header.needs_persistence());
    }

    #[test]
    fn test_atomic_header_children_ops() {
        let header = AtomicNodeHeader::new(4);

        assert_eq!(header.num_children(), 0);

        // Increment
        let prev = header.fetch_add_children(5);
        assert_eq!(prev, 0);
        assert_eq!(header.num_children(), 5);

        // Increment again
        let prev = header.fetch_add_children(3);
        assert_eq!(prev, 5);
        assert_eq!(header.num_children(), 8);

        // Decrement
        let prev = header.fetch_sub_children(2);
        assert_eq!(prev, 8);
        assert_eq!(header.num_children(), 6);
    }

    #[test]
    fn test_atomic_header_children_cas() {
        let header = AtomicNodeHeader::new(4);
        header.fetch_add_children(5);

        // CAS should fail with wrong expected value
        let result = header.compare_exchange_children(3, 10);
        assert_eq!(result, Err(5));
        assert_eq!(header.num_children(), 5);

        // CAS should succeed with correct expected value
        let result = header.compare_exchange_children(5, 10);
        assert_eq!(result, Ok(5));
        assert_eq!(header.num_children(), 10);
    }

    #[test]
    fn test_atomic_header_version_ops() {
        let header = AtomicNodeHeader::new(4);

        assert_eq!(header.version(), 0);
        assert!(header.is_version_stable()); // 0 is even

        // Begin write (increment to odd)
        let prev = header.begin_write();
        assert_eq!(prev, 0);
        assert_eq!(header.version(), 1);
        assert!(!header.is_version_stable()); // 1 is odd

        // End write (increment to even)
        header.end_write();
        assert_eq!(header.version(), 2);
        assert!(header.is_version_stable()); // 2 is even
    }

    #[test]
    fn test_atomic_header_version_check() {
        let header = AtomicNodeHeader::new(4);

        assert!(header.check_version(0));
        assert!(!header.check_version(1));

        header.fetch_add_version();
        assert!(!header.check_version(0));
        assert!(header.check_version(1));
    }

    #[test]
    fn test_atomic_header_clone() {
        let header = AtomicNodeHeader::new(16);
        header.try_set_final();
        header.fetch_add_children(7);
        header.fetch_add_version();

        let cloned = header.clone();

        assert_eq!(cloned.node_type, 16);
        assert!(cloned.is_final());
        assert_eq!(cloned.num_children(), 7);
        assert_eq!(cloned.version(), 1);

        // Modifications to clone don't affect original
        cloned.fetch_add_children(3);
        assert_eq!(header.num_children(), 7);
        assert_eq!(cloned.num_children(), 10);
    }

    #[test]
    fn test_atomic_header_debug() {
        let header = AtomicNodeHeader::new(4);
        let debug_str = format!("{:?}", header);
        assert!(debug_str.contains("AtomicNodeHeader"));
        assert!(debug_str.contains("node_type: 4"));
    }
}