orx_concurrent_iter/implementations/jagged_arrays/owned/
con_iter.rs

1use super::{chunk_puller::ChunkPullerJaggedOwned, into_iter::RawJaggedIterOwned};
2use super::{raw_jagged::RawJagged, slice_iter::RawJaggedSliceIterOwned};
3use crate::implementations::jagged_arrays::indexer::JaggedIndexer;
4use crate::{ConcurrentIter, ExactSizeConcurrentIter};
5use core::sync::atomic::{AtomicUsize, Ordering};
6
7/// Flattened concurrent iterator of a raw jagged array yielding owned elements.
8///
9/// Ensures that all elements are dropped regardless of whether they are iterated over or skipped.
10/// Further, cleans up the allocations of the jagged array.
11pub struct ConIterJaggedOwned<T, X>
12where
13    T: Send,
14    X: JaggedIndexer,
15{
16    jagged: RawJagged<T, X>,
17    counter: AtomicUsize,
18}
19
20unsafe impl<T: Send, X: JaggedIndexer> Sync for ConIterJaggedOwned<T, X> {}
21
22impl<T, X> ConIterJaggedOwned<T, X>
23where
24    T: Send,
25    X: JaggedIndexer,
26{
27    pub(crate) fn new(jagged: RawJagged<T, X>, begin: usize) -> Self {
28        Self {
29            jagged,
30            counter: begin.into(),
31        }
32    }
33
34    fn progress_and_get_begin_idx(&self, number_to_fetch: usize) -> Option<usize> {
35        let begin_idx = self.counter.fetch_add(number_to_fetch, Ordering::Relaxed);
36        match begin_idx < self.jagged.len() {
37            true => Some(begin_idx),
38            false => None,
39        }
40    }
41
42    pub(super) fn progress_and_get_iter(
43        &self,
44        chunk_size: usize,
45    ) -> Option<(usize, RawJaggedSliceIterOwned<'_, T>)> {
46        self.progress_and_get_begin_idx(chunk_size)
47            .map(|begin_idx| {
48                let end_idx = (begin_idx + chunk_size)
49                    .min(self.jagged.len())
50                    .max(begin_idx);
51                let slice = self.jagged.slice(begin_idx, end_idx);
52                let iter = RawJaggedSliceIterOwned::new(slice);
53                (begin_idx, iter)
54            })
55    }
56}
57
58impl<T, X> ConcurrentIter for ConIterJaggedOwned<T, X>
59where
60    T: Send,
61    X: JaggedIndexer,
62{
63    type Item = T;
64
65    type SequentialIter = RawJaggedIterOwned<T, X>;
66
67    type ChunkPuller<'i>
68        = ChunkPullerJaggedOwned<'i, T, X>
69    where
70        Self: 'i;
71
72    fn into_seq_iter(mut self) -> Self::SequentialIter {
73        let num_taken = self.counter.load(Ordering::Acquire).min(self.jagged.len());
74
75        let jagged_to_drop = self.jagged.move_into_new(num_taken);
76
77        RawJaggedIterOwned::new(jagged_to_drop)
78    }
79
80    fn skip_to_end(&self) {
81        let current = self.counter.fetch_max(self.jagged.len(), Ordering::Acquire);
82        let num_taken_before = current.min(self.jagged.len());
83        let slice = self.jagged.slice_from(num_taken_before);
84        let _iter = RawJaggedSliceIterOwned::new(slice);
85    }
86
87    fn next(&self) -> Option<Self::Item> {
88        self.progress_and_get_begin_idx(1).and_then(|idx| {
89            // SAFETY: `counter` ensures that elements from each position is taken only once
90            unsafe { self.jagged.take(idx) }
91        })
92    }
93
94    fn next_with_idx(&self) -> Option<(usize, Self::Item)> {
95        self.progress_and_get_begin_idx(1).and_then(|idx| {
96            // SAFETY: `counter` ensures that elements from each position is taken only once
97            unsafe { self.jagged.take(idx).map(|value| (idx, value)) }
98        })
99    }
100
101    fn size_hint(&self) -> (usize, Option<usize>) {
102        let num_taken = self.counter.load(Ordering::Acquire);
103        let remaining = self.jagged.len().saturating_sub(num_taken);
104        (remaining, Some(remaining))
105    }
106
107    fn is_completed_when_none_returned(&self) -> bool {
108        true
109    }
110
111    fn chunk_puller(&self, chunk_size: usize) -> Self::ChunkPuller<'_> {
112        let chunk_size = chunk_size.min(self.jagged.len());
113        Self::ChunkPuller::new(self, chunk_size)
114    }
115}
116
117impl<T, X> ExactSizeConcurrentIter for ConIterJaggedOwned<T, X>
118where
119    T: Send,
120    T: Send + Sync,
121    X: JaggedIndexer + Send + Sync,
122{
123    fn len(&self) -> usize {
124        let num_taken = self.counter.load(Ordering::Acquire);
125        self.jagged.len().saturating_sub(num_taken)
126    }
127}
128
129impl<T, X> Drop for ConIterJaggedOwned<T, X>
130where
131    T: Send,
132    X: JaggedIndexer,
133{
134    fn drop(&mut self) {
135        if self.jagged.num_taken().is_some() {
136            let num_taken = self.counter.load(Ordering::Acquire);
137            // SAFETY: `num_taken` elements are already taken out by the concurrent iterator.
138            // Before dropping this iterator, this is set as the `num_taken` of the raw
139            // jagged array which is responsible from dropping the elements and allocations.
140            // This will ensure that these `num_taken` elements will not be attempted to be
141            // dropped the second time.
142            unsafe { self.jagged.set_num_taken(Some(num_taken)) };
143        }
144    }
145}
orx_concurrent_iter/implementations/jagged_arrays/owned/con_iter.rs

orx_concurrent_iter/implementations/jagged_arrays/owned/
con_iter.rs
