dirstat/src/main.rs

#![feature(try_blocks)]

#![allow(dead_code)]

#[macro_use] extern crate pin_project;
#[macro_use] extern crate lazy_static;
#[macro_use] extern crate cfg_if;
#[macro_use] extern crate ad_hoc_iter;

#[cfg(feature="inspect")] use serde::{Serialize, Deserialize};

#[cfg(feature="jemalloc")] 
use jemallocator::Jemalloc;

#[cfg(feature="jemalloc")] 
#[global_allocator]
static GLOBAL: Jemalloc = Jemalloc;

use std::convert::{TryFrom, TryInto};
use color_eyre::{
    eyre::{
	self,
	eyre,

	WrapErr as _,
    },
    
    Help as _,
    SectionExt as _,
};

#[macro_use] mod ext;
pub use ext::prelude::*;
use std::sync::Arc;

mod bytes;
mod data;
mod config;
mod state;
mod arg;
mod work;
mod info;
#[cfg(feature="inspect")] mod serial;
#[cfg(feature="defer-drop")] mod defer_drop;

#[cfg(feature="inspect")]
async fn write_graph(graph: Arc<data::HierarchicalINodeGraph>) -> eyre::Result<()>
{
    let cfg = config::get_global();
    match cfg.serialise_output.as_ref().map(|ser_out| {
	cfg_eprintln!(Verbose; cfg, "Writing graph to stream...");

	type BoxedWrite = Box<dyn tokio::io::AsyncWrite + Unpin + Send + Sync>;
	use futures::FutureExt;
	use config::OutputSerialisationMode;
	let should_comp = ser_out.should_compress();
	match ser_out {
	    OutputSerialisationMode::File(output_file) |
	    OutputSerialisationMode::RawFile(output_file) => {
		use tokio::fs::OpenOptions;
		(async move {
		    let stream = OpenOptions::new()
			.write(true)
			.truncate(true)
			.create(true)
			.open(output_file).await
			.wrap_err(eyre!("Failed to open file for writing"))
			.with_section(|| format!("{:?}", output_file).header("File was"))?;

		    Ok::<BoxedWrite, eyre::Report>(Box::new(stream))
		}.boxed(), None, should_comp)
	    },
	    OutputSerialisationMode::Stdout |
	    OutputSerialisationMode::RawStdout => (async move { Ok::<BoxedWrite, _>(Box::new(tokio::io::stdout())) }.boxed(), Option::<&std::path::PathBuf>::None, should_comp),
	    #[cfg(feature="prealloc")] OutputSerialisationMode::PreallocFile(output_file) => {
		(async move {
		    Ok::<BoxedWrite, _>(Box::new(tokio::io::sink())) // we use a sink as a shim stream since it will never be used when tuple item `.1` is Some()
		}.boxed(), Some(output_file), false)
	    },
	}
    }) {
	// We use tuple item `.1` here to indicate if we're in normal write mode.
	// None -> normal
	// Some(path) -> prealloc
	// `.2` indicates if we should compress while in normal write mode.
	Some((stream_fut, None, compress)) => {
	    let stream = stream_fut.await?; 
	    serial::write_async(stream, graph.as_ref(), compress).await
		.wrap_err(eyre!("Failed to serialise graph to stream"))?;
	},
	#[cfg(feature="prealloc")] Some((_task_fut, Some(output_file), _)) => {
	    use tokio::fs::OpenOptions;
	    let file = OpenOptions::new()
		.write(true)
		.read(true) //needed for map I think?
		.truncate(true)
		.create(true)
		.open(&output_file).await
		.wrap_err(eyre!("Failed to open file for mapping"))
		.with_section(|| format!("{:?}", output_file).header("File was"))?;
	    let mut file = file.into_std().await;
            cfg_eprintln!(Verbose; cfg, "Opened file for prealloc-write");
	    tokio::task::spawn_blocking(move || {
		serial::write_sync_map(&mut file, graph.as_ref())
	    }).await.wrap_err(eyre!("Prealloc panicked while dumping"))
		.with_section(|| format!("{:?}", output_file).header("File was"))?
		.wrap_err(eyre!("Prealloc failed to dump graph to file"))
		.with_section(|| format!("{:?}", output_file).header("File was"))?;
	},
	_ => (),
    }
    Ok(())
}

