kedge_cache/
lib.rs

1//! # kedge-cache
2//!
3//! A deterministic, content-addressed cache for **AST compaction** — never for
4//! LLM responses.
5//!
6//! Caching LLM output is a production hazard: cache invalidation is hard, and
7//! returning a stale/hallucinated answer because a prompt *looked* similar breaks
8//! things silently. Tree-sitter compaction, by contrast, is a pure function of the
9//! file's bytes: `sha256(file_contents)` is a perfect key. If a file hasn't
10//! changed, its skeleton hasn't either — return the cached result and skip the
11//! parser entirely.
12
13use std::path::Path;
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::sync::Mutex;
16use std::time::{SystemTime, UNIX_EPOCH};
17
18use kedge_compact::{CompactResult, Compactor};
19use rusqlite::{Connection, OptionalExtension};
20use sha2::{Digest, Sha256};
21use thiserror::Error;
22
23#[derive(Debug, Error)]
24pub enum CacheError {
25    #[error("sqlite error: {0}")]
26    Sqlite(#[from] rusqlite::Error),
27    #[error("compaction error: {0}")]
28    Compact(#[from] kedge_compact::CompactError),
29    #[error("cache lock poisoned")]
30    Poisoned,
31}
32
33pub type Result<T> = std::result::Result<T, CacheError>;
34
35/// Bump when kedge-compact's compaction algorithm changes in a way that alters
36/// output for unchanged input. It's part of the cache key, so a bump invalidates
37/// every prior entry (they simply become misses) — no stale skeletons after an
38/// upgrade.
39const COMPACT_ALGO_VERSION: &str = "1";
40
41// The key is `(hash, lang, version)`, not `hash` alone: identical *bytes* compacted
42// under a different language (a polyglot file, or the same source fed as Rust then
43// Python) must not return the first language's skeleton, and an algorithm-version
44// bump must not serve a stale one. The table is versioned (`_v2`) so an older
45// content-only cache is discarded rather than mis-read.
46const SCHEMA: &str = r#"
47DROP TABLE IF EXISTS ast_cache;
48CREATE TABLE IF NOT EXISTS ast_cache_v2 (
49    hash             TEXT NOT NULL,      -- sha256 of the source bytes
50    lang             TEXT NOT NULL,
51    version          TEXT NOT NULL,      -- COMPACT_ALGO_VERSION
52    text             TEXT NOT NULL,      -- the compacted skeleton
53    original_tokens  INTEGER NOT NULL,
54    compacted_tokens INTEGER NOT NULL,
55    elided_bodies    INTEGER NOT NULL,
56    created_ms       INTEGER NOT NULL,
57    PRIMARY KEY (hash, lang, version)
58);
59"#;
60
61/// The sha256 content hash used as the cache key.
62pub fn content_hash(source: &str) -> String {
63    let mut h = Sha256::new();
64    h.update(source.as_bytes());
65    format!("{:x}", h.finalize())
66}
67
68fn now_ms() -> i64 {
69    SystemTime::now()
70        .duration_since(UNIX_EPOCH)
71        .map(|d| d.as_millis() as i64)
72        .unwrap_or(0)
73}
74
75/// A SQLite-backed cache of compaction results, keyed by content hash.
76pub struct AstCache {
77    conn: Mutex<Connection>,
78    hits: AtomicU64,
79    misses: AtomicU64,
80}
81
82impl AstCache {
83    /// Open (or create) a cache at `path`.
84    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
85        Self::init(Connection::open(path)?)
86    }
87
88    /// An ephemeral in-memory cache (tests, one-shot runs).
89    pub fn in_memory() -> Result<Self> {
90        Self::init(Connection::open_in_memory()?)
91    }
92
93    fn init(conn: Connection) -> Result<Self> {
94        conn.execute_batch(SCHEMA)?;
95        Ok(AstCache {
96            conn: Mutex::new(conn),
97            hits: AtomicU64::new(0),
98            misses: AtomicU64::new(0),
99        })
100    }
101
102    /// Compact `source`, returning a cached result on a content-hash hit and only
103    /// invoking the Tree-sitter parser on a miss.
104    pub fn outline_cached(&self, compactor: &mut Compactor, source: &str) -> Result<CompactResult> {
105        let hash = content_hash(source);
106        let lang = compactor.language().name();
107        if let Some(hit) = self.load(&hash, lang)? {
108            self.hits.fetch_add(1, Ordering::SeqCst);
109            return Ok(hit);
110        }
111        self.misses.fetch_add(1, Ordering::SeqCst);
112        let result = compactor.outline(source)?;
113        self.store(&hash, lang, &result)?;
114        Ok(result)
115    }
116
117    /// Cache hits observed so far.
118    pub fn hits(&self) -> u64 {
119        self.hits.load(Ordering::SeqCst)
120    }
121
122    /// Cache misses (i.e. parses actually run) so far.
123    pub fn misses(&self) -> u64 {
124        self.misses.load(Ordering::SeqCst)
125    }
126
127    fn load(&self, hash: &str, lang: &str) -> Result<Option<CompactResult>> {
128        let conn = self.conn.lock().map_err(|_| CacheError::Poisoned)?;
129        let row = conn
130            .query_row(
131                "SELECT text, original_tokens, compacted_tokens, elided_bodies
132                 FROM ast_cache_v2 WHERE hash = ?1 AND lang = ?2 AND version = ?3",
133                rusqlite::params![hash, lang, COMPACT_ALGO_VERSION],
134                |r| {
135                    Ok(CompactResult {
136                        text: r.get(0)?,
137                        original_tokens: r.get::<_, i64>(1)? as usize,
138                        compacted_tokens: r.get::<_, i64>(2)? as usize,
139                        elided_bodies: r.get::<_, i64>(3)? as usize,
140                    })
141                },
142            )
143            .optional()?;
144        Ok(row)
145    }
146
147    fn store(&self, hash: &str, lang: &str, r: &CompactResult) -> Result<()> {
148        let conn = self.conn.lock().map_err(|_| CacheError::Poisoned)?;
149        conn.execute(
150            "INSERT OR REPLACE INTO ast_cache_v2
151             (hash, lang, version, text, original_tokens, compacted_tokens, elided_bodies, created_ms)
152             VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
153            rusqlite::params![
154                hash,
155                lang,
156                COMPACT_ALGO_VERSION,
157                r.text,
158                r.original_tokens as i64,
159                r.compacted_tokens as i64,
160                r.elided_bodies as i64,
161                now_ms(),
162            ],
163        )?;
164        Ok(())
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    const SRC: &str = r#"
173        /// A greeter.
174        pub fn hello(name: &str) -> String {
175            let mut s = String::new();
176            s.push_str("hi ");
177            s.push_str(name);
178            s
179        }
180    "#;
181
182    #[test]
183    fn hash_is_stable_and_content_sensitive() {
184        assert_eq!(content_hash("abc"), content_hash("abc"));
185        assert_ne!(content_hash("abc"), content_hash("abd"));
186    }
187
188    #[test]
189    fn miss_then_hit_skips_the_parser() {
190        let cache = AstCache::in_memory().unwrap();
191        let mut c = Compactor::rust().unwrap();
192
193        // First call: miss → parses and stores.
194        let first = cache.outline_cached(&mut c, SRC).unwrap();
195        assert_eq!(cache.misses(), 1);
196        assert_eq!(cache.hits(), 0);
197
198        // Second call, identical source: hit → no parse, same skeleton.
199        let second = cache.outline_cached(&mut c, SRC).unwrap();
200        assert_eq!(cache.hits(), 1);
201        assert_eq!(cache.misses(), 1);
202        assert_eq!(first.text, second.text);
203        assert_eq!(first.compacted_tokens, second.compacted_tokens);
204
205        // A changed file: miss again.
206        let changed = format!("{SRC}\n// touched");
207        cache.outline_cached(&mut c, &changed).unwrap();
208        assert_eq!(cache.misses(), 2);
209    }
210
211    #[test]
212    fn same_bytes_different_language_is_not_a_hit() {
213        let cache = AstCache::in_memory().unwrap();
214        let mut rust = Compactor::rust().unwrap();
215        let mut py = Compactor::new(kedge_compact::Language::Python).unwrap();
216        // Store the skeleton under Rust.
217        cache.outline_cached(&mut rust, SRC).unwrap();
218        assert_eq!(cache.misses(), 1);
219        // The exact same bytes compacted as Python must MISS — not return the Rust
220        // skeleton — because the key includes the language.
221        cache.outline_cached(&mut py, SRC).unwrap();
222        assert_eq!(
223            cache.misses(),
224            2,
225            "same bytes under a different language must not be a cross-language hit"
226        );
227        assert_eq!(cache.hits(), 0);
228    }
229
230    #[test]
231    fn cache_survives_reopen_on_disk() {
232        let dir = tempfile::tempdir().unwrap();
233        let path = dir.path().join("ast.sqlite");
234        {
235            let cache = AstCache::open(&path).unwrap();
236            let mut c = Compactor::rust().unwrap();
237            cache.outline_cached(&mut c, SRC).unwrap();
238            assert_eq!(cache.misses(), 1);
239        }
240        // Reopen: the earlier result is a hit, no parse needed.
241        let cache = AstCache::open(&path).unwrap();
242        let mut c = Compactor::rust().unwrap();
243        let r = cache.outline_cached(&mut c, SRC).unwrap();
244        assert_eq!(cache.hits(), 1);
245        assert_eq!(cache.misses(), 0);
246        assert!(r.text.contains("hello"));
247    }
248}
kedge_cache/lib.rs

kedge_cache/
lib.rs
