Skip to main content

rlx_flow/
plugin.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! Type-erased arch blocks — keep model-specific emission out of the core enum.
5
6use crate::blocks::CustomStage;
7use crate::escape::Emit;
8use crate::stage::FlowStage;
9use crate::value::FlowValue;
10
11/// Named plugin stage (alias over tier-2 custom emission).
12pub struct PluginStage(CustomStage);
13
14impl PluginStage {
15    pub fn new<F>(f: F) -> Self
16    where
17        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> anyhow::Result<Option<FlowValue>>
18            + Send
19            + Sync
20            + 'static,
21    {
22        Self(CustomStage::new(f))
23    }
24
25    pub fn named<F>(name: impl Into<String>, f: F) -> Self
26    where
27        F: Fn(&mut Emit<'_>, Option<FlowValue>) -> anyhow::Result<Option<FlowValue>>
28            + Send
29            + Sync
30            + 'static,
31    {
32        Self(CustomStage::named(name, f))
33    }
34
35    pub(crate) fn into_stage(self) -> FlowStage {
36        FlowStage::Custom(self.0)
37    }
38}
39
40pub fn plugin<F>(f: F) -> FlowStage
41where
42    F: Fn(&mut Emit<'_>, Option<FlowValue>) -> anyhow::Result<Option<FlowValue>>
43        + Send
44        + Sync
45        + 'static,
46{
47    PluginStage::new(f).into_stage()
48}
49
50pub fn plugin_named<F>(name: impl Into<String>, f: F) -> FlowStage
51where
52    F: Fn(&mut Emit<'_>, Option<FlowValue>) -> anyhow::Result<Option<FlowValue>>
53        + Send
54        + Sync
55        + 'static,
56{
57    PluginStage::named(name, f).into_stage()
58}