src/vendor/rayon/src/iter/find_first_last/mod.rs

   1 use std::cell::Cell;
   2 use std::sync::atomic::{AtomicUsize, Ordering};
   3 use super::plumbing::*;
   4 use super::*;
   5
   6 #[cfg(test)]
   7 mod test;
   8
   9 // The key optimization for find_first is that a consumer can stop its search if
  10 // some consumer to its left already found a match (and similarly for consumers
  11 // to the right for find_last). To make this work, all consumers need some
  12 // notion of their position in the data relative to other consumers, including
  13 // unindexed consumers that have no built-in notion of position.
  14 //
  15 // To solve this, we assign each consumer a lower and upper bound for an
  16 // imaginary "range" of data that it consumes. The initial consumer starts with
  17 // the range 0..usize::max_value(). The split divides this range in half so that
  18 // one resulting consumer has the range 0..(usize::max_value() / 2), and the
  19 // other has (usize::max_value() / 2)..usize::max_value(). Every subsequent
  20 // split divides the range in half again until it cannot be split anymore
  21 // (i.e. its length is 1), in which case the split returns two consumers with
  22 // the same range. In that case both consumers will continue to consume all
  23 // their data regardless of whether a better match is found, but the reducer
  24 // will still return the correct answer.
  25
  26 #[derive(Copy, Clone)]
  27 enum MatchPosition {
  28     Leftmost,
  29     Rightmost,
  30 }
  31
  32 /// Returns true if pos1 is a better match than pos2 according to MatchPosition
  33 #[inline]
  34 fn better_position(pos1: usize, pos2: usize, mp: MatchPosition) -> bool {
  35     match mp {
  36         MatchPosition::Leftmost => pos1 < pos2,
  37         MatchPosition::Rightmost => pos1 > pos2,
  38     }
  39 }
  40
  41 pub fn find_first<I, P>(pi: I, find_op: P) -> Option<I::Item>
  42     where I: ParallelIterator,
  43           P: Fn(&I::Item) -> bool + Sync
  44 {
  45     let best_found = AtomicUsize::new(usize::max_value());
  46     let consumer = FindConsumer::new(&find_op, MatchPosition::Leftmost, &best_found);
  47     pi.drive_unindexed(consumer)
  48 }
  49
  50 pub fn find_last<I, P>(pi: I, find_op: P) -> Option<I::Item>
  51     where I: ParallelIterator,
  52           P: Fn(&I::Item) -> bool + Sync
  53 {
  54     let best_found = AtomicUsize::new(0);
  55     let consumer = FindConsumer::new(&find_op, MatchPosition::Rightmost, &best_found);
  56     pi.drive_unindexed(consumer)
  57 }
  58
  59 struct FindConsumer<'p, P: 'p> {
  60     find_op: &'p P,
  61     lower_bound: Cell<usize>,
  62     upper_bound: usize,
  63     match_position: MatchPosition,
  64     best_found: &'p AtomicUsize,
  65 }
  66
  67 impl<'p, P> FindConsumer<'p, P> {
  68     fn new(find_op: &'p P, match_position: MatchPosition, best_found: &'p AtomicUsize) -> Self {
  69         FindConsumer {
  70             find_op: find_op,
  71             lower_bound: Cell::new(0),
  72             upper_bound: usize::max_value(),
  73             match_position: match_position,
  74             best_found: best_found,
  75         }
  76     }
  77
  78     fn current_index(&self) -> usize {
  79         match self.match_position {
  80             MatchPosition::Leftmost => self.lower_bound.get(),
  81             MatchPosition::Rightmost => self.upper_bound,
  82         }
  83     }
  84 }
  85
  86 impl<'p, T, P> Consumer<T> for FindConsumer<'p, P>
  87     where T: Send,
  88           P: Fn(&T) -> bool + Sync
  89 {
  90     type Folder = FindFolder<'p, T, P>;
  91     type Reducer = FindReducer;
  92     type Result = Option<T>;
  93
  94     fn split_at(self, _index: usize) -> (Self, Self, Self::Reducer) {
  95         let dir = self.match_position;
  96         (self.split_off_left(), self, FindReducer { match_position: dir })
  97     }
  98
  99     fn into_folder(self) -> Self::Folder {
 100         FindFolder {
 101             find_op: self.find_op,
 102             boundary: self.current_index(),
 103             match_position: self.match_position,
 104             best_found: self.best_found,
 105             item: None,
 106         }
 107     }
 108
 109     fn full(&self) -> bool {
 110         // can stop consuming if the best found index so far is *strictly*
 111         // better than anything this consumer will find
 112         better_position(self.best_found.load(Ordering::Relaxed),
 113                         self.current_index(),
 114                         self.match_position)
 115     }
 116 }
 117
 118 impl<'p, T, P> UnindexedConsumer<T> for FindConsumer<'p, P>
 119     where T: Send,
 120           P: Fn(&T) -> bool + Sync
 121 {
 122     fn split_off_left(&self) -> Self {
 123         // Upper bound for one consumer will be lower bound for the other. This
 124         // overlap is okay, because only one of the bounds will be used for
 125         // comparing against best_found; the other is kept only to be able to
 126         // divide the range in half.
 127         //
 128         // When the resolution of usize has been exhausted (i.e. when
 129         // upper_bound = lower_bound), both results of this split will have the
 130         // same range. When that happens, we lose the ability to tell one
 131         // consumer to stop working when the other finds a better match, but the
 132         // reducer ensures that the best answer is still returned (see the test
 133         // above).
 134         let old_lower_bound = self.lower_bound.get();
 135         let median = old_lower_bound + ((self.upper_bound - old_lower_bound) / 2);
 136         self.lower_bound.set(median);
 137
 138         FindConsumer {
 139             find_op: self.find_op,
 140             lower_bound: Cell::new(old_lower_bound),
 141             upper_bound: median,
 142             match_position: self.match_position,
 143             best_found: self.best_found,
 144         }
 145     }
 146
 147     fn to_reducer(&self) -> Self::Reducer {
 148         FindReducer { match_position: self.match_position }
 149     }
 150 }
 151
 152 struct FindFolder<'p, T, P: 'p> {
 153     find_op: &'p P,
 154     boundary: usize,
 155     match_position: MatchPosition,
 156     best_found: &'p AtomicUsize,
 157     item: Option<T>,
 158 }
 159
 160 impl<'p, P: 'p + Fn(&T) -> bool, T> Folder<T> for FindFolder<'p, T, P> {
 161     type Result = Option<T>;
 162
 163     fn consume(mut self, item: T) -> Self {
 164         let found_best_in_range = match self.match_position {
 165             MatchPosition::Leftmost => self.item.is_some(),
 166             MatchPosition::Rightmost => false,
 167         };
 168
 169         if !found_best_in_range && (self.find_op)(&item) {
 170             // Continuously try to set best_found until we succeed or we
 171             // discover a better match was already found.
 172             let mut current = self.best_found.load(Ordering::Relaxed);
 173             loop {
 174                 if better_position(current, self.boundary, self.match_position) {
 175                     break;
 176                 }
 177                 match self.best_found.compare_exchange_weak(current,
 178                                                             self.boundary,
 179                                                             Ordering::Relaxed,
 180                                                             Ordering::Relaxed) {
 181                     Ok(_) => {
 182                         self.item = Some(item);
 183                         break;
 184                     }
 185                     Err(v) => current = v,
 186                 }
 187             }
 188         }
 189         self
 190     }
 191
 192     fn complete(self) -> Self::Result {
 193         self.item
 194     }
 195
 196     fn full(&self) -> bool {
 197         let found_best_in_range = match self.match_position {
 198             MatchPosition::Leftmost => self.item.is_some(),
 199             MatchPosition::Rightmost => false,
 200         };
 201
 202         found_best_in_range ||
 203         better_position(self.best_found.load(Ordering::Relaxed),
 204                         self.boundary,
 205                         self.match_position)
 206     }
 207 }
 208
 209 struct FindReducer {
 210     match_position: MatchPosition,
 211 }
 212
 213 impl<T> Reducer<Option<T>> for FindReducer {
 214     fn reduce(self, left: Option<T>, right: Option<T>) -> Option<T> {
 215         match self.match_position {
 216             MatchPosition::Leftmost => left.or(right),
 217             MatchPosition::Rightmost => right.or(left),
 218         }
 219     }
 220 }