basic_trie 2.1.0

use std::cmp::Ordering;
use std::fmt::Debug;
use std::{fmt, ops};
use thin_vec::ThinVec;

use crate::child_storage::ChildStorage;
#[cfg(feature = "serde")]
use serde_crate::{Deserialize, Serialize};
use unicode_segmentation::UnicodeSegmentation;

type WordEnd<D> = Option<ThinVec<D>>;

/// Helper struct for returning multiple values for deleting data.
/// It is needed because the 'must_keep' value will at some point change
/// from false to true, but the data stays the same from the beginning of
/// unwinding.
pub(crate) struct RemoveData<D> {
    must_keep: bool,
    pub(crate) data: WordEnd<D>,
}

/// Singular trie node that represents its children and a marker for word ending.
#[cfg_attr(
    feature = "serde",
    derive(Serialize, Deserialize),
    serde(crate = "serde_crate")
)]
pub struct TrieDataNode<D> {
    #[cfg_attr(feature = "serde", serde(rename = "c"))]
    pub(crate) children: ChildStorage<TrieDataNode<D>>,
    #[cfg_attr(feature = "serde", serde(rename = "wed"))]
    word_end_data: WordEnd<D>,
}

impl<D> Default for TrieDataNode<D> {
    fn default() -> Self {
        Self {
            children: ChildStorage::default(),
            word_end_data: None,
        }
    }
}

/// Methods only on nodes that have data.
impl<D> TrieDataNode<D> {
    /// Returns a new instance of a TrieNode.
    pub(crate) fn new() -> Self {
        TrieDataNode {
            children: Default::default(),
            word_end_data: None,
        }
    }

    /// Recursive function that drops all children maps and collects data
    /// regardless of having multiple words branching from them or not.
    pub(crate) fn remove_all_words_collect(&mut self, found_data: &mut Vec<D>) -> usize {
        let num_removed = self
            .children
            .values_mut()
            .map(|child| child.remove_all_words_collect(found_data))
            .sum::<usize>()
            + self.is_associated() as usize;

        if let Some(data_vec) = self.disassociate() {
            found_data.extend(data_vec);
        }

        self.clear_children();

        num_removed
    }

    /// Recursive function that counts the number of words from a starting node.
    pub(crate) fn count_words(&self) -> usize {
        self.children
            .values()
            .map(|child| child.count_words())
            .sum::<usize>()
            + self.is_associated() as usize
    }

    /// Recursive function finds every node that is an end of a word and appends
    /// its data as references to the passed vector.
    pub(crate) fn generate_all_data<'a>(&'a self, found_data: &mut Vec<&'a D>) {
        if let Some(data_vec) = &self.word_end_data {
            found_data.extend(data_vec.iter());
        }

        self.children
            .values()
            .for_each(|x| x.generate_all_data(found_data));
    }

    /// Recursive function finds every node that is an end of a word and appends
    /// its data as mutable references to the passed vector.
    pub(crate) fn generate_all_data_mut<'a>(&'a mut self, found_data: &mut Vec<&'a mut D>) {
        if let Some(data_vec) = &mut self.word_end_data {
            found_data.extend(data_vec.iter_mut());
        }

        self.children
            .values_mut()
            .for_each(|x| x.generate_all_data_mut(found_data));
    }

    /// Function pushes data to the association vector.
    pub(crate) fn push_data(&mut self, data: D) {
        self.get_association_mut().as_mut().unwrap().push(data);
    }

    /// Recursive function for inserting found words from the given node and
    /// given starting substring.
    pub(crate) fn find_words(&self, substring: &mut String, found_words: &mut Vec<String>) {
        if self.is_associated() {
            found_words.push(substring.to_string());
        }

        for (&character, node) in self.children.iter() {
            substring.push(character);
            node.find_words(substring, found_words);
            substring.pop();
        }
    }

    /// The recursive function for finding a vector of shortest and longest words in the TrieNode consists of:
    /// - the DFS tree traversal part for getting to every child node;
    /// - matching lengths of found words in combination with the passed ordering.
    pub(crate) fn words_min_max(
        &self,
        substring: &mut String,
        current_visual_len: usize,
        found_words: &mut Vec<String>,
        current_best_len: &mut Option<usize>,
        #[cfg(feature = "unicode")] is_path_pure_ascii: bool,
        ord: Ordering,
    ) {
        if self.is_associated() {
            match current_best_len {
                Some(best_len) => match current_visual_len.cmp(best_len) {
                    o if o == ord => {
                        *best_len = current_visual_len;
                        found_words.clear();
                        found_words.push(substring.clone());
                    }
                    Ordering::Equal => {
                        found_words.push(substring.clone());
                    }
                    _ => {}
                },
                None => {
                    *current_best_len = Some(current_visual_len);
                    found_words.push(substring.clone());
                }
            }
        }

        for (&character, node) in self.children.iter() {
            substring.push(character);

            #[cfg(feature = "unicode")]
            let (next_visual_len, next_is_ascii) = if is_path_pure_ascii && character.is_ascii() {
                (current_visual_len + 1, true)
            } else {
                (substring.graphemes(true).count(), false)
            };

            #[cfg(not(feature = "unicode"))]
            let next_visual_len = current_visual_len + 1;

            node.words_min_max(
                substring,
                next_visual_len,
                found_words,
                current_best_len,
                #[cfg(feature = "unicode")]
                next_is_ascii,
                ord,
            );
            substring.pop();
        }
    }

    /// Function resets the association of a word and returns the
    /// previous association. If 'keep_word' is true, the association is only
    /// reset.
    pub(crate) fn clear_word_end_association(&mut self, keep_word: bool) -> WordEnd<D> {
        let return_data = self.disassociate();

        if keep_word && return_data.is_some() {
            self.associate();
        }

        return_data
    }

    /// Recursive function for removing and freeing memory of a word that is not needed anymore.
    /// The algorithm first finds the last node of a word given in the form of a character iterator,
    /// then it frees the maps and unwinds to the first node that should not be deleted.
    /// The first node that should not be deleted is either:
    /// - the root node
    /// - the node that has multiple words branching from it
    /// - the node that represents an end to some word with the same prefix
    /// The last node's data is propagated all the way to the final return
    /// with the help of auxiliary 'RemoveData<D>' struct.
    pub(crate) fn remove_one_word<'b>(
        &mut self,
        mut characters: impl Iterator<Item = char>,
    ) -> RemoveData<D> {
        let next_character = match characters.next() {
            None => {
                return RemoveData {
                    must_keep: false,
                    data: self.disassociate(),
                }
            }
            Some(char) => char,
        };

        let next_node = self.children.get_mut(next_character).unwrap();
        let must_keep = next_node.remove_one_word(characters);

        if self.children.len() > 1 || must_keep.must_keep {
            return RemoveData {
                must_keep: true,
                data: must_keep.data,
            };
        }
        self.clear_children();

        RemoveData {
            must_keep: self.is_associated(),
            data: must_keep.data,
        }
    }

    /// Function marks the node as an end of a word.
    pub(crate) fn associate(&mut self) {
        self.word_end_data = Some(ThinVec::new());
    }

    /// Function unmarks the node as an end of a word and returns the data.
    pub(crate) fn disassociate(&mut self) -> WordEnd<D> {
        self.word_end_data.take()
    }

    /// Function returns true if an association is found for the word.
    pub(crate) fn is_associated(&self) -> bool {
        self.word_end_data.is_some()
    }

    /// Function returns the node association.
    pub(crate) fn get_association(&self) -> &WordEnd<D> {
        &self.word_end_data
    }

    /// Function returns the mutable node association.
    pub(crate) fn get_association_mut(&mut self) -> &mut WordEnd<D> {
        &mut self.word_end_data
    }

    /// Function removes all children of a node.
    pub(crate) fn clear_children(&mut self) {
        self.children = Default::default();
    }
}

