leo_passes/loop_unrolling/
range_iterator.rs

1// Copyright (C) 2019-2025 Provable Inc.
2// This file is part of the Leo library.
3
4// The Leo library 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, either version 3 of the License, or
7// (at your option) any later version.
8
9// The Leo library is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with the Leo library. If not, see <https://www.gnu.org/licenses/>.
16
17use num_traits::One;
18use std::{fmt::Display, ops::Add};
19
20use leo_ast::Value;
21use leo_errors::LeoError;
22
23// TODO: Consider the sealing this trait.
24// TODO: Better name.
25/// A trait for whose implementors are concrete values for loop bounds.
26pub(crate) trait LoopBound:
27    Add<Output = Self> + Copy + Display + One + PartialOrd + TryFrom<Value, Error = LeoError>
28{
29}
30
31impl LoopBound for i128 {}
32impl LoopBound for u128 {}
33
34/// Whether or not a bound is inclusive or exclusive.
35pub(crate) enum Clusivity {
36    Inclusive,
37    Exclusive,
38}
39
40/// An iterator over a range of values.
41pub(crate) struct RangeIterator<I: LoopBound> {
42    end: I,
43    current: Option<I>,
44    clusivity: Clusivity,
45}
46
47impl<I: LoopBound> RangeIterator<I> {
48    pub(crate) fn new(start: I, end: I, clusivity: Clusivity) -> Self {
49        Self { end, current: Some(start), clusivity }
50    }
51}
52
53impl<I: LoopBound> Iterator for RangeIterator<I> {
54    type Item = I;
55
56    fn next(&mut self) -> Option<Self::Item> {
57        match self.current {
58            None => None,
59            Some(value) if value < self.end => {
60                self.current = Some(value.add(I::one()));
61                Some(value)
62            }
63            Some(value) => {
64                self.current = None;
65                match self.clusivity {
66                    Clusivity::Exclusive => None,
67                    Clusivity::Inclusive => Some(value),
68                }
69            }
70        }
71    }
72}