flake_edit/lock/
sizes.rs

1//! Store-size estimation for flake inputs.
2//!
3//! `flake-edit follow --stats` reports how much disk follows deduplication saves.
4use std::collections::{BTreeSet, HashMap};
5use std::process::{Command, Stdio};
6use std::time::Duration;
7
8use base64::Engine as _;
9use sha2::{Digest, Sha256};
10
11use crate::follows::{AttrPath, Segment, join_raw};
12use crate::lock::{FlakeLock, NestedInput};
13
14/// Wall-clock cap a single `nix path-info` invocation.
15const PATH_INFO_TIMEOUT: Duration = Duration::from_secs(30);
16
17/// Sizes and store paths for the attribute paths of a lockfile.
18#[derive(Debug, Default)]
19pub struct InputSizes {
20    /// Dotted attribute path to the store path its lock node unpacks to.
21    store_paths: HashMap<String, String>,
22    /// Store path to NAR size in bytes, where known.
23    sizes: HashMap<String, u64>,
24}
25
26impl InputSizes {
27    pub fn query(lock: &FlakeLock, nested_inputs: &[NestedInput]) -> Self {
28        Self::query_with(lock, nested_inputs, query_path_sizes)
29    }
30
31    fn query_with(
32        lock: &FlakeLock,
33        nested_inputs: &[NestedInput],
34        size_lookup: impl FnOnce(&[&str]) -> HashMap<String, u64>,
35    ) -> Self {
36        let store_dir = store_dir();
37        let mut store_paths: HashMap<String, String> = HashMap::new();
38        for (attr, node_key) in attr_to_node_map(lock, nested_inputs) {
39            let nar_hash = lock
40                .nodes
41                .get(&node_key)
42                .and_then(|node| node.locked.as_ref())
43                .and_then(|locked| locked.nar_hash.as_deref());
44            let Some(store_path) =
45                nar_hash.and_then(|hash| store_path_from_nar_hash(&store_dir, hash))
46            else {
47                continue;
48            };
49            store_paths.insert(attr, store_path);
50        }
51        let unique: BTreeSet<&str> = store_paths.values().map(String::as_str).collect();
52        let sizes = if unique.is_empty() {
53            HashMap::new()
54        } else {
55            let paths: Vec<&str> = unique.into_iter().collect();
56            size_lookup(&paths)
57        };
58        Self { store_paths, sizes }
59    }
60
61    pub fn from_attr_entries(
62        entries: impl IntoIterator<Item = (String, String, Option<u64>)>,
63    ) -> Self {
64        let mut store_paths = HashMap::new();
65        let mut sizes = HashMap::new();
66        for (attr, store_path, size) in entries {
67            if let Some(size) = size {
68                sizes.insert(store_path.clone(), size);
69            }
70            store_paths.insert(attr, store_path);
71        }
72        Self { store_paths, sizes }
73    }
74
75    /// NAR size in bytes of the store path `path` resolves to.
76    pub fn size_for(&self, path: &AttrPath) -> Option<u64> {
77        self.sizes.get(self.store_path_for(path)?).copied()
78    }
79
80    /// Store path `path` resolves to.
81    pub fn store_path_for(&self, path: &AttrPath) -> Option<&str> {
82        self.store_paths.get(&dotted(path)).map(String::as_str)
83    }
84
85    /// Shortest attribute path resolving to `store_path`.
86    pub fn label_for_store_path(&self, store_path: &str) -> Option<&str> {
87        self.store_paths
88            .iter()
89            .filter(|(_, sp)| sp.as_str() == store_path)
90            .map(|(attr, _)| attr.as_str())
91            .min_by_key(|attr| (attr.matches('.').count(), attr.len(), *attr))
92    }
93}
94
95/// Dotted raw form of an attribute path (`hls-1.10.nixpkgs`).
96fn dotted(path: &AttrPath) -> String {
97    join_raw(path.segments())
98}
99
100fn store_dir() -> String {
101    std::env::var("NIX_STORE_DIR").unwrap_or_else(|_| "/nix/store".to_string())
102}
103
104fn attr_to_node_map(lock: &FlakeLock, nested_inputs: &[NestedInput]) -> HashMap<String, String> {
105    let mut out = HashMap::new();
106    if let Some(root_inputs) = lock
107        .nodes
108        .get(&lock.root)
109        .and_then(|node| node.inputs.as_ref())
110    {
111        for name in root_inputs.keys() {
112            let Ok(seg) = Segment::from_unquoted(name.clone()) else {
113                continue;
114            };
115            if let Ok(key) = lock.resolve_input_path(&AttrPath::new(seg)) {
116                out.insert(name.clone(), key);
117            }
118        }
119    }
120    for nested in nested_inputs {
121        if let Ok(key) = lock.resolve_input_path(&nested.path) {
122            out.insert(dotted(&nested.path), key);
123        }
124    }
125    out
126}
127
128/// Compute the store path a fixed-output source tree with `nar_hash`
129/// unpacks to.
130///
131/// Returns `None` for hashes that are not SRI sha256.
132fn store_path_from_nar_hash(store_dir: &str, nar_hash: &str) -> Option<String> {
133    const HEX: &[u8; 16] = b"0123456789abcdef";
134    let encoded = nar_hash.strip_prefix("sha256-")?;
135    let raw = base64::engine::general_purpose::STANDARD
136        .decode(encoded)
137        .ok()?;
138    if raw.len() != 32 {
139        return None;
140    }
141    let mut hex = String::with_capacity(64);
142    for byte in &raw {
143        hex.push(HEX[(byte >> 4) as usize] as char);
144        hex.push(HEX[(byte & 0x0f) as usize] as char);
145    }
146    let fingerprint = format!("source:sha256:{hex}:{store_dir}:source");
147    let digest = Sha256::digest(fingerprint.as_bytes());
148    let mut folded = [0u8; 20];
149    for (i, byte) in digest.iter().enumerate() {
150        folded[i % 20] ^= byte;
151    }
152    Some(format!("{store_dir}/{}-source", nix_base32(&folded)))
153}
154
155fn nix_base32(bytes: &[u8; 20]) -> String {
156    const ALPHABET: &[u8; 32] = b"0123456789abcdfghijklmnpqrsvwxyz";
157    let mut out = String::with_capacity(32);
158    for n in (0..32).rev() {
159        let bit = n * 5;
160        let byte = bit / 8;
161        let shift = bit % 8;
162        let mut value = bytes[byte] >> shift;
163        if shift > 0 && byte + 1 < bytes.len() {
164            value |= bytes[byte + 1] << (8 - shift);
165        }
166        out.push(ALPHABET[(value & 0x1f) as usize] as char);
167    }
168    out
169}
170
171fn query_path_sizes(paths: &[&str]) -> HashMap<String, u64> {
172    let mut args: Vec<&str> = vec![
173        "path-info",
174        "--json",
175        "--extra-experimental-features",
176        "nix-command",
177    ];
178    args.extend(paths);
179    let Some(raw) = run_nix(&args, PATH_INFO_TIMEOUT) else {
180        return HashMap::new();
181    };
182    serde_json::from_slice::<serde_json::Value>(&raw)
183        .map(|json| parse_path_info(&json))
184        .unwrap_or_default()
185}
186
187/// Parse `nix path-info --json` output into a `store path -> narSize` map.
188/// Supports both of these versions:
189/// - object keyed by store path: (`{"/nix/store/a": {...}}`).
190/// - array of objects: (`[{"path": "/nix/store/a", "narSize": 1}, ...]`).
191fn parse_path_info(json: &serde_json::Value) -> HashMap<String, u64> {
192    let mut out = HashMap::new();
193    if let Some(entries) = json.as_array() {
194        for entry in entries {
195            if let (Some(path), Some(size)) = (
196                entry.get("path").and_then(|v| v.as_str()),
197                entry.get("narSize").and_then(|v| v.as_u64()),
198            ) {
199                out.insert(path.to_string(), size);
200            }
201        }
202    } else if let Some(entries) = json.as_object() {
203        for (path, entry) in entries {
204            if let Some(size) = entry.get("narSize").and_then(|v| v.as_u64()) {
205                out.insert(path.clone(), size);
206            }
207        }
208    }
209    out
210}
211
212fn run_nix(args: &[&str], timeout: Duration) -> Option<Vec<u8>> {
213    use std::io::Read;
214
215    let mut child = Command::new("nix")
216        .args(args)
217        .stdin(Stdio::null())
218        .stdout(Stdio::piped())
219        .stderr(Stdio::piped())
220        .spawn()
221        .ok()?;
222
223    let mut stdout_pipe = child.stdout.take()?;
224    let stdout_reader = std::thread::spawn(move || {
225        let mut buf = Vec::new();
226        let _ = stdout_pipe.read_to_end(&mut buf);
227        buf
228    });
229    let mut stderr_pipe = child.stderr.take()?;
230    let stderr_reader = std::thread::spawn(move || {
231        let mut buf = Vec::new();
232        let _ = stderr_pipe.read_to_end(&mut buf);
233        buf
234    });
235
236    let start = std::time::Instant::now();
237    let status = loop {
238        match child.try_wait() {
239            Ok(Some(status)) => break status,
240            Ok(None) => {
241                if start.elapsed() >= timeout {
242                    let _ = child.kill();
243                    let _ = child.wait();
244                    tracing::debug!("nix {} timed out after {:?}", args.join(" "), timeout);
245                    return None;
246                }
247                std::thread::sleep(Duration::from_millis(50));
248            }
249            Err(_) => {
250                let _ = child.kill();
251                let _ = child.wait();
252                return None;
253            }
254        }
255    };
256
257    let stdout = stdout_reader.join().unwrap_or_default();
258    if status.success() {
259        Some(stdout)
260    } else {
261        let stderr = stderr_reader.join().unwrap_or_default();
262        tracing::debug!(
263            "nix {} failed ({status}): {}",
264            args.join(" "),
265            String::from_utf8_lossy(&stderr).trim(),
266        );
267        None
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274
275    #[test]
276    fn store_path_from_nar_hash_matches_nix() {
277        assert_eq!(
278            store_path_from_nar_hash(
279                "/nix/store",
280                "sha256-x5UQuRsH3MqI0U9afaXSNqzTPSeZlRLvFAav2Ux1pNw=",
281            )
282            .as_deref(),
283            Some("/nix/store/3a2vdn5i7vd2wl654xs8nb52jf1v6cbh-source"),
284        );
285    }
286
287    #[test]
288    fn store_path_from_nar_hash_rejects_non_sri_sha256() {
289        for hash in [
290            "",
291            "not-a-hash",
292            "sha512-x5UQuRsH3MqI0U9afaXSNqzTPSeZlRLvFAav2Ux1pNw=",
293            "sha256-@@@invalid@@@",
294            // Valid base64 but not 32 bytes of digest.
295            "sha256-YWJj",
296        ] {
297            assert_eq!(
298                store_path_from_nar_hash("/nix/store", hash),
299                None,
300                "hash {hash:?} must not produce a store path",
301            );
302        }
303    }
304
305    #[test]
306    fn parse_path_info_handles_array_shape() {
307        let json: serde_json::Value = serde_json::from_str(
308            r#"[
309                {"path": "/nix/store/a", "narSize": 100},
310                {"path": "/nix/store/b", "narSize": 200}
311            ]"#,
312        )
313        .unwrap();
314        let map = parse_path_info(&json);
315        assert_eq!(map.get("/nix/store/a"), Some(&100));
316        assert_eq!(map.get("/nix/store/b"), Some(&200));
317    }
318
319    #[test]
320    fn parse_path_info_handles_object_shape() {
321        let json: serde_json::Value = serde_json::from_str(
322            r#"{
323                "/nix/store/a": {"narSize": 100},
324                "/nix/store/b": {"narSize": 200}
325            }"#,
326        )
327        .unwrap();
328        let map = parse_path_info(&json);
329        assert_eq!(map.get("/nix/store/a"), Some(&100));
330        assert_eq!(map.get("/nix/store/b"), Some(&200));
331    }
332
333    #[test]
334    fn query_maps_follows_chain_to_shared_store_path() {
335        let lock = FlakeLock::read_from_str(
336            r#"{
337  "nodes": {
338    "nixpkgs": {
339      "locked": {
340        "narHash": "sha256-x5UQuRsH3MqI0U9afaXSNqzTPSeZlRLvFAav2Ux1pNw=",
341        "owner": "nixos", "repo": "nixpkgs", "rev": "aaa", "type": "github"
342      },
343      "original": { "owner": "nixos", "repo": "nixpkgs", "type": "github" }
344    },
345    "home-manager": {
346      "inputs": { "nixpkgs": ["nixpkgs"] },
347      "locked": {
348        "narHash": "sha256-AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA=",
349        "owner": "o", "repo": "r", "rev": "bbb", "type": "github"
350      },
351      "original": { "owner": "o", "repo": "r", "type": "github" }
352    },
353    "root": {
354      "inputs": { "nixpkgs": "nixpkgs", "home-manager": "home-manager" }
355    }
356  },
357  "root": "root",
358  "version": 7
359}"#,
360        )
361        .expect("lock parses");
362        let nested = lock.nested_inputs();
363
364        let nixpkgs_path = "/nix/store/3a2vdn5i7vd2wl654xs8nb52jf1v6cbh-source";
365        let sizes = InputSizes::query_with(&lock, &nested, |paths| {
366            assert!(
367                paths.contains(&nixpkgs_path),
368                "computed store path must reach the size lookup, got {paths:?}",
369            );
370            paths.iter().map(|p| ((*p).to_string(), 42)).collect()
371        });
372
373        let top: AttrPath = "nixpkgs".parse().unwrap();
374        let follower: AttrPath = "home-manager.nixpkgs".parse().unwrap();
375        assert_eq!(sizes.store_path_for(&top), Some(nixpkgs_path));
376        assert_eq!(
377            sizes.store_path_for(&follower),
378            sizes.store_path_for(&top),
379            "follows chain must land on the target's store path",
380        );
381        assert_eq!(sizes.size_for(&follower), Some(42));
382        assert_eq!(
383            sizes.label_for_store_path(nixpkgs_path),
384            Some("nixpkgs"),
385            "shortest attribute path must label the shared store path",
386        );
387    }
388
389    #[test]
390    fn from_attr_entries_shares_sizes_via_store_path() {
391        let sizes = InputSizes::from_attr_entries([
392            ("nixpkgs".to_string(), "/store/a".to_string(), Some(7)),
393            ("crane.nixpkgs".to_string(), "/store/a".to_string(), None),
394            ("crane".to_string(), "/store/b".to_string(), None),
395        ]);
396        let follower: AttrPath = "crane.nixpkgs".parse().unwrap();
397        let unknown: AttrPath = "crane".parse().unwrap();
398        assert_eq!(sizes.size_for(&follower), Some(7));
399        assert_eq!(sizes.size_for(&unknown), None);
400        assert_eq!(sizes.label_for_store_path("/store/a"), Some("nixpkgs"));
401    }
402}
flake_edit/lock/sizes.rs

flake_edit/lock/
sizes.rs
