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.
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::{cell::UnsafeCell, io, path::Path, time::Duration};
13
14use bon::builder;
15use rapidhash::{HashMapExt, RapidHashMap as HashMap};
16use thiserror::Error;
17use tracing::{debug, info, warn};
18use widestring::U16CString;
19use windows::{Win32::Storage::FileSystem::GetDiskFreeSpaceExW, core::PCWSTR};
20
21use crate::{
22    search,
23    wm::{self, EverythingClient, RequestFlags, SearchFlags},
24};
25
26#[derive(Error, Debug)]
27pub enum Error {
28    #[error("path is relative")]
29    RelativePath,
30
31    #[error("folder not found")]
32    NotFound,
33
34    #[error(transparent)]
35    Io(#[from] io::Error),
36
37    #[error(transparent)]
38    Ipc(#[from] wm::IpcError),
39}
40
41thread_local! {
42    static EVERYTHING: UnsafeCell<Option<EverythingClient>> = const { UnsafeCell::new(None) };
43    static LAST_PARENT: UnsafeCell<std::path::PathBuf> = const { UnsafeCell::new(std::path::PathBuf::new()) };
44    static RESULT_MAP: UnsafeCell<Option<HashMap<String, u64>>> = const { UnsafeCell::new(None) };
45}
46
47/// Get the size of a folder.
48///
49/// See [`folder::size`](super::size) for details.
50///
51/// ## Arguments
52/// - `path`: An absolute path to a folder
53/// - `timeout`: Optional timeout for IPC queries.
54///   If `None`, uses the default timeout.
55/// - `parent_max_size`: Optional mutable reference to receive the maximum size of folders in the parent folder.
56///
57/// ## Returns
58/// - `Ok(u64)`: The size in bytes
59/// - `Err(Error)`: If the path is invalid
60#[builder]
61pub fn get_folder_size(
62    #[builder(start_fn)] path: &Path,
63    timeout: Option<Duration>,
64    parent_max_size: Option<&mut u64>,
65) -> Result<u64, Error> {
66    debug_assert_eq!(search::normalize_path_ev(path), path);
67
68    // Get the parent directory
69    let parent = match path.parent() {
70        Some(p) if p.as_os_str().is_empty() => return Err(Error::RelativePath),
71        Some(p) => p,
72        None => {
73            // Handle 3-character paths (e.g., "C:\")
74            if path.as_os_str().len() == 3 {
75                let path_u16 = U16CString::from_os_str(path).unwrap();
76                let mut size = 0u64;
77                if unsafe {
78                    GetDiskFreeSpaceExW(PCWSTR(path_u16.as_ptr()), None, Some(&mut size), None)
79                }
80                .is_ok()
81                {
82                    return Ok(size);
83                }
84            }
85            return Err(Error::RelativePath);
86        }
87    };
88
89    // Get or create the Everything client for this thread
90    let everything = EVERYTHING.with(|cell| -> Result<&EverythingClient, wm::IpcError> {
91        let opt = unsafe { &mut *cell.get() };
92        if opt.is_none() {
93            *opt = Some(EverythingClient::new()?);
94        }
95        Ok(&*opt.as_ref().unwrap())
96    })?;
97
98    // Check if we need to query for a new parent
99    let needs_query = LAST_PARENT.with(|cell| unsafe {
100        let last_path = &mut *cell.get();
101        last_path != parent
102    });
103
104    if needs_query {
105        // Clear and rebuild cache for new parent
106        LAST_PARENT.with(|cell| unsafe {
107            *cell.get() = parent.to_path_buf();
108        });
109
110        RESULT_MAP.with(|cell| unsafe {
111            *cell.get() = None;
112        });
113
114        // Query Everything for files in the folder
115        let search_query = format!(r#"folder:infolder:"{}""#, parent.display());
116        let query_list = everything
117            .query_wait(&search_query)
118            .search_flags(SearchFlags::empty())
119            .request_flags(RequestFlags::FileName | RequestFlags::Size)
120            .maybe_timeout(timeout)
121            .call()
122            .inspect_err(|e| warn!(%e, ?parent, "query failed"))?;
123        info!(len = query_list.len(), "query");
124
125        // Build result map from query results
126        let mut result_map = HashMap::with_capacity(query_list.len());
127        for item in query_list.iter() {
128            if let (Some(filename), Some(file_size)) = (
129                item.get_str(RequestFlags::FileName),
130                item.get_size(RequestFlags::Size),
131            ) {
132                let filename_str = filename.to_string_lossy();
133                result_map.insert(filename_str, file_size);
134            }
135        }
136        RESULT_MAP.with(|cell| unsafe {
137            *cell.get() = Some(result_map);
138        });
139    }
140
141    // Look up the file in the result map
142    // Or PathFindFileNameW()
143    let filename = path
144        .file_name()
145        .and_then(|f| f.to_str())
146        .ok_or(Error::RelativePath)?
147        .to_string();
148
149    match RESULT_MAP.with(|cell| {
150        let map = unsafe { &*cell.get() }.as_ref();
151
152        if let Some(max_size) = parent_max_size {
153            *max_size = map
154                .and_then(|m| m.values().max().copied())
155                .unwrap_or_default();
156        }
157
158        map.and_then(|m| m.get(&filename)).copied()
159    }) {
160        // If size is 0, try with realpath
161        Some(0) => {
162            let realpath = search::canonicalize_path_ev(path)?;
163            if realpath != path {
164                debug!(?realpath);
165                // TODO: pipe?
166                let size = everything
167                    .get_folder_size(&realpath)
168                    .maybe_timeout(timeout)
169                    .call()?;
170
171                // Cache realpath size
172                RESULT_MAP.with(|cell| {
173                    // We got Some(0)
174                    let map = unsafe { &mut *cell.get() }.as_mut().unwrap();
175                    map.insert(filename, size);
176                });
177
178                return Ok(size);
179            }
180        }
181        Some(size) => return Ok(size),
182        None => {
183            RESULT_MAP.with(|cell| {
184                let map = unsafe { &*cell.get() };
185                debug!(filename, ?map);
186            });
187        }
188    }
189
190    // TODO: May be new folder
191    Err(Error::NotFound)
192}
193
194#[cfg(test)]
195mod tests {
196    use super::*;
197
198    #[test_log::test]
199    #[test_log(default_log_filter = "trace")]
200    fn get_folder_size_root() {
201        let r = get_folder_size(Path::new(r"C:\")).call();
202        dbg!(&r);
203        // Just verify it returns something without panicking
204        assert!(r.unwrap() > 0);
205    }
206
207    #[test_log::test]
208    #[test_log(default_log_filter = "trace")]
209    fn get_folder_size_ev() {
210        let r = get_folder_size(Path::new(r"C:\Windows")).call();
211        dbg!(&r);
212        assert!(r.unwrap() > 0);
213
214        let r = get_folder_size(Path::new(r"C:\Users")).call();
215        dbg!(&r);
216        assert!(r.unwrap() > 0);
217    }
218
219    #[test_log::test]
220    #[test_log(default_log_filter = "trace")]
221    fn get_folder_size_ev_max() {
222        let mut max_size: u64 = 0;
223        let r = get_folder_size(Path::new(r"C:\Windows"))
224            .parent_max_size(&mut max_size)
225            .call();
226        dbg!(&r, max_size);
227        assert!(r.unwrap() > 0);
228
229        let mut max_size2: u64 = 0;
230        let r = get_folder_size(Path::new(r"C:\Users"))
231            .parent_max_size(&mut max_size2)
232            .call();
233        dbg!(&r, max_size2);
234        assert!(r.unwrap() > 0);
235
236        assert_eq!(max_size, max_size2);
237    }
238
239    #[test_log::test]
240    #[test_log(default_log_filter = "trace")]
241    fn get_folder_size_ev_realpath() {
242        // Test realpath resolution: "C:\Documents and Settings" -> "C:\Users"
243        let r = get_folder_size(Path::new(r"C:\Documents and Settings"))
244            .call()
245            .unwrap();
246        dbg!(&r);
247        assert!(r > 0);
248        let r1 = get_folder_size(Path::new(r"C:\Documents and Settings"))
249            .call()
250            .unwrap();
251        dbg!(&r1);
252        assert_eq!(r, r1);
253
254        let r2 = get_folder_size(Path::new(r"C:\Users")).call().unwrap();
255        dbg!(&r2);
256        assert!(r2 > 0);
257        assert_eq!(r, r2);
258    }
259}
everything_ipc/folder/size.rs

everything_ipc/folder/
size.rs
