1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

use essential_types::{convert::bool_from_word, Word};

use crate::{
    error::{RepeatError, RepeatResult, StackError},
    OpResult, Stack,
};

#[cfg(test)]
mod tests;

#[derive(Debug, Default, PartialEq)]
/// A stack of repeat counters.
pub struct Repeat {
    stack: Vec<Slot>,
}

#[derive(Debug, PartialEq)]
struct Slot {
    pub counter: Word,
    pub limit: Direction,
    pub repeat_index: usize,
}

#[derive(Debug, PartialEq)]
enum Direction {
    Up(Word),
    Down,
}

/// `Stack::Repeat` implementation.
pub(crate) fn repeat(pc: usize, stack: &mut Stack, repeat: &mut Repeat) -> OpResult<()> {
    let [num_repeats, count_up] = stack.pop2()?;
    let count_up = bool_from_word(count_up).ok_or(RepeatError::InvalidCountDirection)?;
    let pc = pc.checked_add(1).ok_or(StackError::IndexOutOfBounds)?;
    if count_up {
        repeat.repeat_to(pc, num_repeats)?;
    } else {
        repeat.repeat_from(pc, num_repeats)?;
    }
    Ok(())
}

impl Repeat {
    /// Create a new repeat stack.
    pub fn new() -> Self {
        Self::default()
    }

    /// Add a new repeat location and counter to the stack.
    /// Counts down to 0.
    pub fn repeat_from(&mut self, location: usize, amount: Word) -> RepeatResult<()> {
        if self.stack.len() >= super::Stack::SIZE_LIMIT {
            return Err(RepeatError::Overflow);
        }
        self.stack.push(Slot {
            counter: amount,
            limit: Direction::Down,
            repeat_index: location,
        });
        Ok(())
    }

    /// Add a new repeat location and counter to the stack.
    /// Counts up from 0 to limit - 1.
    pub fn repeat_to(&mut self, location: usize, limit: Word) -> RepeatResult<()> {
        if self.stack.len() >= super::Stack::SIZE_LIMIT {
            return Err(RepeatError::Overflow);
        }
        self.stack.push(Slot {
            counter: 0,
            limit: Direction::Up(limit),
            repeat_index: location,
        });
        Ok(())
    }

    /// Get the current repeat counter.
    ///
    /// Returns an error if the stack is empty.
    pub fn counter(&self) -> RepeatResult<Word> {
        self.stack
            .last()
            .map(|s| s.counter)
            .ok_or(RepeatError::NoCounter)
    }

    // TODO: Update this comment.
    /// If there is a counter on the stack and the counter
    /// has greater then 1 repeat left then this will decrement
    /// the counter and return the index to repeat to.
    ///
    /// If the counter is 1 then this will pop the counter and
    /// return None because the repeat is done.
    ///
    /// If called when the stack is empty then this will return
    /// an error.
    ///
    /// Note that because the code has run once before the
    /// `RepeatEnd` is hit then we stop at 1.
    pub fn repeat(&mut self) -> RepeatResult<Option<usize>> {
        let slot = self.stack.last_mut().ok_or(RepeatError::Empty)?;
        match slot.limit {
            Direction::Up(limit) => {
                if slot.counter >= limit.saturating_sub(1) {
                    self.stack.pop();
                    Ok(None)
                } else {
                    slot.counter += 1;
                    Ok(Some(slot.repeat_index))
                }
            }
            Direction::Down => {
                if slot.counter <= 1 {
                    self.stack.pop();
                    Ok(None)
                } else {
                    slot.counter -= 1;
                    Ok(Some(slot.repeat_index))
                }
            }
        }
    }
}