kermit-ds 0.1.1

use {
    super::implementation::{TreeTrie, TrieNode},
    kermit_derive::IntoTrieIter,
    kermit_iters::{LinearIterator, TrieIterable, TrieIterator, TrieIteratorWrapper},
};

/// A [`TrieIterator`] over a [`TreeTrie`].
///
/// # Position model
///
/// `stack` holds `(node, sibling_index)` pairs from root to current depth;
/// `stack.last()` is the node we are currently positioned on. `pos` mirrors
/// the sibling index of the deepest stack entry — kept as a separate field
/// so [`next`](LinearIterator::next) and [`seek`](LinearIterator::seek) can
/// advance it without re-popping. To list the *siblings* of the current
/// node we read the children of the parent: `stack[len - 2]` for depth ≥ 2,
/// or [`TreeTrie::children`] for depth 1.
///
/// Compare to [`ColumnTrieIter`](super::super::column_trie::ColumnTrie):
/// where `ColumnTrieIter` carries three integer coordinates over flat
/// arrays, `TreeTrieIter` walks pointer-linked nodes via the explicit
/// stack.
#[derive(IntoTrieIter)]
struct TreeTrieIter<'a> {
    /// Sibling index of the deepest stack entry (the current position).
    pos: usize,
    /// The trie being iterated.
    trie: &'a TreeTrie,
    /// Path from the root to the current depth.
    stack: Vec<(&'a TrieNode, usize)>,
}

impl<'a> TreeTrieIter<'a> {
    fn new(trie: &'a TreeTrie) -> Self {
        Self {
            pos: 0,
            trie,
            stack: Vec::new(),
        }
    }

    /// Returns the sibling list at the current depth, or `None` if the stack
    /// is empty (iterator not yet opened).
    ///
    /// Stack invariant: each entry is `(current_node, sibling_index)` from
    /// root to the current position. To get the siblings of the current node
    /// we need the *parent's* children list:
    /// - depth 1 → parent is the trie root, so siblings are `trie.children()`
    /// - depth 2+ → parent is at `stack[len - 2]`, so siblings are its children
    fn siblings(&self) -> Option<&'a Vec<TrieNode>> {
        if self.stack.is_empty() {
            None
        } else if self.stack.len() == 1 {
            Some(self.trie.children())
        } else {
            Some(self.stack[self.stack.len() - 2].0.children())
        }
    }
}

impl LinearIterator for TreeTrieIter<'_> {
    fn key(&self) -> Option<usize> { Some(self.siblings()?.get(self.pos)?.key()) }

    fn next(&mut self) -> Option<usize> {
        if let Some(siblings) = self.siblings() {
            if self.at_end() {
                return None;
            }
            self.pos += 1;
            if let Some(node) = siblings.get(self.pos) {
                self.stack.pop();
                self.stack.push((node, self.pos));
                return Some(node.key());
            }
        }
        None
    }

    /// Advances to the least upper bound of `seek_key` among the current
    /// siblings. Returns `false` if no key `≥ seek_key` remains.
    ///
    /// # Panics
    ///
    /// Panics if `seek_key < self.key()`. `seek` is only valid for
    /// **forward** moves — every iterator in the [`LinearIterator`] contract
    /// is positioned at the start of a sorted sequence and may only advance.
    /// The caller is responsible for ensuring this; the join algorithms in
    /// `kermit-algos` enforce it via the leapfrog ring's invariants.
    fn seek(&mut self, seek_key: usize) -> bool {
        if self.at_end() {
            return false;
        }

        if let Some(current_key) = self.key() {
            if current_key > seek_key {
                panic!("seek_key must be ≥ the key at the current position");
            } else {
                let siblings = self
                    .siblings()
                    .expect("If there exists a key, there should ALWAYS be at least one sibling");

                while (!self.at_end()) && seek_key > siblings[self.pos].key() {
                    self.pos += 1;
                }

                if self.at_end() {
                    false
                } else {
                    self.stack.pop();
                    self.stack.push((&siblings[self.pos], self.pos));
                    true
                }
            }
        } else {
            false
        }
    }

    fn at_end(&self) -> bool {
        if let Some(siblings) = self.siblings() {
            self.pos >= siblings.len()
        } else {
            true
        }
    }
}

impl TrieIterator for TreeTrieIter<'_> {
    fn open(&mut self) -> bool {
        if let Some((node, _)) = self.stack.last() {
            if let Some(child) = node.children().first() {
                self.stack.push((child, 0));
                self.pos = 0;
                true
            } else {
                false
            }
        } else if self.trie.children().is_empty() {
            false
        } else {
            self.stack.push((&self.trie.children()[0], 0));
            self.pos = 0;
            true
        }
    }

    fn up(&mut self) -> bool {
        if self.stack.pop().is_some() {
            self.pos = if let Some((_, i)) = self.stack.last() {
                *i
            } else {
                0
            };
            true
        } else {
            false
        }
    }
}

impl TrieIterable for TreeTrie {
    fn trie_iter(&self) -> impl TrieIterator + IntoIterator<Item = Vec<usize>> {
        TreeTrieIter::new(self)
    }
}