essential_constraint_vm/
repeat.rs

1use essential_types::{convert::bool_from_word, Word};
2
3use crate::{
4    error::{RepeatError, RepeatResult, StackError},
5    OpResult, Stack,
6};
7
8#[cfg(test)]
9mod tests;
10
11#[derive(Debug, Default, PartialEq)]
12/// A stack of repeat counters.
13pub 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
30/// `Stack::Repeat` implementation.
31pub(crate) fn repeat(pc: usize, stack: &mut Stack, repeat: &mut Repeat) -> OpResult<()> {
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    /// Create a new repeat stack.
45    pub fn new() -> Self {
46        Self::default()
47    }
48
49    /// Add a new repeat location and counter to the stack.
50    /// Counts down to 0.
51    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    /// Add a new repeat location and counter to the stack.
64    /// Counts up from 0 to limit - 1.
65    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    /// Get the current repeat counter.
78    ///
79    /// Returns an error if the stack is empty.
80    pub fn counter(&self) -> RepeatResult<Word> {
81        self.stack
82            .last()
83            .map(|s| s.counter)
84            .ok_or(RepeatError::NoCounter)
85    }
86
87    // TODO: Update this comment.
88    /// If there is a counter on the stack and the counter
89    /// has greater then 1 repeat left then this will decrement
90    /// the counter and return the index to repeat to.
91    ///
92    /// If the counter is 1 then this will pop the counter and
93    /// return None because the repeat is done.
94    ///
95    /// If called when the stack is empty then this will return
96    /// an error.
97    ///
98    /// Note that because the code has run once before the
99    /// `RepeatEnd` is hit then we stop at 1.
100    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}