@@ -380,6 +380,178 @@ impl<'a, T: 'a, const N: usize> Drop for Drain<'a, T, N> {
380
380
}
381
381
}
382
382
383
+ #[ cfg( feature = "drain_filter" ) ]
384
+ /// An iterator which uses a closure to determine if an element should be removed.
385
+ ///
386
+ /// Returned from [`SmallVec::drain_filter`][1].
387
+ ///
388
+ /// [1]: struct.SmallVec.html#method.drain_filter
389
+ pub struct DrainFilter < ' a , T , const N : usize , F >
390
+ where
391
+ F : FnMut ( & mut T ) -> bool ,
392
+ {
393
+ vec : & ' a mut SmallVec < T , N > ,
394
+ /// The index of the item that will be inspected by the next call to `next`.
395
+ idx : usize ,
396
+ /// The number of items that have been drained (removed) thus far.
397
+ del : usize ,
398
+ /// The original length of `vec` prior to draining.
399
+ old_len : usize ,
400
+ /// The filter test predicate.
401
+ pred : F ,
402
+ /// A flag that indicates a panic has occurred in the filter test predicate.
403
+ /// This is used as a hint in the drop implementation to prevent consumption
404
+ /// of the remainder of the `DrainFilter`. Any unprocessed items will be
405
+ /// backshifted in the `vec`, but no further items will be dropped or
406
+ /// tested by the filter predicate.
407
+ panic_flag : bool ,
408
+ }
409
+
410
+ #[ cfg( feature = "drain_filter" ) ]
411
+ impl < T , F , const N : usize > Iterator for DrainFilter < ' _ , T , N , F >
412
+ where
413
+ F : FnMut ( & mut T ) -> bool ,
414
+ {
415
+ type Item = T ;
416
+
417
+ fn next ( & mut self ) -> Option < T > {
418
+ unsafe {
419
+ while self . idx < self . old_len {
420
+ let i = self . idx ;
421
+ let v = core:: slice:: from_raw_parts_mut ( self . vec . as_mut_ptr ( ) , self . old_len ) ;
422
+ self . panic_flag = true ;
423
+ let drained = ( self . pred ) ( & mut v[ i] ) ;
424
+ self . panic_flag = false ;
425
+ // Update the index *after* the predicate is called. If the index
426
+ // is updated prior and the predicate panics, the element at this
427
+ // index would be leaked.
428
+ self . idx += 1 ;
429
+ if drained {
430
+ self . del += 1 ;
431
+ return Some ( core:: ptr:: read ( & v[ i] ) ) ;
432
+ } else if self . del > 0 {
433
+ let del = self . del ;
434
+ let src: * const Self :: Item = & v[ i] ;
435
+ let dst: * mut Self :: Item = & mut v[ i - del] ;
436
+ core:: ptr:: copy_nonoverlapping ( src, dst, 1 ) ;
437
+ }
438
+ }
439
+ None
440
+ }
441
+ }
442
+
443
+ fn size_hint ( & self ) -> ( usize , Option < usize > ) {
444
+ ( 0 , Some ( self . old_len - self . idx ) )
445
+ }
446
+ }
447
+
448
+ #[ cfg( feature = "drain_filter" ) ]
449
+ impl < T , F , const N : usize > Drop for DrainFilter < ' _ , T , N , F >
450
+ where
451
+ F : FnMut ( & mut T ) -> bool ,
452
+ {
453
+ fn drop ( & mut self ) {
454
+ struct BackshiftOnDrop < ' a , ' b , T , const N : usize , F >
455
+ where
456
+ F : FnMut ( & mut T ) -> bool ,
457
+ {
458
+ drain : & ' b mut DrainFilter < ' a , T , N , F > ,
459
+ }
460
+
461
+ impl < ' a , ' b , T , const N : usize , F > Drop for BackshiftOnDrop < ' a , ' b , T , N , F >
462
+ where
463
+ F : FnMut ( & mut T ) -> bool ,
464
+ {
465
+ fn drop ( & mut self ) {
466
+ unsafe {
467
+ if self . drain . idx < self . drain . old_len && self . drain . del > 0 {
468
+ // This is a pretty messed up state, and there isn't really an
469
+ // obviously right thing to do. We don't want to keep trying
470
+ // to execute `pred`, so we just backshift all the unprocessed
471
+ // elements and tell the vec that they still exist. The backshift
472
+ // is required to prevent a double-drop of the last successfully
473
+ // drained item prior to a panic in the predicate.
474
+ let ptr = self . drain . vec . as_mut_ptr ( ) ;
475
+ let src = ptr. add ( self . drain . idx ) ;
476
+ let dst = src. sub ( self . drain . del ) ;
477
+ let tail_len = self . drain . old_len - self . drain . idx ;
478
+ src. copy_to ( dst, tail_len) ;
479
+ }
480
+ self . drain . vec . set_len ( self . drain . old_len - self . drain . del ) ;
481
+ }
482
+ }
483
+ }
484
+
485
+ let backshift = BackshiftOnDrop { drain : self } ;
486
+
487
+ // Attempt to consume any remaining elements if the filter predicate
488
+ // has not yet panicked. We'll backshift any remaining elements
489
+ // whether we've already panicked or if the consumption here panics.
490
+ if !backshift. drain . panic_flag {
491
+ backshift. drain . for_each ( drop) ;
492
+ }
493
+ }
494
+ }
495
+
496
+ #[ cfg( feature = "drain_keep_rest" ) ]
497
+ impl < T , F , const N : usize > DrainFilter < ' _ , T , N , F >
498
+ where
499
+ F : FnMut ( & mut T ) -> bool ,
500
+ {
501
+ /// Keep unyielded elements in the source `Vec`.
502
+ ///
503
+ /// # Examples
504
+ ///
505
+ /// ```
506
+ /// # use smallvec::{smallvec, SmallVec};
507
+ ///
508
+ /// let mut vec: SmallVec<char, 2> = smallvec!['a', 'b', 'c'];
509
+ /// let mut drain = vec.drain_filter(|_| true);
510
+ ///
511
+ /// assert_eq!(drain.next().unwrap(), 'a');
512
+ ///
513
+ /// // This call keeps 'b' and 'c' in the vec.
514
+ /// drain.keep_rest();
515
+ ///
516
+ /// // If we wouldn't call `keep_rest()`,
517
+ /// // `vec` would be empty.
518
+ /// assert_eq!(vec, SmallVec::<char, 2>::from_slice(&['b', 'c']));
519
+ /// ```
520
+ pub fn keep_rest ( self ) {
521
+ // At this moment layout looks like this:
522
+ //
523
+ // _____________________/-- old_len
524
+ // / \
525
+ // [kept] [yielded] [tail]
526
+ // \_______/ ^-- idx
527
+ // \-- del
528
+ //
529
+ // Normally `Drop` impl would drop [tail] (via .for_each(drop), ie still calling `pred`)
530
+ //
531
+ // 1. Move [tail] after [kept]
532
+ // 2. Update length of the original vec to `old_len - del`
533
+ // a. In case of ZST, this is the only thing we want to do
534
+ // 3. Do *not* drop self, as everything is put in a consistent state already, there is nothing to do
535
+ let mut this = ManuallyDrop :: new ( self ) ;
536
+
537
+ unsafe {
538
+ // ZSTs have no identity, so we don't need to move them around.
539
+ let needs_move = core:: mem:: size_of :: < T > ( ) != 0 ;
540
+
541
+ if needs_move && this. idx < this. old_len && this. del > 0 {
542
+ let ptr = this. vec . as_mut_ptr ( ) ;
543
+ let src = ptr. add ( this. idx ) ;
544
+ let dst = src. sub ( this. del ) ;
545
+ let tail_len = this. old_len - this. idx ;
546
+ src. copy_to ( dst, tail_len) ;
547
+ }
548
+
549
+ let new_len = this. old_len - this. del ;
550
+ this. vec . set_len ( new_len) ;
551
+ }
552
+ }
553
+ }
554
+
383
555
/// An iterator that consumes a `SmallVec` and yields its items by value.
384
556
///
385
557
/// Returned from [`SmallVec::into_iter`][1].
@@ -732,6 +904,71 @@ impl<T, const N: usize> SmallVec<T, N> {
732
904
}
733
905
}
734
906
907
+ #[ cfg( feature = "drain_filter" ) ]
908
+ /// Creates an iterator which uses a closure to determine if an element should be removed.
909
+ ///
910
+ /// If the closure returns true, the element is removed and yielded. If the closure returns
911
+ /// false, the element will remain in the vector and will not be yielded by the iterator.
912
+ ///
913
+ /// Using this method is equivalent to the following code:
914
+ /// ```
915
+ /// # use smallvec::SmallVec;
916
+ /// # let some_predicate = |x: &mut i32| { *x == 2 || *x == 3 || *x == 6 };
917
+ /// # let mut vec: SmallVec<i32, 8> = SmallVec::from_slice(&[1i32, 2, 3, 4, 5, 6]);
918
+ /// let mut i = 0;
919
+ /// while i < vec.len() {
920
+ /// if some_predicate(&mut vec[i]) {
921
+ /// let val = vec.remove(i);
922
+ /// // your code here
923
+ /// } else {
924
+ /// i += 1;
925
+ /// }
926
+ /// }
927
+ ///
928
+ /// # assert_eq!(vec, SmallVec::<i32, 8>::from_slice(&[1i32, 4, 5]));
929
+ /// ```
930
+ /// ///
931
+ /// But `drain_filter` is easier to use. `drain_filter` is also more efficient,
932
+ /// because it can backshift the elements of the array in bulk.
933
+ ///
934
+ /// Note that `drain_filter` also lets you mutate every element in the filter closure,
935
+ /// regardless of whether you choose to keep or remove it.
936
+ ///
937
+ /// # Examples
938
+ ///
939
+ /// Splitting an array into evens and odds, reusing the original allocation:
940
+ ///
941
+ /// ```
942
+ /// # use smallvec::SmallVec;
943
+ /// let mut numbers: SmallVec<i32, 16> = SmallVec::from_slice(&[1i32, 2, 3, 4, 5, 6, 8, 9, 11, 13, 14, 15]);
944
+ ///
945
+ /// let evens = numbers.drain_filter(|x| *x % 2 == 0).collect::<SmallVec<i32, 16>>();
946
+ /// let odds = numbers;
947
+ ///
948
+ /// assert_eq!(evens, SmallVec::<i32, 16>::from_slice(&[2i32, 4, 6, 8, 14]));
949
+ /// assert_eq!(odds, SmallVec::<i32, 16>::from_slice(&[1i32, 3, 5, 9, 11, 13, 15]));
950
+ /// ```
951
+ pub fn drain_filter < F > ( & mut self , filter : F ) -> DrainFilter < ' _ , T , N , F >
952
+ where
953
+ F : FnMut ( & mut T ) -> bool ,
954
+ {
955
+ let old_len = self . len ( ) ;
956
+
957
+ // Guard against us getting leaked (leak amplification)
958
+ unsafe {
959
+ self . set_len ( 0 ) ;
960
+ }
961
+
962
+ DrainFilter {
963
+ vec : self ,
964
+ idx : 0 ,
965
+ del : 0 ,
966
+ old_len,
967
+ pred : filter,
968
+ panic_flag : false ,
969
+ }
970
+ }
971
+
735
972
#[ inline]
736
973
pub fn push ( & mut self , value : T ) {
737
974
let len = self . len ( ) ;
0 commit comments