1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141

use std::collections::HashMap;
use syntax::ast::*;
use syntax::attr::HasAttrs;
use syntax::mut_visit::{self, MutVisitor};
use syntax::ptr::P;
use syntax::visit::{self, Visitor};

use smallvec::SmallVec;

use crate::ast_manip::{GetNodeId, MutVisit, Visit};

struct CollectCfgAttrs {
    node_attrs: HashMap<NodeId, Vec<Attribute>>,
}

impl CollectCfgAttrs {
    fn collect<T: HasAttrs + GetNodeId>(&mut self, x: &T) {
        let attrs = x
            .attrs()
            .iter()
            .filter(|attr| attr.check_name("cfg_attr"))
            .cloned()
            .collect::<Vec<_>>();
        if attrs.len() > 0 {
            self.node_attrs.insert(x.get_node_id(), attrs);
        }
    }
}

macro_rules! collect_cfg_attrs {
    ($($visit_thing:ident($Thing:ty), $walk_thing:ident;)*) => {
        impl<'ast> Visitor<'ast> for CollectCfgAttrs {
            $(
                fn $visit_thing(&mut self, x: &'ast $Thing) {
                    self.collect(x);
                    visit::$walk_thing(self, x);
                }
            )*

            fn visit_mac(&mut self, mac: &'ast Mac) {
                visit::walk_mac(self, mac)
            }
        }
    };
}

collect_cfg_attrs! {
    visit_item(Item), walk_item;
    visit_impl_item(ImplItem), walk_impl_item;
    visit_trait_item(TraitItem), walk_trait_item;
    visit_foreign_item(ForeignItem), walk_foreign_item;
    visit_stmt(Stmt), walk_stmt;
    visit_expr(Expr), walk_expr;
    // TODO: extend this list with the remaining node types
}

pub fn collect_cfg_attrs(krate: &Crate) -> HashMap<NodeId, Vec<Attribute>> {
    let mut v = CollectCfgAttrs {
        node_attrs: HashMap::new(),
    };
    krate.visit(&mut v);
    v.node_attrs
}

struct RestoreCfgAttrs {
    node_attrs: HashMap<NodeId, Vec<Attribute>>,
}

impl RestoreCfgAttrs {
    fn restore<T: HasAttrs + GetNodeId>(&mut self, x: &mut T) {
        if let Some(cfg_attrs) = self.node_attrs.get(&x.get_node_id()) {
            info!(
                "RESTORE ATTRS {:?} onto {:?}",
                cfg_attrs
                    .iter()
                    .map(|a| ::syntax::print::pprust::attr_to_string(a))
                    .collect::<Vec<_>>(),
                x.attrs()
                    .iter()
                    .map(|a| ::syntax::print::pprust::attr_to_string(a))
                    .collect::<Vec<_>>()
            );
            x.visit_attrs(|attrs| {
                // Drop attrs that were produced by evaluation of one of the `#[cfg_attr]`s.
                attrs.retain(|a| !cfg_attrs.iter().any(|ca| ca.span.contains(a.span)));
                // Now put the #[cfg_attr]s themselves back in.
                attrs.extend(cfg_attrs.iter().cloned());
            });
            info!(
                "  attrs changed to {:?}",
                x.attrs()
                    .iter()
                    .map(|a| ::syntax::print::pprust::attr_to_string(a))
                    .collect::<Vec<_>>()
            );
        }
    }
}

impl MutVisitor for RestoreCfgAttrs {
    fn flat_map_item(&mut self, mut i: P<Item>) -> SmallVec<[P<Item>; 1]> {
        self.restore(&mut i);
        mut_visit::noop_flat_map_item(i, self)
    }

    fn flat_map_impl_item(&mut self, mut i: ImplItem) -> SmallVec<[ImplItem; 1]> {
        self.restore(&mut i);
        mut_visit::noop_flat_map_impl_item(i, self)
    }

    fn flat_map_trait_item(&mut self, mut i: TraitItem) -> SmallVec<[TraitItem; 1]> {
        self.restore(&mut i);
        mut_visit::noop_flat_map_trait_item(i, self)
    }

    fn flat_map_foreign_item(&mut self, mut i: ForeignItem) -> SmallVec<[ForeignItem; 1]> {
        self.restore(&mut i);
        mut_visit::noop_flat_map_foreign_item(i, self)
    }

    fn flat_map_stmt(&mut self, mut s: Stmt) -> SmallVec<[Stmt; 1]> {
        self.restore(&mut s);
        mut_visit::noop_flat_map_stmt(s, self)
    }

    fn visit_expr(&mut self, e: &mut P<Expr>) {
        self.restore(e);
        mut_visit::noop_visit_expr(e, self)
    }

    fn visit_mac(&mut self, mac: &mut Mac) {
        mut_visit::noop_visit_mac(mac, self)
    }

    // TODO: extend this impl with the remaining node types
}

pub fn restore_cfg_attrs(krate: &mut Crate, node_attrs: HashMap<NodeId, Vec<Attribute>>) {
    let mut f = RestoreCfgAttrs { node_attrs };
    krate.visit(&mut f)
}