Skip to main content

rig_tools/
overlay.rs

1//! [`Overlay`]: a workspace that keeps its writes to itself, to learn what
2//! a tool would leave on disk without changing it.
3
4use std::ffi::{OsStr, OsString};
5use std::io::{self, ErrorKind};
6use std::path::{Path, PathBuf};
7use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
8
9use ignore::overrides::OverrideBuilder;
10use rig_core::tool::{FileChange, content_hash};
11
12use crate::fs::Previous;
13use crate::process::{Launcher, Process, Spawn};
14use crate::workspace::{Contents, Excerpt, Listing, Stat, Workspace, excerpt_of};
15
16/// A [`Workspace`] over another that reads through to it and keeps every
17/// write, removal and rename to itself: run a file tool on it to learn what
18/// the tool would leave on disk ([`changes`](Self::changes)) while the disk
19/// stays as it is. It runs no commands, since what a command changes cannot
20/// be kept from the disk.
21///
22/// ```
23/// use std::sync::Arc;
24/// use rig_core::tool::ChangeKind;
25/// use rig_tools::overlay::Overlay;
26/// use rig_tools::workspace::{Local, Workspace};
27///
28/// let path = std::env::temp_dir().join(format!("overlay-doc-{}.txt", std::process::id()));
29/// let overlay = Overlay::new(Arc::new(Local));
30/// overlay.write(&path, b"would be written\n")?;
31/// assert!(!path.exists());
32/// let changes = overlay.changes();
33/// assert_eq!(changes.iter().map(|change| change.kind).collect::<Vec<_>>(), [ChangeKind::Created]);
34/// # Ok::<(), std::io::Error>(())
35/// ```
36#[derive(Debug)]
37pub struct Overlay {
38    /// The workspace read through to.
39    inner: Arc<dyn Workspace>,
40    /// Each path written or removed here, in the order first touched: its
41    /// bytes now, or `None` once removed.
42    kept: Mutex<Vec<(PathBuf, Option<Arc<[u8]>>)>>,
43}
44
45impl Overlay {
46    /// An overlay over `inner`, with nothing written yet.
47    pub fn new(inner: Arc<dyn Workspace>) -> Self {
48        Self {
49            inner,
50            kept: Mutex::new(Vec::new()),
51        }
52    }
53
54    /// The paths written or removed here.
55    fn kept(&self) -> MutexGuard<'_, Vec<(PathBuf, Option<Arc<[u8]>>)>> {
56        self.kept.lock().unwrap_or_else(PoisonError::into_inner)
57    }
58
59    /// What this overlay holds for `path`: `Some(Some(bytes))` written,
60    /// `Some(None)` removed, `None` untouched.
61    fn held(&self, path: &Path) -> Option<Option<Arc<[u8]>>> {
62        self.kept()
63            .iter()
64            .find(|(kept, _)| kept == path)
65            .map(|(_, bytes)| bytes.clone())
66    }
67
68    /// Keeps `bytes` (`None`: removed) for `path`.
69    fn keep(&self, path: &Path, bytes: Option<Arc<[u8]>>) {
70        let mut kept = self.kept();
71        match kept.iter_mut().find(|(kept, _)| kept == path) {
72            Some((_, held)) => *held = bytes,
73            None => kept.push((path.to_path_buf(), bytes)),
74        }
75    }
76
77    /// What the writes kept here would do to the files of the workspace
78    /// under it, in the order each file was first touched: created,
79    /// modified or deleted, with its hunks, its text before and the
80    /// [`FileState`] of each side (the [`FileState::unread`] size of a file
81    /// too large to read). A file written back as it was, or created and
82    /// removed again, is no change.
83    ///
84    /// [`FileState`]: rig_core::tool::FileState
85    /// [`FileState::unread`]: rig_core::tool::FileState::unread
86    pub fn changes(&self) -> Vec<FileChange> {
87        let kept = self.kept().clone();
88        kept.into_iter()
89            .filter_map(|(path, after)| {
90                let before = match self.inner.stat(&path) {
91                    Ok(Stat::File { .. }) => self.inner.previous(&path),
92                    _ => None,
93                };
94                let unchanged = |before: &Previous, after: &[u8]| {
95                    let hash = before.state.as_ref().and_then(|it| it.hash.as_ref());
96                    hash.is_some_and(|hash| *hash == content_hash(after))
97                };
98                match (&before, &after) {
99                    (None, None) => None,
100                    (Some(before), Some(after)) if unchanged(before, after) => None,
101                    _ => Some(Previous::change(
102                        before,
103                        &path.to_string_lossy(),
104                        after.as_deref(),
105                    )),
106                }
107            })
108            .collect()
109    }
110
111    /// The bytes of the file `path` as this overlay sees it, whole.
112    fn bytes(&self, path: &Path) -> io::Result<Arc<[u8]>> {
113        match self.held(path) {
114            Some(Some(bytes)) => Ok(bytes),
115            Some(None) => Err(ErrorKind::NotFound.into()),
116            None => self
117                .inner
118                .read(path, u64::MAX)
119                .map(|contents| contents.bytes.into()),
120        }
121    }
122}
123
124impl Workspace for Overlay {
125    fn host(&self) -> Option<&str> {
126        self.inner.host()
127    }
128
129    fn stat(&self, path: &Path) -> io::Result<Stat> {
130        match self.held(path) {
131            Some(Some(bytes)) => Ok(Stat::File {
132                size: bytes.len() as u64,
133                read_only: false,
134                modified: None,
135            }),
136            Some(None) => Ok(Stat::Missing),
137            None => match self.inner.stat(path)? {
138                Stat::Missing
139                    if self
140                        .kept()
141                        .iter()
142                        .any(|(kept, bytes)| bytes.is_some() && kept.starts_with(path)) =>
143                {
144                    Ok(Stat::Dir)
145                }
146                stat => Ok(stat),
147            },
148        }
149    }
150
151    fn read(&self, path: &Path, limit: u64) -> io::Result<Contents> {
152        match self.held(path) {
153            Some(Some(bytes)) => Ok(Contents {
154                size: bytes.len() as u64,
155                bytes: bytes
156                    .get(..limit.try_into().unwrap_or(usize::MAX))
157                    .unwrap_or(&bytes)
158                    .to_vec(),
159            }),
160            Some(None) => Err(ErrorKind::NotFound.into()),
161            None => self.inner.read(path, limit),
162        }
163    }
164
165    fn excerpt(&self, path: &Path, first: usize, count: usize) -> io::Result<Excerpt> {
166        match self.held(path) {
167            None => self.inner.excerpt(path, first, count),
168            Some(_) => excerpt_of(self, path, first, count),
169        }
170    }
171
172    fn write(&self, path: &Path, bytes: &[u8]) -> io::Result<()> {
173        match self.stat(path)? {
174            Stat::Dir | Stat::Other => Err(io::Error::new(
175                ErrorKind::InvalidInput,
176                "not a regular file",
177            )),
178            Stat::File {
179                read_only: true, ..
180            } => Err(io::Error::new(
181                ErrorKind::PermissionDenied,
182                "the file is read-only; leave it unchanged or ask the user",
183            )),
184            Stat::File { .. } | Stat::Missing => {
185                self.keep(path, Some(bytes.into()));
186                Ok(())
187            }
188        }
189    }
190
191    fn remove(&self, path: &Path) -> io::Result<()> {
192        match self.stat(path)? {
193            Stat::Missing => Err(ErrorKind::NotFound.into()),
194            Stat::Dir => Err(io::Error::new(ErrorKind::IsADirectory, "is a directory")),
195            Stat::File { .. } | Stat::Other => {
196                self.keep(path, None);
197                Ok(())
198            }
199        }
200    }
201
202    fn rename(&self, from: &Path, to: &Path) -> io::Result<()> {
203        let bytes = self.bytes(from)?;
204        self.keep(to, Some(bytes));
205        self.keep(from, None);
206        Ok(())
207    }
208
209    fn canonical(&self, path: &Path) -> PathBuf {
210        self.inner.canonical(path)
211    }
212
213    fn join(&self, base: &Path, path: &Path) -> PathBuf {
214        self.inner.join(base, path)
215    }
216
217    /// The files under it, as listed there, without those removed here,
218    /// then the files created here that `listing` keeps by their names
219    /// ([`Listing::keeps`]; ignore files are not read for them).
220    fn files(&self, root: &Path, glob: Option<&str>, listing: Listing) -> io::Result<Vec<PathBuf>> {
221        let listed = self.inner.files(root, glob, listing)?;
222        let matcher = glob
223            .map(|glob| OverrideBuilder::new(root).add(glob)?.build())
224            .transpose()
225            .map_err(|error| {
226                io::Error::new(ErrorKind::InvalidInput, format!("bad glob: {error}"))
227            })?;
228        let kept = self.kept().clone();
229        let removed = |path: &PathBuf| {
230            kept.iter()
231                .any(|(kept, bytes)| kept == path && bytes.is_none())
232        };
233        let created = kept.iter().filter_map(|(path, bytes)| {
234            let visible = listing.keeps(path.strip_prefix(root).ok()?);
235            let matches = matcher
236                .as_ref()
237                .is_none_or(|matcher| matcher.matched(path, false).is_whitelist());
238            (bytes.is_some() && visible && matches && !listed.contains(path)).then(|| path.clone())
239        });
240        let mut files: Vec<PathBuf> = listed
241            .iter()
242            .filter(|path| !removed(path))
243            .cloned()
244            .collect();
245        files.extend(created);
246        Ok(files)
247    }
248
249    fn spawn(
250        &self,
251        _: &Launcher,
252        _: Option<&Path>,
253        _: &OsStr,
254        _: &[OsString],
255        _: Spawn,
256    ) -> io::Result<Process> {
257        Err(io::Error::new(
258            ErrorKind::Unsupported,
259            "an overlay runs no commands: what a command changes cannot be kept from the disk",
260        ))
261    }
262}
263
264#[cfg(test)]
265mod tests;