keymaps 1.2.0

// keymaps - A rust crate which provides a standardized encoding for key codes
//
// Copyright (C) 2025 Benedikt Peetz <benedikt.peetz@b-peetz.de>
// SPDX-License-Identifier: LGPL-3.0-or-later
//
// This file is part of the keymaps crate.
//
// You should have received a copy of the License along with this program.
// If not, see <https://www.gnu.org/licenses/lgpl-3.0.txt>.

use std::collections::HashMap;

use crate::{
    error::{TrieInsert, TrieReplace},
    key_repr::Key,
};

pub mod display;
pub mod iterator;

#[cfg(test)]
mod tests;

/// A prefix tree
///
/// This is the core data structure exposed by this crate.
/// It maps keys to values by utilizing a trie.
///
/// ## [`Node`]s
/// Each [`Node`] contains either further [`Node`]s -- which is called a parent [`Node`] -- or a
/// value, which are called child [`Node`]s.
/// This makes it impossible to associate a value with a [`Node`] whilst also adding further values
/// downstream (i.e., as children to the same [`Node`].)
#[derive(Debug)]
pub struct Trie<K, V>
where
    K: Clone + Eq + std::hash::Hash + std::fmt::Debug,
    V: Clone + PartialEq + std::fmt::Debug,
{
    root: Node<K, V>,
}

