wave_emu/
control_flow.rs

1// Copyright 2026 Ojima Abraham
2// SPDX-License-Identifier: Apache-2.0
3
4//! Structured control flow with divergence support. Maintains a stack of active
5//!
6//! masks for nested if/else/endif and loop/break/endloop constructs. Divergence
7//! occurs when threads disagree on predicates, executing both paths sequentially.
8
9use crate::EmulatorError;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum ControlFrameKind {
13    If,
14    Loop,
15}
16
17#[derive(Debug, Clone)]
18pub struct ControlFrame {
19    pub kind: ControlFrameKind,
20    pub entry_mask: u64,
21    pub then_mask: u64,
22    pub else_mask: u64,
23    pub break_mask: u64,
24    pub loop_start_pc: u32,
25    pub in_else: bool,
26}
27
28impl ControlFrame {
29    pub fn new_if(entry_mask: u64, then_mask: u64, else_mask: u64) -> Self {
30        Self {
31            kind: ControlFrameKind::If,
32            entry_mask,
33            then_mask,
34            else_mask,
35            break_mask: 0,
36            loop_start_pc: 0,
37            in_else: false,
38        }
39    }
40
41    pub fn new_loop(entry_mask: u64, loop_start_pc: u32) -> Self {
42        Self {
43            kind: ControlFrameKind::Loop,
44            entry_mask,
45            then_mask: entry_mask,
46            else_mask: 0,
47            break_mask: 0,
48            loop_start_pc,
49            in_else: false,
50        }
51    }
52}
53
54#[derive(Debug)]
55pub struct ControlFlowStack {
56    frames: Vec<ControlFrame>,
57    max_depth: usize,
58}
59
60impl ControlFlowStack {
61    pub fn new(max_depth: usize) -> Self {
62        Self {
63            frames: Vec::with_capacity(max_depth),
64            max_depth,
65        }
66    }
67
68    pub fn push(&mut self, frame: ControlFrame) -> Result<(), EmulatorError> {
69        if self.frames.len() >= self.max_depth {
70            return Err(EmulatorError::StackOverflow {
71                kind: "control flow".into(),
72            });
73        }
74        self.frames.push(frame);
75        Ok(())
76    }
77
78    pub fn pop(&mut self) -> Option<ControlFrame> {
79        self.frames.pop()
80    }
81
82    pub fn top(&self) -> Option<&ControlFrame> {
83        self.frames.last()
84    }
85
86    pub fn top_mut(&mut self) -> Option<&mut ControlFrame> {
87        self.frames.last_mut()
88    }
89
90    pub fn depth(&self) -> usize {
91        self.frames.len()
92    }
93
94    pub fn is_empty(&self) -> bool {
95        self.frames.is_empty()
96    }
97
98    pub fn find_loop(&self) -> Option<&ControlFrame> {
99        self.frames
100            .iter()
101            .rev()
102            .find(|f| f.kind == ControlFrameKind::Loop)
103    }
104
105    pub fn find_loop_mut(&mut self) -> Option<&mut ControlFrame> {
106        self.frames
107            .iter_mut()
108            .rev()
109            .find(|f| f.kind == ControlFrameKind::Loop)
110    }
111}
112
113#[derive(Debug)]
114pub struct ControlFlowManager {
115    stack: ControlFlowStack,
116}
117
118impl ControlFlowManager {
119    pub fn new() -> Self {
120        Self {
121            stack: ControlFlowStack::new(32),
122        }
123    }
124
125    pub fn handle_if(
126        &mut self,
127        active_mask: u64,
128        predicate_mask: u64,
129    ) -> Result<(u64, Option<u32>), EmulatorError> {
130        let then_mask = active_mask & predicate_mask;
131        let else_mask = active_mask & !predicate_mask;
132
133        let frame = ControlFrame::new_if(active_mask, then_mask, else_mask);
134        self.stack.push(frame)?;
135
136        Ok((then_mask, None))
137    }
138
139    pub fn handle_else(&mut self, active_mask: u64) -> Result<(u64, Option<u32>), EmulatorError> {
140        let frame = self
141            .stack
142            .top_mut()
143            .ok_or_else(|| EmulatorError::ControlFlowError {
144                message: "else without matching if".into(),
145            })?;
146
147        if frame.kind != ControlFrameKind::If {
148            return Err(EmulatorError::ControlFlowError {
149                message: "else without matching if".into(),
150            });
151        }
152
153        frame.in_else = true;
154        let _ = active_mask;
155        Ok((frame.else_mask, None))
156    }
157
158    pub fn handle_endif(&mut self) -> Result<u64, EmulatorError> {
159        let frame = self
160            .stack
161            .pop()
162            .ok_or_else(|| EmulatorError::ControlFlowError {
163                message: "endif without matching if".into(),
164            })?;
165
166        if frame.kind != ControlFrameKind::If {
167            self.stack.push(frame)?;
168            return Err(EmulatorError::ControlFlowError {
169                message: "endif without matching if".into(),
170            });
171        }
172
173        Ok(frame.entry_mask)
174    }
175
176    pub fn handle_loop(&mut self, active_mask: u64, pc: u32) -> Result<u64, EmulatorError> {
177        let frame = ControlFrame::new_loop(active_mask, pc);
178        self.stack.push(frame)?;
179        Ok(active_mask)
180    }
181
182    pub fn handle_break(
183        &mut self,
184        active_mask: u64,
185        predicate_mask: u64,
186    ) -> Result<(u64, Option<u32>), EmulatorError> {
187        let frame = self
188            .stack
189            .find_loop_mut()
190            .ok_or_else(|| EmulatorError::ControlFlowError {
191                message: "break outside of loop".into(),
192            })?;
193
194        let breaking_threads = active_mask & predicate_mask;
195        frame.break_mask |= breaking_threads;
196
197        let new_mask = active_mask & !breaking_threads;
198
199        if new_mask == 0 {
200            Ok((new_mask, Some(0)))
201        } else {
202            Ok((new_mask, None))
203        }
204    }
205
206    pub fn handle_continue(
207        &mut self,
208        active_mask: u64,
209        predicate_mask: u64,
210    ) -> Result<(u64, Option<u32>), EmulatorError> {
211        let frame = self
212            .stack
213            .find_loop()
214            .ok_or_else(|| EmulatorError::ControlFlowError {
215                message: "continue outside of loop".into(),
216            })?;
217
218        let continuing_threads = active_mask & predicate_mask;
219        let new_mask = active_mask & !continuing_threads;
220
221        if new_mask == 0 || continuing_threads != 0 {
222            Ok((new_mask, Some(frame.loop_start_pc)))
223        } else {
224            Ok((new_mask, None))
225        }
226    }
227
228    pub fn handle_endloop(
229        &mut self,
230        active_mask: u64,
231    ) -> Result<(u64, Option<u32>), EmulatorError> {
232        let _ = active_mask;
233
234        let frame = self
235            .stack
236            .top_mut()
237            .ok_or_else(|| EmulatorError::ControlFlowError {
238                message: "endloop without matching loop".into(),
239            })?;
240
241        if frame.kind != ControlFrameKind::Loop {
242            return Err(EmulatorError::ControlFlowError {
243                message: "endloop without matching loop".into(),
244            });
245        }
246
247        let remaining_mask = frame.entry_mask & !frame.break_mask;
248
249        if remaining_mask != 0 {
250            let loop_start = frame.loop_start_pc;
251            Ok((remaining_mask, Some(loop_start)))
252        } else {
253            let frame = self
254                .stack
255                .pop()
256                .ok_or_else(|| EmulatorError::ControlFlowError {
257                    message: "endloop without matching loop".into(),
258                })?;
259            Ok((frame.entry_mask, None))
260        }
261    }
262
263    pub fn pop_loop(&mut self) -> Result<u64, EmulatorError> {
264        let frame = self
265            .stack
266            .pop()
267            .ok_or_else(|| EmulatorError::ControlFlowError {
268                message: "endloop without matching loop".into(),
269            })?;
270
271        if frame.kind != ControlFrameKind::Loop {
272            self.stack.push(frame)?;
273            return Err(EmulatorError::ControlFlowError {
274                message: "endloop without matching loop".into(),
275            });
276        }
277
278        Ok(frame.entry_mask & !frame.break_mask)
279    }
280
281    pub fn depth(&self) -> usize {
282        self.stack.depth()
283    }
284
285    pub fn is_empty(&self) -> bool {
286        self.stack.is_empty()
287    }
288
289    pub fn current_loop_start(&self) -> Option<u32> {
290        self.stack.find_loop().map(|f| f.loop_start_pc)
291    }
292}
293
294impl Default for ControlFlowManager {
295    fn default() -> Self {
296        Self::new()
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::*;
303
304    #[test]
305    fn test_cf_simple_if() {
306        let mut cf = ControlFlowManager::new();
307
308        let (mask, _) = cf.handle_if(0xFFFF_FFFF, 0x0000_FFFF).unwrap();
309        assert_eq!(mask, 0x0000_FFFF);
310
311        let (mask, _) = cf.handle_else(mask).unwrap();
312        assert_eq!(mask, 0xFFFF_0000);
313
314        let mask = cf.handle_endif().unwrap();
315        assert_eq!(mask, 0xFFFF_FFFF);
316    }
317
318    #[test]
319    fn test_cf_if_no_else_needed() {
320        let mut cf = ControlFlowManager::new();
321
322        let (mask, _) = cf.handle_if(0xFFFF_FFFF, 0xFFFF_FFFF).unwrap();
323        assert_eq!(mask, 0xFFFF_FFFF);
324
325        let mask = cf.handle_endif().unwrap();
326        assert_eq!(mask, 0xFFFF_FFFF);
327    }
328
329    #[test]
330    fn test_cf_if_all_take_else() {
331        let mut cf = ControlFlowManager::new();
332
333        let (then_mask, _) = cf.handle_if(0xFFFF_FFFF, 0).unwrap();
334        assert_eq!(then_mask, 0);
335
336        let (else_mask, _) = cf.handle_else(then_mask).unwrap();
337        assert_eq!(else_mask, 0xFFFF_FFFF);
338
339        let mask = cf.handle_endif().unwrap();
340        assert_eq!(mask, 0xFFFF_FFFF);
341    }
342
343    #[test]
344    fn test_cf_simple_loop() {
345        let mut cf = ControlFlowManager::new();
346
347        let mask = cf.handle_loop(0xFFFF_FFFF, 0x100).unwrap();
348        assert_eq!(mask, 0xFFFF_FFFF);
349        assert_eq!(cf.current_loop_start(), Some(0x100));
350
351        let (mask, jump) = cf.handle_break(mask, 0x0000_FFFF).unwrap();
352        assert_eq!(mask, 0xFFFF_0000);
353        assert!(jump.is_none());
354    }
355
356    #[test]
357    fn test_cf_loop_all_break() {
358        let mut cf = ControlFlowManager::new();
359
360        cf.handle_loop(0xFFFF_FFFF, 0x100).unwrap();
361        let (mask, jump) = cf.handle_break(0xFFFF_FFFF, 0xFFFF_FFFF).unwrap();
362        assert_eq!(mask, 0);
363        assert!(jump.is_some());
364    }
365
366    #[test]
367    fn test_cf_nested_if() {
368        let mut cf = ControlFlowManager::new();
369
370        let (mask1, _) = cf.handle_if(0xFFFF_FFFF, 0x00FF_00FF).unwrap();
371        assert_eq!(mask1, 0x00FF_00FF);
372
373        let (mask2, _) = cf.handle_if(mask1, 0x0000_00FF).unwrap();
374        assert_eq!(mask2, 0x0000_00FF);
375
376        let mask2 = cf.handle_endif().unwrap();
377        assert_eq!(mask2, 0x00FF_00FF);
378
379        let mask1 = cf.handle_endif().unwrap();
380        assert_eq!(mask1, 0xFFFF_FFFF);
381    }
382
383    #[test]
384    fn test_cf_stack_overflow() {
385        let mut cf = ControlFlowManager::new();
386
387        for i in 0..32 {
388            cf.handle_if(0xFFFF_FFFF, 0x0000_FFFF).unwrap();
389            assert_eq!(cf.depth(), i + 1);
390        }
391
392        assert!(cf.handle_if(0xFFFF_FFFF, 0x0000_FFFF).is_err());
393    }
394
395    #[test]
396    fn test_cf_mismatched_endif() {
397        let mut cf = ControlFlowManager::new();
398        assert!(cf.handle_endif().is_err());
399    }
400
401    #[test]
402    fn test_cf_break_outside_loop() {
403        let mut cf = ControlFlowManager::new();
404        assert!(cf.handle_break(0xFFFF_FFFF, 0xFFFF_FFFF).is_err());
405    }
406}
wave_emu/control_flow.rs

wave_emu/
control_flow.rs
