Add a MutVisitor for the AST

src/librustc_allocator/expand.rs

@@ -1,10 +1,9 @@
 use rustc::middle::allocator::AllocatorKind;
 use rustc_errors;
-use smallvec::SmallVec;
 use syntax::{
     ast::{
         self, Arg, Attribute, Crate, Expr, FnHeader, Generics, Ident, Item, ItemKind,
-        Mac, Mod, Mutability, Ty, TyKind, Unsafety, VisibilityKind,
+        Mod, Mutability, Ty, TyKind, Unsafety, VisibilityKind, NodeId,
     },
     attr,
     source_map::{
@@ -16,7 +15,8 @@ use syntax::{
         expand::ExpansionConfig,
         hygiene::{self, Mark, SyntaxContext},
     },
-    fold::{self, Folder},
+    fold::Folder,
+    visit_mut::{self, MutVisitor, Action},
     parse::ParseSess,
     ptr::P,
     symbol::Symbol
@@ -28,7 +28,7 @@ use {AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS};
 pub fn modify(
     sess: &ParseSess,
     resolver: &mut dyn Resolver,
-    krate: Crate,
+    mut krate: Crate,
     crate_name: String,
     handler: &rustc_errors::Handler,
 ) -> ast::Crate {
@@ -39,7 +39,8 @@ pub fn modify(
         found: false,
         crate_name: Some(crate_name),
         in_submod: -1, // -1 to account for the "root" module
-    }.fold_crate(krate)
+    }.visit_crate(&mut krate);
+    krate
 }
 
 struct ExpandAllocatorDirectives<'a> {
@@ -54,34 +55,35 @@ struct ExpandAllocatorDirectives<'a> {
     in_submod: isize,
 }
 
