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}