impl<D> ops::AddAssign for TrieDataNode<D> {
    /// Overriding the += operator on nodes.
    /// Function adds two nodes based on the principle:
    /// for every child node and character in the 'rhs' node:
    /// - if the self node doesn't have that character in its children map,
    /// simply move the pointer to the self's children map without any extra cost;
    /// - if the self node has that character, the node of that character (self's child)
    /// is added with the 'rhc's' node.
    /// An edge case exists when the 'rhc's' node has an association but self's node doesn't.
    /// That association is handled based on the result of 'rhc_next_node.word_end_data'.
    /// On Some(data), the self node vector is initialized with the 'rhc' node vector.
    fn add_assign(&mut self, rhs: Self) {
        for (char, mut rhs_next_node) in rhs.children.into_iter() {
            // Does self contain the character?
            match self.children.remove(char) {
                // The whole node is removed, as owned, operated on and returned in self's children.
                Some(mut self_next_node) => {
                    // Edge case: associate self node if the other node is also associated
                    // Example: when adding 'word' to 'word1', 'd' on 'word' needs to be associated
                    if let Some(data_vec_rhs) = rhs_next_node.word_end_data.take() {
                        if let Some(data_vec_self) = &mut self_next_node.word_end_data {
                            data_vec_self.extend(data_vec_rhs);
                        } else {
                            self_next_node.word_end_data = Some(data_vec_rhs);
                        }
                    }

                    self_next_node += rhs_next_node;
                    self.children.insert_direct(char, self_next_node);
                }
                // Self doesn't contain the character, no conflict arises.
                // The whole 'rhs' node is just moved from 'rhs' into self.
                None => {
                    self.children.insert_direct(char, rhs_next_node);
                }
            }
        }
    }
}

impl<D: PartialEq> PartialEq for TrieDataNode<D> {
    /// Operation == can be applied only to TrieNodes whose data implements PartialEq.
    fn eq(&self, other: &Self) -> bool {
        // If keys aren't equal, nodes aren't equal.
        if !self.children.has_same_keys(&other.children) {
            return false;
        }

        // If associations aren't equal, two nodes aren't equal.
        if !match (&self.word_end_data, &other.word_end_data) {
            (Some(self_vec), Some(other_vec)) => {
                // If they both have an association, return true only if the data is identical
                self_vec.len() == other_vec.len() && self_vec.iter().all(|k| other_vec.contains(k))
            }
            // If they both don't have an association, return true
            (None, None) => true,
            _ => false,
        } {
            return false;
        }

        // Every child node that has the same key (character) must be equal.
        self.children
            .iter()
            .map(|(char, self_child)| (self_child, other.children.get(*char).unwrap()))
            .all(|(self_child, other_child)| other_child == self_child)
    }
}

impl<D: Debug> Debug for TrieDataNode<D> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut s = f.debug_struct("Node");

        if let Some(data) = &self.word_end_data {
            s.field("data", &data);
        }

        s.field("children", &self.children);
        s.finish()
    }
}