Skip to main content

wasm_dbms/transaction/overlay/table/
index.rs

1//! Table overlay index tracking for uncommitted transaction changes.
2
3use std::collections::{BTreeMap, HashMap, HashSet};
4
5use wasm_dbms_api::prelude::Value;
6
7/// Overlay for tracking uncommited transaction changes to indexes on a table.
8///
9/// 1 Table has 1 [`IndexOverlay`], which has n BTreeMaps for each index
10#[derive(Debug, Default, Clone)]
11pub struct IndexOverlay(HashMap<Vec<&'static str>, IndexColumnOverlay>);
12
13/// Overlay for tracking uncommited transaction changes to an indexed column.
14#[derive(Debug, Default, Clone)]
15struct IndexColumnOverlay {
16    /// Maps from added indexed values to their corresponding primary key values.
17    added: BTreeMap<Vec<Value>, HashSet<Value>>,
18    /// Maps from removed indexed values to their corresponding primary key values.
19    removed: BTreeMap<Vec<Value>, HashSet<Value>>,
20}
21
22impl IndexOverlay {
23    /// Inserts a record into the index overlay for the specified indexed columns, indexed values, and primary key.
24    ///
25    /// The primary key is added to the set of primary keys associated with the indexed values in the `added` map.
26    /// If the primary key was previously in the `removed` set for these values, it is removed from there.
27    pub fn insert(
28        &mut self,
29        indexed_columns: &[&'static str],
30        indexed_values: Vec<Value>,
31        pk: Value,
32    ) {
33        let column_overlay = self.0.entry(indexed_columns.to_vec()).or_default();
34        column_overlay.insert(indexed_values, pk);
35    }
36
37    /// Removes a record from the index overlay for the specified indexed columns, indexed values, and primary key.
38    ///
39    /// The primary key is added to the set of primary keys associated with the indexed values in the `removed` map.
40    /// If the primary key was previously in the `added` set for these values, it is removed from there.
41    pub fn delete(
42        &mut self,
43        indexed_columns: &[&'static str],
44        indexed_values: Vec<Value>,
45        pk: Value,
46    ) {
47        let column_overlay = self.0.entry(indexed_columns.to_vec()).or_default();
48        column_overlay.delete(indexed_values, pk);
49    }
50
51    /// Returns the set of primary keys added for the given index columns and indexed values.
52    ///
53    /// Returns an empty set if no entries exist.
54    pub fn added_pks(
55        &self,
56        indexed_columns: &[&'static str],
57        indexed_values: &[Value],
58    ) -> HashSet<Value> {
59        self.0
60            .get(indexed_columns)
61            .and_then(|overlay| overlay.added.get(indexed_values))
62            .cloned()
63            .unwrap_or_default()
64    }
65
66    /// Returns the set of primary keys removed for the given index columns and indexed values.
67    ///
68    /// Returns an empty set if no entries exist.
69    pub fn removed_pks(
70        &self,
71        indexed_columns: &[&'static str],
72        indexed_values: &[Value],
73    ) -> HashSet<Value> {
74        self.0
75            .get(indexed_columns)
76            .and_then(|overlay| overlay.removed.get(indexed_values))
77            .cloned()
78            .unwrap_or_default()
79    }
80
81    /// Returns the set of primary keys added within the given inclusive key range.
82    pub fn added_pks_in_range(
83        &self,
84        indexed_columns: &[&'static str],
85        start: Option<&[Value]>,
86        end: Option<&[Value]>,
87    ) -> HashSet<Value> {
88        self.pks_in_range_from_map(indexed_columns, start, end, |overlay| &overlay.added)
89    }
90
91    /// Returns the set of primary keys removed within the given inclusive key range.
92    pub fn removed_pks_in_range(
93        &self,
94        indexed_columns: &[&'static str],
95        start: Option<&[Value]>,
96        end: Option<&[Value]>,
97    ) -> HashSet<Value> {
98        self.pks_in_range_from_map(indexed_columns, start, end, |overlay| &overlay.removed)
99    }
100
101    /// Updates a record in the index overlay by removing the old indexed values and inserting the new ones.
102    ///
103    /// This is equivalent to calling [`delete`](Self::delete) with the old values
104    /// followed by [`insert`](Self::insert) with the new values.
105    pub fn update(
106        &mut self,
107        indexed_columns: &[&'static str],
108        old_indexed_values: Vec<Value>,
109        new_indexed_values: Vec<Value>,
110        pk: Value,
111    ) {
112        let column_overlay = self.0.entry(indexed_columns.to_vec()).or_default();
113        column_overlay.delete(old_indexed_values, pk.clone());
114        column_overlay.insert(new_indexed_values, pk);
115    }
116
117    fn pks_in_range_from_map(
118        &self,
119        indexed_columns: &[&'static str],
120        start: Option<&[Value]>,
121        end: Option<&[Value]>,
122        map_selector: impl Fn(&IndexColumnOverlay) -> &BTreeMap<Vec<Value>, HashSet<Value>>,
123    ) -> HashSet<Value> {
124        use std::ops::Bound;
125
126        let Some(column_overlay) = self.0.get(indexed_columns) else {
127            return HashSet::new();
128        };
129
130        let start_bound = match start {
131            Some(start_key) => Bound::Included(start_key.to_vec()),
132            None => Bound::Unbounded,
133        };
134        let end_bound = match end {
135            Some(end_key) => Bound::Included(end_key.to_vec()),
136            None => Bound::Unbounded,
137        };
138
139        map_selector(column_overlay)
140            .range((start_bound, end_bound))
141            .flat_map(|(_, pks)| pks.iter().cloned())
142            .collect()
143    }
144}
145
146impl IndexColumnOverlay {
147    /// Inserts a record into the index overlay for the specified indexed values and primary key.
148    ///
149    /// The primary key is added to the `added` set and removed from `removed` if present,
150    /// keeping the two sets consistent.
151    fn insert(&mut self, indexed_values: Vec<Value>, pk: Value) {
152        if let Some(removed_pks) = self.removed.get_mut(&indexed_values) {
153            removed_pks.remove(&pk);
154        }
155        self.added.entry(indexed_values).or_default().insert(pk);
156    }
157
158    /// Removes a record from the index overlay for the specified indexed values and primary key.
159    ///
160    /// If the PK was in the `added` set (overlay-only entry), it is simply removed from `added`
161    /// without adding to `removed`, since the base index never had this entry.
162    /// If the PK was **not** in `added`, it is added to `removed` to mark a base index entry
163    /// for exclusion.
164    fn delete(&mut self, indexed_values: Vec<Value>, pk: Value) {
165        let was_in_added = self
166            .added
167            .get_mut(&indexed_values)
168            .is_some_and(|added_pks| added_pks.remove(&pk));
169
170        if !was_in_added {
171            self.removed.entry(indexed_values).or_default().insert(pk);
172        }
173    }
174}
175
176#[cfg(test)]
177mod tests {
178
179    use super::*;
180
181    fn pk(id: u32) -> Value {
182        Value::Uint32(id.into())
183    }
184
185    fn text_val(s: &str) -> Value {
186        Value::Text(s.to_string().into())
187    }
188
189    // -- IndexColumnOverlay tests --
190
191    #[test]
192    fn test_column_overlay_insert_should_add_to_added() {
193        let mut overlay = IndexColumnOverlay::default();
194        overlay.insert(vec![text_val("alice")], pk(1));
195
196        assert!(overlay.added[&vec![text_val("alice")]].contains(&pk(1)));
197        assert!(overlay.removed.is_empty());
198    }
199
200    #[test]
201    fn test_column_overlay_delete_should_add_to_removed() {
202        let mut overlay = IndexColumnOverlay::default();
203        overlay.delete(vec![text_val("alice")], pk(1));
204
205        assert!(overlay.removed[&vec![text_val("alice")]].contains(&pk(1)));
206        assert!(overlay.added.is_empty());
207    }
208
209    #[test]
210    fn test_column_overlay_insert_after_delete_should_move_pk_from_removed_to_added() {
211        let mut overlay = IndexColumnOverlay::default();
212        let key = vec![text_val("alice")];
213
214        overlay.delete(key.clone(), pk(1));
215        assert!(overlay.removed[&key].contains(&pk(1)));
216
217        overlay.insert(key.clone(), pk(1));
218        assert!(overlay.added[&key].contains(&pk(1)));
219        assert!(!overlay.removed[&key].contains(&pk(1)));
220    }
221
222    #[test]
223    fn test_column_overlay_delete_after_insert_should_remove_from_added_without_stale_removed() {
224        let mut overlay = IndexColumnOverlay::default();
225        let key = vec![text_val("alice")];
226
227        overlay.insert(key.clone(), pk(1));
228        assert!(overlay.added[&key].contains(&pk(1)));
229
230        // Delete an overlay-only entry: should remove from added but NOT add to removed,
231        // since this key+pk was never in the base index.
232        overlay.delete(key.clone(), pk(1));
233        assert!(!overlay.added[&key].contains(&pk(1)));
234        assert!(
235            overlay
236                .removed
237                .get(&key)
238                .map_or(true, |pks| !pks.contains(&pk(1)))
239        );
240    }
241
242    #[test]
243    fn test_column_overlay_delete_base_entry_goes_to_removed() {
244        // Simulates deleting a record that exists in the base index (not overlay-inserted).
245        // Since it was never in `added`, it should appear in `removed`.
246        let mut overlay = IndexColumnOverlay::default();
247        let key = vec![text_val("alice")];
248
249        overlay.delete(key.clone(), pk(1));
250
251        assert!(overlay.removed[&key].contains(&pk(1)));
252        assert!(overlay.added.is_empty());
253    }
254
255    #[test]
256    fn test_column_overlay_delete_base_then_reinsert_cancels_removed() {
257        // Delete a base entry, then re-insert the same key+pk.
258        // The re-insert should cancel the removal.
259        let mut overlay = IndexColumnOverlay::default();
260        let key = vec![text_val("alice")];
261
262        overlay.delete(key.clone(), pk(1));
263        assert!(overlay.removed[&key].contains(&pk(1)));
264
265        overlay.insert(key.clone(), pk(1));
266        assert!(overlay.added[&key].contains(&pk(1)));
267        assert!(!overlay.removed[&key].contains(&pk(1)));
268    }
269
270    #[test]
271    fn test_column_overlay_multiple_pks_per_key() {
272        let mut overlay = IndexColumnOverlay::default();
273        let key = vec![text_val("alice")];
274
275        overlay.insert(key.clone(), pk(1));
276        overlay.insert(key.clone(), pk(2));
277        overlay.insert(key.clone(), pk(3));
278
279        let added_pks = &overlay.added[&key];
280        assert_eq!(added_pks.len(), 3);
281        assert!(added_pks.contains(&pk(1)));
282        assert!(added_pks.contains(&pk(2)));
283        assert!(added_pks.contains(&pk(3)));
284    }
285
286    #[test]
287    fn test_column_overlay_delete_one_pk_should_not_affect_others() {
288        let mut overlay = IndexColumnOverlay::default();
289        let key = vec![text_val("alice")];
290
291        overlay.insert(key.clone(), pk(1));
292        overlay.insert(key.clone(), pk(2));
293        overlay.delete(key.clone(), pk(1));
294
295        assert!(!overlay.added[&key].contains(&pk(1)));
296        assert!(overlay.added[&key].contains(&pk(2)));
297        // pk(1) was overlay-only, so it should NOT be in removed
298        assert!(
299            overlay
300                .removed
301                .get(&key)
302                .map_or(true, |pks| !pks.contains(&pk(1)))
303        );
304    }
305
306    // -- IndexOverlay tests --
307
308    #[test]
309    fn test_index_overlay_insert() {
310        let mut overlay = IndexOverlay::default();
311        let cols: &[&str] = &["name"];
312
313        overlay.insert(cols, vec![text_val("alice")], pk(1));
314
315        let column_overlay = &overlay.0[&vec!["name"]];
316        assert!(column_overlay.added[&vec![text_val("alice")]].contains(&pk(1)));
317    }
318
319    #[test]
320    fn test_index_overlay_delete() {
321        let mut overlay = IndexOverlay::default();
322        let cols: &[&str] = &["name"];
323
324        overlay.delete(cols, vec![text_val("alice")], pk(1));
325
326        let column_overlay = &overlay.0[&vec!["name"]];
327        assert!(column_overlay.removed[&vec![text_val("alice")]].contains(&pk(1)));
328    }
329
330    #[test]
331    fn test_index_overlay_update_should_remove_old_and_add_new() {
332        let mut overlay = IndexOverlay::default();
333        let cols: &[&str] = &["name"];
334
335        overlay.update(cols, vec![text_val("alice")], vec![text_val("bob")], pk(1));
336
337        let column_overlay = &overlay.0[&vec!["name"]];
338        assert!(column_overlay.removed[&vec![text_val("alice")]].contains(&pk(1)));
339        assert!(column_overlay.added[&vec![text_val("bob")]].contains(&pk(1)));
340    }
341
342    #[test]
343    fn test_index_overlay_update_same_value_should_be_consistent() {
344        let mut overlay = IndexOverlay::default();
345        let cols: &[&str] = &["name"];
346
347        // Update where old == new: delete removes from added, insert adds back
348        overlay.update(
349            cols,
350            vec![text_val("alice")],
351            vec![text_val("alice")],
352            pk(1),
353        );
354
355        let column_overlay = &overlay.0[&vec!["name"]];
356        // PK should be in added (insert happened after delete)
357        assert!(column_overlay.added[&vec![text_val("alice")]].contains(&pk(1)));
358        // PK should not be in removed (insert cleaned it)
359        assert!(!column_overlay.removed[&vec![text_val("alice")]].contains(&pk(1)));
360    }
361
362    #[test]
363    fn test_index_overlay_multiple_indexes() {
364        let mut overlay = IndexOverlay::default();
365        let name_cols: &[&str] = &["name"];
366        let email_cols: &[&str] = &["email"];
367
368        overlay.insert(name_cols, vec![text_val("alice")], pk(1));
369        overlay.insert(email_cols, vec![text_val("alice@example.com")], pk(1));
370
371        assert!(overlay.0.contains_key(&vec!["name"]));
372        assert!(overlay.0.contains_key(&vec!["email"]));
373
374        let name_overlay = &overlay.0[&vec!["name"]];
375        assert!(name_overlay.added[&vec![text_val("alice")]].contains(&pk(1)));
376
377        let email_overlay = &overlay.0[&vec!["email"]];
378        assert!(email_overlay.added[&vec![text_val("alice@example.com")]].contains(&pk(1)));
379    }
380
381    #[test]
382    fn test_index_overlay_composite_index() {
383        let mut overlay = IndexOverlay::default();
384        let cols: &[&str] = &["first_name", "last_name"];
385
386        overlay.insert(cols, vec![text_val("alice"), text_val("smith")], pk(1));
387        overlay.insert(cols, vec![text_val("alice"), text_val("jones")], pk(2));
388
389        let column_overlay = &overlay.0[&vec!["first_name", "last_name"]];
390        assert!(column_overlay.added[&vec![text_val("alice"), text_val("smith")]].contains(&pk(1)));
391        assert!(column_overlay.added[&vec![text_val("alice"), text_val("jones")]].contains(&pk(2)));
392    }
393
394    #[test]
395    fn test_index_overlay_insert_delete_insert_cycle() {
396        let mut overlay = IndexOverlay::default();
397        let cols: &[&str] = &["name"];
398        let key = vec![text_val("alice")];
399
400        overlay.insert(cols, key.clone(), pk(1));
401        overlay.delete(cols, key.clone(), pk(1));
402        overlay.insert(cols, key.clone(), pk(1));
403
404        let column_overlay = &overlay.0[&vec!["name"]];
405        assert!(column_overlay.added[&key].contains(&pk(1)));
406        // After insert→delete→insert, PK was never in the base index,
407        // so removed should not contain it.
408        assert!(
409            column_overlay
410                .removed
411                .get(&key)
412                .map_or(true, |pks| !pks.contains(&pk(1)))
413        );
414    }
415
416    #[test]
417    fn test_added_pks_in_range_returns_matching_entries() {
418        let mut overlay = IndexOverlay::default();
419        let cols: &[&str] = &["age"];
420        overlay.insert(cols, vec![Value::Uint32(20.into())], pk(1));
421        overlay.insert(cols, vec![Value::Uint32(25.into())], pk(2));
422        overlay.insert(cols, vec![Value::Uint32(30.into())], pk(3));
423        overlay.insert(cols, vec![Value::Uint32(35.into())], pk(4));
424
425        let pks = overlay.added_pks_in_range(
426            cols,
427            Some(&[Value::Uint32(25.into())]),
428            Some(&[Value::Uint32(30.into())]),
429        );
430
431        assert_eq!(pks.len(), 2);
432        assert!(pks.contains(&pk(2)));
433        assert!(pks.contains(&pk(3)));
434    }
435
436    #[test]
437    fn test_removed_pks_in_range_returns_matching_entries() {
438        let mut overlay = IndexOverlay::default();
439        let cols: &[&str] = &["age"];
440        overlay.delete(cols, vec![Value::Uint32(20.into())], pk(1));
441        overlay.delete(cols, vec![Value::Uint32(25.into())], pk(2));
442
443        let pks = overlay.removed_pks_in_range(
444            cols,
445            Some(&[Value::Uint32(20.into())]),
446            Some(&[Value::Uint32(25.into())]),
447        );
448
449        assert_eq!(pks.len(), 2);
450        assert!(pks.contains(&pk(1)));
451        assert!(pks.contains(&pk(2)));
452    }
453
454    #[test]
455    fn test_added_pks_in_range_open_ended() {
456        let mut overlay = IndexOverlay::default();
457        let cols: &[&str] = &["age"];
458        overlay.insert(cols, vec![Value::Uint32(20.into())], pk(1));
459        overlay.insert(cols, vec![Value::Uint32(30.into())], pk(2));
460
461        let pks = overlay.added_pks_in_range(cols, Some(&[Value::Uint32(25.into())]), None);
462        assert_eq!(pks.len(), 1);
463        assert!(pks.contains(&pk(2)));
464
465        let pks = overlay.added_pks_in_range(cols, None, Some(&[Value::Uint32(25.into())]));
466        assert_eq!(pks.len(), 1);
467        assert!(pks.contains(&pk(1)));
468    }
469}