compiler/rustc_data_structures/src/map_in_place.rs

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