async ok

Cargo.toml

@@ -8,10 +8,13 @@ edition = "2018"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-default = ["parallel", "limit-concurrency", "splash"]
+default = ["parallel", "limit-concurrency", "splash", "limit-recursion"]
 
 # Limit the concurrent operations of parallel mode to 4096
 limit-concurrency = ["parallel"]
+# Handle directories recursively
+recursive = []
+limit-recursion = ["recursive"]
 splash = []
 parallel = ["tokio", "futures"]
 threads = ["parallel", "tokio/rt-threaded"]

src/main.rs

@@ -28,6 +28,8 @@ use map::*;
 mod temp;
 mod error;
 mod arg;
+#[cfg(feature="recursive")]
+mod recurse;
 
 cfg_if!{
     if #[cfg(feature="splash")] {
@@ -57,11 +59,20 @@ fn args_or_out<T: ExactSizeIterator>(i: T, low: usize) -> T
 #[cfg(feature="parallel")]
 #[cfg_attr(feature="parallel", tokio::main)]
 async fn main() -> eyre::Result<()> {
-    reyre!(init(), "Failed to initialise")?;
+    use futures::{
+	stream,
+	prelude::*,
+    };
     
-    reyre!(parallel::main(futures::stream::iter(args_or_out(std::env::args(), 2)
-						.skip(1)
-						.dedup())).await,
+    reyre!(init(), "Failed to initialise")?;
+    reyre!(parallel::main(stream::iter(args_or_out(std::env::args(), 2)
+				       .skip(1)
+				       .dedup())
+			  .filter_map(|file| {
+			      async move {
+				  Some(parallel::expand_dir(file).await)
+			      }
+			  }).flatten()).await,
 	   "Jobs failed")
 }
 

src/parallel.rs

@@ -12,9 +12,11 @@ use futures::{
 	Future,
 	OptionFuture,
 	FutureExt,
+	BoxFuture,
 	join_all,
     },
     stream::{
+	self,
 	Stream,
 	StreamExt,
     },
@@ -22,6 +24,7 @@ use futures::{
 use tokio::{
     sync::{
 	Semaphore,
+	mpsc,
     },
     fs::{
 	OpenOptions,
@@ -39,6 +42,8 @@ cfg_if!{
     }
 }
 
+#[cfg(feature="recursive")] use recurse::{MAX_DEPTH, Recursion};
+
 fn gensem() -> Option<Arc<Semaphore>>
 {
     trace!("Limiting concurrency to {:?}", MAX_WORKERS);
@@ -90,7 +95,7 @@ async fn work<P: AsRef<Path>>(apath: P, sem: Option<Arc<Semaphore>>) -> Result<(
 }
 
 async fn join_stream<I: Stream>(stream: I) -> impl Iterator<Item=<I::Item as Future>::Output> + ExactSizeIterator
-    where I::Item: Future
+where I::Item: Future
 {
     //gotta be a better way than heap allocating here, right?
     stream.then(|x| async move { x.await }).collect::<Vec<_>>().await.into_iter()
@@ -108,7 +113,6 @@ pub async fn main<I: Stream<Item=String>>(list: I) -> eyre::Result<()>
     for (i, res) in (0usize..).zip(join_stream(list.map(|file| tokio::spawn(work(file, sem.clone()))))
 				   .map(|x| {trace!("--- {} Finished ---", x.len()); x}).await)
     {
-	//trace!("Done on {:?}", res);
 	match res {
 	    Ok(Ok((path, true))) => info!("<{:?}> OK (processed)", path),
 	    Ok(Ok((path, false))) => info!("<{:?}> OK (skipped)", path),
@@ -156,3 +160,73 @@ pub async fn main<I: Stream<Item=String>>(list: I) -> eyre::Result<()>
 
     Ok(())
 }
+
+#[cfg(feature="recursive")] 
+fn push_dir<'a>(path: &'a Path, depth: usize, to: mpsc::Sender<String>) -> BoxFuture<'a, tokio::io::Result<()>>
+{
+    async move {
+	let mut dir = fs::read_dir(path).await?;
+	let mut workers = match dir.size_hint() {
+	    (0, Some(0)) | (0, None) => Vec::new(),
+	    (x, None) | (_, Some(x)) => Vec::with_capacity(x),
+	};
+	let can_recurse = match MAX_DEPTH {
+	    Recursion::All => true,
+	    Recursion::N(n) if depth < usize::from(n) => true,
+	    _ => false,
+	};
+	while let Some(item) = dir.next_entry().await? {
+	    let mut to = to.clone();
+	    workers.push(async move {
+		match path.join(item.file_name()).into_os_string().into_string() {
+		    Ok(name) => {
+			if item.file_type().await?.is_dir() {
+			    if can_recurse {
+				if let Err(e) = push_dir(name.as_ref(), depth+1, to).await {
+				    error!("Walking dir {:?} failed: {}", item.file_name(), e);
+				}
+			    }
+			} else {
+			    to.send(name).await.unwrap();
+			}
+		    },
+		    Err(err) => {
+			error!("Couldn't process file {:?} because it contains invalid UTF-8", err);
+		    },
+		}
+		Ok::<_, std::io::Error>(())
+	    });
+	}
+	join_all(workers).await;
+	Ok(())
+    }.boxed()
+}
+
+pub async fn expand_dir(p: String) -> impl Stream<Item=String>
+{
+    cfg_if!{
+	if #[cfg(feature="recursive")]  {
+	    let (mut tx, rx) = mpsc::channel(16);
+	    tokio::spawn(async move {
+		let path = Path::new(&p);
+		if path.is_dir() {
+		    if let Err(err) = push_dir(path, 0, tx).await {
+			error!("Walking dir {:?} failed: {}", path, err);
+		    }
+		} else {
+		    tx.send(p).await.unwrap();
+		}
+	    });
+	    rx	
+	} else {
+	    stream::iter(iter::once(p).filter_map(|p| {
+		if Path::new(&p).is_dir() {
+		    warn!("{:?} is a directory, skipping", p);
+		    None
+		} else {
+		    Some(p)
+		}
+	    }))
+	}
+    }
+}

src/recurse.rs

@@ -0,0 +1,33 @@
+//! Recursion stuffs
+use super::*;
+use std::{
+    num::NonZeroUsize,
+    fmt,
+};
+
+#[derive(Debug)]
+pub enum Recursion
+{
+    All,
+    N(NonZeroUsize),
+}
+
+cfg_if!{
+    if #[cfg(feature="limit-recursion")] {
+	pub const MAX_DEPTH: Recursion = Recursion::N(unsafe{NonZeroUsize::new_unchecked(256)});
+    } else {
+	pub const MAX_DEPTH: Recursion = Recursion::All;
+    }
+}
+
+impl fmt::Display for Recursion
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result
+    {
+	match self {
+	    Self::N(n) => write!(f, "{}", n),
+	    Self::All => write!(f, "unlimited"),
+	    _ => write!(f, "no"),
+	}
+    }
+}

src/splash.rs

@@ -39,9 +39,12 @@ For verbose output, set `RUST_LOG` env var to one of the following:
 Made by {} with <3 (Licensed GPL 3.0 or later)"#, arg::program_name(), env!("CARGO_PKG_VERSION"), env!("CARGO_PKG_AUTHORS"));
     println!("\nEnabled extensions: ");
     feature!(in nightly, "\tCompiled with Rust nightly extensions");
-    println!();
+    println!("Features:");
     feature!("parallel", "\tWill run up to {} operations in parallel", parallel::MAX_WORKERS.map(|x| Cow::Owned(x.to_string())).unwrap_or(Cow::Borrowed("unlimited")));
     feature!("limit-concurrency", "Concurrency is capped");
     feature!("threads", "\tUsing thread-pool");
+    feature!("recursive", "\tRecursivly process files up to {} directories deep", recurse::MAX_DEPTH);
+    feature!("limit-recursion", "Concurrency is capped");
+
     std::process::exit(1)
 }