src/libsyntax/util/map_in_place.rs

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