hb/src/directory.rs

use std::fs::read_dir;
use std::hint::unreachable_unchecked;
use std::path::{Path, PathBuf};
use std::io::{Result, ErrorKind};

use crate::unit::Unit;

#[derive(Debug, Clone)]
pub struct Directory {
    name: PathBuf,
    size: u64,
    children: Vec<Directory>,
}
impl Directory {
    #[inline]
    pub const fn size(&self) -> u64 {
        self.size
    }

    #[inline]
    pub fn scale(&self, unit: Unit) -> String {
        unit.convert(self.size)
    }

    #[inline]
    pub fn path(&self) -> &Path {
        self.name.as_ref()
    }

    pub fn new<P>(path: P) -> Result<Self>
    where
        P: AsRef<Path>
    {
        let path = path.as_ref();
        let name = path.canonicalize()?
            .file_name()
            // file_name() returns None if and only if the path ends in "..".
            // due to the call to canonicalize(), this can never be the case,
            // so this can be safely unwrapped
            .unwrap_or_else(|| unsafe { unreachable_unchecked() })
            .into();

        let dir = match read_dir(path) {
            Ok(dir) => dir,
            Err(io_error) => match io_error.kind() {
                ErrorKind::NotADirectory => return Ok(Self {
                    name,
                    size: path.metadata()?.len(),
                    children: Vec::new()
                }),
                _ => return Err(io_error),
            }
        };

        let mut size = 0;
        let mut children = Vec::new();
        for entry in dir {
            let child = Self::new(entry?.path())?;
            size += child.size;
            children.push(child);
        }

        Ok(Self{
            name,
            size,
            children,
        })
    }

    pub fn display(self, unit: Unit) -> String {
        // since self.size is definitionally the greatest value, the tab length
        // is just the length of self.len, plus two for a tab width
        let tab_size = self.size.to_string().len() + 2;
        let mut result = self.vectorise(unit).iter().map(|e| e.stringify_tabbed(tab_size) + "\n").collect::<String>();

        if ! result.is_empty() {
            let final_newline_char_range = result.len()-2 .. result.len();
            result.drain(final_newline_char_range);
        }

        result
    }

    /// TODO: make not recursive, take &self if possible,
    /// and maybe write directly to stdout to not use so much mem
    fn vectorise(self, unit: Unit) -> Vec<TreeEntry> {
        let mut result = Vec::new();

        result.push(TreeEntry::new(
            Vec::new(), self.name.display().to_string(), self.size, unit
        ));

        let mut new_entry_part = TreePart::First;
        let mut continue_part = TreePart::Wait;

        let len = self.children.len();

        for (idx, child) in self.children.into_iter().enumerate() {
            if idx+1 == len {
                new_entry_part = TreePart::Last;
                continue_part = TreePart::Blank;
            }

            let subtree = child.vectorise(unit);

            for mut item in subtree {
                if item.parts.is_empty() {
                    item.parts.push(new_entry_part);
                } else {
                    item.parts.push(continue_part);
                }
                result.push(item);
            }
        }

        result
    }
}

#[derive(Debug)]
struct TreeEntry {
    parts: Vec<TreePart>,
    path: String,
    size: u64,
    unit: Unit
}
impl TreeEntry {
    fn new(parts: Vec<TreePart>, path: String, size: u64, unit: Unit) -> Self {
        Self {
            parts, path, size, unit
        }
    }
    
    fn stringify_tabbed(&self, tab_size: usize) -> String {
        let mut result = format!("{:<tab_size$}", self.unit.convert(self.size));

        for part in self.parts.iter().rev() {
            result += part.display();
        }
        // dont add the space to empty entries
        result += " ";
        result += &self.path;

        result
    }
}

#[derive(PartialEq, Eq, Debug, Clone, Copy)]
enum TreePart {
    First,
    Wait,
    Last,
    Blank
}
impl TreePart {
    /// convert to ascii art
    pub const fn display(&self) -> &str {
        match self {
            Self::First => "├──",
            Self::Wait  => "│  ",
            Self::Last  => "└──",
            Self::Blank => "   ",
        }
    }
}