nix_dev_env/
nix_profile_cache.rs

1use std::{
2    ffi::OsStr,
3    fs,
4    io::{self, Write},
5    path::{Path, PathBuf},
6    process,
7    time::SystemTime,
8};
9
10use serde_json::Value;
11use sha1::{Digest, Sha1};
12
13use crate::nix_command;
14
15#[derive(Debug, Clone)]
16pub struct NixProfileCache {
17    cache_dir: PathBuf,
18    flake_inputs_dir: PathBuf,
19    flake_reference: FlakeReference,
20    files_to_watch: Vec<PathBuf>,
21    profile_symlink: PathBuf,
22    profile_rc_file: PathBuf,
23}
24
25impl NixProfileCache {
26    pub fn new(cache_dir: PathBuf, flake_reference: &str) -> anyhow::Result<Self> {
27        let flake_inputs_dir = cache_dir.join("flake-inputs");
28
29        let flake_reference = FlakeReference::parse(flake_reference)?;
30
31        let mut files_to_watch = vec![];
32        let hash = if let Some(flake_dir) = &flake_reference.flake_dir {
33            files_to_watch.extend_from_slice(&[
34                flake_dir.join("flake.nix"),
35                flake_dir.join("flake.lock"),
36                flake_dir.join("devshell.toml"),
37            ]);
38            hash_files(&files_to_watch)?
39        } else {
40            hash_flake_reference(&flake_reference.flake_reference_string)?
41        };
42
43        let profile_symlink = cache_dir.join(format!("flake-profile-{}", hash));
44        let profile_rc_file = profile_symlink.with_extension("rc");
45        Ok(Self {
46            cache_dir,
47            flake_inputs_dir,
48            flake_reference,
49            files_to_watch,
50            profile_symlink,
51            profile_rc_file,
52        })
53    }
54
55    pub fn needs_update(&self) -> anyhow::Result<bool> {
56        let mut need_update = true;
57
58        if self.profile_rc_file.is_file() && self.profile_symlink.is_file() {
59            let profile_rc_mtime = fs::metadata(&self.profile_rc_file)?.modified()?;
60
61            need_update = self.files_to_watch.iter().any(|file| {
62                fs::metadata(file)
63                    .and_then(|meta| meta.modified())
64                    .unwrap_or(SystemTime::UNIX_EPOCH)
65                    > profile_rc_mtime
66            });
67        }
68
69        Ok(need_update)
70    }
71
72    pub fn update(&self) -> anyhow::Result<()> {
73        clean_old_gcroots(&self.cache_dir, &self.flake_inputs_dir)?;
74
75        let tmp_profile = self
76            .cache_dir
77            .join(format!("flake-tmp-profile.{}", process::id()));
78
79        let stdout_content = nix_command::nix([
80            OsStr::new("print-dev-env"),
81            OsStr::new("--no-write-lock-file"),
82            OsStr::new("--profile"),
83            tmp_profile.as_os_str(),
84            OsStr::new(&self.flake_reference.flake_reference_string),
85        ])?;
86
87        fs::File::create(&self.profile_rc_file)?.write_all(stdout_content.as_bytes())?;
88
89        add_gcroot(&tmp_profile, &self.profile_symlink)?;
90        fs::remove_file(&tmp_profile)?;
91
92        if self.flake_reference.flake_dir.is_some() {
93            for input in get_flake_input_paths(&self.flake_reference.flake_reference_string)? {
94                let store_path = PathBuf::from("/nix/store").join(&input);
95                let symlink_path = self.flake_inputs_dir.join(&input);
96                add_gcroot(&store_path, &symlink_path)?;
97            }
98        }
99
100        Ok(())
101    }
102
103    pub fn profile_rc(&self) -> &Path {
104        &self.profile_rc_file
105    }
106}
107
108#[derive(Debug, Clone)]
109struct FlakeReference {
110    pub flake_reference_string: String,
111    pub flake_dir: Option<PathBuf>,
112}
113
114impl FlakeReference {
115    pub fn parse(flake_reference: &str) -> anyhow::Result<Self> {
116        let mut flake_reference_iter = flake_reference.split('#');
117        let flake_uri = flake_reference_iter
118            .next()
119            .ok_or_else(|| anyhow::anyhow!("Missing flake URI"))?;
120        let flake_specifier = flake_reference_iter.next();
121
122        let expanded_flake_reference_and_flake_dir =
123            if FlakeReference::is_path_type(flake_reference) {
124                let flake_dir_str =
125                    shellexpand::full(flake_uri.strip_prefix("path:").unwrap_or(flake_uri))?;
126                let expanded_flake_reference = format!(
127                    "{}{}{}",
128                    &flake_dir_str,
129                    flake_specifier.map(|_| "#").unwrap_or(""),
130                    flake_specifier.unwrap_or("")
131                );
132                Some((expanded_flake_reference, flake_dir_str))
133            } else {
134                None
135            };
136
137        Ok(Self {
138            flake_dir: expanded_flake_reference_and_flake_dir
139                .as_ref()
140                .map(|x| PathBuf::from(x.1.to_string())),
141            flake_reference_string: expanded_flake_reference_and_flake_dir
142                .map(|x| x.0)
143                .unwrap_or_else(|| String::from(flake_reference)),
144        })
145    }
146
147    fn is_path_type(flake_reference: &str) -> bool {
148        flake_reference.starts_with("path:")
149            || flake_reference.starts_with('~')
150            || flake_reference.starts_with('/')
151            || flake_reference.starts_with("./")
152            || flake_reference.starts_with("../")
153    }
154}
155
156fn hash_files(filenames: impl AsRef<[PathBuf]>) -> anyhow::Result<String> {
157    let (hasher, no_files) = filenames
158        .as_ref()
159        .iter()
160        .filter(|f| {
161            // TODO: figure out what to do if the file doesn't exist
162            f.exists()
163        })
164        .try_fold((Sha1::new(), true), |(mut acc, ..), f| {
165            acc.update(fs::read(f)?);
166            anyhow::Result::<(Sha1, bool)>::Ok((acc, false))
167        })?;
168
169    if no_files {
170        return Err(anyhow::anyhow!("No files found to hash"));
171    }
172
173    Ok(format!("{:x}", hasher.finalize()))
174}
175
176fn hash_flake_reference(flake_reference: &str) -> anyhow::Result<String> {
177    let mut hasher = Sha1::new();
178    hasher.update(flake_reference);
179    Ok(format!("{:x}", hasher.finalize()))
180}
181
182fn clean_old_gcroots(cache_dir: &Path, flake_inputs_dir: &Path) -> anyhow::Result<()> {
183    let res = fs::remove_dir_all(cache_dir);
184    if let Err(e) = &res
185        && e.kind() != io::ErrorKind::NotFound
186    {
187        res?;
188    }
189    fs::create_dir_all(flake_inputs_dir)?;
190    Ok(())
191}
192
193fn add_gcroot(store_path: &Path, symlink: &Path) -> anyhow::Result<()> {
194    nix_command::nix([
195        OsStr::new("build"),
196        OsStr::new("--out-link"),
197        symlink.as_os_str(),
198        store_path.as_os_str(),
199    ])?;
200    Ok(())
201}
202
203fn get_flake_input_paths(flake_reference: &str) -> anyhow::Result<Vec<PathBuf>> {
204    let stdout_content = nix_command::nix([
205        "flake",
206        "archive",
207        "--json",
208        "--no-write-lock-file",
209        flake_reference,
210    ])?;
211    let json = serde_json::from_str::<Value>(&stdout_content)?;
212    Ok(get_paths_from_doc(&json))
213}
214
215fn get_paths_from_doc(doc: &Value) -> Vec<PathBuf> {
216    let mut result = Vec::new();
217
218    if let Some(p) = get_path(doc) {
219        result.push(p);
220    }
221
222    if let Some(inputs) = doc.get("inputs").and_then(|i| i.as_object()) {
223        for (_k, v) in inputs {
224            let sub_paths = get_paths_from_doc(v);
225            result.extend(sub_paths);
226        }
227    }
228
229    result
230}
231
232fn get_path(doc: &Value) -> Option<PathBuf> {
233    doc.get("path")
234        .and_then(|value| value.as_str())
235        .map(|path| {
236            if path.len() > 11 {
237                PathBuf::from(&path[11..])
238            } else {
239                PathBuf::from(path)
240            }
241        })
242}
243
244#[cfg(test)]
245mod tests {
246    use std::{io::Write, path::PathBuf};
247
248    use once_cell::sync::Lazy;
249    use serde_json::json;
250    use tempfile::NamedTempFile;
251
252    use super::{get_path, get_paths_from_doc, hash_files};
253
254    static TEST_FILE: Lazy<NamedTempFile> = Lazy::new(|| {
255        let mut test_file = tempfile::NamedTempFile::new().unwrap();
256        writeln!(test_file.as_file_mut(), r#"echo "1.1.1";"#).unwrap();
257        test_file
258    });
259
260    #[test]
261    fn test_hash_one() {
262        assert_eq!(
263            hash_files([TEST_FILE.path().to_path_buf()]).unwrap(),
264            "6ead949bf4bcae230b9ed9cd11e578e34ce9f9ea"
265        );
266    }
267
268    #[test]
269    fn test_hash_multiple() {
270        assert_eq!(
271            hash_files([
272                TEST_FILE.path().to_path_buf(),
273                TEST_FILE.path().to_path_buf(),
274            ])
275            .unwrap(),
276            "f109b7892a541ed1e3cf39314cd25d21042b984f"
277        );
278    }
279
280    #[test]
281    fn test_hash_filters_nonexistent() {
282        assert_eq!(
283            hash_files([TEST_FILE.path().to_path_buf(), PathBuf::from("FOOBARBAZ"),]).unwrap(),
284            "6ead949bf4bcae230b9ed9cd11e578e34ce9f9ea"
285        );
286    }
287
288    #[test]
289    fn test_get_path_removes_prefix() {
290        let input = json!({
291            "path": "aaaaaaaaaaabbbbb"
292        });
293        let result = get_path(&input);
294        assert_eq!(result, Some(PathBuf::from("bbbbb")));
295    }
296
297    #[test]
298    fn test_get_paths_from_doc() {
299        let input = json!({
300            "path": "aaaaaaaaaaabbbbb",
301            "inputs": {
302                "foo": {
303                    "path": "aaaaaaaaaaaccccc",
304                    "inputs": {
305                        "bar": {
306                            "path": "aaaaaaaaaaaddddd",
307                            "inputs": {}
308                        }
309                    }
310                }
311            }
312        });
313        let result = get_paths_from_doc(&input);
314        assert_eq!(
315            result,
316            vec![
317                "bbbbb".to_string(),
318                "ccccc".to_string(),
319                "ddddd".to_string()
320            ]
321        );
322    }
323}
nix_dev_env/nix_profile_cache.rs

nix_dev_env/
nix_profile_cache.rs
