parade-rs 2.0.0

use std::{
    collections::BTreeMap,
    fmt::{Debug, Display, Formatter},
    hash::Hash,
};

use rapidhash::RapidHashMap as HashMap;

use serde::{Deserialize, Serialize};

/// A looping-enabled nested tree data structure for use in Parade.
/// Keys must be globally unique, and there is no automatic balancing.
/// There is no root node, and nodes can contain themselves or their ancestors. This probably makes "tree" an inaccurate term.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Tree<Q, T>
where
    Q: PartialEq + Eq + Clone + Display + Debug + Hash + Ord,
    T: PartialEq + Eq + Clone,
{
    nodes: HashMap<Q, Node<T, Q>>,
}
impl<Q, T> Default for Tree<Q, T>
where
    Q: PartialEq + Eq + Clone + Display + Debug + Hash + Ord,
    T: PartialEq + Eq + Clone,
{
    fn default() -> Self {
        Self {
            nodes: Default::default(),
        }
    }
}
impl<Q, T> Tree<Q, T>
where
    Q: PartialEq + Eq + Clone + Display + Debug + Hash + Ord,
    T: PartialEq + Eq + Clone,
{
    /// Creates a new empty [Tree].
    pub fn new() -> Self {
        Self::default()
    }

    /// Fetches the parent key for the specified child key, if present.
    ///
    /// # Errors
    ///
    /// This method fails if the specified key is not present in the [Tree].
    pub fn get_parent(&self, key: &Q) -> Result<&Q, Error<Q>> {
        self.nodes
            .get(key)
            .map(|node| &node.parent)
            .ok_or_else(|| Error::NonexistentNode(key.clone()))
    }

    /// Sets the parent key for the specified child key, if present.
    ///
    /// # Errors
    ///
    /// This method fails if either of the specified keys are not present in the [Tree].
    pub fn set_parent(&mut self, child: &Q, parent: Q) -> Result<(), Error<Q>> {
        self.get(&parent)?;
        self.nodes
            .get_mut(child)
            .ok_or_else(|| Error::NonexistentNode(child.clone()))?
            .parent = parent;
        Ok(())
    }

    /// Changes the key of the specified child to the provided one, if available.
    /// All parent-child relationships are preserved.
    ///
    /// # Errors
    ///
    /// This method fails if the specified `from` key is not present in the [Tree], or if the `to` key is.
    pub fn rename(&mut self, from: &Q, to: Q) -> Result<(), Error<Q>> {
        if self.get(&to).is_ok() {
            Err(Error::NonUniqueKey(to))
        } else {
            let node = self.nodes.remove(from);
            match node {
                Some(node) => {
                    self.add(to.clone(), node)?;
                    for child in self.children(from) {
                        self.set_parent(&child, to.clone())?;
                    }
                    Ok(())
                }
                None => Err(Error::NonexistentNode(from.clone())),
            }
        }
    }

    /// Adds a new [Node] to the [Tree] at the specified key, if available.
    ///
    /// # Errors
    ///
    /// This method fails if the specified key is already present in the [Tree].
    pub fn add(&mut self, key: Q, node: Node<T, Q>) -> Result<(), Error<Q>> {
        if self.nodes.contains_key(&key) {
            Err(Error::NonUniqueKey(key))
        } else {
            self.nodes.insert(key.clone(), node);
            Ok(())
        }
    }
    /// Removes the [Node] at the specified key, if present, and returns its held value.
    ///
    /// # Errors
    ///
    /// This method fails if the specified key is not present in the [Tree].
    pub fn remove(&mut self, key: &Q, remove_children: bool) -> Result<T, Error<Q>> {
        let node = self.nodes.remove(key);
        match node {
            Some(node) => {
                if remove_children {
                    for child in self.children(key) {
                        self.remove(&child, remove_children)?;
                    }
                } else {
                    for child in self.children(key) {
                        self.set_parent(&child, node.parent.clone())?;
                    }
                }
                Ok(node.value)
            }
            None => Err(Error::NonexistentNode(key.clone())),
        }
    }
    /// Returns the [Node] at the specified key, if present.
    ///
    /// # Errors
    ///
    /// This method fails if the specified key is not present in the [Tree].
    pub fn get(&self, key: &Q) -> Result<&Node<T, Q>, Error<Q>> {
        self.nodes
            .get(key)
            .ok_or_else(|| Error::NonexistentNode(key.clone()))
    }
    /// Returns a mutable reference to the [Node] at the specified key, if present.
    ///
    /// # Errors
    ///
    /// This method fails if the specified key is not present in the [Tree].
    pub fn get_mut(&mut self, key: &Q) -> Result<&mut Node<T, Q>, Error<Q>> {
        self.nodes
            .get_mut(key)
            .ok_or_else(|| Error::NonexistentNode(key.clone()))
    }
    /// Returns the [Node] at the specified key, if it is present and its parent matches the provided parent key.
    /// Returns [None] if the key is present but has a different parent.
    ///
    /// # Errors
    ///
    /// This method fails if the specified key is not present in the [Tree].
    pub fn get_from_parent(&self, key: &Q, parent: &Q) -> Result<Option<&Node<T, Q>>, Error<Q>> {
        let node = self.get(key)?;
        if self.get_parent(key)? == parent {
            Ok(Some(node))
        } else {
            Ok(None)
        }
    }
    /// Returns a [BTreeMap] containing the [Node]s at the provided keys.
    ///
    /// # Errors
    ///
    /// This method fails if any of the provided keys are not present in the [Tree].
    pub fn get_all(&self, keys: Vec<Q>) -> Result<BTreeMap<Q, Node<T, Q>>, Error<Q>> {
        Ok(keys
            .iter()
            .map(|key| self.get(key))
            .collect::<Result<Vec<&Node<T, Q>>, Error<Q>>>()?
            .into_iter()
            .cloned()
            .enumerate()
            .map_while(|(index, value)| keys.get(index).map(|key| (key.clone(), value)))
            .collect())
    }
    /// Returns a [BTreeMap] containing the values of the [Node]s at the provided keys.
    ///
    /// # Errors
    ///
    /// This method fails if any of the provided keys are not present in the [Tree].
    pub fn get_all_values(&self, keys: Vec<Q>) -> Result<BTreeMap<Q, T>, Error<Q>> {
        Ok(self
            .get_all(keys)?
            .into_iter()
            .map(|(key, node)| (key, node.value))
            .collect())
    }

    /// Returns a [Vec] of the keys of all children of the specified parent key.
    pub fn children(&self, key: &Q) -> Vec<Q> {
        self.nodes
            .iter()
            .filter_map(|(child, node)| {
                if node.parent == *key {
                    Some(child.clone())
                } else {
                    None
                }
            })
            .collect()
    }
    /// Returns a [Vec] of the keys of all siblings of the specified key, optionally including the key itself.
    ///
    /// # Errors
    ///
    /// This method fails if the specified key is not present in the [Tree].
    pub fn siblings(&self, key: &Q, inclusive: bool) -> Result<Vec<Q>, Error<Q>> {
        let nodes = self.children(self.get_parent(key)?);
        if inclusive {
            Ok(nodes)
        } else {
            Ok(nodes.into_iter().filter(|k| k != key).collect())
        }
    }
}

