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

nix_dev_env/
nix_profile_cache.rs
