Support multiple item macros

Closes #4375

1 files changed, 22 insertions, 31 deletions

diff --git a/src/libsyntax/fold.rs b/src/libsyntax/fold.rs
index b268988c3c5..4c0653a3c04 100644
--- a/src/libsyntax/fold.rs
+++ b/src/libsyntax/fold.rs
@@ -13,6 +13,7 @@ use ast;
 use codemap::{respan, Span, Spanned};
 use parse::token;
 use opt_vec::OptVec;
+use util::small_vector::SmallVector;
 
 // We may eventually want to be able to fold over type parameters, too.
 pub trait ast_fold {
@@ -113,7 +114,7 @@ pub trait ast_fold {
         }
     }
 
-    fn fold_item(&self, i: @item) -> Option<@item> {
+    fn fold_item(&self, i: @item) -> SmallVector<@item> {
         noop_fold_item(i, self)
     }
 
@@ -159,7 +160,7 @@ pub trait ast_fold {
         noop_fold_block(b, self)
     }
 
-    fn fold_stmt(&self, s: &Stmt) -> Option<@Stmt> {
+    fn fold_stmt(&self, s: &Stmt) -> SmallVector<@Stmt> {
         noop_fold_stmt(s, self)
     }
 
@@ -216,23 +217,20 @@ pub trait ast_fold {
         }
     }
 
-    fn fold_decl(&self, d: @Decl) -> Option<@Decl> {
+    fn fold_decl(&self, d: @Decl) -> SmallVector<@Decl> {
         let node = match d.node {
-            DeclLocal(ref l) => Some(DeclLocal(self.fold_local(*l))),
+            DeclLocal(ref l) => SmallVector::one(DeclLocal(self.fold_local(*l))),
             DeclItem(it) => {
-                match self.fold_item(it) {
-                    Some(it_folded) => Some(DeclItem(it_folded)),
-                    None => None,
-                }
+                self.fold_item(it).move_iter().map(|i| DeclItem(i)).collect()
             }
         };
 
-        node.map(|node| {
+        node.move_iter().map(|node| {
             @Spanned {
                 node: node,
                 span: d.span,
             }
-        })
+        }).collect()
     }
 
     fn fold_expr(&self, e: @Expr) -> @Expr {
@@ -618,13 +616,7 @@ fn fold_variant_arg_<T:ast_fold>(va: &variant_arg, folder: &T)
 
 pub fn noop_fold_block<T:ast_fold>(b: &Block, folder: &T) -> Block {
     let view_items = b.view_items.map(|x| folder.fold_view_item(x));
-    let mut stmts = ~[];
-    for stmt in b.stmts.iter() {
-        match folder.fold_stmt(*stmt) {
-            None => {}
-            Some(stmt) => stmts.push(stmt)
-        }
-    }
+    let stmts = b.stmts.iter().flat_map(|s| folder.fold_stmt(*s).move_iter()).collect();
     ast::Block {
         view_items: view_items,
         stmts: stmts,
@@ -711,7 +703,7 @@ pub fn noop_fold_mod<T:ast_fold>(m: &_mod, folder: &T) -> _mod {
         view_items: m.view_items
                      .iter()
                      .map(|x| folder.fold_view_item(x)).collect(),
-        items: m.items.iter().filter_map(|x| folder.fold_item(*x)).collect(),
+        items: m.items.iter().flat_map(|x| folder.fold_item(*x).move_iter()).collect(),
     }
 }
 
@@ -728,10 +720,10 @@ pub fn noop_fold_crate<T:ast_fold>(c: Crate, folder: &T) -> Crate {
 }
 
 pub fn noop_fold_item<T:ast_fold>(i: @ast::item, folder: &T)
-                                  -> Option<@ast::item> {
+                                  -> SmallVector<@ast::item> {
     let fold_attribute = |x| fold_attribute_(x, folder);
 
-    Some(@ast::item {
+    SmallVector::one(@ast::item {
         ident: folder.fold_ident(i.ident),
         attrs: i.attrs.map(|e| fold_attribute(*e)),
         id: folder.new_id(i.id),
@@ -867,27 +859,26 @@ pub fn noop_fold_expr<T:ast_fold>(e: @ast::Expr, folder: &T) -> @ast::Expr {
     }
 }
 
-pub fn noop_fold_stmt<T:ast_fold>(s: &Stmt, folder: &T) -> Option<@Stmt> {
-    let node = match s.node {
+pub fn noop_fold_stmt<T:ast_fold>(s: &Stmt, folder: &T) -> SmallVector<@Stmt> {
+    let nodes = match s.node {
         StmtDecl(d, nid) => {
-            match folder.fold_decl(d) {
-                Some(d) => Some(StmtDecl(d, folder.new_id(nid))),
-                None => None,
-            }
+            folder.fold_decl(d).move_iter()
+                    .map(|d| StmtDecl(d, folder.new_id(nid)))
+                    .collect()
         }
         StmtExpr(e, nid) => {
-            Some(StmtExpr(folder.fold_expr(e), folder.new_id(nid)))
+            SmallVector::one(StmtExpr(folder.fold_expr(e), folder.new_id(nid)))
         }
         StmtSemi(e, nid) => {
-            Some(StmtSemi(folder.fold_expr(e), folder.new_id(nid)))
+            SmallVector::one(StmtSemi(folder.fold_expr(e), folder.new_id(nid)))
         }
-        StmtMac(ref mac, semi) => Some(StmtMac(folder.fold_mac(mac), semi))
+        StmtMac(ref mac, semi) => SmallVector::one(StmtMac(folder.fold_mac(mac), semi))
     };
 
-    node.map(|node| @Spanned {
+    nodes.move_iter().map(|node| @Spanned {
         node: node,
         span: folder.new_span(s.span),
-    })
+    }).collect()
 }
 
 #[cfg(test)]