dirstat/src/data/graph.rs

use super::*;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::borrow::Borrow;
use std::iter::FusedIterator;
use std::cell::RefCell;

#[derive(Debug, Clone, PartialEq, Eq)]
/// A reference to an INode within a graph.
pub struct INodeRef<'a>(&'a INodeInfoGraph, INode);

impl<'a> AsRef<Path> for INodeRef<'a>
{
    fn as_ref(&self) -> &Path
    {
	self.0.paths_reverse.get(&self.1).unwrap()
    }
}

impl<'a> Borrow<INode> for INodeRef<'a>
{
    fn borrow(&self) -> &INode
    {
	&self.1
    }
}


impl<'a> INodeRef<'a>
{
    /// The owning graph of this INode reference.
    #[inline] pub fn graph(&self) -> &INodeInfoGraph
    {
	self.0
    }

    /// The `FsInfo` for this INode
    #[inline] pub fn info(&self) -> &FsInfo
    {
	self.0.inodes.get(&self.1).unwrap()
    }

    /// Compute the total size of this INode.
    ///
    /// If this is a file, it is O(1). Otherwise, recursively compute the sizes of all children.
    #[inline] pub fn size(&self) -> u64
    {
	self.0.lookup_size_or(&self.1, move || {
	    match self.info()
	    {
		FsInfo::File(sz, _) => *sz,
		FsInfo::Directory(_) => {
		    self.children().map(|x| x.size()).sum()
		},
	    }	    
	})
    }

    /// The internal INode
    #[inline] pub fn inode(&self) -> INode
    {
	self.1
    }

    /// The path this INode refers to
    #[inline] pub fn path(&self) -> &Path
    {
	self.as_ref()
    }

    /// Get an iterator over the children of this INode, if it has children.
    #[inline] pub fn children(&self) -> Children<'a>
    {
	self.0.children_of(&self.1)
    }
}

/// Contains a graph of all paths and inodes that were successfully stat'd
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct INodeInfoGraph
{
    inodes: HashMap<INode, FsInfo>, // FsInfo contains parent INode that can be used to look up again in this table
    paths: HashMap<PathBuf, INode>, // map absolute paths to INodes to be looked up in `inodes` table.

    paths_reverse: HashMap<INode, PathBuf>, // reverse lookup of INode to PathBuf,
    
    children: HashMap<INode, Vec<INode>>, //reverse lookup for directory INodes and their parent INodes

    total_sizes_cached: RefCell<HashMap<INode, u64>>, // the total sizes of all INodes
}

impl INodeInfoGraph
{
    fn lookup_insert_size<F>(&self, node: INode, with: F) -> u64
    where F: FnOnce() -> u64
    {
	let w = with();

	if let Ok(mut cache) = self.total_sizes_cached.try_borrow_mut()
	{
	    cache.insert(node, w);
	    w
	} else {
	    w
	}
    }
    fn lookup_size_or<F>(&self, node: impl Borrow<INode>, or: F) -> u64
    where F: FnOnce() -> u64
    {
	if let Ok(cache) = self.total_sizes_cached.try_borrow() {
	    let node = node.borrow();
	    if let Some(sz) = cache.get(node) {
		return *sz;
	    }
	}
	self.lookup_insert_size(*node.borrow(), or)
    }
    /// Create a new graph from these linked `HashMap`s
    #[inline] pub fn new(inodes: HashMap<INode, FsInfo>, paths: HashMap<PathBuf, INode>) -> Self
    {
	Self {
	    children: HashMap::with_capacity(inodes.len()),
	    paths_reverse: HashMap::with_capacity(paths.len()),
	    total_sizes_cached: RefCell::new(HashMap::with_capacity(inodes.len())),
	    inodes,
	    paths,
	}
	.compute_child_table()
	    .compute_reverse_path_table()
    }
    #[inline] fn compute_child_table(mut self) -> Self
    {
	for (node, info) in self.inodes.iter()
	{
	    match info {
		FsInfo::Directory(parent_node) |
		FsInfo::File(_, parent_node) => {
		    self.children.entry(*parent_node).or_insert_with(|| Vec::new()).push(*node);
		},
	    }
	}

	self
    }
    #[inline] fn compute_reverse_path_table(mut self) -> Self
    {
	self.paths_reverse.extend(self.paths.iter().map(|(x,y)| (*y,x.clone())));
	self
    }
    /// Get the FsInfo of this `INode`
    #[inline] pub fn get_info(&self, node: impl Borrow<INode>) -> Option<&FsInfo>
    {
	self.inodes.get(node.borrow())
    }
    /// An iterator over top-level children of this node
    pub fn children_of(&self, node: impl Borrow<INode>) -> Children<'_>
    {
	Children(self, match self.children.get(node.borrow()) {
	    Some(slc) => slc.iter(),
	    _ => [].iter(),
	})
    }
    /// An iterator over all the directories in this 
    pub fn directories(&self) -> Directories<'_>
    {
	Directories(self, self.children.keys())
    }
}

/// An iterator over all directories in a graph
#[derive(Debug, Clone)]
pub struct Directories<'a>(&'a INodeInfoGraph, std::collections::hash_map::Keys<'a, INode, Vec<INode>>);

impl<'a> Iterator for Directories<'a>
{
    type Item = INodeRef<'a>;
    #[inline] fn next(&mut self) -> Option<Self::Item>
    {
	self.1.next().map(|x| INodeRef(self.0, *x))
    }
    #[inline] fn size_hint(&self) -> (usize, Option<usize>) {
	self.1.size_hint()
    }
}
impl<'a> ExactSizeIterator for Directories<'a>{}
impl<'a> FusedIterator for Directories<'a>{}

/// The immediate children of a specific INode 
#[derive(Debug, Clone)]
pub struct Children<'a>(&'a INodeInfoGraph, std::slice::Iter<'a, INode>);

impl<'a> Iterator for Children<'a>
{
    type Item = INodeRef<'a>;
    #[inline] fn next(&mut self) -> Option<Self::Item>
    {
	self.1.next().map(|x| INodeRef(self.0, *x))
    }
    fn size_hint(&self) -> (usize, Option<usize>) {
	self.1.size_hint()
    }
}
impl<'a> ExactSizeIterator for Children<'a>{}
impl<'a> FusedIterator for Children<'a>{}
impl<'a> DoubleEndedIterator for Children<'a>
{
    #[inline] fn next_back(&mut self) -> Option<Self::Item> {
	self.1.next_back().map(|x| INodeRef(self.0, *x))
    }
}