perf: double directory size calculation speed (#2683)

yazi-config/src/keymap/key.rs

@@ -198,6 +198,6 @@ impl Display for Key {
 			_ => "Unknown",
 		};
 
-		write!(f, "{}>", code)
+		write!(f, "{code}>")
 	}
 }

yazi-fs/src/calculator.rs

@@ -0,0 +1,96 @@
+use std::{collections::VecDeque, fs::ReadDir, future::poll_fn, mem, path::{Path, PathBuf}, pin::Pin, task::{Poll, ready}, time::{Duration, Instant}};
+
+use tokio::task::JoinHandle;
+use yazi_shared::Either;
+
+type Task = Either<PathBuf, ReadDir>;
+
+pub enum SizeCalculator {
+	Idle((VecDeque<Task>, Option<u64>)),
+	Pending(JoinHandle<(VecDeque<Task>, Option<u64>)>),
+}
+
+impl SizeCalculator {
+	pub async fn new(path: impl AsRef<Path>) -> std::io::Result<Self> {
+		let p = path.as_ref().to_owned();
+		tokio::task::spawn_blocking(|| {
+			let mut buf = VecDeque::from([Either::Right(std::fs::read_dir(p)?)]);
+			let size = Self::next_chunk(&mut buf);
+			Ok(Self::Idle((buf, size)))
+		})
+		.await?
+	}
+
+	pub async fn total(path: impl AsRef<Path>) -> std::io::Result<u64> {
+		let mut it = Self::new(path).await?;
+		let mut total = 0;
+		while let Some(n) = it.next().await? {
+			total += n;
+		}
+		Ok(total)
+	}
+
+	pub async fn next(&mut self) -> std::io::Result<Option<u64>> {
+		poll_fn(|cx| {
+			loop {
+				match self {
+					Self::Idle((buf, size)) => {
+						if let Some(s) = size.take() {
+							return Poll::Ready(Ok(Some(s)));
+						} else if buf.is_empty() {
+							return Poll::Ready(Ok(None));
+						}
+
+						let mut buf = mem::take(buf);
+						*self = Self::Pending(tokio::task::spawn_blocking(move || {
+							let size = Self::next_chunk(&mut buf);
+							(buf, size)
+						}));
+					}
+					Self::Pending(handle) => {
+						*self = Self::Idle(ready!(Pin::new(handle).poll(cx))?);
+					}
+				}
+			}
+		})
+		.await
+	}
+
+	fn next_chunk(buf: &mut VecDeque<Either<PathBuf, ReadDir>>) -> Option<u64> {
+		let (mut i, mut size, now) = (0, 0, Instant::now());
+		macro_rules! pop_and_continue {
+			() => {{
+				buf.pop_front();
+				if buf.is_empty() {
+					return Some(size);
+				}
+				continue;
+			}};
+		}
+
+		while i < 5000 && now.elapsed() < Duration::from_millis(50) {
+			i += 1;
+			let front = buf.front_mut()?;
+
+			if let Either::Left(p) = front {
+				*front = match std::fs::read_dir(p) {
+					Ok(it) => Either::Right(it),
+					Err(_) => pop_and_continue!(),
+				};
+			}
+
+			let Some(next) = front.right_mut()?.next() else {
+				pop_and_continue!();
+			};
+
+			let Ok(ent) = next else { continue };
+			let Ok(ft) = ent.file_type() else { continue };
+			if ft.is_dir() {
+				buf.push_back(Either::Left(ent.path()));
+			} else {
+				size += ent.metadata().map(|m| m.len()).unwrap_or(0);
+			}
+		}
+		Some(size)
+	}
+}

yazi-fs/src/fns.rs

@@ -1,4 +1,4 @@
-use std::{borrow::Cow, collections::{HashMap, HashSet, VecDeque}, ffi::{OsStr, OsString}, path::{Path, PathBuf}};
+use std::{borrow::Cow, collections::{HashMap, HashSet}, ffi::{OsStr, OsString}, path::{Path, PathBuf}};
 
 use anyhow::{Result, bail};
 use tokio::{fs, io::{self, AsyncWriteExt}, select, sync::{mpsc, oneshot}, time};
@@ -181,30 +181,6 @@ pub async fn realname_unchecked<'a>(
 	}
 }
 
