zrx_stream/stream/operator/
sort.rs

1// Copyright (c) 2025-2026 Zensical and contributors
2
3// SPDX-License-Identifier: MIT
4// All contributions are certified under the DCO
5
6// Permission is hereby granted, free of charge, to any person obtaining a copy
7// of this software and associated documentation files (the "Software"), to
8// deal in the Software without restriction, including without limitation the
9// rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
10// sell copies of the Software, and to permit persons to whom the Software is
11// furnished to do so, subject to the following conditions:
12
13// The above copyright notice and this permission notice shall be included in
14// all copies or substantial portions of the Software.
15
16// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
17// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
18// FITNESS FOR A PARTICULAR PURPOSE AND NON-INFRINGEMENT. IN NO EVENT SHALL THE
19// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
20// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
21// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
22// IN THE SOFTWARE.
23
24// ----------------------------------------------------------------------------
25
26//! Sort operator.
27
28use std::ops::Range;
29
30use zrx_scheduler::action::descriptor::Property;
31use zrx_scheduler::action::output::IntoOutputs;
32use zrx_scheduler::action::Descriptor;
33use zrx_scheduler::effect::Item;
34use zrx_scheduler::{Id, Value};
35use zrx_store::decorator::Indexed;
36use zrx_store::Store;
37
38use crate::stream::value::Position;
39use crate::stream::Stream;
40
41use super::{Operator, OperatorExt};
42
43// ----------------------------------------------------------------------------
44// Structs
45// ----------------------------------------------------------------------------
46
47/// Sort operator.
48struct Sort<I, T>
49where
50    I: Id,
51{
52    /// Store of items.
53    store: Indexed<I, T>,
54}
55
56// ----------------------------------------------------------------------------
57// Implementations
58// ----------------------------------------------------------------------------
59
60impl<I, T> Stream<I, T>
61where
62    I: Id,
63    T: Value + Clone + Ord,
64{
65    pub fn sort(&self) -> Stream<I, Position<T>> {
66        self.with_operator(Sort { store: Indexed::default() })
67    }
68}
69
70// ----------------------------------------------------------------------------
71// Trait implementations
72// ----------------------------------------------------------------------------
73
74impl<I, T> Operator<I, T> for Sort<I, T>
75where
76    I: Id,
77    T: Value + Clone + Ord,
78{
79    type Item<'a> = Item<&'a I, Option<&'a T>>;
80
81    /// Handles the given item.
82    ///
83    /// Sorting a stream involves maintaining an internal store of items, and
84    /// updating the positions of items whenever an item is inserted, updated,
85    /// or removed. The returned items are annotated with their new positions,
86    /// which reflect their order in the sorted stream.
87    #[cfg_attr(
88        feature = "tracing",
89        tracing::instrument(level = "debug", skip_all, fields(id = %item.id))
90    )]
91    fn handle(&mut self, item: Self::Item<'_>) -> impl IntoOutputs<I> {
92        let len = self.store.len();
93
94        // After determining the current length of the store, which we need to
95        // discern insertions from in-place updates, we either insert into or
96        // remove the item from the store, which gives us the affected range of
97        // indices. When the range is `None`, it means that nothing changed.
98        match item.data {
99            Some(data) => self.store.insert_if_changed(item.id, data),
100            None => self.store.remove(item.id).map(|n| n..n),
101        }
102        // If nothing changed, we can return early. Otherwise, if the number of
103        // items in the store changed, we must update the positions of all items
104        // that come after the affected range. In this case, we need to extend
105        // the range to cover the end of the store.
106        .map(|Range { start, mut end }| {
107            if len != self.store.len() {
108                end = self.store.len();
109            }
110
111            // If an item was deleted from the store, we need to include the
112            // deletion with all items that changed their positions
113            let mut items = Vec::with_capacity(end - start);
114            if len > self.store.len() {
115                items.push(Item::new(item.id.clone(), None));
116            }
117
118            // Now, we can iterate over the affected range and create new items
119            // with updated positions beginning at the start of the range
120            for (index, pair) in self.store.range(start..end).enumerate() {
121                let item = Item::from(pair).into_owned();
122                items.push(
123                    item.map(|data| Some(Position::new(start + index, data))),
124                );
125            }
126
127            // Return items
128            items
129        })
130        .unwrap_or_default()
131    }
132
133    /// Returns the descriptor.
134    #[inline]
135    fn descriptor(&self) -> Descriptor {
136        Descriptor::builder() // fmt
137            .property(Property::Flush)
138            .build()
139    }
140}