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 }, 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 } = ¤t_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) { let mut current_node = &self.root; let mut current_key = vec![]; for ch in key.iter() { if let NodeValue::Parent { children } = ¤t_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 ¤t_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, has_parent: bool, is_last: bool, f: &mut std::fmt::Formatter<'_>, ) -> std::fmt::Result { let node_value = match &node.value { NodeValue::Parent { children: _ } => "", 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(()) } }