async fn normal(cfg: config::Config) -> eyre::Result<()>
{
    let state = state::State::new(cfg
				  .validate()
				  .wrap_err(eyre!("Invalid config"))
				  .with_suggestion(|| "Try running `--help`")?);
    
    let (graph, cfg) = tokio::select!{
	x = work::work_on_all(state) => {x}
	_ = tokio::signal::ctrl_c() => {
	    return Err(eyre!("Interrupt signalled, exiting"));
	}
    };
    
    let cfg = cfg.make_global();
    let (treemap, graph) = tokio::task::spawn_blocking(move || {
	cfg_println!(Verbose; cfg, "Computing hierarchy...");

	// Create treemap
	#[cfg(feature="treemap")] 
	let treemap = {
	    //Check config for if the treemap flag is present. If it is, get the width and height from there too.
	    if let Some(&(x,y)) = cfg.inspection.treemap.as_ref() {
		Some(info::treemap(cfg, &graph, (
		    x as f64,
		    y as f64,
		)))
	    } else {
		None
	    }
	};
	#[cfg(not(feature="treemap"))]
	let treemap = Option::<std::convert::Infallible>::None;

	// Create recursive graph
	let mut graph = graph.into_hierarchical();
	cfg_println!(Verbose; cfg, "Computing sizes...");

	graph.compute_recursive_sizes();
	(treemap, graph)
    }).await.expect("Failed to compute hierarchy from graph");

    //#[cfg(debug_assertions)] cfg_eprintln!(Verbose; cfg, "{:?}", graph);
    
    #[allow(unused_imports)] use futures::future::OptionFuture;
    let (graph, writer) = {
	cfg_if!{
	    if #[cfg(feature="inspect")] {
		let graph = Arc::new(graph);
		let (g, w) = if cfg.serialise_output.is_some() {
		    (Arc::clone(&graph), Some(tokio::spawn(write_graph(graph))))
		} else {
		    (graph, None)
		};
		(g, OptionFuture::from(w))
	    } else {
		(graph, ())//OptionFuture::from(Option::<futures::future::BoxFuture<'static, ()>>::None))
	    }
	}
    };

    #[cfg(feature="treemap")] 
    if let Some(treemap) = treemap {
	//TODO: Print treemap
	todo!("{:?}",treemap);
    }
    
    info::print_basic_max_info(&cfg, &graph);

    #[cfg(feature="inspect")] 
    match writer.await {
	Some(Ok(Ok(_))) if cfg.serialise_output.is_some() => cfg_eprintln!(Verbose; cfg, "Written successfully"),
	Some(Ok(error @ Err(_))) => return error.wrap_err(eyre!("Failed to write graph to output stream")),
	Some(Err(_)) => cfg_eprintln!(Silent; cfg, "Panic while writing graph to stream"),
	_ => (),
    }
    #[cfg(not(feature="inspect"))] drop(writer); 

    
    Ok(())
}

async fn parse_mode() -> eyre::Result<()>
{
    match arg::parse_args()
	.wrap_err(eyre!("Failed to parse args"))
        .with_suggestion(|| "Try running `--help`")?
    {
	arg::Mode::Normal(cfg) => {
	    #[cfg(debug_assertions)] eprintln!("cfg: {:#?}", cfg);
	    normal(cfg).await
	},
	arg::Mode::Help => arg::help(),
    }    
}

#[tokio::main]
async fn main() -> eyre::Result<()> {
    color_eyre::install()?;

    parse_mode().await
	.wrap_err(eyre!("Fatal error"))?;

    /*    let max_size = graph.directories().map(|x| x.size()).max();
    println!("Max size: {:?}", max_size);*/
    
    //println!("{:?}", graph);
    
    Ok(())
}