compiler/rustc_data_structures/src/map_in_place.rs

   1 use smallvec::{Array, SmallVec};
   2 use std::ptr;
   3 use thin_vec::ThinVec;
   4
   5 pub trait MapInPlace<T>: Sized {
   6     fn map_in_place<F>(&mut self, mut f: F)
   7     where
   8         F: FnMut(T) -> T,
   9     {
  10         self.flat_map_in_place(|e| Some(f(e)))
  11     }
  12
  13     fn flat_map_in_place<F, I>(&mut self, f: F)
  14     where
  15         F: FnMut(T) -> I,
  16         I: IntoIterator<Item = T>;
  17 }
  18
  19 // The implementation of this method is syntactically identical for all the
  20 // different vector types.
  21 macro_rules! flat_map_in_place {
  22     () => {
  23         fn flat_map_in_place<F, I>(&mut self, mut f: F)
  24         where
  25             F: FnMut(T) -> I,
  26             I: IntoIterator<Item = T>,
  27         {
  28             let mut read_i = 0;
  29             let mut write_i = 0;
  30             unsafe {
  31                 let mut old_len = self.len();
  32                 self.set_len(0); // make sure we just leak elements in case of panic
  33
  34                 while read_i < old_len {
  35                     // move the read_i'th item out of the vector and map it
  36                     // to an iterator
  37                     let e = ptr::read(self.as_ptr().add(read_i));
  38                     let iter = f(e).into_iter();
  39                     read_i += 1;
  40
  41                     for e in iter {
  42                         if write_i < read_i {
  43                             ptr::write(self.as_mut_ptr().add(write_i), e);
  44                             write_i += 1;
  45                         } else {
  46                             // If this is reached we ran out of space
  47                             // in the middle of the vector.
  48                             // However, the vector is in a valid state here,
  49                             // so we just do a somewhat inefficient insert.
  50                             self.set_len(old_len);
  51                             self.insert(write_i, e);
  52
  53                             old_len = self.len();
  54                             self.set_len(0);
  55
  56                             read_i += 1;
  57                             write_i += 1;
  58                         }
  59                     }
  60                 }
  61
  62                 // write_i tracks the number of actually written new items.
  63                 self.set_len(write_i);
  64             }
  65         }
  66     };
  67 }
  68
  69 impl<T> MapInPlace<T> for Vec<T> {
  70     flat_map_in_place!();
  71 }
  72
  73 impl<T, A: Array<Item = T>> MapInPlace<T> for SmallVec<A> {
  74     flat_map_in_place!();
  75 }
  76
  77 impl<T> MapInPlace<T> for ThinVec<T> {
  78     flat_map_in_place!();
  79 }