path: root/sys/nixpkgs/pkgs/lf-make-map/src/mapping/map_tree.rs

use std::{
    arch::x86_64::_MM_FROUND_CUR_DIRECTION, cmp::Ordering, collections::HashMap, fmt::Display, mem,
};

use log::{debug, info};

use super::{error, MapKey};

/// A prefix tree
pub struct MappingTree {
    root: Node,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum NodeValue {
    Parent { children: HashMap<MapKey, Node> },
    Child { path: String },
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Node {
    value: NodeValue,
}

impl MappingTree {
    pub fn new() -> Self {
        Self {
            root: Node::new_parent(),
        }
    }

    /// Returns the node at the key, otherwise None
    pub fn get(&self, key: &[MapKey]) -> Option<&Node> {
        let mut current_node = &self.root;
        for ch in key.iter() {
            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: &[MapKey]) -> Option<&mut Node> {
        let mut current_node = &mut self.root;
        for ch in key.iter() {
            if let NodeValue::Parent { children } = &mut current_node.value {
                current_node = children.get_mut(&ch)?
            } else {
                return None;
            }
        }

        Some(current_node)
    }

    /// Returns the node at the key, otherwise the last node that matched.
    pub fn try_get(&self, key: &[MapKey]) -> (&Node, Vec<MapKey>) {
        let mut current_node = &self.root;
        let mut current_key = vec![];

        for ch in key.iter() {
            if let NodeValue::Parent { children } = &current_node.value {
                current_node = if let Some(node) = children.get(&ch) {
                    let (key, _value) = 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)
    }

    pub fn include(&mut self, path: &str) {
        let associated_key = MapKey::new_ones_from_path(path, 1);
        self.insert(&associated_key, path);
    }

    pub fn insert(&mut self, key: &[MapKey], path: &str) {
        let (_node, found_key) = self.try_get(key).clone();

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

            let needed_nodes_length = needed_nodes_key.iter().count();

            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.iter() {
                current_location.push(ch.to_owned());

                let next_node = if counter == needed_nodes_length {
                    Node::new_child(path.to_owned())
                } else {
                    Node::new_parent()
                };

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

                        let children =
                            if let NodeValue::Parent { children } = &mut current_node.value {
                                children
                            } else {
                                unreachable!("This is a parent, we cheched")
                            };

                        children.insert(ch.to_owned(), next_node);
                        children.get_mut(&ch).expect("Was just inserted")
                    }
                    NodeValue::Child { path } => {
                        // A node that should be a parent was classified
                        // as child before:
                        //
                        //  1. Remove the child node and replace it with a parent one.
                        //  2. Add the child node to the parent node as child, but with a '.' as MapKey.
                        //  3. Add the original node also as child to the parent node.

                        let mut children = HashMap::new();
                        let move_child_node = Node::new_child(path.to_owned());

                        children.insert(
                            MapKey {
                                key: ".".to_owned(),
                                part_path: ".".to_owned(),
                                resolution: 1,
                            },
                            move_child_node,
                        );
                        children.insert(ch.to_owned(), next_node);

                        current_node.value = NodeValue::Parent { children };

                        let children =
                            if let NodeValue::Parent { children } = &mut current_node.value {
                                children
                            } else {
                                unreachable!("We just inserted the parent value.")
                            };

                        children.get_mut(&ch).expect("Was just inserted")
                    }
                };

                counter += 1;
            }
        } else {
            // Another node was already inserted with the same key!
            // So we simple increase the resolution of the other node and this node, until their
            // keys are not the same anymore.
            // This only includes the last segment of the `MapKey`
            //
            // 1. Change both keys, until they are not equal any more
            // 2. Move the wrongly placed node to the new place.
            // 3. Insert our node.
            let mut foreign_key = found_key.last().expect("This will exist").clone();
            let mut our_key = key.last().expect("This will exist").clone();

            debug!(
                "'{}' ('{}') and '{}' ('{}') are the same, try to find a better combination!",
                our_key, our_key.part_path, foreign_key, foreign_key.part_path,
            );

            while our_key == foreign_key {
                our_key.increment();
                foreign_key.increment();
            }

            debug!(
                "Found a better one: '{}' ('{}') and '{}' ('{}')",
                our_key, our_key.part_path, foreign_key, foreign_key.part_path,
            );

            let parent = self
                .get_mut(&found_key[..found_key.len() - 1])
                .expect("This node will exist");

            if let NodeValue::Parent { children } = &mut parent.value {
                let old = children
                    .remove(found_key.last().expect("This will exist"))
                    .expect("This will be there");
                children.insert(foreign_key, old);
                children.insert(our_key, Node::new_child(path.to_owned()));
            } else {
                unreachable!("This node will be a parent");
            }
        }
    }
}

impl Node {
    pub fn new_child(path: String) -> Self {
        Self {
            value: NodeValue::Child { path },
        }
    }
    pub fn new_parent() -> Self {
        Self {
            value: NodeValue::Parent {
                children: HashMap::new(),
            },
        }
    }
}

impl Display for MappingTree {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        fn write_node(
            node: &Node,
            indention: String,
            location: Vec<MapKey>,
            has_parent: bool,
            is_last: bool,
            f: &mut std::fmt::Formatter<'_>,
        ) -> std::fmt::Result {
            let node_value = match &node.value {
                NodeValue::Parent { children: _ } => "<Parent>",
                NodeValue::Child { path } => path,
            };

            let new_idention = if !has_parent {
                String::new()
            } else if is_last {
                indention.replace('│', " ") + "    "
            } else {
                indention.clone() + "│   "
            };

            let bullet = if has_parent {
                if is_last {
                    String::from("└── ")
                } else {
                    String::from("├── ")
                }
            } else {
                String::new()
            };

            write!(
                f,
                "{}{}\x1b[1;33m{}\x1b[0m: {}\n",
                indention,
                bullet,
                MapKey::display(&location),
                node_value,
            )?;

            match &node.value {
                NodeValue::Parent { children } => {
                    let value_length = children.len();
                    let mut counter = 1;

                    let mut children_vec: Vec<(&MapKey, &Node)> = children.iter().collect();
                    children_vec.sort_by(|(a, _), (b, _)| a.key.cmp(&b.key));

                    for (key, child) in children_vec {
                        let mut new_location = location.clone();
                        new_location.push(key.to_owned());

                        if counter == value_length {
                            write_node(
                                child,
                                new_idention.clone(),
                                new_location.clone(),
                                true,
                                true,
                                f,
                            )?;
                        } else {
                            write_node(
                                child,
                                new_idention.clone(),
                                new_location.clone(),
                                true,
                                false,
                                f,
                            )?;
                        };
                        counter += 1;
                    }
                }
                NodeValue::Child { path: _ } => {
                    // Do nothing and stop the recursion
                }
            }

            Ok(())
        }

        write_node(&self.root, String::new(), vec![], false, false, f)?;

        Ok(())
    }
}