Some TLC for the MoveMap trait

8 files changed, 120 insertions, 56 deletions

diff --git a/src/libsyntax/ext/deriving/generic/mod.rs b/src/libsyntax/ext/deriving/generic/mod.rs
index 625abf706ca..bd89430d81d 100644
--- a/src/libsyntax/ext/deriving/generic/mod.rs
+++ b/src/libsyntax/ext/deriving/generic/mod.rs
@@ -204,7 +204,7 @@ use ext::build::AstBuilder;
 use codemap::{self, DUMMY_SP};
 use codemap::Span;
 use diagnostic::SpanHandler;
-use fold::MoveMap;
+use util::move_map::MoveMap;
 use owned_slice::OwnedSlice;
 use parse::token::{intern, InternedString};
 use parse::token::special_idents;
diff --git a/src/libsyntax/ext/expand.rs b/src/libsyntax/ext/expand.rs
index 44dab257726..908b2bd2429 100644
--- a/src/libsyntax/ext/expand.rs
+++ b/src/libsyntax/ext/expand.rs
@@ -24,6 +24,7 @@ use ext::base::*;
 use feature_gate::{self, Features, GatedCfg};
 use fold;
 use fold::*;
+use util::move_map::MoveMap;
 use parse;
 use parse::token::{fresh_mark, fresh_name, intern};
 use ptr::P;
diff --git a/src/libsyntax/fold.rs b/src/libsyntax/fold.rs
index 0d5e80eab17..73f2c51b246 100644
--- a/src/libsyntax/fold.rs
+++ b/src/libsyntax/fold.rs
@@ -26,34 +26,11 @@ use codemap::{respan, Span, Spanned};
 use owned_slice::OwnedSlice;
 use parse::token;
 use ptr::P;
-use std::ptr;
 use util::small_vector::SmallVector;
+use util::move_map::MoveMap;
 
 use std::rc::Rc;
 
