Skip to main content

cadmpeg_core/decode/
unique.rs

1// SPDX-License-Identifier: Apache-2.0
2//! One-table indexes that exclude every repeated key.
3
4use super::{u64_from_index, DecodeContext, ScopedReservation};
5use crate::CodecError;
6use std::collections::HashMap;
7use std::hash::Hash;
8
9impl DecodeContext<'_> {
10    /// Builds a scoped table. Repeated keys become tombstones until the final
11    /// in-place removal, so a third occurrence cannot restore a duplicate.
12    /// Surviving values are Some; no second table is allocated.
13    pub fn unique_index<K: Eq + Hash, V>(
14        &self,
15        entries: impl IntoIterator<Item = (K, V)>,
16        key_work: impl Fn(&K) -> Result<u64, CodecError>,
17        operation: &'static str,
18    ) -> Result<(HashMap<K, Option<V>>, ScopedReservation<'_>), CodecError> {
19        let mut table = HashMap::<K, Option<V>>::new();
20        let mut storage = self.reserve_scoped(0, operation)?;
21        for (key, value) in entries {
22            self.charge_work(1, operation)?;
23            self.charge_work(key_work(&key)?, operation)?;
24            if let Some(previous) = table.get_mut(&key) {
25                *previous = None;
26                continue;
27            }
28            self.charge_collection_items(1, operation)?;
29            if table.len() == table.capacity() {
30                self.charge_work(u64_from_index(table.capacity()), operation)?;
31                for stored in table.keys() {
32                    self.charge_work(key_work(stored)?, operation)?;
33                }
34                // A hash table reserves at most four buckets per requested
35                // entry, including load-factor rounding. Each bucket holds
36                // its pair and control bytes. The extra 32 bytes cover
37                // alignment and control-group padding. Keep the old and new
38                // tables admitted together while fallible growth rehashes.
39                let required = table
40                    .len()
41                    .checked_add(1)
42                    .and_then(|count| count.checked_mul(4))
43                    .and_then(|count| {
44                        count.checked_mul(
45                            std::mem::size_of::<(K, Option<V>)>()
46                                .max(1)
47                                .checked_add(32)?,
48                        )
49                    })
50                    .ok_or_else(|| {
51                        CodecError::from(self.budget.scoped_size_overflow_limit(operation))
52                    })?;
53                storage.grow(u64_from_index(required))?;
54                table.try_reserve(1).map_err(|_| {
55                    self.budget
56                        .scoped_allocation_failed(u64_from_index(required), operation)
57                })?;
58            }
59            self.charge_work(key_work(&key)?, operation)?;
60            // discarded-value: the key was absent before the admitted insertion.
61            let _ = table.insert(key, Some(value));
62        }
63        self.charge_work(u64_from_index(table.capacity()), operation)?;
64        table.retain(|_, value| value.is_some());
65        Ok((table, storage))
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use crate::decode::{DecodeArena, DecodeContext, DecodePolicy, ResourceDimension};
72    use crate::CodecError;
73
74    #[test]
75    fn unique_index_removes_all_duplicate_occurrences_in_one_table() {
76        let arena = DecodeArena::new();
77        let mut policy = DecodePolicy::service();
78        policy.limits.max_collection_items = 2;
79        policy.limits.max_materialized_bytes =
80            super::u64_from_index(4 * (std::mem::size_of::<(u32, Option<i32>)>() + 32));
81        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("context");
82        let (table, storage) = ctx
83            .unique_index(
84                [(1_u32, 2), (1, 3), (1, 4), (2, 5)],
85                |_| Ok(4),
86                "unique test",
87            )
88            .expect("two slots");
89        assert_eq!(table.len(), 1);
90        assert_eq!(table.get(&2), Some(&Some(5)));
91        assert!(!table.contains_key(&1));
92        drop((table, storage));
93        ctx.reserve_scoped(policy.limits.max_materialized_bytes, "released table")
94            .expect("storage released");
95    }
96
97    #[test]
98    fn unique_index_one_slot_keeps_one_record() {
99        let arena = DecodeArena::new();
100        let mut policy = DecodePolicy::service();
101        policy.limits.max_collection_items = 1;
102        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("context");
103        let (table, _storage) = ctx
104            .unique_index([(7_u32, 7)], |_| Ok(4), "unique test")
105            .expect("one table slot");
106        assert_eq!(table.get(&7), Some(&Some(7)));
107    }
108
109    #[test]
110    fn unique_index_propagates_slot_workspace_and_key_work_refusals() {
111        for dimension in [
112            ResourceDimension::CollectionItems,
113            ResourceDimension::MaterializedBytes,
114            ResourceDimension::WorkUnits,
115        ] {
116            let arena = DecodeArena::new();
117            let mut policy = DecodePolicy::service();
118            match dimension {
119                ResourceDimension::CollectionItems => policy.limits.max_collection_items = 0,
120                ResourceDimension::MaterializedBytes => policy.limits.max_materialized_bytes = 0,
121                ResourceDimension::WorkUnits => policy.limits.max_work_units = 3,
122                _ => panic!("test dimension"),
123            }
124            let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &policy).expect("context");
125            assert!(
126                matches!(ctx.unique_index([("long key", 1)], |key| Ok(super::u64_from_index(key.len())), "unique test"),
127                Err(CodecError::ResourceLimit(limit)) if limit.dimension == dimension)
128            );
129        }
130    }
131}