use std::collections::HashMap;
use std::path::{Component, Path, PathBuf};
pub fn generate_tree(paths: &[PathBuf]) -> String {
if paths.is_empty() {
return String::new();
}
let total_path_len: usize = paths.iter().map(|p| p.to_string_lossy().len()).sum();
let mut output = String::with_capacity(total_path_len + paths.len() * 8);
let mut tree = TreeNode::new();
for path in paths {
add_path_to_tree(&mut tree, path);
}
output.push_str("Directory structure:\n");
render_tree(&tree, &mut output, "", true);
output.push('\n');
output
}
#[derive(Debug)]
struct TreeNode {
name: String,
children: HashMap<String, TreeNode>,
is_file: bool,
}
impl TreeNode {
fn new() -> Self {
TreeNode {
name: String::new(),
children: HashMap::new(),
is_file: false,
}
}
fn new_with_name(name: String, is_file: bool) -> Self {
TreeNode {
name,
children: HashMap::new(),
is_file,
}
}
}
pub fn clean_path_components(path: &Path) -> Vec<String> {
path.components()
.filter_map(|component| match component {
Component::Prefix(_) | Component::RootDir => None,
Component::CurDir => None, Component::ParentDir => Some("..".to_string()), Component::Normal(os_str) => Some(os_str.to_string_lossy().to_string()),
})
.collect()
}
fn add_path_to_tree(root: &mut TreeNode, path: &Path) {
add_path_to_tree_with_type(root, path, true)
}
fn add_path_to_tree_with_type(root: &mut TreeNode, path: &Path, final_is_file: bool) {
let components = clean_path_components(path);
if components.is_empty() {
return;
}
let mut current = root;
for (i, name) in components.iter().enumerate() {
let is_last = i == components.len() - 1;
if is_last {
match current.children.get_mut(name) {
Some(existing_entry) => {
if existing_entry.is_file && !final_is_file {
existing_entry.is_file = false;
} else if !existing_entry.is_file && final_is_file {
if existing_entry.children.is_empty() {
existing_entry.is_file = true;
}
}
}
None => {
current.children.insert(
name.clone(),
TreeNode::new_with_name(name.clone(), final_is_file),
);
}
}
} else {
let entry = current
.children
.entry(name.clone())
.or_insert_with(|| TreeNode::new_with_name(name.clone(), false));
if entry.is_file {
entry.is_file = false;
}
current = entry;
}
}
}
fn render_child(
child: &TreeNode,
output: &mut String,
current_prefix: &str,
is_last: bool,
is_root: bool,
) {
if !is_root {
output.push_str(current_prefix);
}
let child_prefix = if is_last { "└── " } else { "├── " };
output.push_str(child_prefix);
output.push_str(&child.name);
if !child.is_file {
output.push('/');
}
output.push('\n');
let next_prefix = if is_root {
if is_last { " " } else { "│ " }.to_string()
} else {
let mut next = String::with_capacity(current_prefix.len() + 4);
next.push_str(current_prefix);
next.push_str(if is_last { " " } else { "│ " });
next
};
render_tree(child, output, &next_prefix, false);
}
fn render_tree(node: &TreeNode, output: &mut String, prefix: &str, is_root: bool) {
let mut children: Vec<_> = node.children.values().collect();
children.sort_by(|a, b| {
match (a.is_file, b.is_file) {
(false, true) => std::cmp::Ordering::Less,
(true, false) => std::cmp::Ordering::Greater,
_ => a.name.cmp(&b.name),
}
});
for (i, child) in children.iter().enumerate() {
let is_last = i == children.len() - 1;
render_child(child, output, prefix, is_last, is_root);
}
}