-// This could have a better place to live.
-pub trait MoveMap<T> {
-    fn move_map<F>(self, f: F) -> Self where F: FnMut(T) -> T;
-}
-
-impl<T> MoveMap<T> for Vec<T> {
-    fn move_map<F>(mut self, mut f: F) -> Vec<T> where F: FnMut(T) -> T {
-        for p in &mut self {
-            unsafe {
-                // FIXME(#5016) this shouldn't need to zero to be safe.
-                ptr::write(p, f(ptr::read_and_drop(p)));
-            }
-        }
-        self
-    }
-}
-
-impl<T> MoveMap<T> for OwnedSlice<T> {
-    fn move_map<F>(self, f: F) -> OwnedSlice<T> where F: FnMut(T) -> T {
-        OwnedSlice::from_vec(self.into_vec().move_map(f))
-    }
-}
-
 pub trait Folder : Sized {
     // Any additions to this trait should happen in form
     // of a call to a public `noop_*` function that only calls
@@ -362,7 +339,7 @@ pub fn noop_fold_view_path<T: Folder>(view_path: P<ViewPath>, fld: &mut T) -> P<
 }
 
 pub fn fold_attrs<T: Folder>(attrs: Vec<Attribute>, fld: &mut T) -> Vec<Attribute> {
-    attrs.into_iter().flat_map(|x| fld.fold_attribute(x)).collect()
+    attrs.move_flat_map(|x| fld.fold_attribute(x))
 }
 
 pub fn fold_thin_attrs<T: Folder>(attrs: ThinAttributes, fld: &mut T) -> ThinAttributes {
@@ -623,6 +600,8 @@ pub fn noop_fold_tt<T: Folder>(tt: &TokenTree, fld: &mut T) -> TokenTree {
 }
 
 pub fn noop_fold_tts<T: Folder>(tts: &[TokenTree], fld: &mut T) -> Vec<TokenTree> {
+    // FIXME: Does this have to take a tts slice?
+    // Could use move_map otherwise...
     tts.iter().map(|tt| fld.fold_tt(tt)).collect()
 }
 
@@ -904,7 +883,7 @@ fn noop_fold_bounds<T: Folder>(bounds: TyParamBounds, folder: &mut T)
 pub fn noop_fold_block<T: Folder>(b: P<Block>, folder: &mut T) -> P<Block> {
     b.map(|Block {id, stmts, expr, rules, span}| Block {
         id: folder.new_id(id),
-        stmts: stmts.into_iter().flat_map(|s| folder.fold_stmt(s).into_iter()).collect(),
+        stmts: stmts.move_flat_map(|s| folder.fold_stmt(s).into_iter()),
         expr: expr.and_then(|x| folder.fold_opt_expr(x)),
         rules: rules,
         span: folder.new_span(span),
@@ -953,9 +932,9 @@ pub fn noop_fold_item_underscore<T: Folder>(i: Item_, folder: &mut T) -> Item_ {
             ItemDefaultImpl(unsafety, folder.fold_trait_ref((*trait_ref).clone()))
         }
         ItemImpl(unsafety, polarity, generics, ifce, ty, impl_items) => {
-            let new_impl_items = impl_items.into_iter().flat_map(|item| {
-                folder.fold_impl_item(item).into_iter()
-            }).collect();
+            let new_impl_items = impl_items.move_flat_map(|item| {
+                folder.fold_impl_item(item)
+            });
             let ifce = match ifce {
                 None => None,
                 Some(ref trait_ref) => {
@@ -971,9 +950,9 @@ pub fn noop_fold_item_underscore<T: Folder>(i: Item_, folder: &mut T) -> Item_ {
         }
         ItemTrait(unsafety, generics, bounds, items) => {
             let bounds = folder.fold_bounds(bounds);
-            let items = items.into_iter().flat_map(|item| {
-                folder.fold_trait_item(item).into_iter()
-            }).collect();
+            let items = items.move_flat_map(|item| {
+                folder.fold_trait_item(item)
+            });
             ItemTrait(unsafety,
                       folder.fold_generics(generics),
                       bounds,
@@ -1032,7 +1011,7 @@ pub fn noop_fold_impl_item<T: Folder>(i: P<ImplItem>, folder: &mut T)
 pub fn noop_fold_mod<T: Folder>(Mod {inner, items}: Mod, folder: &mut T) -> Mod {
     Mod {
         inner: folder.new_span(inner),
-        items: items.into_iter().flat_map(|x| folder.fold_item(x).into_iter()).collect(),
+        items: items.move_flat_map(|x| folder.fold_item(x)),
     }
 }
 
@@ -1353,8 +1332,7 @@ pub fn noop_fold_opt_expr<T: Folder>(e: P<Expr>, folder: &mut T) -> Option<P<Exp
 }
 
 pub fn noop_fold_exprs<T: Folder>(es: Vec<P<Expr>>, folder: &mut T) -> Vec<P<Expr>> {
-    // FIXME: Needs a efficient in-place flat_map
-    es.into_iter().flat_map(|e| folder.fold_opt_expr(e)).collect()
+    es.move_flat_map(|e| folder.fold_opt_expr(e))
 }
 
 pub fn noop_fold_stmt<T: Folder>(Spanned {node, span}: Stmt, folder: &mut T)
diff --git a/src/libsyntax/lib.rs b/src/libsyntax/lib.rs
index 475408472ee..da62c253d8f 100644
--- a/src/libsyntax/lib.rs
+++ b/src/libsyntax/lib.rs
@@ -70,6 +70,7 @@ pub mod util {
     #[cfg(test)]
     pub mod parser_testing;
     pub mod small_vector;
+    pub mod move_map;
 }
 
 pub mod diagnostics {
diff --git a/src/libsyntax/print/pprust.rs b/src/libsyntax/print/pprust.rs
index 8ea93e5b8da..aa55cb847fa 100644
--- a/src/libsyntax/print/pprust.rs
+++ b/src/libsyntax/print/pprust.rs
@@ -1653,8 +1653,6 @@ impl<'a> State<'a> {
         }
         if parse::classify::stmt_ends_with_semi(&st.node) {
             try!(word(&mut self.s, ";"));
-        } else {
-            //try!(word(&mut self.s, ""));
         }
         self.maybe_print_trailing_comment(st.span, None)
     }
diff --git a/src/libsyntax/test.rs b/src/libsyntax/test.rs
index aab106b7a67..6fd3833a3cd 100644
--- a/src/libsyntax/test.rs
+++ b/src/libsyntax/test.rs
@@ -29,7 +29,8 @@ use entry::{self, EntryPointType};
 use ext::base::ExtCtxt;
 use ext::build::AstBuilder;
 use ext::expand::ExpansionConfig;
-use fold::{Folder, MoveMap};
+use fold::Folder;
+use util::move_map::MoveMap;
 use fold;
 use owned_slice::OwnedSlice;
 use parse::token::{intern, InternedString};
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)) },
+        }
     }
 }