zsh/extensions/
cow_map.rs

1//! `cow_map` — copy-on-write `HashMap` wrapper for cheap subshell
2//! snapshot/restore.
3//!
4//! !!! WARNING: RUST-ONLY HELPER !!!
5//! C zsh has no counterpart. `$( … )` and `( … )` are forks there
6//! (`c:Src/exec.c:4783` `getoutput` → `entersubsh`), so the child gets
7//! the parent's whole address space by page-table copy and the kernel
8//! does the copy-on-write. zshrs runs both forms IN PROCESS and has to
9//! snapshot the mutable globals by hand, which turned "enter a
10//! substitution" into O(total shell state). This type is the userspace
11//! stand-in for the page-table trick: sharing until someone writes.
12//!
13//! Every read goes through `Deref` and touches the shared map. Every
14//! `&mut` method goes through `DerefMut`, which calls `Arc::make_mut` —
15//! a no-op when the map is unshared (the ordinary case), and a single
16//! deep copy the first time a subshell body writes while the parent
17//! still holds a snapshot. After that copy the subshell owns its map
18//! and the parent's snapshot is frozen at the pre-write contents, which
19//! is exactly what the fork gave C.
20//!
21//! `clone()` is deliberately NOT a deep copy — it is the snapshot
22//! operation, so it must stay O(1). Independence is still guaranteed:
23//! whichever side writes first is the one that pays for the split.
24
25use std::collections::HashMap;
26use std::hash::Hash;
27use std::ops::{Deref, DerefMut};
28use std::sync::Arc;
29
30/// Copy-on-write associative store used for subshell snapshot/restore.
31///
32/// Drop-in for `HashMap<K, V>` at read and write call sites via
33/// `Deref`/`DerefMut`; the difference is only visible in the cost of
34/// `clone()`.
35#[derive(Debug)]
36pub struct CowHashMap<K, V> {
37    inner: Arc<HashMap<K, V>>,
38}
39
40impl<K, V> CowHashMap<K, V> {
41    /// Empty map, sharing nothing.
42    pub fn new() -> Self {
43        Self {
44            inner: Arc::new(HashMap::new()),
45        }
46    }
47
48    /// True while another handle (a live subshell snapshot) shares this
49    /// map, i.e. while the next write will pay for a deep copy.
50    pub fn is_shared(&self) -> bool {
51        Arc::strong_count(&self.inner) > 1
52    }
53}
54
55impl<K, V> Default for CowHashMap<K, V> {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61/// O(1) — a refcount bump, NOT a deep copy. This is the snapshot
62/// operation; see the module docs.
63impl<K, V> Clone for CowHashMap<K, V> {
64    fn clone(&self) -> Self {
65        Self {
66            inner: Arc::clone(&self.inner),
67        }
68    }
69}
70
71impl<K, V> Deref for CowHashMap<K, V> {
72    type Target = HashMap<K, V>;
73
74    fn deref(&self) -> &HashMap<K, V> {
75        &self.inner
76    }
77}
78
79impl<K: Clone + Eq + Hash, V: Clone> DerefMut for CowHashMap<K, V> {
80    /// Splits the map away from any snapshot sharing it, then hands out
81    /// the `&mut`. A caller that takes `&mut` only to read still pays
82    /// the split — that is a cost bug, never a correctness one.
83    fn deref_mut(&mut self) -> &mut HashMap<K, V> {
84        Arc::make_mut(&mut self.inner)
85    }
86}
87
88impl<K, V> From<HashMap<K, V>> for CowHashMap<K, V> {
89    fn from(map: HashMap<K, V>) -> Self {
90        Self {
91            inner: Arc::new(map),
92        }
93    }
94}
95
96impl<K: Clone + Eq + Hash, V: Clone> FromIterator<(K, V)> for CowHashMap<K, V> {
97    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
98        Self::from(HashMap::from_iter(iter))
99    }
100}
101
102impl<'a, K, V> IntoIterator for &'a CowHashMap<K, V> {
103    type Item = (&'a K, &'a V);
104    type IntoIter = std::collections::hash_map::Iter<'a, K, V>;
105
106    fn into_iter(self) -> Self::IntoIter {
107        self.inner.iter()
108    }
109}
110
111impl<K: Eq + Hash, V: PartialEq> PartialEq for CowHashMap<K, V> {
112    fn eq(&self, other: &Self) -> bool {
113        // Sharing the same allocation is equality without a walk.
114        Arc::ptr_eq(&self.inner, &other.inner) || *self.inner == *other.inner
115    }
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    /// The whole point: a snapshot must not walk the map, and must not
123    /// see writes made after it was taken.
124    #[test]
125    fn snapshot_is_shared_until_a_write_splits_it() {
126        let mut live: CowHashMap<String, Vec<String>> = CowHashMap::new();
127        live.insert("_comps".into(), vec!["git".into()]);
128
129        let snap = live.clone();
130        assert!(live.is_shared(), "clone must share, not copy");
131
132        // Subshell body writes: the split happens here, not at clone.
133        live.insert("added_in_subshell".into(), vec![]);
134        assert!(!live.is_shared(), "make_mut must have split the map");
135
136        assert!(snap.get("added_in_subshell").is_none());
137        assert_eq!(snap.get("_comps").map(Vec::len), Some(1));
138        assert_eq!(live.len(), 2);
139    }
140
141    /// Restoring the snapshot must undo the subshell's writes, which is
142    /// what `$( … )` does at its tail.
143    #[test]
144    fn restoring_a_snapshot_drops_the_subshell_writes() {
145        let mut live: CowHashMap<String, String> = CowHashMap::new();
146        live.insert("keep".into(), "outer".into());
147        let snap = live.clone();
148
149        live.insert("keep".into(), "inner".into());
150        live.insert("leaked".into(), "inner".into());
151        live.remove("nothing");
152
153        live = snap;
154        assert_eq!(live.get("keep").map(String::as_str), Some("outer"));
155        assert!(live.get("leaked").is_none());
156    }
157
158    /// Mutating through the snapshot handle must not reach back into the
159    /// live map either — the split is symmetric.
160    #[test]
161    fn writing_through_the_snapshot_does_not_reach_the_live_map() {
162        let mut live: CowHashMap<String, String> = CowHashMap::new();
163        live.insert("k".into(), "live".into());
164        let mut snap = live.clone();
165
166        snap.insert("k".into(), "snap".into());
167
168        assert_eq!(live.get("k").map(String::as_str), Some("live"));
169        assert_eq!(snap.get("k").map(String::as_str), Some("snap"));
170    }
171}
zsh/extensions/cow_map.rs

zsh/extensions/
cow_map.rs
