rlx_flow/blocks/
repeat.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 anyhow::Result;
17
18use crate::context::FlowCtx;
19use crate::stage::FlowStage;
20use crate::value::FlowValue;
21pub struct RepeatStage {
22    pub count: usize,
23    pub stage_for_index: std::sync::Arc<dyn Fn(usize) -> FlowStage + Send + Sync>,
24}
25
26impl std::fmt::Debug for RepeatStage {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        f.debug_struct("RepeatStage")
29            .field("count", &self.count)
30            .finish_non_exhaustive()
31    }
32}
33
34impl RepeatStage {
35    pub fn new(
36        count: usize,
37        stage_for_index: impl Fn(usize) -> FlowStage + Send + Sync + 'static,
38    ) -> Self {
39        Self {
40            count,
41            stage_for_index: std::sync::Arc::new(stage_for_index),
42        }
43    }
44}
45
46impl Clone for RepeatStage {
47    fn clone(&self) -> Self {
48        Self {
49            count: self.count,
50            stage_for_index: std::sync::Arc::clone(&self.stage_for_index),
51        }
52    }
53}
54
55impl RepeatStage {
56    pub fn emit(
57        &self,
58        ctx: &mut FlowCtx<'_>,
59        mut input: Option<FlowValue>,
60    ) -> Result<Option<FlowValue>> {
61        for i in 0..self.count {
62            let stage = (self.stage_for_index)(i);
63            input = stage.emit(ctx, input)?;
64        }
65        Ok(input)
66    }
67}
rlx_flow/blocks/repeat.rs

rlx_flow/blocks/
repeat.rs
