rlx_flow/blocks/
repeat.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3
4use anyhow::Result;
5
6use crate::context::FlowCtx;
7use crate::stage::FlowStage;
8use crate::value::FlowValue;
9pub struct RepeatStage {
10    pub count: usize,
11    pub stage_for_index: std::sync::Arc<dyn Fn(usize) -> FlowStage + Send + Sync>,
12}
13
14impl std::fmt::Debug for RepeatStage {
15    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
16        f.debug_struct("RepeatStage")
17            .field("count", &self.count)
18            .finish_non_exhaustive()
19    }
20}
21
22impl RepeatStage {
23    pub fn new(
24        count: usize,
25        stage_for_index: impl Fn(usize) -> FlowStage + Send + Sync + 'static,
26    ) -> Self {
27        Self {
28            count,
29            stage_for_index: std::sync::Arc::new(stage_for_index),
30        }
31    }
32}
33
34impl Clone for RepeatStage {
35    fn clone(&self) -> Self {
36        Self {
37            count: self.count,
38            stage_for_index: std::sync::Arc::clone(&self.stage_for_index),
39        }
40    }
41}
42
43impl RepeatStage {
44    pub fn emit(
45        &self,
46        ctx: &mut FlowCtx<'_>,
47        mut input: Option<FlowValue>,
48    ) -> Result<Option<FlowValue>> {
49        for i in 0..self.count {
50            let stage = (self.stage_for_index)(i);
51            input = stage.emit(ctx, input)?;
52        }
53        Ok(input)
54    }
55}
rlx_flow/blocks/repeat.rs

rlx_flow/blocks/
repeat.rs
