1use essential_types::{convert::bool_from_word, Word};
2
3use crate::{
4    error::{OpSyncResult, RepeatError, RepeatResult, StackError},
5    Stack,
6};
7
8#[cfg(test)]
9mod tests;
10
11#[derive(Debug, Default, PartialEq)]
12pub struct Repeat {
14    stack: Vec<Slot>,
15}
16
17#[derive(Debug, PartialEq)]
18struct Slot {
19    pub counter: Word,
20    pub limit: Direction,
21    pub repeat_index: usize,
22}
23
24#[derive(Debug, PartialEq)]
25enum Direction {
26    Up(Word),
27    Down,
28}
29
30pub(crate) fn repeat(pc: usize, stack: &mut Stack, repeat: &mut Repeat) -> OpSyncResult<()> {
32    let [num_repeats, count_up] = stack.pop2()?;
33    let count_up = bool_from_word(count_up).ok_or(RepeatError::InvalidCountDirection)?;
34    let pc = pc.checked_add(1).ok_or(StackError::IndexOutOfBounds)?;
35    if count_up {
36        repeat.repeat_to(pc, num_repeats)?;
37    } else {
38        repeat.repeat_from(pc, num_repeats)?;
39    }
40    Ok(())
41}
42
43impl Repeat {
44    pub fn new() -> Self {
46        Self::default()
47    }
48
49    pub fn repeat_from(&mut self, location: usize, amount: Word) -> RepeatResult<()> {
52        if self.stack.len() >= super::Stack::SIZE_LIMIT {
53            return Err(RepeatError::Overflow);
54        }
55        self.stack.push(Slot {
56            counter: amount,
57            limit: Direction::Down,
58            repeat_index: location,
59        });
60        Ok(())
61    }
62
63    pub fn repeat_to(&mut self, location: usize, limit: Word) -> RepeatResult<()> {
66        if self.stack.len() >= super::Stack::SIZE_LIMIT {
67            return Err(RepeatError::Overflow);
68        }
69        self.stack.push(Slot {
70            counter: 0,
71            limit: Direction::Up(limit),
72            repeat_index: location,
73        });
74        Ok(())
75    }
76
77    pub fn counter(&self) -> RepeatResult<Word> {
81        self.stack
82            .last()
83            .map(|s| s.counter)
84            .ok_or(RepeatError::NoCounter)
85    }
86
87    pub fn repeat(&mut self) -> RepeatResult<Option<usize>> {
101        let slot = self.stack.last_mut().ok_or(RepeatError::Empty)?;
102        match slot.limit {
103            Direction::Up(limit) => {
104                if slot.counter >= limit.saturating_sub(1) {
105                    self.stack.pop();
106                    Ok(None)
107                } else {
108                    slot.counter += 1;
109                    Ok(Some(slot.repeat_index))
110                }
111            }
112            Direction::Down => {
113                if slot.counter <= 1 {
114                    self.stack.pop();
115                    Ok(None)
116                } else {
117                    slot.counter -= 1;
118                    Ok(Some(slot.repeat_index))
119                }
120            }
121        }
122    }
123}