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}