update fold_method to return a smallvector

This is nice for macros, to allow them to expand into multiple methods

3 files changed, 38 insertions, 15 deletions

diff --git a/src/librustc/middle/astencode.rs b/src/librustc/middle/astencode.rs
index 8fa7d7c2aaa..d58023a4875 100644
--- a/src/librustc/middle/astencode.rs
+++ b/src/librustc/middle/astencode.rs
@@ -345,7 +345,9 @@ fn simplify_ast(ii: e::InlinedItemRef) -> ast::InlinedItem {
         // HACK we're not dropping items.
         e::IIItemRef(i) => ast::IIItem(fold::noop_fold_item(i, &mut fld)
                                        .expect_one("expected one item")),
-        e::IIMethodRef(d, p, m) => ast::IIMethod(d, p, fold::noop_fold_method(m, &mut fld)),
+        e::IIMethodRef(d, p, m) => ast::IIMethod(d, p, fold::noop_fold_method(m, &mut fld)
+                                                 .expect_one(
+                "noop_fold_method must produce exactly one method")),
         e::IIForeignRef(i) => ast::IIForeign(fold::noop_fold_foreign_item(i, &mut fld))
     }
 }
@@ -387,7 +389,8 @@ fn renumber_and_map_ast(xcx: &ExtendedDecodeContext,
                 ast::IIItem(fld.fold_item(i).expect_one("expected one item"))
             }
             ast::IIMethod(d, is_provided, m) => {
-                ast::IIMethod(xcx.tr_def_id(d), is_provided, fld.fold_method(m))
+                ast::IIMethod(xcx.tr_def_id(d), is_provided, fld.fold_method(m)
+                              .expect_one("expected one method"))
             }
             ast::IIForeign(i) => ast::IIForeign(fld.fold_foreign_item(i))
         }
diff --git a/src/libsyntax/ast_map.rs b/src/libsyntax/ast_map.rs
index d0a0c4fe380..b8a0a31f9c3 100644
--- a/src/libsyntax/ast_map.rs
+++ b/src/libsyntax/ast_map.rs
@@ -571,13 +571,14 @@ impl<'a, F: FoldOps> Folder for Ctx<'a, F> {
         m
     }
 
-    fn fold_method(&mut self, m: Gc<Method>) -> Gc<Method> {
+    fn fold_method(&mut self, m: Gc<Method>) -> SmallVector<Gc<Method>> {
         let parent = self.parent;
         self.parent = DUMMY_NODE_ID;
-        let m = fold::noop_fold_method(&*m, self);
+        let m = fold::noop_fold_method(&*m, self).expect_one(
+            "noop_fold_method must produce exactly one method");
         assert_eq!(self.parent, m.id);
         self.parent = parent;
-        m
+        SmallVector::one(m)
     }
 
     fn fold_fn_decl(&mut self, decl: &FnDecl) -> P<FnDecl> {
diff --git a/src/libsyntax/fold.rs b/src/libsyntax/fold.rs
index 88c8318e1b7..fd786192cb4 100644
--- a/src/libsyntax/fold.rs
+++ b/src/libsyntax/fold.rs
@@ -114,7 +114,7 @@ pub trait Folder {
         noop_fold_type_method(m, self)
     }
 
-    fn fold_method(&mut self, m: Gc<Method>) -> Gc<Method>  {
+    fn fold_method(&mut self, m: Gc<Method>) -> SmallVector<Gc<Method>>  {
         noop_fold_method(&*m, self)
     }
 
@@ -465,10 +465,16 @@ fn fold_interpolated<T: Folder>(nt : &token::Nonterminal, fld: &mut T) -> token:
     match *nt {
         token::NtItem(item) =>
             token::NtItem(fld.fold_item(item)
+                          // this is probably okay, because the only folds likely
+                          // to peek inside interpolated nodes will be renamings/markings,
+                          // which map single items to single items
                           .expect_one("expected fold to produce exactly one item")),
         token::NtBlock(block) => token::NtBlock(fld.fold_block(block)),
         token::NtStmt(stmt) =>
             token::NtStmt(fld.fold_stmt(stmt)
+                          // this is probably okay, because the only folds likely
+                          // to peek inside interpolated nodes will be renamings/markings,
+                          // which map single items to single items
                           .expect_one("expected fold to produce exactly one statement")),
         token::NtPat(pat) => token::NtPat(fld.fold_pat(pat)),
         token::NtExpr(expr) => token::NtExpr(fld.fold_expr(expr)),
@@ -683,15 +689,26 @@ pub fn noop_fold_item_underscore<T: Folder>(i: &Item_, folder: &mut T) -> Item_
             ItemImpl(fold_generics(generics, folder),
                      ifce.as_ref().map(|p| fold_trait_ref(p, folder)),
                      folder.fold_ty(ty),
-                     methods.iter().map(|x| folder.fold_method(*x)).collect()
+                     methods.iter().flat_map(|x| folder.fold_method(*x).move_iter()).collect()
             )
         }
         ItemTrait(ref generics, ref unbound, ref traits, ref methods) => {
-            let methods = methods.iter().map(|method| {
-                match *method {
-                    Required(ref m) => Required(folder.fold_type_method(m)),
-                    Provided(method) => Provided(folder.fold_method(method))
-                }
+            let methods = methods.iter().flat_map(|method| {
+                let r = match *method {
+                    Required(ref m) =>
+                            SmallVector::one(Required(folder.fold_type_method(m))).move_iter(),
+                    Provided(method) => {
+                            // the awkward collect/iter idiom here is because
+                            // even though an iter and a map satisfy the same trait bound,
+                            // they're not actually the same type, so the method arms
+                            // don't unify.
+                            let methods : SmallVector<ast::TraitMethod> =
+                                folder.fold_method(method).move_iter()
+                                .map(|m| Provided(m)).collect();
+                            methods.move_iter()
+                        }
+                };
+                r
             }).collect();
             ItemTrait(fold_generics(generics, folder),
                       unbound.clone(),
@@ -791,9 +808,11 @@ pub fn noop_fold_foreign_item<T: Folder>(ni: &ForeignItem,
     }
 }
 
-pub fn noop_fold_method<T: Folder>(m: &Method, folder: &mut T) -> Gc<Method> {
+// Default fold over a method.
+// Invariant: produces exactly one method.
+pub fn noop_fold_method<T: Folder>(m: &Method, folder: &mut T) -> SmallVector<Gc<Method>> {
     let id = folder.new_id(m.id); // Needs to be first, for ast_map.
-    box(GC) Method {
+    SmallVector::one(box(GC) Method {
         attrs: m.attrs.iter().map(|a| folder.fold_attribute(*a)).collect(),
         id: id,
         span: folder.new_span(m.span),
@@ -809,7 +828,7 @@ pub fn noop_fold_method<T: Folder>(m: &Method, folder: &mut T) -> Gc<Method> {
             },
             MethMac(ref mac) => MethMac(folder.fold_mac(mac)),
         }
-    }
+    })
 }
 
 pub fn noop_fold_pat<T: Folder>(p: Gc<Pat>, folder: &mut T) -> Gc<Pat> {