Some TLC for the MoveMap trait

2 files changed, 102 insertions, 17 deletions

diff --git a/src/libsyntax/util/move_map.rs b/src/libsyntax/util/move_map.rs
new file mode 100644
index 00000000000..95c24c66630
--- /dev/null
+++ b/src/libsyntax/util/move_map.rs
@@ -0,0 +1,79 @@
+// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+use owned_slice::OwnedSlice;
+
+use std::ptr;
+
+pub trait MoveMap<T>: Sized {
+    fn move_map<F>(self, mut f: F) -> Self where F: FnMut(T) -> T {
+        self.move_flat_map(|e| Some(f(e)))
+    }
+
+    fn move_flat_map<F, I>(self, f: F) -> Self
+        where F: FnMut(T) -> I,
+              I: IntoIterator<Item=T>;
+}
+
+impl<T> MoveMap<T> for Vec<T> {
+    fn move_flat_map<F, I>(mut self, mut f: F) -> Self
+        where F: FnMut(T) -> I,
+              I: IntoIterator<Item=T>
+    {
+        let mut read_i = 0;
+        let mut write_i = 0;
+        unsafe {
+            let mut old_len = self.len();
+            self.set_len(0); // make sure we just leak elements in case of panic
+
+            while read_i < old_len {
+                // move the read_i'th item out of the vector and map it
+                // to an iterator
+                let e = ptr::read(self.get_unchecked(read_i));
+                let mut iter = f(e).into_iter();
+                read_i += 1;
+
+                while let Some(e) = iter.next() {
+                    if write_i < read_i {
+                        ptr::write(self.get_unchecked_mut(write_i), e);
+                        write_i += 1;
+                    } else {
+                        // If this is reached we ran out of space
+                        // in the middle of the vector.
+                        // However, the vector is in a valid state here,
+                        // so we just do a somewhat inefficient insert.
+                        self.set_len(old_len);
+                        self.insert(write_i, e);
+
+                        old_len = self.len();
+                        self.set_len(0);
+
+                        read_i += 1;
+                        write_i += 1;
+                    }
+                }
+            }
+
+            // write_i tracks the number of actually written new items.
+            self.set_len(write_i);
+        }
+
+        self
+    }
+}
+
+impl<T> MoveMap<T> for OwnedSlice<T> {
+    fn move_flat_map<F, I>(self, f: F) -> Self
+        where F: FnMut(T) -> I,
+              I: IntoIterator<Item=T>
+    {
+        OwnedSlice::from_vec(self.into_vec().move_flat_map(f))
+    }
+}
diff --git a/src/libsyntax/util/small_vector.rs b/src/libsyntax/util/small_vector.rs
index dac3b054165..ee183d7f3e9 100644
--- a/src/libsyntax/util/small_vector.rs
+++ b/src/libsyntax/util/small_vector.rs
@@ -16,7 +16,7 @@ use std::mem;
 use std::slice;
 use std::vec;
 
-use fold::MoveMap;
+use util::move_map::MoveMap;
 
 /// A vector type optimized for cases where the size is almost always 0 or 1
 pub struct SmallVector<T> {
@@ -134,15 +134,6 @@ impl<T> SmallVector<T> {
         self.into_iter()
     }
 
-    pub fn into_iter(self) -> IntoIter<T> {
-        let repr = match self.repr {
-            Zero => ZeroIterator,
-            One(v) => OneIterator(v),
-            Many(vs) => ManyIterator(vs.into_iter())
-        };
-        IntoIter { repr: repr }
-    }
-
     pub fn len(&self) -> usize {
         match self.repr {
             Zero => 0,
@@ -154,6 +145,19 @@ impl<T> SmallVector<T> {
     pub fn is_empty(&self) -> bool { self.len() == 0 }
 }
 
+impl<T> IntoIterator for SmallVector<T> {
+    type Item = T;
+    type IntoIter = IntoIter<T>;
+    fn into_iter(self) -> Self::IntoIter {
+        let repr = match self.repr {
+            Zero => ZeroIterator,
+            One(v) => OneIterator(v),
+            Many(vs) => ManyIterator(vs.into_iter())
+        };
+        IntoIter { repr: repr }
+    }
+}
+
 pub struct IntoIter<T> {
     repr: IntoIterRepr<T>,
 }
@@ -192,13 +196,15 @@ impl<T> Iterator for IntoIter<T> {
 }
 
 impl<T> MoveMap<T> for SmallVector<T> {
-    fn move_map<F>(self, mut f: F) -> SmallVector<T> where F: FnMut(T) -> T {
-        let repr = match self.repr {
-            Zero => Zero,
-            One(v) => One(f(v)),
-            Many(vs) => Many(vs.move_map(f))
-        };
-        SmallVector { repr: repr }
+    fn move_flat_map<F, I>(self, mut f: F) -> Self
+        where F: FnMut(T) -> I,
+              I: IntoIterator<Item=T>
+    {
+        match self.repr {
+            Zero => Self::zero(),
+            One(v) => f(v).into_iter().collect(),
+            Many(vs) => SmallVector { repr: Many(vs.move_flat_map(f)) },
+        }
     }
 }