-pub async fn calculate_size(path: &Path) -> u64 {
-	let mut total = 0;
-	let mut stack = VecDeque::from([path.to_path_buf()]);
-	while let Some(path) = stack.pop_front() {
-		let Ok(meta) = fs::symlink_metadata(&path).await else { continue };
-		if !meta.is_dir() {
-			total += meta.len();
-			continue;
-		}
-
-		let Ok(mut it) = fs::read_dir(path).await else { continue };
-		while let Ok(Some(entry)) = it.next_entry().await {
-			let Ok(meta) = entry.metadata().await else { continue };
-
-			if meta.is_dir() {
-				stack.push_back(entry.path());
-			} else {
-				total += meta.len();
-			}
-		}
-	}
-	total
-}
-
 pub fn copy_with_progress(
 	from: &Path,
 	to: &Path,

yazi-fs/src/lib.rs

@@ -2,7 +2,7 @@
 
 yazi_macro::mod_pub!(cha mounts);
 
-yazi_macro::mod_flat!(cwd file files filter fns op path sorter sorting stage step xdg);
+yazi_macro::mod_flat!(calculator cwd file files filter fns op path sorter sorting stage step xdg);
 
 pub fn init() {
 	CWD.init(<_>::default());

yazi-scheduler/src/file/file.rs

@@ -4,7 +4,7 @@ use anyhow::{Result, anyhow};
 use tokio::{fs::{self, DirEntry}, io::{self, ErrorKind::{AlreadyExists, NotFound}}, sync::mpsc};
 use tracing::warn;
 use yazi_config::YAZI;
-use yazi_fs::{calculate_size, cha::Cha, copy_with_progress, maybe_exists, ok_or_not_found, path_relative_to, skip_path};
+use yazi_fs::{SizeCalculator, cha::Cha, copy_with_progress, maybe_exists, ok_or_not_found, path_relative_to, skip_path};
 use yazi_shared::url::Url;
 
 use super::{FileOp, FileOpDelete, FileOpHardlink, FileOpLink, FileOpPaste, FileOpTrash};
@@ -49,7 +49,7 @@ impl File {
 								&& matches!(e.raw_os_error(), Some(1) | Some(93)) =>
 						{
 							task.retry += 1;
-							self.log(task.id, format!("Paste task retry: {:?}", task))?;
+							self.log(task.id, format!("Paste task retry: {task:?}"))?;
 							self.queue(FileOp::Paste(task), LOW).await?;
 							return Ok(());
 						}
@@ -122,7 +122,7 @@ impl File {
 			FileOp::Delete(task) => {
 				if let Err(e) = fs::remove_file(&task.target).await {
 					if e.kind() != NotFound && maybe_exists(&task.target).await {
-						self.fail(task.id, format!("Delete task failed: {:?}, {e}", task))?;
+						self.fail(task.id, format!("Delete task failed: {task:?}, {e}"))?;
 						Err(e)?
 					}
 				}
@@ -318,7 +318,7 @@ impl File {
 
 	pub async fn trash(&self, mut task: FileOpTrash) -> Result<()> {
 		let id = task.id;
-		task.length = calculate_size(&task.target).await;
+		task.length = SizeCalculator::total(&task.target).await?;
 
 		self.prog.send(TaskProg::New(id, task.length))?;
 		self.queue(FileOp::Trash(task), LOW).await?;

yazi-scheduler/src/prework/prework.rs

@@ -6,7 +6,7 @@ use parking_lot::{Mutex, RwLock};
 use tokio::sync::mpsc;
 use tracing::error;
 use yazi_config::Priority;
-use yazi_fs::{FilesOp, calculate_size};
+use yazi_fs::{FilesOp, SizeCalculator};
 use yazi_plugin::isolate;
 use yazi_shared::{event::CmdCow, url::Url};
 
@@ -73,7 +73,7 @@ impl Prework {
 				self.prog.send(TaskProg::Adv(task.id, 1, 0))?;
 			}
 			PreworkOp::Size(task) => {
-				let length = calculate_size(&task.target).await;
+				let length = SizeCalculator::total(&task.target).await?;
 				task.throttle.done((task.target, length), |buf| {
 					{
 						let mut loading = self.size_loading.write();