Skip to main content

paralight/iter/source/
range.rs

1// Copyright 2024-2026 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9use super::{
10    ExactParallelSource, ExactSourceDescriptor, IntoExactParallelSource,
11    SimpleExactSourceDescriptor, SourceCleanup,
12};
13#[cfg(feature = "nightly")]
14use std::iter::Step;
15use std::ops::{Range, RangeInclusive};
16
17pub struct RangeSourceDescriptor<T> {
18    start: T,
19    len: usize,
20}
21
22impl<T> SourceCleanup for RangeSourceDescriptor<T> {
23    const NEEDS_CLEANUP: bool = false;
24
25    fn len(&self) -> usize {
26        self.len
27    }
28
29    unsafe fn cleanup_item_range(&self, _range: Range<usize>) {
30        // Nothing to cleanup
31    }
32}
33
34#[cfg(feature = "nightly")]
35impl<T: Step + Copy + Send + Sync> SimpleExactSourceDescriptor for RangeSourceDescriptor<T> {
36    type Item = T;
37
38    unsafe fn simple_exact_fetch_item(&self, index: usize) -> Self::Item {
39        debug_assert!(index < self.len);
40        T::forward(self.start, index)
41    }
42}
43
44#[cfg(not(feature = "nightly"))]
45impl SimpleExactSourceDescriptor for RangeSourceDescriptor<usize> {
46    type Item = usize;
47
48    unsafe fn simple_exact_fetch_item(&self, index: usize) -> Self::Item {
49        debug_assert!(index < self.len);
50        self.start + index
51    }
52}
53
54/// A parallel source over a [`Range`]. This struct is created by the
55/// [`into_par_iter()`](IntoExactParallelSource::into_par_iter) method on
56/// [`IntoExactParallelSource`].
57///
58/// You most likely won't need to interact with this struct directly, as it
59/// implements the [`ExactParallelSource`] and
60/// [`ExactParallelSourceExt`](super::ExactParallelSourceExt) traits, but it is
61/// nonetheless public because of the `must_use` annotation.
62///
63/// See also [`RangeInclusiveParallelSource`].
64///
65/// ### Stability blockers
66///
67/// On stable Rust, this struct is currently only implemented for ranges of
68/// [`usize`]. Ranges over any [`Step`] type are only available on Rust nightly
69/// with the `nightly` feature of Paralight enabled. This is because the
70/// implementation depends on the
71/// [`step_trait`](https://github.com/rust-lang/rust/issues/42168) nightly Rust
72/// feature.
73///
74/// ```
75/// # use paralight::iter::RangeParallelSource;
76/// # use paralight::prelude::*;
77/// # let mut thread_pool = ThreadPoolBuilder {
78/// #     num_threads: ThreadCount::AvailableParallelism,
79/// #     range_strategy: RangeStrategy::WorkStealing,
80/// #     cpu_pinning: CpuPinningPolicy::No,
81/// # }
82/// # .build();
83/// let iter: RangeParallelSource<usize> = (1..10).into_par_iter();
84/// let sum = iter.with_thread_pool(&mut thread_pool).sum::<usize>();
85/// assert_eq!(sum, 5 * 9);
86/// ```
87#[must_use = "iterator adaptors are lazy"]
88pub struct RangeParallelSource<T> {
89    range: Range<T>,
90}
91
92#[cfg(feature = "nightly")]
93impl<T: Step + Copy + Send + Sync> IntoExactParallelSource for Range<T> {
94    type Item = T;
95    type Source = RangeParallelSource<T>;
96
97    fn into_par_iter(self) -> Self::Source {
98        RangeParallelSource { range: self }
99    }
100}
101
102#[cfg(feature = "nightly")]
103impl<T: Step + Copy + Send + Sync> ExactParallelSource for RangeParallelSource<T> {
104    type Item = T;
105
106    fn exact_descriptor(self) -> impl ExactSourceDescriptor<Item = Self::Item> + Sync {
107        let range = self.range;
108        let (len_hint, len) = T::steps_between(&range.start, &range.end);
109        let len = len.unwrap_or_else(|| {
110            if len_hint == 0 {
111                panic!("cannot iterate over a backward range");
112            } else {
113                panic!(
114                    "cannot iterate over a range with more than usize::MAX items ({})",
115                    usize::MAX
116                );
117            }
118        });
119        RangeSourceDescriptor {
120            start: range.start,
121            len,
122        }
123    }
124}
125
126#[cfg(not(feature = "nightly"))]
127impl IntoExactParallelSource for Range<usize> {
128    type Item = usize;
129    type Source = RangeParallelSource<usize>;
130
131    fn into_par_iter(self) -> Self::Source {
132        RangeParallelSource { range: self }
133    }
134}
135
136#[cfg(not(feature = "nightly"))]
137impl ExactParallelSource for RangeParallelSource<usize> {
138    type Item = usize;
139
140    fn exact_descriptor(self) -> impl ExactSourceDescriptor<Item = Self::Item> + Sync {
141        let range = self.range;
142        let len = range
143            .end
144            .checked_sub(range.start)
145            .expect("cannot iterate over a backward range");
146        RangeSourceDescriptor {
147            start: range.start,
148            len,
149        }
150    }
151}
152
153/// A parallel source over a [`RangeInclusive`]. This struct is created by the
154/// [`into_par_iter()`](IntoExactParallelSource::into_par_iter) method on
155/// [`IntoExactParallelSource`].
156///
157/// You most likely won't need to interact with this struct directly, as it
158/// implements the [`ExactParallelSource`] and
159/// [`ExactParallelSourceExt`](super::ExactParallelSourceExt) traits, but it is
160/// nonetheless public because of the `must_use` annotation.
161///
162/// See also [`RangeParallelSource`].
163///
164/// ### Stability blockers
165///
166/// On stable Rust, this struct is currently only implemented for ranges of
167/// [`usize`]. Ranges over any [`Step`] type are only available on Rust nightly
168/// with the `nightly` feature of Paralight enabled. This is because the
169/// implementation depends on the
170/// [`step_trait`](https://github.com/rust-lang/rust/issues/42168) nightly Rust
171/// feature.
172///
173/// ```
174/// # use paralight::iter::RangeInclusiveParallelSource;
175/// # use paralight::prelude::*;
176/// # let mut thread_pool = ThreadPoolBuilder {
177/// #     num_threads: ThreadCount::AvailableParallelism,
178/// #     range_strategy: RangeStrategy::WorkStealing,
179/// #     cpu_pinning: CpuPinningPolicy::No,
180/// # }
181/// # .build();
182/// let iter: RangeInclusiveParallelSource<usize> = (1..=10).into_par_iter();
183/// let sum = iter.with_thread_pool(&mut thread_pool).sum::<usize>();
184/// assert_eq!(sum, 5 * 11);
185/// ```
186#[must_use = "iterator adaptors are lazy"]
187pub struct RangeInclusiveParallelSource<T> {
188    range: RangeInclusive<T>,
189}
190
191#[cfg(feature = "nightly")]
192impl<T: Step + Copy + Send + Sync> IntoExactParallelSource for RangeInclusive<T> {
193    type Item = T;
194    type Source = RangeInclusiveParallelSource<T>;
195
196    fn into_par_iter(self) -> Self::Source {
197        RangeInclusiveParallelSource { range: self }
198    }
199}
200
201#[cfg(feature = "nightly")]
202impl<T: Step + Copy + Send + Sync> ExactParallelSource for RangeInclusiveParallelSource<T> {
203    type Item = T;
204
205    fn exact_descriptor(self) -> impl ExactSourceDescriptor<Item = Self::Item> + Sync {
206        let (start, end) = self.range.into_inner();
207        let (len_hint, len) = T::steps_between(&start, &end);
208        let len = len.unwrap_or_else(|| {
209            if len_hint == 0 {
210                panic!("cannot iterate over a backward range");
211            } else {
212                panic!(
213                    "cannot iterate over a range with more than usize::MAX items ({})",
214                    usize::MAX
215                );
216            }
217        });
218        let len = len.checked_add(1).unwrap_or_else(|| {
219            panic!(
220                "cannot iterate over a range with more than usize::MAX items ({})",
221                usize::MAX
222            );
223        });
224        RangeSourceDescriptor { start, len }
225    }
226}
227
228#[cfg(not(feature = "nightly"))]
229impl IntoExactParallelSource for RangeInclusive<usize> {
230    type Item = usize;
231    type Source = RangeInclusiveParallelSource<usize>;
232
233    fn into_par_iter(self) -> Self::Source {
234        RangeInclusiveParallelSource { range: self }
235    }
236}
237
238#[cfg(not(feature = "nightly"))]
239impl ExactParallelSource for RangeInclusiveParallelSource<usize> {
240    type Item = usize;
241
242    fn exact_descriptor(self) -> impl ExactSourceDescriptor<Item = Self::Item> + Sync {
243        let (start, end) = self.range.into_inner();
244        let len = end
245            .checked_sub(start)
246            .expect("cannot iterate over a backward range");
247        let len = len.checked_add(1).unwrap_or_else(|| {
248            panic!(
249                "cannot iterate over a range with more than usize::MAX items ({})",
250                usize::MAX
251            );
252        });
253        RangeSourceDescriptor { start, len }
254    }
255}