introduce `LintTable`

4 files changed, 75 insertions, 25 deletions

diff --git a/src/librustc/lint/context.rs b/src/librustc/lint/context.rs
index 3506a9c067c..edf5666a3a5 100644
--- a/src/librustc/lint/context.rs
+++ b/src/librustc/lint/context.rs
@@ -773,11 +773,10 @@ impl<'a, 'tcx> hir_visit::Visitor<'tcx> for LateContext<'a, 'tcx> {
 
     // Output any lints that were previously added to the session.
     fn visit_id(&mut self, id: ast::NodeId) {
-        if let Some(lints) = self.sess().lints.borrow_mut().remove(&id) {
-            debug!("LateContext::visit_id: id={:?} lints={:?}", id, lints);
-            for early_lint in lints {
-                self.early_lint(early_lint);
-            }
+        let lints = self.sess().lints.borrow_mut().take(id);
+        for early_lint in lints {
+            debug!("LateContext::visit_id: id={:?} early_lint={:?}", id, early_lint);
+            self.early_lint(early_lint);
         }
     }
 
@@ -1232,7 +1231,7 @@ pub fn check_crate<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
 
     // If we missed any lints added to the session, then there's a bug somewhere
     // in the iteration code.
-    for (id, v) in tcx.sess.lints.borrow().iter() {
+    if let Some((id, v)) = tcx.sess.lints.borrow().get_any() {
         for early_lint in v {
             span_bug!(early_lint.diagnostic.span.clone(),
                       "unprocessed lint {:?} at {}",
@@ -1250,10 +1249,9 @@ pub fn check_ast_crate(sess: &Session, krate: &ast::Crate) {
     // Visit the whole crate.
     cx.with_lint_attrs(&krate.attrs, |cx| {
         // Lints may be assigned to the whole crate.
-        if let Some(lints) = cx.sess.lints.borrow_mut().remove(&ast::CRATE_NODE_ID) {
-            for early_lint in lints {
-                cx.early_lint(early_lint);
-            }
+        let lints = cx.sess.lints.borrow_mut().take(ast::CRATE_NODE_ID);
+        for early_lint in lints {
+            cx.early_lint(early_lint);
         }
 
         // since the root module isn't visited as an item (because it isn't an
@@ -1270,7 +1268,7 @@ pub fn check_ast_crate(sess: &Session, krate: &ast::Crate) {
 
     // If we missed any lints added to the session, then there's a bug somewhere
     // in the iteration code.
-    for (_, v) in sess.lints.borrow().iter() {
+    for (_, v) in sess.lints.borrow().get_any() {
         for early_lint in v {
             span_bug!(early_lint.diagnostic.span.clone(), "unprocessed lint {:?}", early_lint);
         }
diff --git a/src/librustc/lint/mod.rs b/src/librustc/lint/mod.rs
index 7e0da00694c..704e32e2d0c 100644
--- a/src/librustc/lint/mod.rs
+++ b/src/librustc/lint/mod.rs
@@ -43,6 +43,8 @@ pub use lint::context::{LateContext, EarlyContext, LintContext, LintStore,
                         raw_emit_lint, check_crate, check_ast_crate, gather_attrs,
                         raw_struct_lint, FutureIncompatibleInfo, EarlyLint, IntoEarlyLint};
 
+pub use lint::table::LintTable;
+
 /// Specification of a single lint.
 #[derive(Copy, Clone, Debug)]
 pub struct Lint {
@@ -346,3 +348,4 @@ pub type LevelSource = (Level, LintSource);
 
 pub mod builtin;
 mod context;
+mod table;
diff --git a/src/librustc/lint/table.rs b/src/librustc/lint/table.rs
new file mode 100644
index 00000000000..3b6d268b08f
--- /dev/null
+++ b/src/librustc/lint/table.rs
@@ -0,0 +1,56 @@
+use syntax::ast;
+use syntax_pos::MultiSpan;
+use util::nodemap::NodeMap;
+
+use super::{Lint, LintId, EarlyLint, IntoEarlyLint};
+
+pub struct LintTable {
+    map: NodeMap<Vec<EarlyLint>>
+}
+
+impl LintTable {
+    pub fn new() -> Self {
+        LintTable { map: NodeMap() }
+    }
+
+    pub fn add_lint<S: Into<MultiSpan>>(&mut self,
+                                        lint: &'static Lint,
+                                        id: ast::NodeId,
+                                        sp: S,
+                                        msg: String)
+    {
+        self.add_lint_diagnostic(lint, id, (sp, &msg[..]))
+    }
+
+    pub fn add_lint_diagnostic<M>(&mut self,
+                                  lint: &'static Lint,
+                                  id: ast::NodeId,
+                                  msg: M)
+        where M: IntoEarlyLint,
+    {
+        let lint_id = LintId::of(lint);
+        let early_lint = msg.into_early_lint(lint_id);
+        let arr = self.map.entry(id).or_insert(vec![]);
+        if !arr.contains(&early_lint) {
+            arr.push(early_lint);
+        }
+    }
+
+    pub fn get(&self, id: ast::NodeId) -> &[EarlyLint] {
+        self.map.get(&id).map(|v| &v[..]).unwrap_or(&[])
+    }
+
+    pub fn take(&mut self, id: ast::NodeId) -> Vec<EarlyLint> {
+        self.map.remove(&id).unwrap_or(vec![])
+    }
+
+    /// Returns the first (id, lint) pair that is non-empty. Used to
+    /// implement a sanity check in lints that all node-ids are
+    /// visited.
+    pub fn get_any(&self) -> Option<(&ast::NodeId, &Vec<EarlyLint>)> {
+        self.map.iter()
+                .filter(|&(_, v)| !v.is_empty())
+                .next()
+    }
+}
+
diff --git a/src/librustc/session/mod.rs b/src/librustc/session/mod.rs
index 36a887e0622..f10ea3544f2 100644
--- a/src/librustc/session/mod.rs
+++ b/src/librustc/session/mod.rs
@@ -20,7 +20,7 @@ use middle::dependency_format;
 use session::search_paths::PathKind;
 use session::config::DebugInfoLevel;
 use ty::tls;
-use util::nodemap::{NodeMap, FxHashMap, FxHashSet};
+use util::nodemap::{FxHashMap, FxHashSet};
 use util::common::duration_to_secs_str;
 use mir::transform as mir_pass;
 
@@ -78,7 +78,7 @@ pub struct Session {
     pub local_crate_source_file: Option<PathBuf>,
     pub working_dir: PathBuf,
     pub lint_store: RefCell<lint::LintStore>,
-    pub lints: RefCell<NodeMap<Vec<lint::EarlyLint>>>,
+    pub lints: RefCell<lint::LintTable>,
     /// Set of (LintId, span, message) tuples tracking lint (sub)diagnostics
     /// that have been set once, but should not be set again, in order to avoid
     /// redundantly verbose output (Issue #24690).
@@ -270,13 +270,14 @@ impl Session {
     pub fn unimpl(&self, msg: &str) -> ! {
         self.diagnostic().unimpl(msg)
     }
+
     pub fn add_lint<S: Into<MultiSpan>>(&self,
                                         lint: &'static lint::Lint,
                                         id: ast::NodeId,
                                         sp: S,
                                         msg: String)
     {
-        self.add_lint_diagnostic(lint, id, (sp, &msg[..]))
+        self.lints.borrow_mut().add_lint(lint, id, sp, msg);
     }
 
     pub fn add_lint_diagnostic<M>(&self,
@@ -285,17 +286,9 @@ impl Session {
                                   msg: M)
         where M: lint::IntoEarlyLint,
     {
-        let lint_id = lint::LintId::of(lint);
-        let mut lints = self.lints.borrow_mut();
-        let early_lint = msg.into_early_lint(lint_id);
-        if let Some(arr) = lints.get_mut(&id) {
-            if !arr.contains(&early_lint) {
-                arr.push(early_lint);
-            }
-            return;
-        }
-        lints.insert(id, vec![early_lint]);
+        self.lints.borrow_mut().add_lint_diagnostic(lint, id, msg);
     }
+
     pub fn reserve_node_ids(&self, count: usize) -> ast::NodeId {
         let id = self.next_node_id.get();
 
@@ -617,7 +610,7 @@ pub fn build_session_(sopts: config::Options,
         local_crate_source_file: local_crate_source_file,
         working_dir: env::current_dir().unwrap(),
         lint_store: RefCell::new(lint::LintStore::new()),
-        lints: RefCell::new(NodeMap()),
+        lints: RefCell::new(lint::LintTable::new()),
         one_time_diagnostics: RefCell::new(FxHashSet()),
         plugin_llvm_passes: RefCell::new(Vec::new()),
         mir_passes: RefCell::new(mir_pass::Passes::new()),