-impl<'a> Folder for ExpandAllocatorDirectives<'a> {
-    fn fold_item(&mut self, item: P<Item>) -> SmallVec<[P<Item>; 1]> {
+impl<'a> MutVisitor for ExpandAllocatorDirectives<'a> {
+    fn visit_item(&mut self, item: &mut P<Item>) -> Action<P<Item>> {
         debug!("in submodule {}", self.in_submod);
 
         let name = if attr::contains_name(&item.attrs, "global_allocator") {
             "global_allocator"
         } else {
-            return fold::noop_fold_item(item, self);
+            visit_mut::walk_item(self, item);
+            return Action::Reuse;
         };
         match item.node {
             ItemKind::Static(..) => {}
             _ => {
                 self.handler
                     .span_err(item.span, "allocators must be statics");
-                return smallvec![item];
+                return Action::Reuse;
             }
         }
 
         if self.in_submod > 0 {
             self.handler
                 .span_err(item.span, "`global_allocator` cannot be used in submodules");
-            return smallvec![item];
+            return Action::Reuse;
         }
 
         if self.found {
             self.handler
                 .span_err(item.span, "cannot define more than one #[global_allocator]");
-            return smallvec![item];
+            return Action::Reuse;
         }
         self.found = true;
 
@@ -142,23 +144,18 @@ impl<'a> Folder for ExpandAllocatorDirectives<'a> {
         let module = f.cx.monotonic_expander().fold_item(module).pop().unwrap();
 
         // Return the item and new submodule
-        smallvec![item, module]
+        Action::Add(vec![module])
     }
 
     // If we enter a submodule, take note.
-    fn fold_mod(&mut self, m: Mod) -> Mod {
+    fn visit_mod(&mut self, m: &mut Mod, _s: Span, _attrs: &[Attribute], _n: NodeId) {
         debug!("enter submodule");
         self.in_submod += 1;
-        let ret = fold::noop_fold_mod(m, self);
+        let ret = visit_mut::walk_mod(self, m);
         self.in_submod -= 1;
         debug!("exit submodule");
         ret
     }
-
-    // `fold_mac` is disabled by default. Enable it here.
-    fn fold_mac(&mut self, mac: Mac) -> Mac {
-        fold::noop_fold_mac(mac, self)
-    }
 }
 
 struct AllocFnFactory<'a> {

src/librustc_allocator/lib.rs

@@ -8,7 +8,6 @@ extern crate rustc_errors;
 extern crate rustc_target;
 extern crate syntax;
 extern crate syntax_pos;
-#[macro_use]
 extern crate smallvec;
 
 pub mod expand;

src/librustc_data_structures/thin_vec.rs

@@ -8,6 +8,14 @@ impl<T> ThinVec<T> {
     pub fn new() -> Self {
         ThinVec(None)
     }
+
+    #[inline]
+    pub fn as_vec(&mut self) -> Option<&mut Vec<T>> {
+        match *self {
+            ThinVec(None) => None,
+            ThinVec(Some(ref mut vec)) => Some(vec),
+        }
+    }
 }
 
 impl<T> From<Vec<T>> for ThinVec<T> {
@@ -39,6 +47,16 @@ impl<T> ::std::ops::Deref for ThinVec<T> {
     }
 }
 
+impl<T> ::std::ops::DerefMut for ThinVec<T> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut [T] {
+        match *self {
+            ThinVec(None) => &mut [],
+            ThinVec(Some(ref mut vec)) => vec,
+        }
+    }
+}
+
 impl<T> Extend<T> for ThinVec<T> {
     fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
         match *self {

src/libsyntax/attr/mod.rs

@@ -695,11 +695,17 @@ impl LitKind {
 
 pub trait HasAttrs: Sized {
     fn attrs(&self) -> &[ast::Attribute];
+    fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>>;
     fn map_attrs<F: FnOnce(Vec<ast::Attribute>) -> Vec<ast::Attribute>>(self, f: F) -> Self;
 }
 
 impl<T: HasAttrs> HasAttrs for Spanned<T> {
-    fn attrs(&self) -> &[ast::Attribute] { self.node.attrs() }
+    fn attrs(&self) -> &[ast::Attribute] {
+        self.node.attrs()
+    }
+    fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>> {
+        self.node.attrs_mut()
+    }
     fn map_attrs<F: FnOnce(Vec<ast::Attribute>) -> Vec<ast::Attribute>>(self, f: F) -> Self {
         respan(self.span, self.node.map_attrs(f))
     }
@@ -709,6 +715,9 @@ impl HasAttrs for Vec<Attribute> {
     fn attrs(&self) -> &[Attribute] {
         self
     }
+    fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>> {
+        Some(self)
+    }
     fn map_attrs<F: FnOnce(Vec<Attribute>) -> Vec<Attribute>>(self, f: F) -> Self {
         f(self)
     }
@@ -718,6 +727,9 @@ impl HasAttrs for ThinVec<Attribute> {
     fn attrs(&self) -> &[Attribute] {
         self
     }
+    fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>> {
+        self.as_vec()
+    }
     fn map_attrs<F: FnOnce(Vec<Attribute>) -> Vec<Attribute>>(self, f: F) -> Self {
         f(self.into()).into()
     }
@@ -727,6 +739,9 @@ impl<T: HasAttrs + 'static> HasAttrs for P<T> {
     fn attrs(&self) -> &[Attribute] {
         (**self).attrs()
     }
+    fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>> {
+        (**self).attrs_mut()
+    }
     fn map_attrs<F: FnOnce(Vec<Attribute>) -> Vec<Attribute>>(self, f: F) -> Self {
         self.map(|t| t.map_attrs(f))
     }
@@ -745,6 +760,18 @@ impl HasAttrs for StmtKind {
         }
     }
 
+    fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>> {
+        match *self {
+            StmtKind::Local(ref mut local) => local.attrs_mut(),
+            StmtKind::Item(..) => None,
+            StmtKind::Expr(ref mut expr) | StmtKind::Semi(ref mut expr) => expr.attrs_mut(),
+            StmtKind::Mac(ref mut mac) => {
+                let (_, _, ref mut attrs) = **mac;
+                attrs.attrs_mut()
+            }
+        }
+    }
+
     fn map_attrs<F: FnOnce(Vec<Attribute>) -> Vec<Attribute>>(self, f: F) -> Self {
         match self {
             StmtKind::Local(local) => StmtKind::Local(local.map_attrs(f)),
@@ -760,6 +787,9 @@ impl HasAttrs for StmtKind {
 
 impl HasAttrs for Stmt {
     fn attrs(&self) -> &[ast::Attribute] { self.node.attrs() }
+    fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>> {
+        self.node.attrs_mut()
+    }
     fn map_attrs<F: FnOnce(Vec<ast::Attribute>) -> Vec<ast::Attribute>>(self, f: F) -> Self {
         Stmt { id: self.id, node: self.node.map_attrs(f), span: self.span }
     }
@@ -770,6 +800,10 @@ impl HasAttrs for GenericParam {
         &self.attrs
     }
 
+    fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>> {
+        self.attrs.attrs_mut()
+    }
+
     fn map_attrs<F: FnOnce(Vec<Attribute>) -> Vec<Attribute>>(mut self, f: F) -> Self {
         self.attrs = self.attrs.map_attrs(f);
         self
@@ -783,6 +817,10 @@ macro_rules! derive_has_attrs {
                 &self.attrs
             }
 
+            fn attrs_mut(&mut self) -> Option<&mut Vec<ast::Attribute>> {
+                self.attrs.attrs_mut()
+            }
+
             fn map_attrs<F>(mut self, f: F) -> Self
                 where F: FnOnce(Vec<Attribute>) -> Vec<Attribute>,
             {