cadmpeg_core/decode/
unique.rs1use super::{u64_from_index, DecodeContext, ScopedReservation};
5use crate::CodecError;
6use std::collections::HashMap;
7use std::hash::Hash;
8
9impl DecodeContext<'_> {
10 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 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 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}