3 files changed, 75 insertions, 59 deletions

diff --git a/crates/ra_hir/src/code_model_api.rs b/crates/ra_hir/src/code_model_api.rs
index bc0f74c89..5437133b8 100644
--- a/crates/ra_hir/src/code_model_api.rs
+++ b/crates/ra_hir/src/code_model_api.rs
@@ -168,6 +168,22 @@ impl Module {
 
     pub fn diagnostics(&self, db: &impl HirDatabase, sink: &mut DiagnosticSink) {
         db.crate_def_map(self.krate).add_diagnostics(db, self.module_id, sink);
+        for decl in self.declarations(db) {
+            match decl {
+                crate::ModuleDef::Function(f) => f.diagnostics(db, sink),
+                crate::ModuleDef::Module(f) => f.diagnostics(db, sink),
+                _ => (),
+            }
+        }
+
+        for impl_block in self.impl_blocks(db) {
+            for item in impl_block.items(db) {
+                match item {
+                    crate::ImplItem::Method(f) => f.diagnostics(db, sink),
+                    _ => (),
+                }
+            }
+        }
     }
 
     pub fn resolver(&self, db: &impl HirDatabase) -> Resolver {
diff --git a/crates/ra_hir/src/diagnostics.rs b/crates/ra_hir/src/diagnostics.rs
index d6b28159e..a6ca68d86 100644
--- a/crates/ra_hir/src/diagnostics.rs
+++ b/crates/ra_hir/src/diagnostics.rs
@@ -1,9 +1,9 @@
 use std::{fmt, any::Any};
 
 use ra_syntax::{SyntaxNodePtr, AstPtr, ast};
+use relative_path::RelativePathBuf;
 
 use crate::HirFileId;
-use relative_path::RelativePathBuf;
 
 /// Diagnostic defines hir API for errors and warnings.
 ///
@@ -21,7 +21,7 @@ pub trait Diagnostic: Any + Send + Sync + fmt::Debug + 'static {
     fn message(&self) -> String;
     fn file(&self) -> HirFileId;
     fn syntax_node(&self) -> SyntaxNodePtr;
-    fn as_any(&self) -> &(Any + Send + 'static);
+    fn as_any(&self) -> &(dyn Any + Send + 'static);
 }
 
 impl dyn Diagnostic {
@@ -30,18 +30,37 @@ impl dyn Diagnostic {
     }
 }
 
-#[derive(Debug, Default)]
-pub struct DiagnosticSink {
-    data: Vec<Box<dyn Diagnostic>>,
+pub struct DiagnosticSink<'a> {
+    callbacks: Vec<Box<dyn FnMut(&dyn Diagnostic) -> Result<(), ()> + 'a>>,
+    default_callback: Box<dyn FnMut(&dyn Diagnostic) + 'a>,
 }
 
-impl DiagnosticSink {
-    pub fn push(&mut self, d: impl Diagnostic) {
-        self.data.push(Box::new(d))
+impl<'a> DiagnosticSink<'a> {
+    pub fn new(cb: impl FnMut(&dyn Diagnostic) + 'a) -> DiagnosticSink<'a> {
+        DiagnosticSink { callbacks: Vec::new(), default_callback: Box::new(cb) }
+    }
+
+    pub fn on<D: Diagnostic, F: FnMut(&D) + 'a>(mut self, mut cb: F) -> DiagnosticSink<'a> {
+        let cb = move |diag: &dyn Diagnostic| match diag.downcast_ref::<D>() {
+            Some(d) => {
+                cb(d);
+                Ok(())
+            }
+            None => Err(()),
+        };
+        self.callbacks.push(Box::new(cb));
+        self
     }
 
-    pub fn into_diagnostics(self) -> Vec<Box<dyn Diagnostic>> {
-        self.data
+    pub(crate) fn push(&mut self, d: impl Diagnostic) {
+        let d: &dyn Diagnostic = &d;
+        for cb in self.callbacks.iter_mut() {
+            match cb(d) {
+                Ok(()) => return,
+                Err(()) => (),
+            }
+        }
+        (self.default_callback)(d)
     }
 }
 
diff --git a/crates/ra_ide_api/src/diagnostics.rs b/crates/ra_ide_api/src/diagnostics.rs
index 5b1a90c82..bf77c9ab1 100644
--- a/crates/ra_ide_api/src/diagnostics.rs
+++ b/crates/ra_ide_api/src/diagnostics.rs
@@ -1,3 +1,5 @@
+use std::cell::RefCell;
+
 use itertools::Itertools;
 use hir::{source_binder, diagnostics::{Diagnostic as _, DiagnosticSink}};
 use ra_db::SourceDatabase;
@@ -25,11 +27,36 @@ pub(crate) fn diagnostics(db: &RootDatabase, file_id: FileId) -> Vec<Diagnostic>
         check_unnecessary_braces_in_use_statement(&mut res, file_id, node);
         check_struct_shorthand_initialization(&mut res, file_id, node);
     }
-
+    let res = RefCell::new(res);
+    let mut sink = DiagnosticSink::new(|d| {
+        res.borrow_mut().push(Diagnostic {
+            message: d.message(),
+            range: d.syntax_node().range(),
+            severity: Severity::Error,
+            fix: None,
+        })
+    })
+    .on::<hir::diagnostics::UnresolvedModule, _>(|d| {
+        let source_root = db.file_source_root(d.file().original_file(db));
+        let create_file = FileSystemEdit::CreateFile { source_root, path: d.candidate.clone() };
+        let fix = SourceChange {
+            label: "create module".to_string(),
+            source_file_edits: Vec::new(),
+            file_system_edits: vec![create_file],
+            cursor_position: None,
+        };
+        res.borrow_mut().push(Diagnostic {
+            range: d.syntax_node().range(),
+            message: d.message(),
+            severity: Severity::Error,
+            fix: Some(fix),
+        })
+    });
     if let Some(m) = source_binder::module_from_file_id(db, file_id) {
-        check_module(&mut res, db, m);
+        m.diagnostics(db, &mut sink);
     };
-    res
+    drop(sink);
+    res.into_inner()
 }
 
 fn syntax_errors(acc: &mut Vec<Diagnostic>, source_file: &SourceFile) {
@@ -127,52 +154,6 @@ fn check_struct_shorthand_initialization(
     Some(())
 }
 
-fn check_module(acc: &mut Vec<Diagnostic>, db: &RootDatabase, module: hir::Module) {
-    let mut diagnostics = DiagnosticSink::default();
-    module.diagnostics(db, &mut diagnostics);
-    for decl in module.declarations(db) {
-        match decl {
-            hir::ModuleDef::Function(f) => f.diagnostics(db, &mut diagnostics),
-            _ => (),
-        }
-    }
-
-    for impl_block in module.impl_blocks(db) {
-        for item in impl_block.items(db) {
-            match item {
-                hir::ImplItem::Method(f) => f.diagnostics(db, &mut diagnostics),
-                _ => (),
-            }
-        }
-    }
-
-    for d in diagnostics.into_diagnostics().iter() {
-        if let Some(d) = d.downcast_ref::<hir::diagnostics::UnresolvedModule>() {
-            let source_root = db.file_source_root(d.file().original_file(db));
-            let create_file = FileSystemEdit::CreateFile { source_root, path: d.candidate.clone() };
-            let fix = SourceChange {
-                label: "create module".to_string(),
-                source_file_edits: Vec::new(),
-                file_system_edits: vec![create_file],
-                cursor_position: None,
-            };
-            acc.push(Diagnostic {
-                range: d.syntax_node().range(),
-                message: d.message(),
-                severity: Severity::Error,
-                fix: Some(fix),
-            })
-        } else {
-            acc.push(Diagnostic {
-                message: d.message(),
-                range: d.syntax_node().range(),
-                severity: Severity::Error,
-                fix: None,
-            })
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use test_utils::assert_eq_text;