use std::collections::HashMap;
use std::error::Error as StdError;
use std::path::Path;
#[derive(Debug, Clone)]
pub enum FileTreeNode {
File {
contents: Vec<u8>,
},
Directory {
files: HashMap<String, FileTreeNode>,
},
}
impl FileTreeNode {
pub fn get_node<P: AsRef<Path>>(&self, path: P) -> Option<&FileTreeNode> {
let path = path.as_ref();
if path == Path::new("") {
return Some(self);
}
let components: Vec<_> = path
.components()
.filter_map(|c| {
if let std::path::Component::Normal(s) = c {
s.to_str()
} else {
None
}
})
.collect();
if components.is_empty() {
return Some(self);
}
let mut current_node = self;
for component in components {
match current_node {
FileTreeNode::Directory { files } => {
current_node = files.get(component)?;
}
FileTreeNode::File { .. } => {
return None; }
}
}
Some(current_node)
}
pub fn get_file<P: AsRef<Path>>(&self, path: P) -> Option<&[u8]> {
match self.get_node(path)? {
FileTreeNode::File { contents } => Some(contents.as_slice()),
FileTreeNode::Directory { .. } => None,
}
}
pub fn file_exists<P: AsRef<Path>>(&self, path: P) -> bool {
matches!(self.get_node(path), Some(FileTreeNode::File { .. }))
}
pub fn dir_exists<P: AsRef<Path>>(&self, path: P) -> bool {
matches!(self.get_node(path), Some(FileTreeNode::Directory { .. }))
}
pub fn list_files<P: AsRef<Path>>(&self, dir_path: P) -> Vec<String> {
match self.get_node(dir_path) {
Some(FileTreeNode::Directory { files }) => files
.iter()
.filter_map(|(name, node)| {
if matches!(node, FileTreeNode::File { .. }) {
Some(name.clone())
} else {
None
}
})
.collect(),
_ => Vec::new(),
}
}
pub fn list_subdirectories<P: AsRef<Path>>(&self, dir_path: P) -> Vec<String> {
match self.get_node(dir_path) {
Some(FileTreeNode::Directory { files }) => files
.iter()
.filter_map(|(name, node)| {
if matches!(node, FileTreeNode::Directory { .. }) {
Some(name.clone())
} else {
None
}
})
.collect(),
_ => Vec::new(),
}
}
pub fn insert<P: AsRef<Path>>(
&mut self,
path: P,
node: FileTreeNode,
) -> Result<(), Box<dyn StdError + Send + Sync>> {
let path = path.as_ref();
let components: Vec<_> = path
.components()
.filter_map(|c| {
if let std::path::Component::Normal(s) = c {
s.to_str().map(|s| s.to_string())
} else {
None
}
})
.collect();
if components.is_empty() {
return Err("Cannot insert at root path".into());
}
let mut current_node = self;
for component in &components[..components.len() - 1] {
match current_node {
FileTreeNode::Directory { files } => {
current_node =
files
.entry(component.clone())
.or_insert_with(|| FileTreeNode::Directory {
files: HashMap::new(),
});
}
FileTreeNode::File { .. } => {
return Err("Cannot traverse into a file".into());
}
}
}
let filename = &components[components.len() - 1];
match current_node {
FileTreeNode::Directory { files } => {
files.insert(filename.clone(), node);
Ok(())
}
FileTreeNode::File { .. } => Err("Cannot insert into a file".into()),
}
}
pub(crate) fn from_json_value(
value: &serde_json::Value,
) -> Result<Self, Box<dyn StdError + Send + Sync>> {
if let Some(contents_str) = value.get("contents").and_then(|v| v.as_str()) {
Ok(FileTreeNode::File {
contents: contents_str.as_bytes().to_vec(),
})
} else if let Some(bytes_array) = value.get("contents").and_then(|v| v.as_array()) {
let contents: Vec<u8> = bytes_array
.iter()
.filter_map(|v| v.as_u64().and_then(|n| u8::try_from(n).ok()))
.collect();
Ok(FileTreeNode::File { contents })
} else if let Some(obj) = value.as_object() {
let mut files = HashMap::new();
for (name, child_value) in obj {
files.insert(name.clone(), Self::from_json_value(child_value)?);
}
Ok(FileTreeNode::Directory { files })
} else {
Err(format!("Invalid file tree node: {:?}", value).into())
}
}
pub fn print_tree(&self) -> String {
self.print_tree_recursive("", "", true)
}
fn print_tree_recursive(&self, name: &str, prefix: &str, is_last: bool) -> String {
let mut result = String::new();
let connector = if is_last { "└── " } else { "├── " };
let extension = if is_last { " " } else { "│ " };
match self {
FileTreeNode::File { .. } => {
result.push_str(&format!("{}{}{}\n", prefix, connector, name));
}
FileTreeNode::Directory { files } => {
result.push_str(&format!("{}{}{}/\n", prefix, connector, name));
let child_prefix = format!("{}{}", prefix, extension);
let count = files.len();
for (i, (child_name, node)) in files.iter().enumerate() {
let is_last_child = i == count - 1;
result.push_str(&node.print_tree_recursive(
child_name,
&child_prefix,
is_last_child,
));
}
}
}
result
}
}