rlx_flow/blocks/
custom.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4use std::fmt;
5use std::sync::Arc;
6
7use anyhow::Result;
8
9use crate::context::FlowCtx;
10use crate::escape::Emit;
11use crate::value::FlowValue;
12type CustomFn =
13    Arc<dyn Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>> + Send + Sync>;
14
15/// User-defined stage — tier-2 escape hatch for novel subgraphs.
16#[derive(Clone)]
17pub struct CustomStage {
18    pub name: Option<String>,
19    f: CustomFn,
20}
21
22impl fmt::Debug for CustomStage {
23    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
24        f.debug_struct("CustomStage")
25            .field("name", &self.name)
26            .finish_non_exhaustive()
27    }
28}
29
30impl CustomStage {
31    pub fn new<F>(f: F) -> Self
32    where
33        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>>
34            + Send
35            + Sync
36            + 'static,
37    {
38        Self {
39            name: None,
40            f: Arc::new(f),
41        }
42    }
43
44    pub fn named<F>(name: impl Into<String>, f: F) -> Self
45    where
46        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> Result<Option<FlowValue>>
47            + Send
48            + Sync
49            + 'static,
50    {
51        Self {
52            name: Some(name.into()),
53            f: Arc::new(f),
54        }
55    }
56
57    pub fn emit(
58        &self,
59        ctx: &mut FlowCtx<'_>,
60        input: Option<FlowValue>,
61    ) -> Result<Option<FlowValue>> {
62        let mut emit = Emit::from_ctx(ctx);
63        (self.f)(&mut emit, input)
64    }
65}
rlx_flow/blocks/custom.rs

rlx_flow/blocks/
custom.rs
