rlevtree/src/trie.rs

use super::keychecker::KeyChecker;
use super::trienode::TrieKey;
use super::trienode::TrieNode;
use std::collections::BTreeSet;
use std::iter::Iterator;
use std::marker::PhantomData;

pub enum VisitOutcome {
    Continue,
    Skip,
    EarlyExit,
}

pub struct Trie<KEY, KEYCHECKER>
where
    KEY: TrieKey,
    KEYCHECKER: KeyChecker<KEY>,
{
    pub(crate) nodes: Vec<TrieNode<KEY>>,
    tails: BTreeSet<usize>,
    checker: PhantomData<KEYCHECKER>,
}
impl<KEY, KEYCHECKER> Default for Trie<KEY, KEYCHECKER>
where
    KEY: TrieKey,
    KEYCHECKER: KeyChecker<KEY>,
{
    fn default() -> Self {
        Trie {
            nodes: vec![TrieNode::new0(None)],
            tails: BTreeSet::new(),
            checker: PhantomData,
        }
    }
}
impl<KEY, KEYCHECKER> Trie<KEY, KEYCHECKER>
where
    KEY: TrieKey,
    KEYCHECKER: KeyChecker<KEY>,
{
    pub fn trie_from_words<T, U>(wordlist: U) -> Trie<KEY, KEYCHECKER>
    where
        T: IntoIterator<Item = KEY>,
        U: IntoIterator<Item = T>,
    {
        let mut result = Trie::default();
        for word in wordlist {
            result.add(word);
        }
        result
    }

    pub(crate) fn get_node_mut(&mut self, index: usize) -> &mut TrieNode<KEY> {
        &mut self.nodes[index]
    }

    pub(crate) fn get_node(&self, index: usize) -> &TrieNode<KEY> {
        &self.nodes[index]
    }

    pub(crate) fn nodes(&self) -> usize {
        self.nodes.len()
    }

    fn add_node(&mut self, key: Option<KEY>, parent: usize, prev: Option<usize>) -> usize {
        let mut result = TrieNode::new0(key);
        let result_index = self.nodes();
        result.parent = Some(parent);
        match prev {
            Some(prev_node) => {
                self.get_node_mut(prev_node).next = Some(result_index);
                result.prev = prev;
            }
            None => {
                let parent_node = self.get_node_mut(parent);
                match parent_node.child {
                    None => {
                        parent_node.child = Some(result_index);
                    }
                    Some(parent_child) => {
                        let mut node = parent_child;
                        loop {
                            let next = self.get_node(node).next;
                            match next {
                                Some(next_node) => {
                                    node = next_node;
                                }
                                None => {
                                    break;
                                }
                            }
                        }
                        self.get_node_mut(node).next = Some(result_index);
                        result.prev = Some(node)
                    }
                }
            }
        }
        self.nodes.push(result);
        result_index
    }

    pub fn add<T>(&mut self, path: T) -> (bool, usize)
    where
        T: IntoIterator<Item = KEY>,
    {
        let mut result = false;
        let mut pnode = 0;
        'wordLoop: for key in path {
            let mut cnode = self.get_node(pnode).child;
            while let Some(cnode_index) = cnode {
                let cnode_node = self.get_node(cnode_index);
                if KEYCHECKER::check(cnode_node.key, Some(key)) {
                    pnode = cnode_index;
                    continue 'wordLoop;
                } else if self.get_node(cnode_index).next.is_none() {
                    break;
                } else {
                    cnode = self.get_node(cnode_index).next;
                }
            }
            pnode = self.add_node(Some(key), pnode, cnode);
            result = true;
        }
        if result {
            let tail = self.add_node(None, pnode, None);
            self.tails.insert(tail);
            let mut node = Some(tail);
            while let Some(n) = node {
                let current_node = self.get_node_mut(n);
                current_node.ref_count += 1;
                node = current_node.parent;
            }
            (true, tail)
        } else {
            (false, pnode)
        }
    }

    pub fn search(&mut self, path: Vec<KEY>) -> Option<usize> {
        let mut result: Option<usize> = None;
        let visit_pre = |stack: &Vec<usize>| -> VisitOutcome {
            if stack.len() == 1 {
                VisitOutcome::Continue
            } else {
                let last = *stack.last().expect("");
                let index = stack.len() - 2;
                let node = self.get_node(last);
                if index < path.len() {
                    if KEYCHECKER::check(node.key, Some(path[index])) {
                        VisitOutcome::Continue
                    } else {
                        VisitOutcome::Skip
                    }
                } else {
                    if node.key.is_none() {
                        result = Some(last);
                    }
                    VisitOutcome::EarlyExit
                }
            }
        };
        let visit_post = |_: &Vec<usize>| {};
        self.walk(visit_pre, visit_post);
        result
    }

    pub fn lineal_descendant(&self, start: usize) -> impl Iterator<Item = &KEY> {
        let mut nodes: Vec<usize> = vec![];
        let mut node_option = Some(start);
        while let Some(node) = node_option {
            let key = &self.get_node(node).key;
            if key.is_some() {
                nodes.push(node);
            }
            node_option = self.get_node(node).parent;
        }
        nodes
            .into_iter()
            .rev()
            .map(|node_index| self.get_node(node_index).key.as_ref().unwrap())
    }

    pub(crate) fn walk<CB1, CB2>(&self, mut visit_pre: CB1, mut visit_post: CB2)
    where
        CB1: FnMut(&Vec<usize>) -> VisitOutcome,
        CB2: FnMut(&Vec<usize>),
    {
        let mut stack: Vec<(usize, Option<usize>)> = vec![];
        let mut public_stack: Vec<usize> = vec![];
        let root_node = self.get_node(0);
        stack.push((0, root_node.child));
        public_stack.push(0);
        visit_pre(&public_stack);
        while !stack.is_empty() {
            let last = &mut stack.last_mut().unwrap();
            match last.1 {
                Some(child_node_id) => {
                    let child_node = self.get_node(child_node_id);
                    last.1 = child_node.next;
                    public_stack.push(child_node_id);
                    let visit_pre_outcome = visit_pre(&public_stack);
                    match visit_pre_outcome {
                        VisitOutcome::Continue => {
                            stack.push((child_node_id, child_node.child));
                        }
                        VisitOutcome::Skip => {
                            stack.push((child_node_id, None));
                        }
                        VisitOutcome::EarlyExit => return,
                    }
                }
                None => {
                    visit_post(&public_stack);
                    stack.pop();
                    public_stack.pop();
                }
            }
        }
    }

    pub fn tails(&self) -> &BTreeSet<usize> {
        &self.tails
    }
}