cadmpeg_core/decode/
sort.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Work admission for in-place unstable sorting.
3
4use std::cmp::Ordering;
5
6use super::{u64_from_index, DecodeContext};
7use crate::CodecError;
8
9pub(super) const INLINE_SORT_LIMIT: usize = 20;
10
11/// Admit a linear check for a large run before allocating a sort permutation.
12pub(super) fn admit_ordered_run<T>(
13    ctx: &DecodeContext<'_>,
14    values: &[T],
15    compare: &mut impl FnMut(&T, &T) -> Ordering,
16    key_bytes: &impl Fn(&T) -> usize,
17    operation: &'static str,
18) -> Result<bool, CodecError> {
19    // Small unstable runs keep their fixed admission estimate, including
20    // runs collected from hash tables.
21    if values.len() <= INLINE_SORT_LIMIT {
22        return Ok(false);
23    }
24    for pair in values.windows(2) {
25        let work = u64_from_index(key_bytes(&pair[0]))
26            .checked_add(u64_from_index(key_bytes(&pair[1])))
27            .and_then(|work| {
28                work.checked_add(u64_from_index(std::mem::size_of::<T>()).checked_mul(2)?)
29            })
30            .and_then(|work| work.checked_add(1))
31            .ok_or_else(|| ctx.refuse_codec_limit(operation, u64::MAX - 1, u64::MAX))?;
32        ctx.charge_work(work, operation)?;
33        if compare(&pair[0], &pair[1]).is_gt() {
34            return Ok(false);
35        }
36    }
37    Ok(true)
38}
39
40impl DecodeContext<'_> {
41    /// Sorts admitted values in place without allocating scratch.
42    ///
43    /// `key_bytes` states the external bytes a comparison can read from each value.
44    pub fn sort_unstable_by<T>(
45        &self,
46        values: &mut [T],
47        mut compare: impl FnMut(&T, &T) -> Ordering,
48        key_bytes: impl Fn(&T) -> usize,
49        operation: &'static str,
50    ) -> Result<(), CodecError> {
51        let count = u64_from_index(values.len());
52        self.charge_work(count, operation)?;
53        if admit_ordered_run(self, values, &mut compare, &key_bytes, operation)? {
54            return Ok(());
55        }
56        let bytes = values
57            .iter()
58            .try_fold(0u64, |bytes, value| {
59                bytes.checked_add(u64_from_index(key_bytes(value)))
60            })
61            .ok_or_else(|| self.refuse_codec_limit(operation, u64::MAX - 1, u64::MAX))?;
62        let levels = u64::from(u64::BITS - count.leading_zeros()) + 1;
63        let work = count
64            .checked_mul(u64_from_index(std::mem::size_of::<T>()))
65            .and_then(|storage| storage.checked_add(bytes.checked_mul(2)?))
66            .and_then(|work| work.checked_mul(levels))
67            .and_then(|work| work.checked_mul(8))
68            .ok_or_else(|| self.refuse_codec_limit(operation, u64::MAX - 1, u64::MAX))?;
69        self.charge_work(work, operation)?;
70        values.sort_unstable_by(compare);
71        Ok(())
72    }
73}
74
75#[cfg(test)]
76mod tests {
77    use super::DecodeContext;
78    use crate::decode::{DecodeArena, DecodePolicy, ResourceDimension};
79    use crate::CodecError;
80
81    #[test]
82    fn charged_unstable_sort_work_refusal_preserves_input() {
83        let mut values = [3u64, 1, 2];
84        let expected = values;
85        let arena = DecodeArena::new();
86        let mut policy = DecodePolicy::service();
87        // Three elements, eight bytes each, three levels, eight work units per byte.
88        policy.limits.max_work_units = 3 + 3 * 8 * 3 * 8 - 1;
89        let (ctx, _) =
90            DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root is admitted");
91        assert!(matches!(
92            ctx.sort_unstable_by(&mut values, Ord::cmp, |_| 0, "test unstable sort"),
93            Err(CodecError::ResourceLimit(limit))
94                if limit.dimension == ResourceDimension::WorkUnits
95        ));
96        assert_eq!(values, expected);
97    }
98
99    #[test]
100    fn charged_unstable_sort_orders_values_without_scratch() {
101        let mut values = [3u64, 1, 2, 1];
102        let arena = DecodeArena::new();
103        let mut policy = DecodePolicy::service();
104        policy.limits.max_materialized_bytes = 0;
105        let (ctx, _) =
106            DecodeContext::from_root_bytes(&[], &arena, &policy).expect("empty root is admitted");
107        ctx.sort_unstable_by(&mut values, Ord::cmp, |_| 0, "test unstable sort")
108            .expect("sort is admitted");
109        assert_eq!(values, [1, 1, 2, 3]);
110    }
111
112    #[test]
113    fn charged_unstable_sort_key_byte_overflow_preserves_input() {
114        let mut values = [2u8, 1];
115        let expected = values;
116        let arena = DecodeArena::new();
117        let (ctx, _) = DecodeContext::from_root_bytes(&[], &arena, &DecodePolicy::service())
118            .expect("empty root is admitted");
119        assert!(matches!(
120            ctx.sort_unstable_by(&mut values, Ord::cmp, |_| usize::MAX, "test unstable sort"),
121            Err(CodecError::ResourceLimit(_))
122        ));
123        assert_eq!(values, expected);
124    }
125}
cadmpeg_core/decode/sort.rs

cadmpeg_core/decode/
sort.rs
