vendor/rustc-ap-rustc_data_structures/src/binary_search_util/mod.rs

   1 #[cfg(test)]
   2 mod tests;
   3
   4 /// Uses a sorted slice `data: &[E]` as a kind of "multi-map". The
   5 /// `key_fn` extracts a key of type `K` from the data, and this
   6 /// function finds the range of elements that match the key. `data`
   7 /// must have been sorted as if by a call to `sort_by_key` for this to
   8 /// work.
   9 pub fn binary_search_slice<E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E]
  10 where
  11     K: Ord,
  12 {
  13     let mid = match data.binary_search_by_key(key, &key_fn) {
  14         Ok(mid) => mid,
  15         Err(_) => return &[],
  16     };
  17     let size = data.len();
  18
  19     // We get back *some* element with the given key -- so do
  20     // a galloping search backwards to find the *first* one.
  21     let mut start = mid;
  22     let mut previous = mid;
  23     let mut step = 1;
  24     loop {
  25         start = start.saturating_sub(step);
  26         if start == 0 || key_fn(&data[start]) != *key {
  27             break;
  28         }
  29         previous = start;
  30         step *= 2;
  31     }
  32     step = previous - start;
  33     while step > 1 {
  34         let half = step / 2;
  35         let mid = start + half;
  36         if key_fn(&data[mid]) != *key {
  37             start = mid;
  38         }
  39         step -= half;
  40     }
  41     // adjust by one if we have overshot
  42     if start < size && key_fn(&data[start]) != *key {
  43         start += 1;
  44     }
  45
  46     // Now search forward to find the *last* one.
  47     let mut end = mid;
  48     let mut previous = mid;
  49     let mut step = 1;
  50     loop {
  51         end = end.saturating_add(step).min(size);
  52         if end == size || key_fn(&data[end]) != *key {
  53             break;
  54         }
  55         previous = end;
  56         step *= 2;
  57     }
  58     step = end - previous;
  59     while step > 1 {
  60         let half = step / 2;
  61         let mid = end - half;
  62         if key_fn(&data[mid]) != *key {
  63             end = mid;
  64         }
  65         step -= half;
  66     }
  67
  68     &data[start..end]
  69 }