flake_edit/
cache.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3use std::sync::OnceLock;
4
5use directories::ProjectDirs;
6use serde::{Deserialize, Serialize};
7
8static CACHE_FILE_NAME: &str = "flake_edit.json";
9
10fn cache_dir() -> &'static PathBuf {
11    static CACHE_DIR: OnceLock<PathBuf> = OnceLock::new();
12    CACHE_DIR.get_or_init(|| {
13        let project_dir = ProjectDirs::from("com", "a-kenji", "flake-edit").unwrap();
14        project_dir.data_dir().to_path_buf()
15    })
16}
17
18fn cache_file() -> &'static PathBuf {
19    static CACHE_FILE: OnceLock<PathBuf> = OnceLock::new();
20    CACHE_FILE.get_or_init(|| cache_dir().join(CACHE_FILE_NAME))
21}
22
23#[derive(Debug, Default, Clone, Serialize, Deserialize)]
24struct CacheEntry {
25    id: String,
26    uri: String,
27    hit: u32,
28}
29
30/// Generate the cache entry key from an id and uri.
31///
32/// The key format is `{id}.{uri}` which allows multiple URIs per input ID
33/// (e.g., both a github: and path: URI for the same input).
34fn entry_key(id: &str, uri: &str) -> String {
35    format!("{}.{}", id, uri)
36}
37
38/// Cache for storing previously used flake URIs.
39///
40/// Used for shell completions to suggest frequently used inputs.
41#[derive(Debug, Default, Clone, Serialize, Deserialize)]
42pub struct Cache {
43    entries: HashMap<String, CacheEntry>,
44}
45
46impl Cache {
47    /// Save the cache to disk.
48    pub fn commit(&self) -> std::io::Result<()> {
49        let cache_dir = cache_dir();
50        if !cache_dir.exists() {
51            std::fs::create_dir_all(cache_dir)?;
52        }
53        let cache_file_location = cache_file();
54        let cache_file = std::fs::File::create(cache_file_location)?;
55        serde_json::to_writer(cache_file, self)
56            .map_err(|e| std::io::Error::other(e.to_string()))?;
57        Ok(())
58    }
59
60    /// Load the cache from the default location, or return a default empty cache.
61    pub fn load() -> Self {
62        Self::from_path(cache_file())
63    }
64
65    /// Load the cache from a specific file path, or return a default empty cache.
66    pub fn from_path(path: &std::path::Path) -> Self {
67        Self::try_from_path(path).unwrap_or_else(|e| {
68            tracing::warn!("Could not read cache file {:?}: {}", path, e);
69            Self::default()
70        })
71    }
72
73    /// Try to load the cache from a specific file path.
74    pub fn try_from_path(path: &std::path::Path) -> std::io::Result<Self> {
75        let file = std::fs::File::open(path)?;
76        serde_json::from_reader(file).map_err(|e| std::io::Error::other(e.to_string()))
77    }
78
79    /// Add or update a cache entry.
80    pub fn add_entry(&mut self, id: String, uri: String) {
81        let key = entry_key(&id, &uri);
82        match self.entries.get_mut(&key) {
83            Some(entry) => entry.hit += 1,
84            None => {
85                let entry = CacheEntry { id, uri, hit: 0 };
86                self.entries.insert(key, entry);
87            }
88        }
89    }
90
91    /// List cached URIs sorted by hit count (most used first).
92    pub fn list_uris(&self) -> Vec<String> {
93        let mut entries: Vec<_> = self.entries.values().collect();
94        entries.sort_by(|a, b| b.hit.cmp(&a.hit));
95        entries.iter().map(|e| e.uri.clone()).collect()
96    }
97
98    /// List cached URIs for a specific input ID, sorted by hit count (most used first).
99    ///
100    /// This is useful for the "change" workflow where we want to suggest URIs
101    /// that were previously used for the same input ID (e.g., both the remote
102    /// github: URI and a local path: URI for testing).
103    pub fn list_uris_for_id(&self, id: &str) -> Vec<String> {
104        let mut entries: Vec<_> = self.entries.values().filter(|e| e.id == id).collect();
105        entries.sort_by(|a, b| b.hit.cmp(&a.hit));
106        entries.iter().map(|e| e.uri.clone()).collect()
107    }
108
109    /// Add entries for all inputs without incrementing hit counts.
110    ///
111    /// This is used to populate the cache with inputs discovered while
112    /// running any command (list, change, update, etc.), not just add.
113    /// Unlike `add_entry`, this does NOT increment hit counts for existing
114    /// entries - it only adds new entries that don't exist yet.
115    pub fn populate_from_inputs<'a>(&mut self, inputs: impl Iterator<Item = (&'a str, &'a str)>) {
116        for (id, uri) in inputs {
117            let key = entry_key(id, uri);
118            // Only add if not already present (don't increment hit count)
119            self.entries.entry(key).or_insert_with(|| CacheEntry {
120                id: id.to_string(),
121                uri: uri.to_string(),
122                hit: 0,
123            });
124        }
125    }
126}
127
128/// Populate the cache with inputs from a flake.
129///
130/// This is a convenience function that loads the cache, adds any new inputs,
131/// and commits the changes. It's designed to be called from any command that
132/// reads inputs, helping build up the cache over time.
133///
134/// Errors are logged but don't cause failures - caching is best-effort.
135/// If `no_cache` is true, this function does nothing.
136pub fn populate_cache_from_inputs<'a>(
137    inputs: impl Iterator<Item = (&'a str, &'a str)>,
138    no_cache: bool,
139) {
140    if no_cache {
141        return;
142    }
143
144    let mut cache = Cache::load();
145    let initial_len = cache.entries.len();
146    cache.populate_from_inputs(inputs);
147
148    // Only write if we added new entries
149    if cache.entries.len() > initial_len
150        && let Err(e) = cache.commit()
151    {
152        tracing::debug!("Could not write to cache: {}", e);
153    }
154}
155
156/// Populate the cache from an InputMap.
157///
158/// Convenience wrapper around `populate_cache_from_inputs` for use with
159/// the `FlakeEdit::list()` result. URIs are trimmed of surrounding quotes
160/// since the raw syntax representation includes them.
161/// If `no_cache` is true, this function does nothing.
162pub fn populate_cache_from_input_map(inputs: &crate::edit::InputMap, no_cache: bool) {
163    populate_cache_from_inputs(
164        inputs
165            .iter()
166            .map(|(id, input)| (id.as_str(), input.url().trim_matches('"'))),
167        no_cache,
168    );
169}
170
171/// Default flake URI type prefixes for completion.
172pub const DEFAULT_URI_TYPES: [&str; 14] = [
173    "github:",
174    "gitlab:",
175    "sourcehut:",
176    "git+https://",
177    "git+ssh://",
178    "git+http://",
179    "git+file://",
180    "git://",
181    "path:",
182    "file://",
183    "tarball:",
184    "https://",
185    "http://",
186    "flake:",
187];
188
189/// Configuration for cache usage.
190///
191/// Controls whether and where to read/write the URI completion cache.
192#[derive(Debug, Clone, Default)]
193pub enum CacheConfig {
194    /// Use the default cache location (~/.local/share/flake-edit/)
195    #[default]
196    Default,
197    /// Don't use any cache (for --no-cache flag)
198    None,
199    /// Use a custom cache file path (for --cache flag or testing)
200    Custom(std::path::PathBuf),
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    #[test]
208    fn test_cache_add_and_list() {
209        let mut cache = Cache::default();
210        cache.add_entry("nixpkgs".into(), "github:NixOS/nixpkgs".into());
211        cache.add_entry(
212            "home-manager".into(),
213            "github:nix-community/home-manager".into(),
214        );
215        cache.add_entry("nixpkgs".into(), "github:NixOS/nixpkgs".into()); // Increment hit
216
217        let uris = cache.list_uris();
218        assert_eq!(uris.len(), 2);
219        // nixpkgs should be first due to higher hit count
220        assert_eq!(uris[0], "github:NixOS/nixpkgs");
221    }
222
223    #[test]
224    fn test_list_uris_for_id() {
225        let mut cache = Cache::default();
226        // Add multiple URIs for the same ID (simulating local/remote toggle workflow)
227        cache.add_entry("treefmt-nix".into(), "github:numtide/treefmt-nix".into());
228        cache.add_entry(
229            "treefmt-nix".into(),
230            "path:/home/user/dev/treefmt-nix".into(),
231        );
232        // Add unrelated entry
233        cache.add_entry("nixpkgs".into(), "github:NixOS/nixpkgs".into());
234        // Increment hit on the github one
235        cache.add_entry("treefmt-nix".into(), "github:numtide/treefmt-nix".into());
236
237        let uris = cache.list_uris_for_id("treefmt-nix");
238        assert_eq!(uris.len(), 2);
239        // github should be first due to higher hit count
240        assert_eq!(uris[0], "github:numtide/treefmt-nix");
241        assert_eq!(uris[1], "path:/home/user/dev/treefmt-nix");
242
243        // Should not include nixpkgs
244        assert!(!uris.contains(&"github:NixOS/nixpkgs".to_string()));
245    }
246
247    #[test]
248    fn test_list_uris_for_id_empty() {
249        let cache = Cache::default();
250        let uris = cache.list_uris_for_id("nonexistent");
251        assert!(uris.is_empty());
252    }
253
254    #[test]
255    fn test_populate_from_inputs() {
256        let mut cache = Cache::default();
257
258        // Add some initial entries
259        cache.add_entry("nixpkgs".into(), "github:NixOS/nixpkgs".into());
260        cache.add_entry("nixpkgs".into(), "github:NixOS/nixpkgs".into()); // hit = 1
261
262        // Populate with inputs (simulating what happens when running a command)
263        let inputs = vec![
264            ("nixpkgs", "github:NixOS/nixpkgs"),           // Already exists
265            ("flake-utils", "github:numtide/flake-utils"), // New
266            ("home-manager", "github:nix-community/home-manager"), // New
267        ];
268        cache.populate_from_inputs(inputs.into_iter());
269
270        // Should have 3 entries total
271        let uris = cache.list_uris();
272        assert_eq!(uris.len(), 3);
273
274        // nixpkgs should still be first (hit=1, others hit=0)
275        assert_eq!(uris[0], "github:NixOS/nixpkgs");
276
277        // New entries should exist
278        assert!(uris.contains(&"github:numtide/flake-utils".to_string()));
279        assert!(uris.contains(&"github:nix-community/home-manager".to_string()));
280    }
281
282    #[test]
283    fn test_populate_does_not_increment_hits() {
284        let mut cache = Cache::default();
285
286        // Add entry with hit count
287        cache.add_entry("nixpkgs".into(), "github:NixOS/nixpkgs".into());
288        cache.add_entry("nixpkgs".into(), "github:NixOS/nixpkgs".into()); // hit = 1
289
290        // Populate with same entry
291        let inputs = vec![("nixpkgs", "github:NixOS/nixpkgs")];
292        cache.populate_from_inputs(inputs.into_iter());
293
294        // Hit count should still be 1, not 2
295        let entry = cache.entries.get("nixpkgs.github:NixOS/nixpkgs").unwrap();
296        assert_eq!(entry.hit, 1);
297    }
298}
flake_edit/cache.rs

flake_edit/
cache.rs
