Skip to main content

zrx_stream/stream/operator/
lift.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//! Lift operator.
27
28use ahash::HashMap;
29
30use zrx_scheduler::action::descriptor::Property;
31use zrx_scheduler::action::output::IntoOutputs;
32use zrx_scheduler::action::{Descriptor, Report};
33use zrx_scheduler::effect::Item;
34use zrx_scheduler::{Id, Value};
35use zrx_store::behavior::StoreDelta;
36use zrx_store::StoreMutRef;
37
38use crate::stream::function::LiftFn;
39use crate::stream::value::Delta;
40use crate::stream::Stream;
41
42use super::{Operator, OperatorExt};
43
44// ----------------------------------------------------------------------------
45// Structs
46// ----------------------------------------------------------------------------
47
48/// Lift operator.
49struct Lift<F, I, U> {
50    /// Operator function.
51    function: F,
52    /// Store of change sets.
53    store: HashMap<I, HashMap<I, U>>,
54}
55
56// ----------------------------------------------------------------------------
57// Implementations
58// ----------------------------------------------------------------------------
59
60impl<I, T> Stream<I, T>
61where
62    I: Id,
63    T: Value,
64{
65    pub fn lift<F, U>(&self, f: F) -> Stream<I, Delta<I, U>>
66    where
67        F: LiftFn<I, T, U>,
68        U: Value + Clone + Eq,
69    {
70        self.with_operator(Lift {
71            function: f,
72            store: HashMap::default(),
73        })
74    }
75}
76
77// ----------------------------------------------------------------------------
78// Trait implementations
79// ----------------------------------------------------------------------------
80
81impl<I, T, F, U> Operator<I, T> for Lift<F, I, U>
82where
83    I: Id,
84    T: Value,
85    F: LiftFn<I, T, U>,
86    U: Value + Clone + Eq,
87{
88    type Item<'a> = Item<&'a I, Option<&'a T>>;
89
90    /// Handles the given item.
91    ///
92    /// Lifting is an essential operation in stream processing, as it allows an
93    /// input item to be transformed into multiple related output items with a
94    /// user-defined function. This operator always returns a delta of items,
95    /// ensuring that the related items can be processed differentially.
96    #[cfg_attr(
97        feature = "tracing",
98        tracing::instrument(level = "debug", skip_all, fields(id = %item.id))
99    )]
100    fn handle(&mut self, item: Self::Item<'_>) -> impl IntoOutputs<I> {
101        if let Some(data) = item.data {
102            // When new data arrives, we apply the operator function to generate
103            // related items. We then compute the delta between the previous set
104            // of related items and the returned set, ensuring that only changes
105            // are propagated downstream.
106            self.function.execute(item.id, data).map(|report| {
107                report.map(|data| {
108                    let store = self.store.get_or_insert_default(item.id);
109                    let delta = store
110                        .changes(data.into_iter().map(Item::into_parts))
111                        .map(|(id, data)| Item::new(id, data))
112                        .collect();
113
114                    // Return delta of items
115                    Item::new(item.id.clone(), Some(delta))
116                })
117            })
118        } else {
119            // If the incoming item has no data, interpret this as a deletion,
120            // removing any previously stored related items. By emitting a delta
121            // of items instead of a deletion, we can ensure that all downstream
122            // operators can update their internal state accordingly.
123            let store = self.store.remove(item.id);
124            let delta = store.map_or_else(Delta::default, |inner| {
125                inner.into_keys().map(|id| Item::new(id, None)).collect()
126            });
127
128            // Return delta of items
129            Ok(Report::new(Item::new(item.id.clone(), Some(delta))))
130        }
131    }
132
133    /// Returns the descriptor.
134    #[inline]
135    fn descriptor(&self) -> Descriptor {
136        Descriptor::builder() // fmt
137            .property(Property::Stable)
138            .build()
139    }
140}