rlx_flow/blocks/
custom.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16use std::fmt;
17use std::sync::Arc;
18
19use anyhow::Result;
20
21use crate::context::FlowCtx;
22use crate::escape::Emit;
23use crate::value::FlowValue;
24type CustomFn =
25    Arc<dyn Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>> + Send + Sync>;
26
27/// User-defined stage — tier-2 escape hatch for novel subgraphs.
28#[derive(Clone)]
29pub struct CustomStage {
30    pub name: Option<String>,
31    f: CustomFn,
32}
33
34impl fmt::Debug for CustomStage {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        f.debug_struct("CustomStage")
37            .field("name", &self.name)
38            .finish_non_exhaustive()
39    }
40}
41
42impl CustomStage {
43    pub fn new<F>(f: F) -> Self
44    where
45        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>>
46            + Send
47            + Sync
48            + 'static,
49    {
50        Self {
51            name: None,
52            f: Arc::new(f),
53        }
54    }
55
56    pub fn named<F>(name: impl Into<String>, f: F) -> Self
57    where
58        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>>
59            + Send
60            + Sync
61            + 'static,
62    {
63        Self {
64            name: Some(name.into()),
65            f: Arc::new(f),
66        }
67    }
68
69    pub fn emit(
70        &self,
71        ctx: &mut FlowCtx<'_>,
72        input: Option<FlowValue>,
73    ) -> Result<Option<FlowValue>> {
74        let mut emit = Emit::from_ctx(ctx);
75        (self.f)(&mut emit, input)
76    }
77}
rlx_flow/blocks/custom.rs

rlx_flow/blocks/
custom.rs
