Skip to main content

imbl/vector/
focus.rs

1// This Source Code Form is subject to the terms of the Mozilla Public
2// License, v. 2.0. If a copy of the MPL was not distributed with this
3// file, You can obtain one at http://mozilla.org/MPL/2.0/.
4
5use std::mem::{replace, swap};
6use std::ops::{Range, RangeBounds};
7use std::ptr::null;
8use std::sync::atomic::{AtomicPtr, Ordering};
9
10use archery::{SharedPointer, SharedPointerKind};
11
12use crate::nodes::chunk::Chunk;
13use crate::sync::Lock;
14use crate::util::to_range;
15use crate::vector::{
16    GenericVector, Iter, IterMut,
17    VectorInner::{Full, Inline, Single},
18    RRB,
19};
20
21fn check_indices<const N: usize>(len: usize, indices: &[usize; N]) -> Option<()> {
22    let mut seen = [None; N];
23    for idx in indices {
24        if *idx > len || seen.contains(&Some(*idx)) {
25            return None;
26        }
27        let empty = seen.iter_mut().find(|a| a.is_none()).unwrap();
28        *empty = Some(*idx);
29    }
30    Some(())
31}
32
33/// Focused indexing over a [`Vector`][Vector].
34///
35/// By remembering the last tree node accessed through an index lookup and the
36/// path we took to get there, we can speed up lookups for adjacent indices
37/// tremendously. Lookups on indices in the same node are instantaneous, and
38/// lookups on sibling nodes are also very fast.
39///
40/// A `Focus` can also be used as a restricted view into a vector, using the
41/// [`narrow`][narrow] and [`split_at`][split_at] methods.
42///
43/// # When should I use a `Focus` for better performance?
44///
45/// `Focus` is useful when you need to perform a large number of index lookups
46/// that are more likely than not to be close to each other. It's usually worth
47/// using a `Focus` in any situation where you're batching a lot of index
48/// lookups together, even if they're not obviously adjacent - there's likely
49/// to be some performance gain for even completely random access.
50///
51/// If you're just iterating forwards or backwards over the [`Vector`][Vector]
52/// in order, you're better off with a regular iterator, which, in fact, is
53/// implemented using a `Focus`, but provides a simpler interface.
54///
55/// If you're just doing a very small number of index lookups, the setup cost
56/// for the `Focus` is probably not worth it.
57///
58/// A `Focus` is never faster than an index lookup on a small [`Vector`][Vector]
59/// with a length below the internal RRB tree's branching factor of 64.
60///
61/// # Examples
62///
63/// This example is contrived, as the better way to iterate forwards or
64/// backwards over a vector is with an actual iterator. Even so, the version
65/// using a `Focus` should run nearly an order of magnitude faster than the
66/// version using index lookups at a length of 1000. It should also be noted
67/// that [`vector::Iter`][Iter] is actually implemented using a `Focus` behind
68/// the scenes, so the performance of the two should be identical.
69///
70/// ```rust
71/// # #[macro_use] extern crate imbl;
72/// # use imbl::vector::Vector;
73/// # use std::iter::FromIterator;
74/// let mut vec: Vector<i64> = Vector::from_iter(0..1000);
75///
76/// // Summing a vector, the slow way:
77/// let mut sum = 0;
78/// for i in 0..1000 {
79///     sum += *vec.get(i).unwrap();
80/// }
81/// assert_eq!(499500, sum);
82///
83/// // Summing a vector faster using a Focus:
84/// let mut sum = 0;
85/// let mut focus = vec.focus();
86/// for i in 0..1000 {
87///     sum += *focus.get(i).unwrap();
88/// }
89/// assert_eq!(499500, sum);
90///
91/// // And the easy way, for completeness:
92/// let sum: i64 = vec.iter().sum();
93/// assert_eq!(499500, sum);
94/// ```
95///
96/// [Vector]: type.Vector.html
97/// [Iter]: struct.Iter.html
98/// [narrow]: #method.narrow
99/// [split_at]: #method.split_at
100pub enum Focus<'a, A, P: SharedPointerKind> {
101    #[doc(hidden)]
102    /// The Single variant is a focus of a simple Vector that can be represented as a single slice.
103    Single(&'a [A]),
104    #[doc(hidden)]
105    /// The Full variant is a focus of a more complex Vector that cannot be represented as a single slice.
106    Full(TreeFocus<A, P>),
107}
108
109impl<'a, A, P: SharedPointerKind> Focus<'a, A, P>
110where
111    A: 'a,
112{
113    /// Construct a `Focus` for a [`Vector`][Vector].
114    ///
115    /// [Vector]: type.Vector.html
116    pub fn new(vector: &'a GenericVector<A, P>) -> Self {
117        match &vector.vector {
118            Inline(chunk) => Focus::Single(chunk),
119            Single(chunk) => Focus::Single(chunk),
120            Full(tree) => Focus::Full(TreeFocus::new(tree)),
121        }
122    }
123
124    /// Get the length of the focused [`Vector`][Vector].
125    ///
126    /// [Vector]: type.Vector.html
127    pub fn len(&self) -> usize {
128        match self {
129            Focus::Single(chunk) => chunk.len(),
130            Focus::Full(tree) => tree.len(),
131        }
132    }
133
134    /// Test if the focused [`Vector`][Vector] is empty.
135    ///
136    /// [Vector]: type.Vector.html
137    pub fn is_empty(&self) -> bool {
138        self.len() == 0
139    }
140
141    /// Get a reference to the value at a given index.
142    pub fn get(&mut self, index: usize) -> Option<&A> {
143        match self {
144            Focus::Single(chunk) => chunk.get(index),
145            Focus::Full(tree) => tree.get(index),
146        }
147    }
148
149    /// Get a reference to the value at a given index.
150    ///
151    /// Panics if the index is out of bounds.
152    pub fn index(&mut self, index: usize) -> &A {
153        self.get(index).expect("index out of bounds")
154    }
155
156    /// Get the chunk for the given index.
157    ///
158    /// This gives you a reference to the leaf node that contains the index,
159    /// along with its start and end indices.
160    pub fn chunk_at(&mut self, index: usize) -> (Range<usize>, &[A]) {
161        let len = self.len();
162        if index >= len {
163            panic!("vector::Focus::chunk_at: index out of bounds");
164        }
165        match self {
166            Focus::Single(chunk) => (0..len, chunk),
167            Focus::Full(tree) => tree.get_chunk(index),
168        }
169    }
170
171    /// Narrow the focus onto a subslice of the vector.
172    ///
173    /// `Focus::narrow(range)` has the same effect as `&slice[range]`, without
174    /// actually modifying the underlying vector.
175    ///
176    /// Panics if the range isn't fully inside the current focus.
177    ///
178    /// ## Examples
179    ///
180    /// ```rust
181    /// # #[macro_use] extern crate imbl;
182    /// # use imbl::vector::Vector;
183    /// # use std::iter::FromIterator;
184    /// let vec: Vector<i64> = Vector::from_iter(0..1000);
185    /// let narrowed = vec.focus().narrow(100..200);
186    /// let narrowed_vec: Vector<i64> = narrowed.into_iter().cloned().collect();
187    /// assert_eq!(Vector::from_iter(100..200), narrowed_vec);
188    /// ```
189    ///
190    /// [slice::split_at]: https://doc.rust-lang.org/std/primitive.slice.html#method.split_at
191    /// [Vector::split_at]: type.Vector.html#method.split_at
192    #[must_use]
193    pub fn narrow<R>(self, range: R) -> Self
194    where
195        R: RangeBounds<usize>,
196    {
197        let r = to_range(&range, self.len());
198        if r.start > r.end || r.end > self.len() {
199            panic!("vector::Focus::narrow: range out of bounds");
200        }
201        match self {
202            Focus::Single(chunk) => Focus::Single(&chunk[r]),
203            Focus::Full(tree) => Focus::Full(tree.narrow(r)),
204        }
205    }
206
207    /// Split the focus into two.
208    ///
209    /// Given an index `index`, consume the focus and produce two new foci, the
210    /// left onto indices `0..index`, and the right onto indices `index..N`
211    /// where `N` is the length of the current focus.
212    ///
213    /// Panics if the index is out of bounds.
214    ///
215    /// This is the moral equivalent of [`slice::split_at`][slice::split_at], in
216    /// that it leaves the underlying data structure unchanged, unlike
217    /// [`Vector::split_at`][Vector::split_at].
218    ///
219    /// ## Examples
220    ///
221    /// ```rust
222    /// # #[macro_use] extern crate imbl;
223    /// # use imbl::vector::Vector;
224    /// # use std::iter::FromIterator;
225    /// let vec: Vector<i64> = Vector::from_iter(0..1000);
226    /// let (left, right) = vec.focus().split_at(500);
227    /// let left_vec: Vector<i64> = left.into_iter().cloned().collect();
228    /// let right_vec: Vector<i64> = right.into_iter().cloned().collect();
229    /// assert_eq!(Vector::from_iter(0..500), left_vec);
230    /// assert_eq!(Vector::from_iter(500..1000), right_vec);
231    /// ```
232    ///
233    /// [slice::split_at]: https://doc.rust-lang.org/std/primitive.slice.html#method.split_at
234    /// [Vector::split_at]: type.Vector.html#method.split_at
235    pub fn split_at(self, index: usize) -> (Self, Self) {
236        if index > self.len() {
237            panic!("vector::Focus::split_at: index out of bounds");
238        }
239        match self {
240            Focus::Single(chunk) => {
241                let (left, right) = chunk.split_at(index);
242                (Focus::Single(left), Focus::Single(right))
243            }
244            Focus::Full(tree) => {
245                let (left, right) = tree.split_at(index);
246                (Focus::Full(left), Focus::Full(right))
247            }
248        }
249    }
250}
251
252impl<'a, A, P: SharedPointerKind + 'a> IntoIterator for Focus<'a, A, P>
253where
254    A: Clone + 'a,
255{
256    type Item = &'a A;
257    type IntoIter = Iter<'a, A, P>;
258
259    fn into_iter(self) -> Self::IntoIter {
260        Iter::from_focus(self)
261    }
262}
263
264impl<'a, A, P: SharedPointerKind> Clone for Focus<'a, A, P>
265where
266    A: Clone + 'a,
267{
268    fn clone(&self) -> Self {
269        match self {
270            Focus::Single(chunk) => Focus::Single(chunk),
271            Focus::Full(tree) => Focus::Full(tree.clone()),
272        }
273    }
274}
275
276pub struct TreeFocus<A, P: SharedPointerKind> {
277    /// A clone of the Vector's internal tree that this focus points to. A clone ensures that we don't require a
278    /// reference to the original tree.
279    tree: RRB<A, P>,
280    /// The view represents the range of the tree that this TreeFocus can see. The view can be narrowed by calling
281    /// either the narrow or split_at methods.
282    view: Range<usize>,
283    /// The tree version of the Vector is represented as the concatenation of 2 chunks, followed by the tree root,
284    /// followed by 2 chunks. The middle_range refers to the range of the Vector that the tree covers.
285    middle_range: Range<usize>,
286    /// This implementation of a focus stores only a single chunk for the Vector. This chunk can refer to one of the 4
287    /// chunks front/back chunks or one of the leaves of the tree. The target_ptr is the pointer to the actual chunk
288    /// in question. The target_range is the range that the chunk represents.
289    target_range: Range<usize>,
290    target_ptr: *const Chunk<A>,
291}
292
293impl<A, P: SharedPointerKind> Clone for TreeFocus<A, P> {
294    fn clone(&self) -> Self {
295        let tree = self.tree.clone();
296        TreeFocus {
297            view: self.view.clone(),
298            middle_range: self.middle_range.clone(),
299            target_range: 0..0,
300            target_ptr: null(),
301            tree,
302        }
303    }
304}
305
306unsafe impl<A: Send, P: SharedPointerKind + Send> Send for TreeFocus<A, P> {}
307unsafe impl<A: Sync, P: SharedPointerKind + Sync> Sync for TreeFocus<A, P> {}
308
309#[inline]
310fn contains<A: Ord>(range: &Range<A>, index: &A) -> bool {
311    *index >= range.start && *index < range.end
312}
313
314impl<A, P: SharedPointerKind> TreeFocus<A, P> {
315    /// Creates a new TreeFocus for a Vector's RRB tree.
316    fn new(tree: &RRB<A, P>) -> Self {
317        let middle_start = tree.outer_f.len() + tree.inner_f.len();
318        let middle_end = middle_start + tree.middle.len();
319        TreeFocus {
320            tree: tree.clone(),
321            view: 0..tree.length,
322            middle_range: middle_start..middle_end,
323            target_range: 0..0,
324            target_ptr: null(),
325        }
326    }
327
328    /// Returns the number of elements that the TreeFocus is valid for.
329    fn len(&self) -> usize {
330        self.view.end - self.view.start
331    }
332
333    /// Restricts the TreeFocus to a subrange of itself.
334    fn narrow(self, mut view: Range<usize>) -> Self {
335        view.start += self.view.start;
336        view.end += self.view.start;
337        TreeFocus {
338            view,
339            middle_range: self.middle_range.clone(),
340            target_range: 0..0,
341            target_ptr: null(),
342            tree: self.tree,
343        }
344    }
345
346    /// Splits the TreeFocus into two disjoint foci. The first TreeFocus is valid for ..index while the
347    /// second is valid for index.. .
348    fn split_at(self, index: usize) -> (Self, Self) {
349        let len = self.len();
350        let left = self.clone().narrow(0..index);
351        let right = self.narrow(index..len);
352        (left, right)
353    }
354
355    /// Computes an absolute index in the RRBTree for the given index relative to the start of this TreeFocus.
356    fn physical_index(&self, index: usize) -> usize {
357        debug_assert!(index < self.view.end);
358        self.view.start + index
359    }
360
361    /// Computes a range relative to the TreeFocus given one that is absolute in the RRBTree.
362    fn logical_range(&self, range: &Range<usize>) -> Range<usize> {
363        (range.start - self.view.start)..(range.end - self.view.start)
364    }
365
366    /// Sets the internal chunk to the one that contains the given absolute index.
367    fn set_focus(&mut self, index: usize) {
368        if index < self.middle_range.start {
369            let outer_len = self.tree.outer_f.len();
370            if index < outer_len {
371                self.target_range = 0..outer_len;
372                self.target_ptr = &*self.tree.outer_f;
373            } else {
374                self.target_range = outer_len..self.middle_range.start;
375                self.target_ptr = &*self.tree.inner_f;
376            }
377        } else if index >= self.middle_range.end {
378            let outer_start = self.middle_range.end + self.tree.inner_b.len();
379            if index < outer_start {
380                self.target_range = self.middle_range.end..outer_start;
381                self.target_ptr = &*self.tree.inner_b;
382            } else {
383                self.target_range = outer_start..self.tree.length;
384                self.target_ptr = &*self.tree.outer_b;
385            }
386        } else {
387            let tree_index = index - self.middle_range.start;
388            let (range, ptr) = self
389                .tree
390                .middle
391                .lookup_chunk(self.tree.middle_level, 0, tree_index);
392            self.target_range =
393                (range.start + self.middle_range.start)..(range.end + self.middle_range.start);
394            self.target_ptr = ptr;
395        }
396    }
397
398    /// Gets the chunk that this TreeFocus is focused on.
399    fn get_focus(&self) -> &Chunk<A> {
400        unsafe { &*self.target_ptr }
401    }
402
403    /// Gets the value at the given index relative to the TreeFocus.
404    pub fn get(&mut self, index: usize) -> Option<&A> {
405        if index >= self.len() {
406            return None;
407        }
408        let phys_index = self.physical_index(index);
409        if !contains(&self.target_range, &phys_index) {
410            self.set_focus(phys_index);
411        }
412        let target_phys_index = phys_index - self.target_range.start;
413        Some(&self.get_focus()[target_phys_index])
414    }
415
416    /// Gets the chunk for an index as a slice and its corresponding range within the TreeFocus.
417    pub fn get_chunk(&mut self, index: usize) -> (Range<usize>, &[A]) {
418        let phys_index = self.physical_index(index);
419        if !contains(&self.target_range, &phys_index) {
420            self.set_focus(phys_index);
421        }
422        let mut slice: &[A] = self.get_focus();
423        let mut left = 0;
424        let mut right = 0;
425        if self.target_range.start < self.view.start {
426            left = self.view.start - self.target_range.start;
427        }
428        if self.target_range.end > self.view.end {
429            right = self.target_range.end - self.view.end;
430        }
431        slice = &slice[left..(slice.len() - right)];
432        let phys_range = (self.target_range.start + left)..(self.target_range.end - right);
433        (self.logical_range(&phys_range), slice)
434    }
435}
436
437/// A mutable version of [`Focus`][Focus].
438///
439/// See [`Focus`][Focus] for more details.
440///
441/// You can only build one `FocusMut` at a time for a vector, effectively
442/// keeping a lock on the vector until you're done with the focus, which relies
443/// on the structure of the vector not changing while it exists.
444///
445/// ```rust,compile_fail
446/// # #[macro_use] extern crate imbl;
447/// # use imbl::vector::Vector;
448/// # use std::iter::FromIterator;
449/// let mut vec: Vector<i64> = Vector::from_iter(0..1000);
450/// let focus1 = vec.focus_mut();
451/// // Fails here in 2015 edition because you're creating
452/// // two mutable references to the same thing.
453/// let focus2 = vec.focus_mut();
454/// // Fails here in 2018 edition because creating focus2
455/// // made focus1's lifetime go out of scope.
456/// assert_eq!(Some(&0), focus1.get(0));
457/// ```
458///
459/// On the other hand, you can split that one focus into multiple sub-focuses,
460/// which is safe because they can't overlap:
461///
462/// ```rust
463/// # #[macro_use] extern crate imbl;
464/// # use imbl::vector::Vector;
465/// # use std::iter::FromIterator;
466/// let mut vec: Vector<i64> = Vector::from_iter(0..1000);
467/// let focus = vec.focus_mut();
468/// let (mut left, mut right) = focus.split_at(500);
469/// assert_eq!(Some(&0), left.get(0));
470/// assert_eq!(Some(&500), right.get(0));
471/// ```
472///
473/// These sub-foci also work as a lock on the vector, even if the focus they
474/// were created from goes out of scope.
475///
476/// ```rust,compile_fail
477/// # #[macro_use] extern crate imbl;
478/// # use imbl::vector::Vector;
479/// # use std::iter::FromIterator;
480/// let mut vec: Vector<i64> = Vector::from_iter(0..1000);
481/// let (left, right) = {
482///     let focus = vec.focus_mut();
483///     focus.split_at(500)
484/// };
485/// // `left` and `right` are still in scope even if `focus` isn't, so we can't
486/// // create another focus:
487/// let focus2 = vec.focus_mut();
488/// assert_eq!(Some(&0), left.get(0));
489/// ```
490///
491/// [Focus]: enum.Focus.html
492pub enum FocusMut<'a, A, P: SharedPointerKind> {
493    #[doc(hidden)]
494    /// The Single variant is a focusmut of a simple Vector that can be represented as a single slice.
495    Single(&'a mut [A]),
496    #[doc(hidden)]
497    /// The Full variant is a focus of a more complex Vector that cannot be represented as a single slice.
498    Full(TreeFocusMut<'a, A, P>),
499}
500
501impl<'a, A, P: SharedPointerKind> FocusMut<'a, A, P>
502where
503    A: 'a,
504{
505    /// Get the length of the focused `Vector`.
506    pub fn len(&self) -> usize {
507        match self {
508            FocusMut::Single(chunk) => chunk.len(),
509            FocusMut::Full(tree) => tree.len(),
510        }
511    }
512
513    /// Test if the focused `Vector` is empty.
514    pub fn is_empty(&self) -> bool {
515        self.len() == 0
516    }
517
518    /// Narrow the focus onto a subslice of the vector.
519    ///
520    /// `FocusMut::narrow(range)` has the same effect as `&slice[range]`, without
521    /// actually modifying the underlying vector.
522    ///
523    /// Panics if the range isn't fully inside the current focus.
524    ///
525    /// ## Examples
526    ///
527    /// ```rust
528    /// # #[macro_use] extern crate imbl;
529    /// # use imbl::vector::Vector;
530    /// # use std::iter::FromIterator;
531    /// let mut vec: Vector<i64> = Vector::from_iter(0..1000);
532    /// let narrowed = vec.focus_mut().narrow(100..200);
533    /// let narrowed_vec: Vector<i64> = narrowed.unmut().into_iter().cloned().collect();
534    /// assert_eq!(Vector::from_iter(100..200), narrowed_vec);
535    /// ```
536    ///
537    /// [slice::split_at]: https://doc.rust-lang.org/std/primitive.slice.html#method.split_at
538    /// [Vector::split_at]: type.Vector.html#method.split_at
539    #[must_use]
540    pub fn narrow<R>(self, range: R) -> Self
541    where
542        R: RangeBounds<usize>,
543    {
544        let r = to_range(&range, self.len());
545        if r.start > r.end || r.start > self.len() {
546            panic!("vector::FocusMut::narrow: range out of bounds");
547        }
548        match self {
549            FocusMut::Single(chunk) => FocusMut::Single(&mut chunk[r]),
550            FocusMut::Full(tree) => FocusMut::Full(tree.narrow(r)),
551        }
552    }
553
554    /// Split the focus into two.
555    ///
556    /// Given an index `index`, consume the focus and produce two new foci, the
557    /// left onto indices `0..index`, and the right onto indices `index..N`
558    /// where `N` is the length of the current focus.
559    ///
560    /// Panics if the index is out of bounds.
561    ///
562    /// This is the moral equivalent of [`slice::split_at`][slice::split_at], in
563    /// that it leaves the underlying data structure unchanged, unlike
564    /// [`Vector::split_at`][Vector::split_at].
565    ///
566    /// ## Examples
567    ///
568    /// ```rust
569    /// # #[macro_use] extern crate imbl;
570    /// # use imbl::vector::Vector;
571    /// # use std::iter::FromIterator;
572    /// let mut vec: Vector<i64> = Vector::from_iter(0..1000);
573    /// {
574    ///     let (left, right) = vec.focus_mut().split_at(500);
575    ///     for ptr in left {
576    ///         *ptr += 100;
577    ///     }
578    ///     for ptr in right {
579    ///         *ptr -= 100;
580    ///     }
581    /// }
582    /// let expected = Vector::from_iter(100..600)
583    ///              + Vector::from_iter(400..900);
584    /// assert_eq!(expected, vec);
585    /// ```
586    ///
587    /// [slice::split_at]: https://doc.rust-lang.org/std/primitive.slice.html#method.split_at
588    /// [Vector::split_at]: type.Vector.html#method.split_at
589    #[allow(clippy::redundant_clone)]
590    pub fn split_at(self, index: usize) -> (Self, Self) {
591        if index > self.len() {
592            panic!("vector::FocusMut::split_at: index out of bounds");
593        }
594        match self {
595            FocusMut::Single(chunk) => {
596                let (left, right) = chunk.split_at_mut(index);
597                (FocusMut::Single(left), FocusMut::Single(right))
598            }
599            FocusMut::Full(tree) => {
600                let (left, right) = tree.split_at(index);
601                (FocusMut::Full(left), FocusMut::Full(right))
602            }
603        }
604    }
605
606    /// Convert a `FocusMut` into a `Focus`.
607    pub fn unmut(self) -> Focus<'a, A, P> {
608        match self {
609            FocusMut::Single(chunk) => Focus::Single(chunk),
610            FocusMut::Full(mut tree) => Focus::Full(TreeFocus {
611                tree: {
612                    let t = tree.tree.lock().unwrap();
613                    (*t).clone()
614                },
615                view: tree.view.clone(),
616                middle_range: tree.middle_range.clone(),
617                target_range: 0..0,
618                target_ptr: null(),
619            }),
620        }
621    }
622}
623
624impl<'a, A, P: SharedPointerKind> FocusMut<'a, A, P>
625where
626    A: Clone + 'a,
627{
628    /// Construct a `FocusMut` for a `Vector`.
629    pub fn new(vector: &'a mut GenericVector<A, P>) -> Self {
630        match &mut vector.vector {
631            Inline(chunk) => FocusMut::Single(chunk),
632            Single(chunk) => FocusMut::Single(SharedPointer::make_mut(chunk).as_mut_slice()),
633            Full(tree) => FocusMut::Full(TreeFocusMut::new(tree)),
634        }
635    }
636
637    /// Get a reference to the value at a given index.
638    pub fn get(&mut self, index: usize) -> Option<&A> {
639        self.get_mut(index).map(|r| &*r)
640    }
641
642    /// Get a mutable reference to the value at a given index.
643    pub fn get_mut(&mut self, index: usize) -> Option<&mut A> {
644        match self {
645            FocusMut::Single(chunk) => chunk.get_mut(index),
646            FocusMut::Full(tree) => tree.get(index),
647        }
648    }
649
650    fn get_many_mut<const N: usize>(&mut self, indices: [usize; N]) -> Option<[&mut A; N]> {
651        check_indices(self.len(), &indices)?;
652        match self {
653            FocusMut::Single(chunk) => {
654                // FIXME: Stable polyfill for std `get_many_mut`
655                let chunk: *mut A = (*chunk).as_mut_ptr();
656                Some(indices.map(|index| {
657                    // Safety:
658                    // - `check_indices` ensures each index is `< self.len()`, which for `FocusMut::Single` is `chunk.len()`
659                    // - `check_indices` ensures the indexes do not overlap
660                    unsafe { &mut *chunk.add(index) }
661                }))
662            }
663            FocusMut::Full(tree) => tree.get_many(indices),
664        }
665    }
666
667    /// Get a reference to the value at a given index.
668    ///
669    /// Panics if the index is out of bounds.
670    pub fn index(&mut self, index: usize) -> &A {
671        &*self.index_mut(index)
672    }
673
674    /// Get a mutable reference to the value at a given index.
675    ///
676    /// Panics if the index is out of bounds.
677    #[allow(clippy::should_implement_trait)] // would if I could
678    pub fn index_mut(&mut self, index: usize) -> &mut A {
679        self.get_mut(index).expect("index out of bounds")
680    }
681
682    /// Gets mutable references for a non-overlapping collection of indices.
683    ///
684    /// Panics if any indices are non-unique.
685    fn index_many_mut<const N: usize>(&mut self, indices: [usize; N]) -> [&mut A; N] {
686        self.get_many_mut(indices)
687            .expect("index out of bounds or overlapping")
688    }
689
690    /// Update the value at a given index.
691    ///
692    /// Returns `None` if the index is out of bounds, or the replaced value
693    /// otherwise.
694    pub fn set(&mut self, index: usize, value: A) -> Option<A> {
695        self.get_mut(index).map(|pos| replace(pos, value))
696    }
697
698    /// Swap the values at two given indices.
699    ///
700    /// Panics if either index is out of bounds.
701    ///
702    /// If the indices are equal, this function returns without doing anything.
703    pub fn swap(&mut self, a: usize, b: usize) {
704        if a == b {
705            return;
706        }
707        self.pair(a, b, |left, right| swap(left, right));
708    }
709
710    /// Lookup two indices simultaneously and run a function over them.
711    ///
712    /// Useful because the borrow checker won't let you have more than one
713    /// mutable reference into the same data structure at any given time.
714    ///
715    /// Panics if either index is out of bounds, or if they are the same index.
716    ///
717    /// # Examples
718    ///
719    /// ```rust
720    /// # #[macro_use] extern crate imbl;
721    /// # use imbl::vector::Vector;
722    /// # use std::iter::FromIterator;
723    /// let mut vec: Vector<i64> = vector![1, 2, 3, 4, 5];
724    /// vec.focus_mut().pair(1, 3, |a, b| *a += *b);
725    /// assert_eq!(vector![1, 6, 3, 4, 5], vec);
726    /// ```
727    pub fn pair<F, B>(&mut self, a: usize, b: usize, mut f: F) -> B
728    where
729        F: FnMut(&mut A, &mut A) -> B,
730    {
731        if a == b {
732            panic!("vector::FocusMut::pair: indices cannot be equal!");
733        }
734        let [pa, pb] = self.index_many_mut([a, b]);
735        f(pa, pb)
736    }
737
738    /// Lookup three indices simultaneously and run a function over them.
739    ///
740    /// Useful because the borrow checker won't let you have more than one
741    /// mutable reference into the same data structure at any given time.
742    ///
743    /// Panics if any index is out of bounds, or if any indices are equal.
744    ///
745    /// # Examples
746    ///
747    /// ```rust
748    /// # #[macro_use] extern crate imbl;
749    /// # use imbl::vector::Vector;
750    /// # use std::iter::FromIterator;
751    /// let mut vec: Vector<i64> = vector![1, 2, 3, 4, 5];
752    /// vec.focus_mut().triplet(0, 2, 4, |a, b, c| *a += *b + *c);
753    /// assert_eq!(vector![9, 2, 3, 4, 5], vec);
754    /// ```
755    pub fn triplet<F, B>(&mut self, a: usize, b: usize, c: usize, mut f: F) -> B
756    where
757        F: FnMut(&mut A, &mut A, &mut A) -> B,
758    {
759        if a == b || b == c || a == c {
760            panic!("vector::FocusMut::triplet: indices cannot be equal!");
761        }
762        let [pa, pb, pc] = self.index_many_mut([a, b, c]);
763        f(pa, pb, pc)
764    }
765
766    /// Get the chunk for the given index.
767    ///
768    /// This gives you a reference to the leaf node that contains the index,
769    /// along with its start and end indices.
770    pub fn chunk_at(&mut self, index: usize) -> (Range<usize>, &mut [A]) {
771        let len = self.len();
772        if index >= len {
773            panic!("vector::FocusMut::chunk_at: index out of bounds");
774        }
775        match self {
776            FocusMut::Single(chunk) => (0..len, chunk),
777            FocusMut::Full(tree) => {
778                let (range, chunk) = tree.get_chunk(index);
779                (range, chunk)
780            }
781        }
782    }
783}
784
785impl<'a, A, P: SharedPointerKind> IntoIterator for FocusMut<'a, A, P>
786where
787    A: Clone + 'a,
788{
789    type Item = &'a mut A;
790    type IntoIter = IterMut<'a, A, P>;
791
792    fn into_iter(self) -> Self::IntoIter {
793        IterMut::from_focus(self)
794    }
795}
796
797impl<'a, A, P: SharedPointerKind> From<FocusMut<'a, A, P>> for Focus<'a, A, P>
798where
799    A: Clone + 'a,
800{
801    fn from(f: FocusMut<'a, A, P>) -> Focus<'a, A, P> {
802        f.unmut()
803    }
804}
805
806// NOTE: The documentation the mutable version is similar to the non-mutable version. I will comment for the places
807// where there are differences, otherwise the documentation is copied directly.
808pub struct TreeFocusMut<'a, A, P: SharedPointerKind> {
809    /// The tree that this TreeFocusMut refers to. Unlike the non-mutable version, TreeFocusMut needs to store a
810    /// mutable reference. Additionally, there may be multiple TreeFocusMuts that refer to the same tree so we need a
811    /// Lock to synchronise the changes.
812    tree: Lock<&'a mut RRB<A, P>>,
813    /// The view represents the range of the tree that this TreeFocusMut can see. The view can be narrowed by calling
814    /// either the narrow or split_at methods.
815    view: Range<usize>,
816    /// The tree version of the Vector is represented as the concatenation of 2 chunks, followed by the tree root,
817    /// followed by 2 chunks. The middle_range refers to the range of the Vector that the tree covers.
818    middle_range: Range<usize>,
819    /// This implementation of a focusmut stores only a single chunk for the Vector. This chunk can refer to one of the
820    /// 4 chunks front/back chunks or one of the leaves of the tree. The target_ptr is the pointer to the actual chunk
821    /// in question. The target_range is the range that the chunk represents.
822    target_range: Range<usize>,
823    /// Not actually sure why this needs to be an atomic, it seems like it is unneccessary. This is just a pointer to
824    /// the chunk referred to above.
825    target_ptr: AtomicPtr<Chunk<A>>,
826}
827
828impl<'a, A, P: SharedPointerKind> TreeFocusMut<'a, A, P>
829where
830    A: 'a,
831{
832    /// Creates a new TreeFocusMut for a Vector's RRB tree.
833    fn new(tree: &'a mut RRB<A, P>) -> Self {
834        let middle_start = tree.outer_f.len() + tree.inner_f.len();
835        let middle_end = middle_start + tree.middle.len();
836        TreeFocusMut {
837            view: 0..tree.length,
838            tree: Lock::new(tree),
839            middle_range: middle_start..middle_end,
840            target_range: 0..0,
841            target_ptr: AtomicPtr::default(),
842        }
843    }
844
845    /// Returns the number of elements that the TreeFocusMut is valid for.
846    fn len(&self) -> usize {
847        self.view.end - self.view.start
848    }
849
850    /// Restricts the TreeFocusMut to a subrange of itself.
851    fn narrow(self, mut view: Range<usize>) -> Self {
852        view.start += self.view.start;
853        view.end += self.view.start;
854        TreeFocusMut {
855            view,
856            middle_range: self.middle_range.clone(),
857            target_range: 0..0,
858            target_ptr: AtomicPtr::default(),
859            tree: self.tree,
860        }
861    }
862
863    /// Splits the TreeFocusMut into two disjoint foci. The first TreeFocusMut is valid for ..index while the
864    /// second is valid for index.. .
865    fn split_at(self, index: usize) -> (Self, Self) {
866        let len = self.len();
867        debug_assert!(index <= len);
868        let left = TreeFocusMut {
869            view: self.view.start..(self.view.start + index),
870            middle_range: self.middle_range.clone(),
871            target_range: 0..0,
872            target_ptr: AtomicPtr::default(),
873            tree: self.tree.clone(),
874        };
875        let right = TreeFocusMut {
876            view: (self.view.start + index)..(self.view.start + len),
877            middle_range: self.middle_range.clone(),
878            target_range: 0..0,
879            target_ptr: AtomicPtr::default(),
880            tree: self.tree,
881        };
882        (left, right)
883    }
884
885    /// Computes an absolute index in the RRBTree for the given index relative to the start of this TreeFocusMut.
886    fn physical_index(&self, index: usize) -> usize {
887        debug_assert!(index < self.view.end);
888        self.view.start + index
889    }
890
891    /// Computes a range relative to the TreeFocusMut given one that is absolute in the RRBTree.
892    fn logical_range(&self, range: &Range<usize>) -> Range<usize> {
893        (range.start - self.view.start)..(range.end - self.view.start)
894    }
895
896    /// Gets the chunk for an index and its corresponding range within the TreeFocusMut.
897    fn get_focus(&mut self) -> &mut Chunk<A> {
898        unsafe { &mut *self.target_ptr.load(Ordering::Relaxed) }
899    }
900
901    fn get_focus_ptr(&mut self) -> *mut Chunk<A> {
902        self.target_ptr.load(Ordering::Relaxed)
903    }
904}
905
906impl<'a, A, P: SharedPointerKind> TreeFocusMut<'a, A, P>
907where
908    A: Clone + 'a,
909{
910    /// Sets the internal chunk to the one that contains the given absolute index.
911    fn set_focus(&mut self, index: usize) {
912        let mut tree = self
913            .tree
914            .lock()
915            .expect("imbl::vector::Focus::set_focus: unable to acquire exclusive lock on Vector");
916        if index < self.middle_range.start {
917            let outer_len = tree.outer_f.len();
918            if index < outer_len {
919                self.target_range = 0..outer_len;
920                self.target_ptr.store(
921                    SharedPointer::make_mut(&mut tree.outer_f),
922                    Ordering::Relaxed,
923                );
924            } else {
925                self.target_range = outer_len..self.middle_range.start;
926                self.target_ptr.store(
927                    SharedPointer::make_mut(&mut tree.inner_f),
928                    Ordering::Relaxed,
929                );
930            }
931        } else if index >= self.middle_range.end {
932            let outer_start = self.middle_range.end + tree.inner_b.len();
933            if index < outer_start {
934                self.target_range = self.middle_range.end..outer_start;
935                self.target_ptr.store(
936                    SharedPointer::make_mut(&mut tree.inner_b),
937                    Ordering::Relaxed,
938                );
939            } else {
940                self.target_range = outer_start..tree.length;
941                self.target_ptr.store(
942                    SharedPointer::make_mut(&mut tree.outer_b),
943                    Ordering::Relaxed,
944                );
945            }
946        } else {
947            let tree_index = index - self.middle_range.start;
948            let level = tree.middle_level;
949            let middle = SharedPointer::make_mut(&mut tree.middle);
950            let (range, ptr) = middle.lookup_chunk_mut(level, 0, tree_index);
951            self.target_range =
952                (range.start + self.middle_range.start)..(range.end + self.middle_range.start);
953            self.target_ptr.store(ptr, Ordering::Relaxed);
954        }
955    }
956
957    /// Gets the value at the given index relative to the TreeFocusMut.
958    pub fn get(&mut self, index: usize) -> Option<&mut A> {
959        if index >= self.len() {
960            return None;
961        }
962        let phys_index = self.physical_index(index);
963        if !contains(&self.target_range, &phys_index) {
964            self.set_focus(phys_index);
965        }
966        let target_phys_index = phys_index - self.target_range.start;
967        Some(&mut self.get_focus()[target_phys_index])
968    }
969
970    fn get_many<const N: usize>(&mut self, indices: [usize; N]) -> Option<[&mut A; N]> {
971        check_indices(self.len(), &indices)?;
972        Some(indices.map(|phys_idx| {
973            if !contains(&self.target_range, &phys_idx) {
974                self.set_focus(phys_idx);
975            }
976            let target_idx = phys_idx - self.target_range.start;
977            // Safety: we have called `set_focus` to get a valid chunk pointer
978            // and `target_idx` lies within it
979            unsafe {
980                let chunk = self.get_focus_ptr();
981                let ptr: *mut [A] = Chunk::as_mut_slice_ptr(chunk);
982                &mut *ptr.cast::<A>().add(target_idx)
983            }
984        }))
985    }
986
987    /// Gets the chunk for an index as a slice and its corresponding range within the TreeFocusMut.
988    pub fn get_chunk(&mut self, index: usize) -> (Range<usize>, &mut [A]) {
989        let phys_index = self.physical_index(index);
990        if !contains(&self.target_range, &phys_index) {
991            self.set_focus(phys_index);
992        }
993        let mut left = 0;
994        let mut right = 0;
995        if self.target_range.start < self.view.start {
996            left = self.view.start - self.target_range.start;
997        }
998        if self.target_range.end > self.view.end {
999            right = self.target_range.end - self.view.end;
1000        }
1001        let phys_range = (self.target_range.start + left)..(self.target_range.end - right);
1002        let log_range = self.logical_range(&phys_range);
1003        let slice_len = self.get_focus().len();
1004        let slice = &mut self.get_focus().as_mut_slice()[left..(slice_len - right)];
1005        (log_range, slice)
1006    }
1007}