zsh/extensions/
deparse_cache.rs

1//! rkyv-backed cache of deparsed function bodies — the text `$functions`
2//! and `${functions[name]}` print.
3//!
4//! **zshrs-original — no C counterpart.** C stores a function as wordcode
5//! and `getpermtext` walks it (c:Src/Modules/parameter.c:495), which is
6//! cheap. zshrs stores the body as text, so producing the same output means
7//! re-parsing and deparsing it. `modules/parameter.rs` memoizes that per
8//! process; this file keeps the results across processes, so a new shell's
9//! first `${functions}` read (zpwr's `_parameters` does one on every
10//! command-position `<TAB>`) does not re-parse every function it has.
11//!
12//! Single-file shard at `$ZSHRS_HOME/deparse.rkyv`, same discipline as
13//! [`crate::autoload_cache`]: versioned header, exact
14//! producing-binary identity, atomic-rename writes, a `pending` buffer
15//! flushed once per prompt and at exit.
16//!
17//! # What identifies an entry
18//!
19//! The deparse is a function of everything the re-lex sees, so the key is a
20//! SHA-256 over all of it, taken while `funcdef_lex_pin` is in force:
21//!
22//!   * the function name and the exact body text;
23//!   * every option bit (`isset(1..OPT_SIZE)`), which covers the pinned
24//!     RCQUOTES and a sticky emulation's option set;
25//!   * the `emulation` value the parser consults directly.
26//!
27//! A body re-lexed with aliases ON also depends on the whole alias table.
28//! Those bodies are not stored here ([`key`] returns `None`) and keep only
29//! the per-process memo. Bodies the parser captured and `autoload -U`
30//! bodies re-lex with aliases off, and they are the bulk.
31//!
32//! The whole shard is stamped with the producing binary's `(mtime, len)`:
33//! deparse output is a property of the build, so a shard from any other
34//! binary is discarded rather than read.
35
36use std::collections::HashMap;
37use std::fs::File;
38use std::path::{Path, PathBuf};
39
40use memmap2::Mmap;
41use parking_lot::Mutex;
42use rkyv::{Archive, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
43
44/// "ZRDP" little-endian.
45pub const SHARD_MAGIC: u32 = 0x5A52_4450;
46/// Bump when the key recipe or the entry layout changes.
47/// v2: token chars moved to the Private Use Area (crate::token_char).
48pub const SHARD_FORMAT_VERSION: u32 = 3;
49/// Entries buffered before a flush is forced. A memory bound for a process
50/// that never reaches a prompt; the real batch boundary is the prompt.
51const PENDING_FLUSH_MAX: usize = 4096;
52/// A shard that grows past this is started over. Old keys are never removed
53/// otherwise (an edited function leaves its previous body's entry behind),
54/// and a rebuild costs one re-deparse per live function.
55const SHARD_MAX_ENTRIES: usize = 100_000;
56
57/// On-disk shard. `entries` maps a [`key`] digest to the deparsed body.
58#[derive(Archive, RkyvDeserialize, RkyvSerialize, Debug, Clone)]
59#[archive(check_bytes)]
60pub struct DeparseShard {
61    /// [`SHARD_MAGIC`].
62    pub magic: u32,
63    /// [`SHARD_FORMAT_VERSION`].
64    pub format_version: u32,
65    /// `size_of::<usize>()` of the writer.
66    pub pointer_width: u32,
67    /// mtime of the binary that produced every entry.
68    pub binary_mtime: i64,
69    /// Byte length of that binary.
70    pub binary_len: u64,
71    /// key digest → deparsed body.
72    pub entries: HashMap<[u8; 32], String>,
73}
74
75struct MmappedShard {
76    _mmap: Mmap,
77    archived: *const ArchivedDeparseShard,
78}
79
80// The mapping is read-only and lives as long as the struct.
81unsafe impl Send for MmappedShard {}
82unsafe impl Sync for MmappedShard {}
83
84impl MmappedShard {
85    fn open(path: &Path) -> Option<Self> {
86        let file = File::open(path).ok()?;
87        let mmap = unsafe { Mmap::map(&file).ok()? };
88        let archived = rkyv::check_archived_root::<DeparseShard>(&mmap[..]).ok()?;
89        let archived = archived as *const ArchivedDeparseShard;
90        let shard = Self { _mmap: mmap, archived };
91        shard.header_ok().then_some(shard)
92    }
93
94    fn shard(&self) -> &ArchivedDeparseShard {
95        unsafe { &*self.archived }
96    }
97
98    fn header_ok(&self) -> bool {
99        let s = self.shard();
100        let Some((mtime, len)) = current_binary_identity() else {
101            return false;
102        };
103        let magic: u32 = s.magic.into();
104        let fv: u32 = s.format_version.into();
105        let pw: u32 = s.pointer_width.into();
106        let bm: i64 = s.binary_mtime.into();
107        let bl: u64 = s.binary_len.into();
108        magic == SHARD_MAGIC
109            && fv == SHARD_FORMAT_VERSION
110            && pw as usize == std::mem::size_of::<usize>()
111            && bm == mtime
112            && bl == len
113    }
114
115    fn get(&self, key: &[u8; 32]) -> Option<String> {
116        self.shard().entries.get(key).map(|s| s.as_str().to_string())
117    }
118}
119
120/// The per-process handle: the mapped shard plus writes not yet on disk.
121pub struct DeparseCache {
122    path: PathBuf,
123    lock_path: PathBuf,
124    /// Opened on first lookup; `Some(None)` = no usable shard on disk.
125    mmap: Mutex<Option<Option<MmappedShard>>>,
126    pending: Mutex<HashMap<[u8; 32], String>>,
127}
128
129impl DeparseCache {
130    fn open(path: &Path) -> Self {
131        let lock_path = path.with_extension("rkyv.lock");
132        Self {
133            path: path.to_path_buf(),
134            lock_path,
135            mmap: Mutex::new(None),
136            pending: Mutex::new(HashMap::new()),
137        }
138    }
139
140    fn get(&self, key: &[u8; 32]) -> Option<String> {
141        if let Some(hit) = self.pending.lock().get(key) {
142            return Some(hit.clone());
143        }
144        let mut guard = self.mmap.lock();
145        let shard = guard.get_or_insert_with(|| MmappedShard::open(&self.path));
146        shard.as_ref()?.get(key)
147    }
148
149    fn put(&self, key: [u8; 32], deparsed: String) {
150        let full = {
151            let mut pending = self.pending.lock();
152            pending.insert(key, deparsed);
153            pending.len() >= PENDING_FLUSH_MAX
154        };
155        if full {
156            let _ = self.flush_pending();
157        }
158    }
159
160    /// Fold every buffered entry into the shard in one rewrite. A no-op when
161    /// nothing is buffered. A failed write drops the batch instead of
162    /// retrying at every prompt; the cost is one re-deparse per entry.
163    ///
164    /// The lock is taken without blocking: this runs before every prompt,
165    /// and waiting on another process's rewrite left a new shell with no
166    /// prompt. A busy lock puts the batch back for the next prompt.
167    fn flush_pending(&self) -> Result<(), String> {
168        let batch = {
169            let mut pending = self.pending.lock();
170            if pending.is_empty() {
171                return Ok(());
172            }
173            std::mem::take(&mut *pending)
174        };
175        if let Some(parent) = self.path.parent() {
176            std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
177        }
178        let Some(_lock) = try_acquire_lock(&self.lock_path) else {
179            // Entries buffered since the take are newer; keep them. Bounded
180            // by PENDING_FLUSH_MAX, past which the rest cost a re-deparse.
181            let mut pending = self.pending.lock();
182            for (key, deparsed) in batch {
183                if pending.len() >= PENDING_FLUSH_MAX {
184                    break;
185                }
186                pending.entry(key).or_insert(deparsed);
187            }
188            return Ok(());
189        };
190        let Some((mtime, len)) = current_binary_identity() else {
191            return Ok(());
192        };
193        let mut shard = read_owned_shard(&self.path)
194            .filter(|s| {
195                s.magic == SHARD_MAGIC
196                    && s.format_version == SHARD_FORMAT_VERSION
197                    && s.pointer_width as usize == std::mem::size_of::<usize>()
198                    && s.binary_mtime == mtime
199                    && s.binary_len == len
200                    && s.entries.len() < SHARD_MAX_ENTRIES
201            })
202            .unwrap_or_else(|| DeparseShard {
203                magic: SHARD_MAGIC,
204                format_version: SHARD_FORMAT_VERSION,
205                pointer_width: std::mem::size_of::<usize>() as u32,
206                binary_mtime: mtime,
207                binary_len: len,
208                entries: HashMap::new(),
209            });
210        shard.entries.extend(batch);
211        let bytes =
212            rkyv::to_bytes::<_, 4096>(&shard).map_err(|e| format!("rkyv serialize: {}", e))?;
213        crate::atomic_write::write_bytes_atomic(&self.path, &bytes)?;
214        // The next lookup maps the file just written.
215        *self.mmap.lock() = None;
216        Ok(())
217    }
218}
219
220/// Exclusive lock on the shard's side file, or `None` when another
221/// process holds it.
222fn try_acquire_lock(path: &Path) -> Option<nix::fcntl::Flock<File>> {
223    let f = File::options()
224        .read(true)
225        .write(true)
226        .create(true)
227        .truncate(false)
228        .open(path)
229        .ok()?;
230    nix::fcntl::Flock::lock(f, nix::fcntl::FlockArg::LockExclusiveNonblock).ok()
231}
232
233fn read_owned_shard(path: &Path) -> Option<DeparseShard> {
234    let bytes = std::fs::read(path).ok()?;
235    let archived = rkyv::check_archived_root::<DeparseShard>(&bytes[..]).ok()?;
236    archived.deserialize(&mut rkyv::Infallible).ok()
237}
238
239fn current_binary_identity() -> Option<(i64, u64)> {
240    use std::os::unix::fs::MetadataExt;
241    static BIN_ID: std::sync::OnceLock<Option<(i64, u64)>> = std::sync::OnceLock::new();
242    *BIN_ID.get_or_init(|| {
243        let meta = std::fs::metadata(std::env::current_exe().ok()?).ok()?;
244        Some((meta.mtime(), meta.len()))
245    })
246}
247
248fn default_cache_path() -> PathBuf {
249    let root = match std::env::var_os("ZSHRS_HOME") {
250        Some(custom) => PathBuf::from(custom),
251        None => dirs::home_dir()
252            .unwrap_or_else(|| PathBuf::from("/tmp"))
253            .join(".zshrs"),
254    };
255    root.join("deparse.rkyv")
256}
257
258static CACHE: once_cell::sync::Lazy<Option<DeparseCache>> = once_cell::sync::Lazy::new(|| {
259    crate::autoload_cache::cache_enabled().then(|| DeparseCache::open(&default_cache_path()))
260});
261
262/// The cache key for deparsing `text` as the body of `name`, or `None` when
263/// the result also depends on the live alias table (see the module doc).
264/// Call it while `funcdef_lex_pin` is held, so the pinned options are the
265/// ones hashed.
266pub fn key(name: &str, text: &str) -> Option<[u8; 32]> {
267    use sha2::Digest;
268    if !crate::ported::lex::noaliases() {
269        return None;
270    }
271    let mut bits = [0u8; (crate::ported::zsh_h::OPT_SIZE as usize).div_ceil(8)];
272    for optno in 1..crate::ported::zsh_h::OPT_SIZE {
273        if crate::ported::zsh_h::isset(optno) {
274            bits[optno as usize / 8] |= 1 << (optno % 8);
275        }
276    }
277    let emulation = crate::ported::options::emulation.load(std::sync::atomic::Ordering::Relaxed);
278    let mut h = sha2::Sha256::new();
279    h.update(name.as_bytes());
280    h.update([0]);
281    h.update(text.as_bytes());
282    h.update([0]);
283    h.update(bits);
284    h.update(emulation.to_le_bytes());
285    Some(h.finalize().into())
286}
287
288/// A previously stored deparse for `key`.
289pub fn lookup(key: &[u8; 32]) -> Option<String> {
290    CACHE.as_ref()?.get(key)
291}
292
293/// Remember `deparsed` for `key`; it reaches disk at the next flush.
294pub fn record(key: [u8; 32], deparsed: &str) {
295    if let Some(cache) = CACHE.as_ref() {
296        cache.put(key, deparsed.to_string());
297    }
298}
299
300/// Write out everything [`record`] buffered. Called beside
301/// `autoload_cache::try_flush_pending` at the prompt, at `zexit`, and from
302/// the `atexit` hook, for the same reasons given there.
303pub fn try_flush_pending() {
304    if let Some(cache) = CACHE.as_ref() {
305        if let Err(e) = cache.flush_pending() {
306            if !crate::atexit_teardown::active() {
307                tracing::warn!(error = %e, "deparse cache: could not flush");
308            }
309        }
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316
317    /// A flushed entry is served by a fresh handle on the same file, which is
318    /// what a new shell sees.
319    #[test]
320    fn flushed_entry_survives_into_a_new_handle() {
321        let dir = tempfile::tempdir().unwrap();
322        let path = dir.path().join("deparse.rkyv");
323        let key = [7u8; 32];
324        let writer = DeparseCache::open(&path);
325        writer.put(key, "print hi".to_string());
326        assert_eq!(writer.get(&key).as_deref(), Some("print hi"), "pending read");
327        writer.flush_pending().unwrap();
328        let reader = DeparseCache::open(&path);
329        assert_eq!(reader.get(&key).as_deref(), Some("print hi"));
330        assert_eq!(reader.get(&[8u8; 32]), None);
331    }
332
333    /// A second flush keeps what the first one wrote.
334    #[test]
335    fn flush_merges_with_the_shard_on_disk() {
336        let dir = tempfile::tempdir().unwrap();
337        let path = dir.path().join("deparse.rkyv");
338        let a = DeparseCache::open(&path);
339        a.put([1u8; 32], "one".to_string());
340        a.flush_pending().unwrap();
341        let b = DeparseCache::open(&path);
342        b.put([2u8; 32], "two".to_string());
343        b.flush_pending().unwrap();
344        let c = DeparseCache::open(&path);
345        assert_eq!(c.get(&[1u8; 32]).as_deref(), Some("one"));
346        assert_eq!(c.get(&[2u8; 32]).as_deref(), Some("two"));
347    }
348
349    /// A shard written under a different binary identity is not read.
350    #[test]
351    fn shard_from_another_binary_is_ignored() {
352        let dir = tempfile::tempdir().unwrap();
353        let path = dir.path().join("deparse.rkyv");
354        let (mtime, len) = current_binary_identity().unwrap();
355        let mut entries = HashMap::new();
356        entries.insert([3u8; 32], "stale".to_string());
357        let shard = DeparseShard {
358            magic: SHARD_MAGIC,
359            format_version: SHARD_FORMAT_VERSION,
360            pointer_width: std::mem::size_of::<usize>() as u32,
361            binary_mtime: mtime,
362            binary_len: len + 1,
363            entries,
364        };
365        let bytes = rkyv::to_bytes::<_, 4096>(&shard).unwrap();
366        std::fs::write(&path, &bytes).unwrap();
367        assert_eq!(DeparseCache::open(&path).get(&[3u8; 32]), None);
368    }
369}
zsh/extensions/deparse_cache.rs

zsh/extensions/
deparse_cache.rs