/// A node in a [Tree].
/// Contains the key of its parent node (which can potentially be its own key), as well as a value.
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct Node<T, Q>
where
    Q: PartialEq + Eq + Clone + Display + Debug + Hash + Ord,
    T: PartialEq + Eq + Clone,
{
    /// The value contained in the node.
    pub value: T,
    /// The key of the parent node.
    pub parent: Q,
}
impl<T, Q> Node<T, Q>
where
    Q: PartialEq + Eq + Clone + Display + Debug + Hash + Ord,
    T: PartialEq + Eq + Clone,
{
    /// Creates a new [Node] with the specified value and parent key.
    pub fn new(value: T, parent: Q) -> Self {
        Self { value, parent }
    }
}

/// The errors that can occur when interacting with a [Tree].
#[derive(Debug)]
pub enum Error<Q>
where
    Q: PartialEq + Eq + Clone + Display + Debug + Hash + Ord,
{
    /// Occurs when an attempt is made to insert a [Node] at a key which is already in use in the given [Tree].
    NonUniqueKey(Q),
    /// Occurs when an attempt is made to fetch a [Node] at a key which is not in use in the given [Tree].
    NonexistentNode(Q),
}
impl<Q> Display for Error<Q>
where
    Q: PartialEq + Eq + Clone + Display + Debug + Hash + Ord,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::NonUniqueKey(key) => write!(f, "The key is already in use: {:?}", key),
            Self::NonexistentNode(key) => {
                write!(f, "Specified node or nodes do not exist: {:?}", key)
            }
        }
    }
}
impl<Q> std::error::Error for Error<Q> where Q: PartialEq + Eq + Clone + Display + Debug + Hash + Ord
{}