Skip to main content

orx_split_vec/concurrent_pinned_vec/
con_pinvec.rs

1use crate::{
2    Doubling, Fragment, GrowthWithConstantTimeAccess, SplitVec,
3    common_traits::iterator::{IterOfSlicesOfCon, SliceBorrowAsMut, SliceBorrowAsRef},
4    concurrent_pinned_vec::{into_iter::ConcurrentSplitVecIntoIter, iter_ptr::IterPtrOfCon},
5    fragment::transformations::{fragment_from_raw, fragment_into_raw},
6};
7use alloc::vec::Vec;
8use core::ops::RangeBounds;
9use core::sync::atomic::{AtomicUsize, Ordering};
10use core::{cell::UnsafeCell, ops::Range};
11use orx_pinned_vec::ConcurrentPinnedVec;
12
13pub struct FragmentData {
14    pub f: usize,
15    pub len: usize,
16    pub capacity: usize,
17}
18
19/// Concurrent wrapper ([`orx_pinned_vec::ConcurrentPinnedVec`]) for the `SplitVec`.
20pub struct ConcurrentSplitVec<T, G: GrowthWithConstantTimeAccess = Doubling> {
21    growth: G,
22    data: Vec<UnsafeCell<*mut T>>,
23    capacity: AtomicUsize,
24    maximum_capacity: usize,
25    max_num_fragments: usize,
26    pinned_vec_len: usize,
27}
28
29impl<T, G: GrowthWithConstantTimeAccess> Drop for ConcurrentSplitVec<T, G> {
30    fn drop(&mut self) {
31        fn take_fragment<T>(_fragment: Fragment<T>) {}
32        unsafe { self.process_into_fragments(self.pinned_vec_len, &mut take_fragment) };
33        self.zero();
34    }
35}
36
37impl<T, G: GrowthWithConstantTimeAccess> ConcurrentSplitVec<T, G> {
38    pub(super) fn destruct(mut self) -> (G, Vec<UnsafeCell<*mut T>>, usize) {
39        let mut data = Vec::new();
40        core::mem::swap(&mut self.data, &mut data);
41        let capacity = self.capacity.load(Ordering::Relaxed);
42        let growth = self.growth.clone();
43        self.zero();
44        (growth, data, capacity)
45    }
46
47    unsafe fn get_raw_mut_unchecked_fi(&self, f: usize, i: usize) -> *mut T {
48        let p = unsafe { *self.data[f].get() };
49        unsafe { p.add(i) }
50    }
51
52    unsafe fn get_raw_mut_unchecked_idx(&self, idx: usize) -> *mut T {
53        let (f, i) = self.growth.get_fragment_and_inner_indices_unchecked(idx);
54        unsafe { self.get_raw_mut_unchecked_fi(f, i) }
55    }
56
57    fn capacity_of(&self, f: usize) -> usize {
58        self.growth.fragment_capacity_of(f)
59    }
60
61    fn layout(len: usize) -> alloc::alloc::Layout {
62        alloc::alloc::Layout::array::<T>(len).expect("len must not overflow")
63    }
64
65    unsafe fn to_fragment(&self, data: FragmentData) -> Fragment<T> {
66        let ptr = unsafe { *self.data[data.f].get() };
67        unsafe { fragment_from_raw(ptr, data.len, data.capacity) }
68    }
69
70    unsafe fn process_into_fragments<F>(&mut self, len: usize, take_fragment: &mut F)
71    where
72        F: FnMut(Fragment<T>),
73    {
74        let mut process_in_cap = |x: FragmentData| {
75            let _fragment_to_drop = unsafe { self.to_fragment(x) };
76        };
77        let mut process_in_len = |x: FragmentData| {
78            let fragment = unsafe { self.to_fragment(x) };
79            take_fragment(fragment);
80        };
81
82        unsafe { self.process_fragments(len, &mut process_in_len, &mut process_in_cap) };
83    }
84
85    unsafe fn process_fragments<P, Q>(
86        &self,
87        len: usize,
88        process_in_len: &mut P,
89        process_in_cap: &mut Q,
90    ) where
91        P: FnMut(FragmentData),
92        Q: FnMut(FragmentData),
93    {
94        let capacity = self.capacity();
95        assert!(capacity >= len);
96
97        let mut remaining_len = len;
98        let mut f = 0;
99        let mut taken_out_capacity = 0;
100
101        while remaining_len > 0 {
102            let capacity = self.capacity_of(f);
103            taken_out_capacity += capacity;
104
105            let len = match remaining_len <= capacity {
106                true => remaining_len,
107                false => capacity,
108            };
109
110            let fragment = FragmentData { f, len, capacity };
111            process_in_len(fragment);
112            remaining_len -= len;
113            f += 1;
114        }
115
116        while capacity > taken_out_capacity {
117            let capacity = self.capacity_of(f);
118            taken_out_capacity += capacity;
119            let len = 0;
120            let fragment = FragmentData { f, len, capacity };
121            process_in_cap(fragment);
122            f += 1;
123        }
124    }
125
126    fn zero(&mut self) {
127        self.capacity = 0.into();
128        self.maximum_capacity = 0;
129        self.max_num_fragments = 0;
130        self.pinned_vec_len = 0;
131    }
132
133    fn num_fragments_for_capacity(&self, capacity: usize) -> usize {
134        match capacity {
135            0 => 0,
136            _ => {
137                self.growth
138                    .get_fragment_and_inner_indices_unchecked(capacity - 1)
139                    .0
140                    + 1
141            }
142        }
143    }
144}
145
146impl<T, G: GrowthWithConstantTimeAccess> From<SplitVec<T, G>> for ConcurrentSplitVec<T, G> {
147    fn from(value: SplitVec<T, G>) -> Self {
148        let (fragments, growth, pinned_vec_len) = (value.fragments, value.growth, value.len);
149
150        let num_fragments = fragments.len();
151        let max_num_fragments = fragments.capacity();
152
153        let mut data = Vec::with_capacity(max_num_fragments);
154        let mut total_len = 0;
155        let mut maximum_capacity = 0;
156
157        for (f, fragment) in fragments.into_iter().enumerate() {
158            let (p, len, cap) = fragment_into_raw(fragment);
159
160            let expected_cap = growth.fragment_capacity_of(f);
161            if core::mem::size_of::<T>() > 0 {
162                assert_eq!(cap, expected_cap);
163            }
164
165            total_len += len;
166            maximum_capacity += cap;
167
168            data.push(UnsafeCell::new(p));
169        }
170        assert_eq!(total_len, pinned_vec_len);
171        let capacity = maximum_capacity;
172
173        for f in num_fragments..data.capacity() {
174            let expected_cap = growth.fragment_capacity_of(f);
175            maximum_capacity += expected_cap;
176
177            data.push(UnsafeCell::new(core::ptr::null_mut()));
178        }
179
180        Self {
181            growth,
182            data,
183            capacity: capacity.into(),
184            maximum_capacity,
185            max_num_fragments,
186            pinned_vec_len,
187        }
188    }
189}
190
191impl<T, G: GrowthWithConstantTimeAccess> ConcurrentPinnedVec<T> for ConcurrentSplitVec<T, G> {
192    type P = SplitVec<T, G>;
193
194    type SliceIter<'a>
195        = IterOfSlicesOfCon<'a, T, G, SliceBorrowAsRef>
196    where
197        Self: 'a;
198
199    type SliceMutIter<'a>
200        = IterOfSlicesOfCon<'a, T, G, SliceBorrowAsMut>
201    where
202        Self: 'a;
203
204    type PtrIter<'a>
205        = IterPtrOfCon<'a, T, G>
206    where
207        Self: 'a;
208
209    type IntoIter = ConcurrentSplitVecIntoIter<T, G>;
210
211    unsafe fn into_inner(mut self, len: usize) -> Self::P {
212        let mut fragments = Vec::with_capacity(self.max_num_fragments);
213        let mut take_fragment = |fragment| fragments.push(fragment);
214        unsafe { self.process_into_fragments(len, &mut take_fragment) };
215
216        self.zero();
217        SplitVec::from_raw_parts(len, fragments, self.growth.clone())
218    }
219
220    unsafe fn clone_with_len(&self, len: usize) -> Self
221    where
222        T: Clone,
223    {
224        let mut fragments = Vec::with_capacity(self.max_num_fragments);
225        let mut clone_fragment = |x: FragmentData| {
226            let mut fragment = Fragment::new_empty(x.capacity);
227            let dst: *mut T = fragment.as_mut_ptr();
228            let src = unsafe { *self.data[x.f].get() };
229            for i in 0..x.len {
230                let value = unsafe { src.add(i).as_ref() }.expect("must be some");
231                unsafe { dst.add(i).write(value.clone()) };
232            }
233            unsafe { fragment.set_len(x.len) };
234            fragments.push(fragment);
235        };
236
237        unsafe { self.process_fragments(len, &mut clone_fragment, &mut |_| {}) };
238
239        let split_vec = SplitVec::from_raw_parts(len, fragments, self.growth.clone());
240        split_vec.into()
241    }
242
243    fn slices<R: RangeBounds<usize>>(&self, range: R) -> Self::SliceIter<'_> {
244        Self::SliceIter::new(self.capacity(), &self.data, self.growth.clone(), range)
245    }
246
247    unsafe fn iter<'a>(&'a self, len: usize) -> impl Iterator<Item = &'a T> + 'a
248    where
249        T: 'a,
250    {
251        self.slices(0..len).flat_map(|x| x.iter())
252    }
253
254    unsafe fn iter_over_range<'a, R: RangeBounds<usize>>(
255        &'a self,
256        range: R,
257    ) -> impl Iterator<Item = &'a T> + 'a
258    where
259        T: 'a,
260    {
261        let [a, b] = orx_pinned_vec::utils::slice::vec_range_limits(&range, None);
262        self.slices(a..b).flat_map(|x| x.iter())
263    }
264
265    unsafe fn slices_mut<R: RangeBounds<usize>>(&self, range: R) -> Self::SliceMutIter<'_> {
266        Self::SliceMutIter::new(self.capacity(), &self.data, self.growth.clone(), range)
267    }
268
269    unsafe fn iter_mut<'a>(&'a mut self, len: usize) -> impl Iterator<Item = &'a mut T> + 'a
270    where
271        T: 'a,
272    {
273        unsafe { self.slices_mut(0..len) }.flat_map(|x| x.iter_mut())
274    }
275
276    unsafe fn get(&self, index: usize) -> Option<&T> {
277        match index < self.capacity() {
278            true => {
279                let p = unsafe { self.get_raw_mut_unchecked_idx(index) };
280                Some(unsafe { &*p })
281            }
282            false => None,
283        }
284    }
285
286    unsafe fn get_mut(&mut self, index: usize) -> Option<&mut T> {
287        match index < self.capacity() {
288            true => {
289                let p = unsafe { self.get_raw_mut_unchecked_idx(index) };
290                Some(unsafe { &mut *p })
291            }
292            false => None,
293        }
294    }
295
296    unsafe fn get_ptr_mut(&self, index: usize) -> *mut T {
297        unsafe { self.get_raw_mut_unchecked_idx(index) }
298    }
299
300    fn max_capacity(&self) -> usize {
301        self.maximum_capacity
302    }
303
304    fn capacity(&self) -> usize {
305        self.capacity.load(Ordering::Acquire)
306    }
307
308    fn grow_to(&self, new_capacity: usize) -> Result<usize, orx_pinned_vec::PinnedVecGrowthError> {
309        let capacity = self.capacity.load(Ordering::Acquire);
310        match new_capacity <= capacity {
311            true => Ok(capacity),
312            false => {
313                let mut f = self.num_fragments_for_capacity(capacity);
314                let mut current_capacity = capacity;
315
316                while new_capacity > current_capacity {
317                    let new_fragment_capacity = self.capacity_of(f);
318                    let layout = Self::layout(new_fragment_capacity);
319                    let ptr = unsafe { alloc::alloc::alloc(layout) } as *mut T;
320                    unsafe { *self.data[f].get() = ptr };
321
322                    f += 1;
323                    current_capacity += new_fragment_capacity;
324                }
325
326                self.capacity.store(current_capacity, Ordering::Release);
327
328                Ok(current_capacity)
329            }
330        }
331    }
332
333    fn grow_to_and_fill_with<F>(
334        &self,
335        new_capacity: usize,
336        fill_with: F,
337    ) -> Result<usize, orx_pinned_vec::PinnedVecGrowthError>
338    where
339        F: Fn() -> T,
340    {
341        let capacity = self.capacity.load(Ordering::Acquire);
342        match new_capacity <= capacity {
343            true => Ok(capacity),
344            false => {
345                let mut f = self.num_fragments_for_capacity(capacity);
346
347                let mut current_capacity = capacity;
348
349                while new_capacity > current_capacity {
350                    let new_fragment_capacity = self.capacity_of(f);
351                    let layout = Self::layout(new_fragment_capacity);
352                    let ptr = unsafe { alloc::alloc::alloc(layout) } as *mut T;
353
354                    for i in 0..new_fragment_capacity {
355                        unsafe { ptr.add(i).write(fill_with()) };
356                    }
357
358                    unsafe { *self.data[f].get() = ptr };
359
360                    f += 1;
361                    current_capacity += new_fragment_capacity;
362                }
363
364                self.capacity.store(current_capacity, Ordering::Release);
365
366                Ok(current_capacity)
367            }
368        }
369    }
370
371    fn fill_with<F>(&self, range: core::ops::Range<usize>, fill_with: F)
372    where
373        F: Fn() -> T,
374    {
375        for i in range {
376            unsafe { self.get_ptr_mut(i).write(fill_with()) };
377        }
378    }
379
380    unsafe fn reserve_maximum_concurrent_capacity(
381        &mut self,
382        _current_len: usize,
383        new_maximum_capacity: usize,
384    ) -> usize {
385        assert_eq!(self.max_num_fragments, self.data.len());
386        assert_eq!(self.max_num_fragments, self.data.capacity());
387
388        let mut num_required_fragments = 0;
389        let mut max_cap = self.maximum_capacity;
390        let mut f = self.data.len();
391
392        while max_cap < new_maximum_capacity {
393            max_cap += self.capacity_of(f);
394            num_required_fragments += 1;
395            f += 1;
396        }
397
398        if num_required_fragments > 0 {
399            self.data.reserve_exact(num_required_fragments);
400        }
401
402        for _ in self.max_num_fragments..self.data.capacity() {
403            self.data.push(UnsafeCell::new(core::ptr::null_mut()));
404        }
405
406        self.maximum_capacity = (0..self.data.len()).map(|f| self.capacity_of(f)).sum();
407        self.max_num_fragments = self.data.len();
408
409        assert_eq!(self.max_num_fragments, self.data.len());
410        assert_eq!(self.max_num_fragments, self.data.capacity());
411
412        self.maximum_capacity
413    }
414
415    unsafe fn reserve_maximum_concurrent_capacity_fill_with<F>(
416        &mut self,
417        current_len: usize,
418        new_maximum_capacity: usize,
419        _fill_with: F,
420    ) -> usize
421    where
422        F: Fn() -> T,
423    {
424        unsafe { self.reserve_maximum_concurrent_capacity(current_len, new_maximum_capacity) }
425    }
426
427    unsafe fn set_pinned_vec_len(&mut self, len: usize) {
428        self.pinned_vec_len = len;
429    }
430
431    unsafe fn clear(&mut self, len: usize) {
432        let mut take_fragment = |_fragment: Fragment<T>| {};
433        unsafe { self.process_into_fragments(len, &mut take_fragment) };
434        self.zero();
435
436        let max_num_fragments = self.data.len();
437        self.data.clear();
438
439        for _ in 0..max_num_fragments {
440            self.data.push(UnsafeCell::new(core::ptr::null_mut()));
441        }
442
443        self.maximum_capacity = (0..self.data.len()).map(|f| self.capacity_of(f)).sum();
444        self.pinned_vec_len = 0;
445    }
446
447    unsafe fn ptr_iter_unchecked(&self, range: Range<usize>) -> Self::PtrIter<'_> {
448        IterPtrOfCon::new(self.capacity(), &self.data, self.growth.clone(), range)
449    }
450
451    unsafe fn into_iter(self, range: Range<usize>) -> Self::IntoIter {
452        let (growth, data, capacity) = self.destruct();
453        ConcurrentSplitVecIntoIter::new(capacity, data, growth, range)
454    }
455}