@@ -125,6 +125,9 @@ use core::marker::PhantomData;
125
125
#[ cfg( feature = "write" ) ]
126
126
use std:: io;
127
127
128
+ #[ cfg( feature = "drain_keep_rest" ) ]
129
+ use core:: mem:: ManuallyDrop ;
130
+
128
131
/// Creates a [`SmallVec`] containing the arguments.
129
132
///
130
133
/// `smallvec!` allows `SmallVec`s to be defined with the same syntax as array expressions.
@@ -410,6 +413,198 @@ impl<'a, T: 'a + Array> Drop for Drain<'a, T> {
410
413
}
411
414
}
412
415
416
+ #[ cfg( feature = "drain_filter" ) ]
417
+ /// An iterator which uses a closure to determine if an element should be removed.
418
+ ///
419
+ /// Returned from [`SmallVec::drain_filter`][1].
420
+ ///
421
+ /// [1]: struct.SmallVec.html#method.drain_filter
422
+ pub struct DrainFilter < ' a , T , F >
423
+ where
424
+ F : FnMut ( & mut T :: Item ) -> bool ,
425
+ T : Array ,
426
+ {
427
+ vec : & ' a mut SmallVec < T > ,
428
+ /// The index of the item that will be inspected by the next call to `next`.
429
+ idx : usize ,
430
+ /// The number of items that have been drained (removed) thus far.
431
+ del : usize ,
432
+ /// The original length of `vec` prior to draining.
433
+ old_len : usize ,
434
+ /// The filter test predicate.
435
+ pred : F ,
436
+ /// A flag that indicates a panic has occurred in the filter test predicate.
437
+ /// This is used as a hint in the drop implementation to prevent consumption
438
+ /// of the remainder of the `DrainFilter`. Any unprocessed items will be
439
+ /// backshifted in the `vec`, but no further items will be dropped or
440
+ /// tested by the filter predicate.
441
+ panic_flag : bool ,
442
+ }
443
+
444
+ #[ cfg( feature = "drain_filter" ) ]
445
+ impl < T , F > fmt:: Debug for DrainFilter < ' _ , T , F >
446
+ where
447
+ F : FnMut ( & mut T :: Item ) -> bool ,
448
+ T : Array ,
449
+ T :: Item : fmt:: Debug ,
450
+ {
451
+ fn fmt ( & self , f : & mut fmt:: Formatter < ' _ > ) -> fmt:: Result {
452
+ f. debug_tuple ( "DrainFilter" ) . field ( & self . vec . as_slice ( ) ) . finish ( )
453
+ }
454
+ }
455
+
456
+ #[ cfg( feature = "drain_filter" ) ]
457
+ impl < T , F > Iterator for DrainFilter < ' _ , T , F >
458
+ where
459
+ F : FnMut ( & mut T :: Item ) -> bool ,
460
+ T : Array ,
461
+ {
462
+ type Item = T :: Item ;
463
+
464
+ fn next ( & mut self ) -> Option < T :: Item >
465
+ {
466
+ unsafe {
467
+ while self . idx < self . old_len {
468
+ let i = self . idx ;
469
+ let v = slice:: from_raw_parts_mut ( self . vec . as_mut_ptr ( ) , self . old_len ) ;
470
+ self . panic_flag = true ;
471
+ let drained = ( self . pred ) ( & mut v[ i] ) ;
472
+ self . panic_flag = false ;
473
+ // Update the index *after* the predicate is called. If the index
474
+ // is updated prior and the predicate panics, the element at this
475
+ // index would be leaked.
476
+ self . idx += 1 ;
477
+ if drained {
478
+ self . del += 1 ;
479
+ return Some ( ptr:: read ( & v[ i] ) ) ;
480
+ } else if self . del > 0 {
481
+ let del = self . del ;
482
+ let src: * const Self :: Item = & v[ i] ;
483
+ let dst: * mut Self :: Item = & mut v[ i - del] ;
484
+ ptr:: copy_nonoverlapping ( src, dst, 1 ) ;
485
+ }
486
+ }
487
+ None
488
+ }
489
+ }
490
+
491
+ fn size_hint ( & self ) -> ( usize , Option < usize > ) {
492
+ ( 0 , Some ( self . old_len - self . idx ) )
493
+ }
494
+ }
495
+
496
+ #[ cfg( feature = "drain_filter" ) ]
497
+ impl < T , F > Drop for DrainFilter < ' _ , T , F >
498
+ where
499
+ F : FnMut ( & mut T :: Item ) -> bool ,
500
+ T : Array ,
501
+ {
502
+ fn drop ( & mut self ) {
503
+ struct BackshiftOnDrop < ' a , ' b , T , F >
504
+ where
505
+ F : FnMut ( & mut T :: Item ) -> bool ,
506
+ T : Array
507
+ {
508
+ drain : & ' b mut DrainFilter < ' a , T , F > ,
509
+ }
510
+
511
+ impl < ' a , ' b , T , F > Drop for BackshiftOnDrop < ' a , ' b , T , F >
512
+ where
513
+ F : FnMut ( & mut T :: Item ) -> bool ,
514
+ T : Array
515
+ {
516
+ fn drop ( & mut self ) {
517
+ unsafe {
518
+ if self . drain . idx < self . drain . old_len && self . drain . del > 0 {
519
+ // This is a pretty messed up state, and there isn't really an
520
+ // obviously right thing to do. We don't want to keep trying
521
+ // to execute `pred`, so we just backshift all the unprocessed
522
+ // elements and tell the vec that they still exist. The backshift
523
+ // is required to prevent a double-drop of the last successfully
524
+ // drained item prior to a panic in the predicate.
525
+ let ptr = self . drain . vec . as_mut_ptr ( ) ;
526
+ let src = ptr. add ( self . drain . idx ) ;
527
+ let dst = src. sub ( self . drain . del ) ;
528
+ let tail_len = self . drain . old_len - self . drain . idx ;
529
+ src. copy_to ( dst, tail_len) ;
530
+ }
531
+ self . drain . vec . set_len ( self . drain . old_len - self . drain . del ) ;
532
+ }
533
+ }
534
+ }
535
+
536
+ let backshift = BackshiftOnDrop { drain : self } ;
537
+
538
+ // Attempt to consume any remaining elements if the filter predicate
539
+ // has not yet panicked. We'll backshift any remaining elements
540
+ // whether we've already panicked or if the consumption here panics.
541
+ if !backshift. drain . panic_flag {
542
+ backshift. drain . for_each ( drop) ;
543
+ }
544
+ }
545
+ }
546
+
547
+ #[ cfg( feature = "drain_keep_rest" ) ]
548
+ impl < T , F > DrainFilter < ' _ , T , F >
549
+ where
550
+ F : FnMut ( & mut T :: Item ) -> bool ,
551
+ T : Array
552
+ {
553
+ /// Keep unyielded elements in the source `Vec`.
554
+ ///
555
+ /// # Examples
556
+ ///
557
+ /// ```
558
+ /// # use smallvec::{smallvec, SmallVec};
559
+ ///
560
+ /// let mut vec: SmallVec<[char; 2]> = smallvec!['a', 'b', 'c'];
561
+ /// let mut drain = vec.drain_filter(|_| true);
562
+ ///
563
+ /// assert_eq!(drain.next().unwrap(), 'a');
564
+ ///
565
+ /// // This call keeps 'b' and 'c' in the vec.
566
+ /// drain.keep_rest();
567
+ ///
568
+ /// // If we wouldn't call `keep_rest()`,
569
+ /// // `vec` would be empty.
570
+ /// assert_eq!(vec, SmallVec::<[char; 2]>::from_slice(&['b', 'c']));
571
+ /// ```
572
+ pub fn keep_rest ( self )
573
+ {
574
+ // At this moment layout looks like this:
575
+ //
576
+ // _____________________/-- old_len
577
+ // / \
578
+ // [kept] [yielded] [tail]
579
+ // \_______/ ^-- idx
580
+ // \-- del
581
+ //
582
+ // Normally `Drop` impl would drop [tail] (via .for_each(drop), ie still calling `pred`)
583
+ //
584
+ // 1. Move [tail] after [kept]
585
+ // 2. Update length of the original vec to `old_len - del`
586
+ // a. In case of ZST, this is the only thing we want to do
587
+ // 3. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
588
+ let mut this = ManuallyDrop :: new ( self ) ;
589
+
590
+ unsafe {
591
+ // ZSTs have no identity, so we don't need to move them around.
592
+ let needs_move = mem:: size_of :: < T > ( ) != 0 ;
593
+
594
+ if needs_move && this. idx < this. old_len && this. del > 0 {
595
+ let ptr = this. vec . as_mut_ptr ( ) ;
596
+ let src = ptr. add ( this. idx ) ;
597
+ let dst = src. sub ( this. del ) ;
598
+ let tail_len = this. old_len - this. idx ;
599
+ src. copy_to ( dst, tail_len) ;
600
+ }
601
+
602
+ let new_len = this. old_len - this. del ;
603
+ this. vec . set_len ( new_len) ;
604
+ }
605
+ }
606
+ }
607
+
413
608
#[ cfg( feature = "union" ) ]
414
609
union SmallVecData < A : Array > {
415
610
inline : core:: mem:: ManuallyDrop < MaybeUninit < A > > ,
@@ -847,6 +1042,65 @@ impl<A: Array> SmallVec<A> {
847
1042
}
848
1043
}
849
1044
1045
+
1046
+ #[ cfg( feature = "drain_filter" ) ]
1047
+ /// Creates an iterator which uses a closure to determine if an element should be removed.
1048
+ ///
1049
+ /// If the closure returns true, the element is removed and yielded. If the closure returns
1050
+ /// false, the element will remain in the vector and will not be yielded by the iterator.
1051
+ ///
1052
+ /// Using this method is equivalent to the following code:
1053
+ /// ```
1054
+ /// # use smallvec::SmallVec;
1055
+ /// # let some_predicate = |x: &mut i32| { *x == 2 || *x == 3 || *x == 6 };
1056
+ /// # let mut vec: SmallVec<[i32; 8]> = SmallVec::from_slice(&[1i32, 2, 3, 4, 5, 6]);
1057
+ /// let mut i = 0;
1058
+ /// while i < vec.len() {
1059
+ /// if some_predicate(&mut vec[i]) {
1060
+ /// let val = vec.remove(i);
1061
+ /// // your code here
1062
+ /// } else {
1063
+ /// i += 1;
1064
+ /// }
1065
+ /// }
1066
+ ///
1067
+ /// # assert_eq!(vec, SmallVec::<[i32; 8]>::from_slice(&[1i32, 4, 5]));
1068
+ /// ```
1069
+ /// ///
1070
+ /// But `drain_filter` is easier to use. `drain_filter` is also more efficient,
1071
+ /// because it can backshift the elements of the array in bulk.
1072
+ ///
1073
+ /// Note that `drain_filter` also lets you mutate every element in the filter closure,
1074
+ /// regardless of whether you choose to keep or remove it.
1075
+ ///
1076
+ /// # Examples
1077
+ ///
1078
+ /// Splitting an array into evens and odds, reusing the original allocation:
1079
+ ///
1080
+ /// ```
1081
+ /// # use smallvec::SmallVec;
1082
+ /// let mut numbers: SmallVec<[i32; 16]> = SmallVec::from_slice(&[1i32, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15]);
1083
+ ///
1084
+ /// let evens = numbers.drain_filter(|x| *x % 2 == 0).collect::<SmallVec<[i32; 16]>>();
1085
+ /// let odds = numbers;
1086
+ ///
1087
+ /// assert_eq!(evens, SmallVec::<[i32; 16]>::from_slice(&[2i32, 4, 6, 8, 14]));
1088
+ /// assert_eq!(odds, SmallVec::<[i32; 16]>::from_slice(&[1i32, 3, 5, 9, 11, 13, 15]));
1089
+ /// ```
1090
+ pub fn drain_filter < F > ( & mut self , filter : F ) -> DrainFilter < ' _ , A , F , >
1091
+ where
1092
+ F : FnMut ( & mut A :: Item ) -> bool ,
1093
+ {
1094
+ let old_len = self . len ( ) ;
1095
+
1096
+ // Guard against us getting leaked (leak amplification)
1097
+ unsafe {
1098
+ self . set_len ( 0 ) ;
1099
+ }
1100
+
1101
+ DrainFilter { vec : self , idx : 0 , del : 0 , old_len, pred : filter, panic_flag : false }
1102
+ }
1103
+
850
1104
/// Append an item to the vector.
851
1105
#[ inline]
852
1106
pub fn push ( & mut self , value : A :: Item ) {
0 commit comments