luminol_filesystem/
trie.rs

1// Copyright (C) 2024 Melody Madeline Lyons
2//
3// This file is part of Luminol.
4//
5// Luminol is free software: you can redistribute it and/or modify
6// it under the terms of the GNU General Public License as published by
7// the Free Software Foundation, either version 3 of the License, or
8// (at your option) any later version.
9//
10// Luminol is distributed in the hope that it will be useful,
11// but WITHOUT ANY WARRANTY; without even the implied warranty of
12// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13// GNU General Public License for more details.
14//
15// You should have received a copy of the GNU General Public License
16// along with Luminol.  If not, see <http://www.gnu.org/licenses/>.
17
18use qp_trie::wrapper::{BStr, BString};
19
20type Trie<T> = qp_trie::Trie<BString, DirTrie<T>>;
21type DirTrie<T> = qp_trie::Trie<BString, Option<T>>;
22
23pub struct FileSystemTrieDirIter<'a, T>(FileSystemTrieDirIterInner<'a, T>);
24
25enum FileSystemTrieDirIterInner<'a, T> {
26    Direct(qp_trie::Iter<'a, BString, Option<T>>, usize),
27    Prefix(std::iter::Once<(&'a str, Option<&'a T>)>),
28}
29
30impl<'a, T> std::iter::FusedIterator for FileSystemTrieDirIter<'a, T> {}
31
32impl<'a, T> std::iter::ExactSizeIterator for FileSystemTrieDirIter<'a, T> {
33    fn len(&self) -> usize {
34        match &self.0 {
35            FileSystemTrieDirIterInner::Direct(_, len) => *len,
36            FileSystemTrieDirIterInner::Prefix(iter) => iter.len(),
37        }
38    }
39}
40
41impl<'a, T> Iterator for FileSystemTrieDirIter<'a, T> {
42    type Item = (&'a str, Option<&'a T>);
43    fn next(&mut self) -> Option<Self::Item> {
44        match &mut self.0 {
45            FileSystemTrieDirIterInner::Direct(iter, len) => {
46                *len = len.saturating_sub(1);
47                iter.next()
48                    .map(|(key, value)| (key.as_str(), value.as_ref()))
49            }
50            FileSystemTrieDirIterInner::Prefix(iter) => iter.next(),
51        }
52    }
53}
54
55pub struct FileSystemTrieIter<'a, T> {
56    path: camino::Utf8PathBuf,
57    trie: &'a Trie<T>,
58    root_done: bool,
59    trie_iter: Option<qp_trie::Iter<'a, BString, DirTrie<T>>>,
60    dir_iter: Option<(camino::Utf8PathBuf, qp_trie::Iter<'a, BString, Option<T>>)>,
61}
62
63impl<'a, T> std::iter::FusedIterator for FileSystemTrieIter<'a, T> {}
64
65impl<'a, T> Iterator for FileSystemTrieIter<'a, T> {
66    type Item = (camino::Utf8PathBuf, &'a T);
67    fn next(&mut self) -> Option<Self::Item> {
68        loop {
69            if let Some((prefix, dir_iter)) = &mut self.dir_iter {
70                match dir_iter.next() {
71                    Some((filename, Some(value))) => {
72                        return Some((format!("{prefix}/{}", filename.as_str()).into(), value));
73                    }
74                    None => {
75                        self.dir_iter = None;
76                        self.root_done = true;
77                    }
78                    _ => {}
79                }
80            } else if let Some(trie_iter) = &mut self.trie_iter {
81                let (prefix, dir_trie) = trie_iter.next()?;
82                self.dir_iter = Some((prefix.as_str().to_owned().into(), dir_trie.iter()));
83            } else if self.path.as_str() == "" {
84                self.root_done = true;
85                self.trie_iter = Some(self.trie.iter())
86            } else if self.root_done {
87                self.trie_iter = Some(
88                    self.trie
89                        .iter_prefix::<BStr>(format!("{}/", self.path).as_str().into()),
90                );
91            } else {
92                self.dir_iter = self
93                    .trie
94                    .get_str(self.path.as_str())
95                    .map(|t| (self.path.clone(), t.iter()));
96            }
97        }
98    }
99}
100
101/// A container for a directory tree-like cache where the "files" are a data type of your choice.
102#[derive(Debug, Clone)]
103pub struct FileSystemTrie<T>(Trie<T>);
104
105impl<T> Default for FileSystemTrie<T> {
106    fn default() -> Self {
107        Self(Default::default())
108    }
109}
110
111impl<T> FileSystemTrie<T> {
112    pub fn new() -> Self {
113        Default::default()
114    }
115
116    /// Adds a directory to the trie. If parent directories do not exist, they will be created.
117    pub fn create_dir(&mut self, path: impl AsRef<camino::Utf8Path>) {
118        let path = path.as_ref();
119
120        // Nothing to do if the path is already in the trie
121        if self.0.contains_key_str(path.as_str()) {
122            return;
123        }
124
125        // Otherwise, find the longest prefix in the trie that is shared with some directory
126        // path that is in the trie
127        let prefix = self.get_dir_prefix(path).as_str().to_string();
128
129        // Check if the trie already contains an entry for this prefix, and if not, create one
130        if !self.0.contains_key_str(&prefix) {
131            let mut dir_trie = DirTrie::new();
132            let prefix_with_trailing_slash = format!("{}/", &prefix);
133            if let Some((child_path, _)) = self
134                .0
135                .iter_prefix_str(if prefix.is_empty() {
136                    ""
137                } else {
138                    &prefix_with_trailing_slash
139                })
140                .next()
141            {
142                // If there is a different path in the trie that has this as a prefix (there can be
143                // at most one), add it to this directory entry
144                dir_trie.insert_str(
145                    camino::Utf8Path::new(child_path.as_str())
146                        .strip_prefix(&prefix)
147                        .unwrap()
148                        .iter()
149                        .next()
150                        .unwrap(),
151                    None,
152                );
153            }
154            self.0.insert_str(&prefix, dir_trie);
155        }
156
157        // Add the new path to the entry at this prefix
158        if let Some(dirname) = path.strip_prefix(&prefix).unwrap().iter().next() {
159            let prefix_trie = self.0.get_mut_str(&prefix).unwrap();
160            prefix_trie.insert_str(dirname, None);
161        }
162
163        // Add the actual entry for the new path
164        self.0.insert_str(path.as_str(), DirTrie::new());
165    }
166
167    /// Adds a file to the trie with the given value. If parent directories do not exist, they will
168    /// be created. If there already was a file at the given path, the original value will be
169    /// returned.
170    pub fn create_file(&mut self, path: impl AsRef<camino::Utf8Path>, value: T) -> Option<T> {
171        let path = path.as_ref().to_owned();
172
173        let dir = path.parent().unwrap_or(camino::Utf8Path::new(""));
174        let filename = path.iter().next_back().unwrap();
175
176        // Add the parent directory to the trie
177        self.create_dir(dir);
178        let dir_trie = self.0.get_mut_str(dir.as_str()).unwrap();
179
180        if let Some(option) = dir_trie.get_mut_str(filename) {
181            // If the path is already in the trie, replace the value and return the original
182            option.replace(value)
183        } else {
184            // Add the file to the parent directory's entry
185            dir_trie.insert_str(filename, Some(value));
186            None
187        }
188    }
189
190    /// Returns whether or not a directory exists at the given path.
191    pub fn contains_dir(&self, path: impl AsRef<camino::Utf8Path>) -> bool {
192        let path = path.as_ref().as_str();
193        if path.is_empty() {
194            return true;
195        }
196        self.0.contains_key_str(path)
197            || (path.is_empty() && self.0.count() != 0)
198            || self
199                .0
200                .longest_common_prefix::<BStr>(format!("{path}/").as_str().into())
201                .as_str()
202                .len()
203                == path.len() + 1
204    }
205
206    /// Returns whether or not a file exists at the given path.
207    pub fn contains_file(&self, path: impl AsRef<camino::Utf8Path>) -> bool {
208        let path = path.as_ref();
209        let Some(filename) = path.iter().next_back() else {
210            return false;
211        };
212        let dir = path.parent().unwrap_or(camino::Utf8Path::new(""));
213        self.0.get_str(dir.as_str()).map_or(false, |dir_trie| {
214            dir_trie
215                .get_str(filename)
216                .and_then(|o| o.as_ref())
217                .is_some()
218        })
219    }
220
221    /// Returns whether or not a file or directory exists at the given path.
222    pub fn contains(&self, path: impl AsRef<camino::Utf8Path>) -> bool {
223        self.contains_file(&path) || self.contains_dir(&path)
224    }
225
226    /// Gets the number of direct children in the directory at the given path, if it exists.
227    pub fn get_dir_size(&self, path: impl AsRef<camino::Utf8Path>) -> Option<usize> {
228        let path = path.as_ref().as_str();
229        if let Some(dir_trie) = self.0.get_str(path) {
230            Some(dir_trie.count())
231        } else if (path.is_empty() && self.0.count() != 0)
232            || self
233                .0
234                .longest_common_prefix::<BStr>(format!("{path}/").as_str().into())
235                .as_str()
236                .len()
237                == path.len() + 1
238        {
239            Some(1)
240        } else {
241            None
242        }
243    }
244
245    /// Gets an immutable reference to the value in the file at the given path, if it exists.
246    pub fn get_file(&self, path: impl AsRef<camino::Utf8Path>) -> Option<&T> {
247        let path = path.as_ref();
248        let filename = path.iter().next_back()?;
249        let dir = path.parent().unwrap_or(camino::Utf8Path::new(""));
250        self.0
251            .get_str(dir.as_str())
252            .and_then(|dir_trie| dir_trie.get_str(filename))
253            .and_then(|o| o.as_ref())
254    }
255
256    /// Gets a mutable reference to the value in the file at the given path, if it exists.
257    pub fn get_file_mut(&mut self, path: impl AsRef<camino::Utf8Path>) -> Option<&mut T> {
258        let path = path.as_ref();
259        let filename = path.iter().next_back()?;
260        let dir = path.parent().unwrap_or(camino::Utf8Path::new(""));
261        self.0
262            .get_mut_str(dir.as_str())
263            .and_then(|dir_trie| dir_trie.get_mut_str(filename))
264            .and_then(|o| o.as_mut())
265    }
266
267    /// Gets an immutable reference to the value in the file at the given path or creates the file
268    /// with the given value if it doesn't.
269    pub fn get_or_create_file(&mut self, path: impl AsRef<camino::Utf8Path>, value: T) -> &T {
270        let path = path.as_ref();
271        if !self.contains_file(path) {
272            self.create_file(path, value);
273        }
274        self.get_file(path).unwrap()
275    }
276
277    /// Gets an immutable reference to the value in the file at the given path or creates the file
278    /// with the given value if it doesn't.
279    pub fn get_or_create_file_with(
280        &mut self,
281        path: impl AsRef<camino::Utf8Path>,
282        f: impl FnOnce() -> T,
283    ) -> &T {
284        let path = path.as_ref();
285        if !self.contains_file(path) {
286            self.create_file(path, f());
287        }
288        self.get_file(path).unwrap()
289    }
290
291    /// Gets a mutable reference to the value in the file at the given path or creates the file
292    /// with the given value if it doesn't.
293    pub fn get_or_create_file_mut(
294        &mut self,
295        path: impl AsRef<camino::Utf8Path>,
296        value: T,
297    ) -> &mut T {
298        let path = path.as_ref();
299        if !self.contains_file(path) {
300            self.create_file(path, value);
301        }
302        self.get_file_mut(path).unwrap()
303    }
304
305    /// Gets a mutable reference to the value in the file at the given path or creates the file
306    /// with the given value if it doesn't.
307    pub fn get_or_create_file_with_mut(
308        &mut self,
309        path: impl AsRef<camino::Utf8Path>,
310        f: impl FnOnce() -> T,
311    ) -> &mut T {
312        let path = path.as_ref();
313        if !self.contains_file(path) {
314            self.create_file(path, f());
315        }
316        self.get_file_mut(path).unwrap()
317    }
318
319    /// Gets the longest prefix of the given path that is the path of a directory in the trie.
320    pub fn get_dir_prefix(&self, path: impl AsRef<camino::Utf8Path>) -> &camino::Utf8Path {
321        let path = path.as_ref();
322        if self.0.contains_key_str(path.as_str()) {
323            return self
324                .0
325                .longest_common_prefix::<BStr>(path.as_str().into())
326                .as_str()
327                .into();
328        }
329        let prefix = self
330            .0
331            .longest_common_prefix::<BStr>(format!("{path}/").as_str().into())
332            .as_str();
333        let prefix = if !self.0.contains_key_str(prefix) || !path.starts_with(prefix) {
334            prefix.rsplit_once('/').map(|(s, _)| s).unwrap_or_default()
335        } else {
336            prefix
337        };
338        prefix.into()
339    }
340
341    /// Removes the file at the given path if it exists, and returns the original value.
342    pub fn remove_file(&mut self, path: impl AsRef<camino::Utf8Path>) -> Option<T> {
343        let path = path.as_ref();
344        let filename = path.iter().next_back()?;
345        let dir = path.parent().unwrap_or(camino::Utf8Path::new(""));
346        self.0
347            .get_mut_str(dir.as_str())
348            .and_then(|dir_trie| dir_trie.remove_str(filename))
349            .flatten()
350    }
351
352    /// Removes the directory at the given path and all of its contents if it exists, and returns
353    /// whether or not it existed.
354    pub fn remove_dir(&mut self, path: impl AsRef<camino::Utf8Path>) -> bool {
355        let path = path.as_ref();
356        let path_str = path.as_str();
357        if path_str.is_empty() {
358            self.0.clear();
359            true
360        } else if self.0.contains_key_str(path_str) {
361            self.0.remove_prefix_str(&format!("{path_str}/"));
362            self.0.remove_str(path_str);
363            if let Some(parent) = path.parent() {
364                self.create_dir(parent);
365                if let (Some(parent_trie), Some(dirname)) =
366                    (self.0.get_mut_str(parent.as_str()), path.iter().next_back())
367                {
368                    parent_trie.remove_str(dirname);
369                }
370            }
371            true
372        } else if self
373            .0
374            .longest_common_prefix::<BStr>(format!("{path_str}/").as_str().into())
375            .as_str()
376            .len()
377            == path_str.len() + 1
378        {
379            self.0.remove_prefix_str(&format!("{path_str}/"));
380            if let Some(parent) = path.parent() {
381                self.create_dir(parent);
382                if let (Some(parent_trie), Some(dirname)) =
383                    (self.0.get_mut_str(parent.as_str()), path.iter().next_back())
384                {
385                    parent_trie.remove_str(dirname);
386                }
387            }
388            true
389        } else {
390            false
391        }
392    }
393
394    /// Given the path to a directory, returns an iterator over its children if it exists.
395    /// The iterator's items are of the form `(key, value)` where `key` is the name of the child as
396    /// `&str` and `value` is the data of the child if it's a file, as `Option<&T>`.
397    pub fn iter_dir(
398        &self,
399        path: impl AsRef<camino::Utf8Path>,
400    ) -> Option<FileSystemTrieDirIter<'_, T>> {
401        let path = path.as_ref();
402        let prefix_with_trailing_slash = format!("{}/", path.as_str());
403        if let Some(dir_trie) = self.0.get_str(path.as_str()) {
404            Some(FileSystemTrieDirIter(FileSystemTrieDirIterInner::Direct(
405                dir_trie.iter(),
406                dir_trie.count(),
407            )))
408        } else if let Some((key, _)) = self
409            .0
410            .iter_prefix_str(if path.as_str().is_empty() {
411                ""
412            } else {
413                &prefix_with_trailing_slash
414            })
415            .next()
416        {
417            Some(FileSystemTrieDirIter(FileSystemTrieDirIterInner::Prefix(
418                std::iter::once((
419                    camino::Utf8Path::new(key.as_str())
420                        .strip_prefix(path)
421                        .unwrap()
422                        .iter()
423                        .next()
424                        .unwrap(),
425                    None,
426                )),
427            )))
428        } else {
429            None
430        }
431    }
432
433    /// Given the path to a directory, returns an iterator over immutable references to its
434    /// descendant files if it exists.
435    pub fn iter_prefix(
436        &self,
437        path: impl AsRef<camino::Utf8Path>,
438    ) -> Option<FileSystemTrieIter<'_, T>> {
439        let path = path.as_ref();
440        if self.contains_dir(path) {
441            Some(FileSystemTrieIter {
442                path: path.to_owned(),
443                trie: &self.0,
444                root_done: false,
445                trie_iter: None,
446                dir_iter: None,
447            })
448        } else {
449            None
450        }
451    }
452}
luminol_filesystem/trie.rs

luminol_filesystem/
trie.rs
