archival/
file_system_memory.rs

1use serde::{Deserialize, Serialize};
2use std::{
3    collections::{BTreeMap, BTreeSet},
4    error::Error,
5    fmt::Debug,
6    hash::Hash,
7    ops::Deref,
8    path::{Path, PathBuf},
9};
10#[cfg(feature = "verbose-logging")]
11use tracing::debug;
12use tracing::error;
13
14use crate::ArchivalError;
15
16use super::FileSystemAPI;
17
18#[derive(Debug, Clone, Eq, Serialize, Deserialize)]
19pub struct DirEntry {
20    path: PathBuf,
21    is_file: bool,
22}
23
24impl PartialOrd for DirEntry {
25    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
26        Some(self.cmp(other))
27    }
28}
29
30impl Ord for DirEntry {
31    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
32        self.path.cmp(&other.path)
33    }
34}
35
36impl PartialEq for DirEntry {
37    fn eq(&self, other: &Self) -> bool {
38        self.path == other.path
39    }
40}
41
42impl Hash for DirEntry {
43    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
44        self.path.hash(state);
45    }
46}
47
48impl DirEntry {
49    fn new(path: &Path, is_file: bool) -> Self {
50        Self {
51            path: path.to_path_buf(),
52            is_file,
53        }
54    }
55    fn copy(&self) -> Self {
56        Self {
57            path: self.path.to_owned(),
58            is_file: self.is_file,
59        }
60    }
61}
62
63impl Deref for DirEntry {
64    type Target = PathBuf;
65    fn deref(&self) -> &Self::Target {
66        &self.path
67    }
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
71pub struct FileGraphNode {
72    path: PathBuf,
73    pub(crate) files: BTreeSet<DirEntry>,
74}
75
76impl FileGraphNode {
77    pub fn new(path: &Path) -> Self {
78        Self {
79            path: path.to_path_buf(),
80            files: BTreeSet::new(),
81        }
82    }
83    pub fn key(path: &Path) -> String {
84        path.to_string_lossy().to_lowercase()
85    }
86    pub fn is_empty(&self) -> bool {
87        self.files.is_empty()
88    }
89    pub fn add(&mut self, path: &Path, is_file: bool) {
90        self.files.insert(DirEntry::new(path, is_file));
91    }
92    pub fn remove(&mut self, path: &Path) {
93        // Since we impl Ord (which is what BTreeSet uses for comparison),
94        // is_file doesn't matter for comparison
95        self.files.remove(&DirEntry::new(path, false));
96    }
97    pub fn copy(&self) -> Self {
98        Self {
99            path: self.path.to_path_buf(),
100            files: self.files.iter().map(|r| r.copy()).collect(),
101        }
102    }
103}
104
105#[derive(Default, Clone, Serialize, Deserialize, PartialEq)]
106pub struct MemoryFileSystem {
107    fs: BTreeMap<String, Vec<u8>>,
108    tree: BTreeMap<String, FileGraphNode>,
109}
110impl Debug for MemoryFileSystem {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        #[derive(Debug)]
113        #[allow(dead_code)]
114        struct DisplayMemoryFileSystem<'a> {
115            files: Vec<&'a String>,
116            tree: &'a BTreeMap<String, FileGraphNode>,
117        }
118        let Self { fs, tree } = self;
119
120        Debug::fmt(
121            &DisplayMemoryFileSystem {
122                files: fs.keys().collect(),
123                tree,
124            },
125            f,
126        )
127    }
128}
129
130impl FileSystemAPI for MemoryFileSystem {
131    fn root_dir(&self) -> &Path {
132        Path::new("<memory>")
133    }
134    fn exists(&self, path: &Path) -> Result<bool, Box<dyn Error>> {
135        if self.fs.contains_key(&path.to_string_lossy().to_lowercase()) || self.is_dir(path)? {
136            Ok(true)
137        } else {
138            Ok(false)
139        }
140    }
141    fn is_dir(&self, path: &Path) -> Result<bool, Box<dyn Error>> {
142        Ok(self.tree.contains_key(&FileGraphNode::key(path)))
143    }
144    fn remove_dir_all(&mut self, path: &Path) -> Result<(), Box<dyn Error>> {
145        self.remove_from_graph(path);
146        Ok(())
147    }
148    fn create_dir_all(&mut self, _path: &Path) -> Result<(), Box<dyn Error>> {
149        // dirs are implicitly created when files are created in them
150        Ok(())
151    }
152    fn read(&self, path: &Path) -> Result<Option<Vec<u8>>, Box<dyn Error>> {
153        Ok(self.read_file(path))
154    }
155    fn read_to_string(&self, path: &Path) -> Result<Option<String>, Box<dyn Error>> {
156        if let Some(file) = self.read_file(path) {
157            Ok(Some(std::str::from_utf8(&file)?.to_string()))
158        } else {
159            Ok(None)
160        }
161    }
162    fn delete(&mut self, path: &Path) -> Result<(), Box<dyn Error>> {
163        if self.is_dir(path)? {
164            return Err(ArchivalError::new("use remove_dir_all to delete directories").into());
165        }
166        self.delete_file(path);
167        Ok(())
168    }
169    fn write(&mut self, path: &Path, contents: Vec<u8>) -> Result<(), Box<dyn Error>> {
170        if self.is_dir(path)? {
171            return Err(ArchivalError::new("cannot write to a folder").into());
172        }
173        self.write_file(path, contents);
174        Ok(())
175    }
176    fn write_str(&mut self, path: &Path, contents: String) -> Result<(), Box<dyn Error>> {
177        self.write(path, contents.as_bytes().to_vec())
178    }
179    fn walk_dir(
180        &self,
181        path: &Path,
182        include_dirs: bool,
183    ) -> Result<Box<dyn Iterator<Item = PathBuf>>, Box<dyn Error>> {
184        let path = path.to_path_buf();
185        let children = self.all_children(&path);
186        let mut all_files: Vec<PathBuf> = vec![];
187        for child in children {
188            let node = self.get_node(&child);
189            all_files.append(
190                &mut node
191                    .files
192                    .into_iter()
193                    .filter_map(|de| {
194                        if !include_dirs && !de.is_file {
195                            return None;
196                        }
197                        match de.path.strip_prefix(&path) {
198                            Ok(path) => Some(path.to_owned()),
199                            Err(e) => {
200                                error!("Ignorning invalid path {} ({})", path.display(), e);
201                                None
202                            }
203                        }
204                    })
205                    .collect(),
206            );
207        }
208        Ok(Box::new(all_files.into_iter()))
209    }
210}
211
212impl MemoryFileSystem {
213    fn write_file(&mut self, path: &Path, data: Vec<u8>) {
214        #[cfg(feature = "verbose-logging")]
215        debug!("write: {}", path.display());
216        self.fs.insert(path.to_string_lossy().to_lowercase(), data);
217        self.write_to_graph(path, true);
218    }
219
220    fn write_to_graph(&mut self, path: &Path, is_file: bool) {
221        // Traverse up the path and add each file to its parent's node
222        let mut last_path: PathBuf = PathBuf::new();
223        let mut is_file = is_file;
224        for ancestor in path.ancestors() {
225            let a_path = ancestor.to_path_buf();
226            // Skip the actual file path, since only directories are nodes
227            if a_path.to_string_lossy() != path.to_string_lossy() {
228                let mut node = self.get_node(&a_path);
229                node.add(&last_path, is_file);
230                // After we add the first file, everything else will be directories.
231                is_file = false;
232                self.tree.insert(FileGraphNode::key(&a_path), node);
233            }
234            a_path.clone_into(&mut last_path);
235        }
236    }
237
238    fn get_node(&self, path: &Path) -> FileGraphNode {
239        match self.tree.get(&FileGraphNode::key(path)) {
240            Some(n) => n.copy(),
241            None => FileGraphNode::new(path),
242        }
243    }
244
245    fn all_children(&self, path: &Path) -> Vec<PathBuf> {
246        // All children of this directory will use keys prefixed with this
247        // directory's key.
248        let node_key = FileGraphNode::key(path);
249        self.tree
250            .keys()
251            .filter(|k| k.starts_with(&node_key))
252            .map(PathBuf::from)
253            .collect()
254    }
255
256    fn remove_from_graph(&mut self, path: &Path) {
257        // If this is a directory, remove it and its children from the graph
258        let node = self.get_node(path);
259        if !node.is_empty() {
260            for path in self.all_children(path) {
261                // delete the node and object
262                self.tree.remove(&FileGraphNode::key(&path));
263                self.delete_file_only(&path);
264            }
265        }
266
267        let mut last_path: PathBuf = path.to_path_buf();
268        for ancestor in path.ancestors() {
269            let a_path = ancestor.to_path_buf();
270            if a_path.to_string_lossy() == path.to_string_lossy() {
271                // We've handled the leaf above
272                continue;
273            }
274            let mut node = self.get_node(&a_path);
275            node.remove(&last_path);
276            if node.is_empty() {
277                // If this is the last item in a node, remove the node entirely.
278                self.tree.remove(&FileGraphNode::key(&a_path));
279            } else {
280                // If the file has siblings, just write the updated value and
281                // stop traversing.
282                self.tree.insert(FileGraphNode::key(&a_path), node);
283                break;
284            }
285            a_path.clone_into(&mut last_path);
286        }
287    }
288
289    fn delete_file_only(&mut self, path: &Path) {
290        self.fs.remove(&path.to_string_lossy().to_lowercase());
291    }
292
293    fn delete_file(&mut self, path: &Path) {
294        self.delete_file_only(path);
295        self.remove_from_graph(path);
296    }
297
298    fn read_file(&self, path: &Path) -> Option<Vec<u8>> {
299        #[cfg(feature = "verbose-logging")]
300        debug!("read {}", path.display());
301        self.fs
302            .get(&path.to_string_lossy().to_lowercase())
303            .map(|v| v.to_vec())
304    }
305}