Skip to main content

rlx_flow/
extension.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4//! Flow-level extension wiring — applies [`rlx_ir::hir_extension`] after build.
5
6use rlx_ir::hir::HirModule;
7use rlx_ir::{apply_hir_extensions, apply_hir_extensions_named};
8
9/// Names of HIR extensions to apply when building this flow (empty = all registered).
10#[derive(Debug, Clone, Default, PartialEq, Eq)]
11pub struct FlowExtensionPlan {
12    pub enabled: Vec<String>,
13    pub apply_all: bool,
14}
15
16impl FlowExtensionPlan {
17    pub fn all() -> Self {
18        Self {
19            enabled: Vec::new(),
20            apply_all: true,
21        }
22    }
23
24    pub fn only(names: impl IntoIterator<Item = impl Into<String>>) -> Self {
25        Self {
26            enabled: names.into_iter().map(Into::into).collect(),
27            apply_all: false,
28        }
29    }
30
31    pub fn apply(&self, hir: &mut HirModule) {
32        if self.apply_all {
33            apply_hir_extensions(hir);
34        } else {
35            let refs: Vec<&str> = self.enabled.iter().map(String::as_str).collect();
36            apply_hir_extensions_named(hir, &refs);
37        }
38    }
39}