everything_ipc/folder/
size.rs

1/*!
2Folder size batch lookup and cache.
3
4- Batch lookup of all folder sizes in the parent folder.
5- Thread-local cache with [`get_folder_size()`].
6- For drivers (like `C:\`), it uses Windows API directly.
7
8## References
9- [`IbDOpusExt/ViewerPlugin/DOpusExt.cpp`](https://github.com/Chaoses-Ib/IbDOpusExt/blob/421397f1f73d49b1351ec6cebdf35a74dddb9019/ViewerPlugin/DOpusExt.cpp#L40-L113)
10*/
11
12use std::{
13    cell::UnsafeCell,
14    io,
15    path::{Path, PathBuf},
16    time::Duration,
17};
18
19use bon::{bon, builder};
20use rapidhash::{HashMapExt, RapidHashMap as HashMap};
21use thiserror::Error;
22use tracing::{debug, info, warn};
23use widestring::U16CString;
24use windows::{Win32::Storage::FileSystem::GetDiskFreeSpaceExW, core::PCWSTR};
25
26use crate::{
27    search,
28    wm::{self, EverythingClient, RequestFlags, SearchFlags},
29};
30
31#[derive(Error, Debug)]
32pub enum Error {
33    #[error("path is relative")]
34    RelativePath,
35
36    #[error("folder not found")]
37    NotFound,
38
39    #[error(transparent)]
40    Io(#[from] io::Error),
41
42    #[error(transparent)]
43    Ipc(#[from] wm::IpcError),
44}
45
46#[cfg(any(doc, not(feature = "drop-join-thread")))]
47thread_local! {
48    static CLIENT: UnsafeCell<FolderSizeClient> = const { UnsafeCell::new(FolderSizeClient::new()) };
49}
50
51/// Get the size of a folder.
52///
53/// Uses a thread-local [`FolderSizeClient`] for caching.
54///
55/// See [`folder::size`](super::size) for details.
56///
57/// ## Arguments
58/// - `path`: An absolute path to a folder
59/// - `timeout`: Optional timeout for IPC queries.
60///   If `None`, uses the default timeout.
61/// - `parent_max_size`: Optional mutable reference to receive the maximum size of folders in the parent folder.
62///
63///   You may also want to set `eager_get_links`.
64///
65/// - `eager_get_links`: If `true`, for folders in the parent folder with size 0,
66///   eagerly resolve symlinks/junctions and query the resolved path's size.
67///
68/// ## Returns
69/// - `Ok(u64)`: The size in bytes
70/// - `Err(Error)`: If the path is invalid
71#[cfg(any(doc, not(feature = "drop-join-thread")))]
72#[builder]
73pub fn get_folder_size(
74    #[builder(start_fn)] path: &Path,
75    timeout: Option<Duration>,
76    parent_max_size: Option<&mut u64>,
77    #[builder(default)] eager_get_links: bool,
78) -> Result<u64, Error> {
79    CLIENT.with(|cell| {
80        let client = unsafe { &mut *cell.get() };
81        client
82            .get_folder_size(path)
83            .maybe_timeout(timeout)
84            .maybe_parent_max_size(parent_max_size)
85            .eager_get_links(eager_get_links)
86            .call()
87    })
88}
89
90/// Folder size client with parent directory cache.
91#[derive(Default)]
92pub struct FolderSizeClient {
93    everything: Option<EverythingClient>,
94    last_parent: PathBuf,
95    /// `HashMap::new()` is not const.
96    result_map: Option<HashMap<String, u64>>,
97}
98
99#[bon]
100impl FolderSizeClient {
101    pub const fn new() -> Self {
102        Self {
103            everything: None,
104            last_parent: PathBuf::new(),
105            result_map: None,
106        }
107    }
108
109    /// Get the size of a folder.
110    ///
111    /// See [`folder::size`](super::size) for details.
112    ///
113    /// ## Arguments
114    /// - `path`: An absolute path to a folder
115    /// - `timeout`: Optional timeout for IPC queries.
116    ///   If `None`, uses the default timeout.
117    /// - `parent_max_size`: Optional mutable reference to receive the maximum size of folders in the parent folder.
118    ///
119    ///   You may also want to set `eager_get_links`.
120    ///
121    /// - `eager_get_links`: If `true`, for folders in the parent folder with size 0,
122    ///   eagerly resolve symlinks/junctions and query the resolved path's size.
123    ///
124    /// ## Returns
125    /// - `Ok(u64)`: The size in bytes
126    /// - `Err(Error)`: If the path is invalid
127    #[builder]
128    pub fn get_folder_size(
129        &mut self,
130        #[builder(start_fn)] path: &Path,
131        timeout: Option<Duration>,
132        parent_max_size: Option<&mut u64>,
133        #[builder(default)] eager_get_links: bool,
134    ) -> Result<u64, Error> {
135        debug_assert_eq!(search::normalize_path_ev(path), path);
136
137        // Get the parent directory
138        let parent = match path.parent() {
139            Some(p) if p.as_os_str().is_empty() => return Err(Error::RelativePath),
140            Some(p) => p,
141            None => {
142                // Handle 3-character paths (e.g., "C:\")
143                if path.as_os_str().len() == 3 {
144                    let path_u16 = U16CString::from_os_str(path).unwrap();
145                    let mut size = 0u64;
146                    if unsafe {
147                        GetDiskFreeSpaceExW(PCWSTR(path_u16.as_ptr()), None, Some(&mut size), None)
148                    }
149                    .is_ok()
150                    {
151                        return Ok(size);
152                    }
153                }
154                return Err(Error::RelativePath);
155            }
156        };
157
158        // Get or create the Everything client
159        // TODO: get_or_try_insert_with()
160        let everything = match self.everything.as_mut() {
161            Some(everything) => everything,
162            None => self.everything.insert(EverythingClient::new()?),
163        };
164
165        // Check if we need to query for a new parent
166        let needs_query = self.last_parent != parent;
167
168        if needs_query {
169            // Clear and rebuild cache for new parent
170            self.last_parent = parent.to_path_buf();
171            self.result_map = None;
172
173            // Query Everything for files in the folder
174            let search_query = format!(r#"folder:infolder:"{}""#, parent.display());
175            let query_list = everything
176                .query_wait(&search_query)
177                .search_flags(SearchFlags::empty())
178                .request_flags(RequestFlags::FileName | RequestFlags::Size)
179                .maybe_timeout(timeout)
180                .call()
181                .inspect_err(|e| warn!(%e, ?parent, "query failed"))?;
182            info!(len = query_list.len(), "query");
183
184            // Build result map from query results
185            let mut result_map = HashMap::with_capacity(query_list.len());
186            for item in query_list.iter() {
187                if let (Some(filename), Some(mut file_size)) = (
188                    item.get_str(RequestFlags::FileName),
189                    item.get_size(RequestFlags::Size),
190                ) {
191                    let filename_str = filename.to_string_lossy();
192
193                    // For folders with size 0, check if it's a symlink/junction when eager_get_links is true
194                    if eager_get_links && file_size == 0 {
195                        let path = parent.join(&filename_str);
196                        match search::canonicalize_path_ev(&path) {
197                            // Resolve path is different, query for the resolved path size
198                            Ok(realpath) if realpath != path => {
199                                debug!(dir = filename_str, ?realpath);
200                                match everything
201                                    .get_folder_size(&realpath)
202                                    .maybe_timeout(timeout)
203                                    .call()
204                                {
205                                    Ok(size) => {
206                                        file_size = size;
207                                    }
208                                    e => warn!(?e, ?realpath, "query realpath failed"),
209                                }
210                            }
211                            Ok(_) => (),
212                            Err(e) => warn!(%e, ?path, "realpath failed"),
213                        }
214                    }
215
216                    result_map.insert(filename_str, file_size);
217                }
218            }
219            self.result_map = Some(result_map);
220        }
221
222        // Look up the file in the result map
223        // Or PathFindFileNameW()
224        let filename = path
225            .file_name()
226            .and_then(|f| f.to_str())
227            .ok_or(Error::RelativePath)?
228            .to_string();
229
230        match self.result_map.as_ref().and_then(|m| {
231            if let Some(max_size) = parent_max_size {
232                *max_size = m.values().max().copied().unwrap_or_default();
233            }
234
235            m.get(&filename).copied()
236        }) {
237            // Empty folder
238            Some(0) if eager_get_links => {
239                return Ok(0);
240            }
241            // If size is 0, try with realpath
242            Some(0) if !eager_get_links => {
243                let realpath = search::canonicalize_path_ev(path)?;
244                if realpath != path {
245                    debug!(?realpath);
246                    // TODO: pipe?
247                    let size = everything
248                        .get_folder_size(&realpath)
249                        .maybe_timeout(timeout)
250                        .call()?;
251
252                    // Cache realpath size
253                    // We got Some(0)
254                    self.result_map.as_mut().unwrap().insert(filename, size);
255
256                    return Ok(size);
257                }
258                // Empty folder
259                return Ok(0);
260            }
261            Some(size) => return Ok(size),
262            None => {
263                debug!(filename, map = ?self.result_map);
264            }
265        }
266
267        // TODO: May be new folder
268        Err(Error::NotFound)
269    }
270}
271
272#[cfg(not(feature = "drop-join-thread"))]
273#[cfg(test)]
274mod tests {
275    use super::*;
276
277    #[test_log::test]
278    #[test_log(default_log_filter = "trace")]
279    fn get_folder_size_root() {
280        let r = get_folder_size(Path::new(r"C:\")).call();
281        dbg!(&r);
282        // Just verify it returns something without panicking
283        assert!(r.unwrap() > 0);
284    }
285
286    #[test_log::test]
287    #[test_log(default_log_filter = "trace")]
288    fn get_folder_size_ev() {
289        let r = get_folder_size(Path::new(r"C:\Windows")).call();
290        dbg!(&r);
291        assert!(r.unwrap() > 0);
292
293        let r = get_folder_size(Path::new(r"C:\Users")).call();
294        dbg!(&r);
295        assert!(r.unwrap() > 0);
296    }
297
298    #[test_log::test]
299    #[test_log(default_log_filter = "trace")]
300    fn get_folder_size_ev_max() {
301        let mut max_size: u64 = 0;
302        let r = get_folder_size(Path::new(r"C:\Windows"))
303            .parent_max_size(&mut max_size)
304            .call();
305        dbg!(&r, max_size);
306        assert!(r.unwrap() > 0);
307
308        let mut max_size2: u64 = 0;
309        let r = get_folder_size(Path::new(r"C:\Users"))
310            .parent_max_size(&mut max_size2)
311            .call();
312        dbg!(&r, max_size2);
313        assert!(r.unwrap() > 0);
314
315        assert_eq!(max_size, max_size2);
316    }
317
318    #[test_log::test]
319    #[test_log(default_log_filter = "trace")]
320    fn get_folder_size_ev_realpath() {
321        // Test realpath resolution: "C:\Documents and Settings" -> "C:\Users"
322        let r = get_folder_size(Path::new(r"C:\Documents and Settings"))
323            .call()
324            .unwrap();
325        dbg!(&r);
326        assert!(r > 0);
327        let r1 = get_folder_size(Path::new(r"C:\Documents and Settings"))
328            .call()
329            .unwrap();
330        dbg!(&r1);
331        assert_eq!(r, r1);
332
333        let r2 = get_folder_size(Path::new(r"C:\Users")).call().unwrap();
334        dbg!(&r2);
335        assert!(r2 > 0);
336        assert_eq!(r, r2);
337    }
338
339    #[test_log::test]
340    #[test_log(default_log_filter = "trace")]
341    fn get_folder_size_ev_realpath_eager() {
342        // Test realpath resolution: "C:\Documents and Settings" -> "C:\Users"
343        let mut max_size: u64 = 0;
344        let r = get_folder_size(Path::new(r"C:\Documents and Settings"))
345            .parent_max_size(&mut max_size)
346            .eager_get_links(true)
347            .call()
348            .unwrap();
349        info!(r, max_size);
350        assert!(r > 0);
351        let r1 = get_folder_size(Path::new(r"C:\Documents and Settings"))
352            .parent_max_size(&mut max_size)
353            .eager_get_links(true)
354            .call()
355            .unwrap();
356        info!(r1, max_size);
357        assert_eq!(r, r1);
358
359        let r2 = get_folder_size(Path::new(r"C:\Users"))
360            .parent_max_size(&mut max_size)
361            .eager_get_links(true)
362            .call()
363            .unwrap();
364        info!(r2, max_size);
365        assert!(r2 > 0);
366        assert_eq!(r, r2);
367    }
368}
everything_ipc/folder/size.rs

everything_ipc/folder/
size.rs