/// This is an type alias for the common use-case of simply using the [`Key`]s directly.
pub type MapTrie<V> = Trie<Key, V>;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum NodeValue<K, V>
where
    K: Clone + Eq + std::hash::Hash + std::fmt::Debug,
    V: Clone + PartialEq + std::fmt::Debug,
{
    Parent { children: HashMap<K, Node<K, V>> },
    Child { value: V },
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Node<K, V>
where
    K: Clone + Eq + std::hash::Hash + std::fmt::Debug,
    V: Clone + PartialEq + std::fmt::Debug,
{
    value: NodeValue<K, V>,
}

impl<K, V> Default for Trie<K, V>
where
    K: Clone + Eq + std::hash::Hash + std::fmt::Debug,
    V: Clone + PartialEq + std::fmt::Debug,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<K, V> Trie<K, V>
where
    K: Clone + Eq + std::hash::Hash + std::fmt::Debug,
    V: Clone + PartialEq + std::fmt::Debug,
{
    #[must_use]
    pub fn new() -> Self {
        Self {
            root: Node::new_parent(),
        }
    }

    /// Returns the root [`Node`] of this [`Trie`].
    pub fn root_node(&self) -> &Node<K, V> {
        &self.root
    }

    /// Returns the [`Node`] at the `key`, otherwise [`None`].
    pub fn get(&self, key: &[K]) -> Option<&Node<K, V>> {
        self.root.get(key)
    }

    /// Returns the [`Node`] at the `key`, otherwise [`None`].
    /// The [`Node`] can be changed
    pub fn get_mut(&mut self, key: &[K]) -> Option<&mut Node<K, V>> {
        self.root.get_mut(key)
    }

    /// Returns the [`Node`] specified by the `key`, otherwise the [`Node`]
    /// that matched the most part of the `key`.
    ///
    /// To differentiate between these options, it also returns the `key` of the matched [`Node`].
    #[allow(clippy::missing_panics_doc)]
    pub fn try_get(&self, key: &[K]) -> (&Node<K, V>, Vec<K>) {
        self.root.try_get(key)
    }

    /// Insert a Value into the [`Trie`].
    /// This is effectively a convenient wrapper around [`Trie::insert_node`] that inserts a
    /// child [`Node`] with the specified value.
    ///
    /// # Errors
    /// See [`Trie::insert_node`] for possible error conditions.
    pub fn insert(&mut self, key: &[K], value: V) -> Result<(), TrieInsert<K, V>> {
        self.root.insert(key, value)
    }

    /// Replace the [`Node`] specified by the `key`.
    ///
    /// # Errors
    /// This will return an Error if the `key` does not match a [`Node`].
    pub fn replace_node(
        &mut self,
        key: &[K],
        node: Node<K, V>,
    ) -> Result<Node<K, V>, TrieReplace<K>> {
        self.root.replace_node(key, node)
    }

    /// Insert a [`Node`] into the [`Trie`].
    ///
    /// # Errors
    /// This will return an Error when
    ///   - a [`Node`] with the `key` is already inserted (use [`Trie::replace_node`] to replace it)
    ///   - the [`Node`] referenced by a part of the `key` is a child [`Node`]
    #[allow(clippy::missing_panics_doc)]
    pub fn insert_node(&mut self, key: &[K], node: &Node<K, V>) -> Result<(), TrieInsert<K, V>> {
        self.root.insert_node(key, node)
    }
}

impl<K, V> Node<K, V>
where
    K: Clone + Eq + std::hash::Hash + std::fmt::Debug,
    V: Clone + PartialEq + std::fmt::Debug,
{
    /// Returns the [`Node`] at the `key`, otherwise [`None`].
    pub fn get(&self, key: &[K]) -> Option<&Node<K, V>> {
        let mut current_node = self;
        for ch in key {
            if let NodeValue::Parent { children } = &current_node.value {
                current_node = children.get(ch)?;
            } else {
                return None;
            }
        }

        Some(current_node)
    }

    /// Returns the [`Node`] at the `key`, otherwise [`None`].
    /// The [`Node`] can be changed
    pub fn get_mut(&mut self, key: &[K]) -> Option<&mut Node<K, V>> {
        let mut current_node = self;
        for ch in key {
            if let NodeValue::Parent { children } = &mut current_node.value {
                current_node = children.get_mut(ch)?;
            } else {
                return None;
            }
        }

        Some(current_node)
    }

    /// Returns the [`Node`] specified by the `key`, otherwise the [`Node`]
    /// that matched the most part of the `key`.
    ///
    /// To differentiate between these options, it also returns the `key` of the matched [`Node`].
    #[allow(clippy::missing_panics_doc)]
    pub fn try_get(&self, key: &[K]) -> (&Node<K, V>, Vec<K>) {
        let mut current_node = self;
        let mut current_key = vec![];

        for ch in key {
            if let NodeValue::Parent { children } = &current_node.value {
                current_node = if let Some(node) = children.get(ch) {
                    let (key, _) = children.get_key_value(ch).expect("This exists, we checked");
                    current_key.push(key.clone());

                    node
                } else {
                    return (current_node, current_key);
                };
            } else {
                return (current_node, current_key);
            }
        }

        (current_node, current_key)
    }

    /// Insert a Value into this [`Node`].
    /// This is effectively a convenient wrapper around [`Node::insert_node`] that inserts a
    /// child [`Node`] with the specified value.
    ///
    /// # Errors
    /// See [`Node::insert_node`] for possible error conditions.
    pub fn insert(&mut self, key: &[K], value: V) -> Result<(), TrieInsert<K, V>> {
        self.insert_node(key, &Node::new_child(value))
    }

    /// Replace the [`Node`] specified by the `key`.
    ///
    /// # Errors
    /// This will return an Error if the `key` does not match a [`Node`].
    pub fn replace_node(
        &mut self,
        key: &[K],
        node: Node<K, V>,
    ) -> Result<Node<K, V>, TrieReplace<K>> {
        if let Some(old_node) = self.get_mut(key) {
            Ok(std::mem::replace(old_node, node))
        } else {
            Err(TrieReplace::NoNodeWithKey(key.to_vec()))
        }
    }

    /// Insert a child [`Node`] into this [`Node`].
    ///
    /// # Errors
    /// This will return an Error when
    ///   - a [`Node`] with the `key` is already inserted (use [`Node::replace_node`] to replace it)
    ///   - the [`Node`] referenced by a part of the `key` is a child [`Node`]
    #[allow(clippy::missing_panics_doc)]
    pub fn insert_node(&mut self, key: &[K], node: &Node<K, V>) -> Result<(), TrieInsert<K, V>> {
        let (_, found_key) = self.try_get(key).clone();

        if found_key == key {
            // Another node was already inserted with the same key!
            // NOTE(@bpeetz): We cannot use `key` here, because “equal” does not mean truly equal. <2025-04-27>
            return Err(TrieInsert::KeyAlreadySet(found_key.clone()));
        }

        let needed_nodes_key = key
            .strip_prefix(&found_key[..])
            .expect("The node's location is a prefix");

        let needed_nodes_length = needed_nodes_key.len();

        let mut current_node = self
            .get_mut(&found_key[..])
            .expect("This should always exists");
        let mut current_location = found_key.clone();
        let mut counter = 1;

        for ch in needed_nodes_key {
            current_location.push(ch.to_owned());

            let next_node = if counter == needed_nodes_length {
                node.clone()
            } else {
                Node::new_parent()
            };

            current_node = match &mut current_node.value {
                NodeValue::Parent { children } => {
                    assert_eq!(children.get(ch), None);

                    children.insert(ch.to_owned(), next_node);
                    children.get_mut(ch).expect("Was just inserted")
                }
                NodeValue::Child { value } => {
                    // A node that should be a parent was added
                    // as a child before.

                    return Err(TrieInsert::KeyIncludesChild {
                        child_key: key.to_vec(),
                        child_value: value.to_owned(),
                    });
                }
            };

            counter += 1;
        }
        Ok(())
    }
}

impl<K, V> Node<K, V>
where
    K: Clone + Eq + std::hash::Hash + std::fmt::Debug,
    V: Clone + PartialEq + std::fmt::Debug,
{
    /// Construct a new child Node.
    #[must_use]
    pub fn new_child(value: V) -> Self {
        Self {
            value: NodeValue::Child { value },
        }
    }

    /// Construct a new parent Node.
    #[must_use]
    pub fn new_parent() -> Self {
        Self {
            value: NodeValue::Parent {
                children: HashMap::new(),
            },
        }
    }

    /// Return the value of this node.
    /// This returns [`None`], if this Node is a Parent.
    pub fn value(&self) -> Option<&V> {
        match &self.value {
            NodeValue::Parent { children: _ } => None,
            NodeValue::Child { value } => Some(value),
        }
    }

    /// Return the children of this node.
    /// This returns [`None`], if this Node is a Child.
    pub fn children(&self) -> Option<&HashMap<K, Self>> {
        match &self.value {
            NodeValue::Parent { children } => Some(children),
            NodeValue::Child { value: _ } => None,
        }
    }
}