1 files changed, 17 insertions, 5 deletions

diff --git a/crates/ra_hir_ty/src/test_db.rs b/crates/ra_hir_ty/src/test_db.rs
index 0481a7b12..fddf0604d 100644
--- a/crates/ra_hir_ty/src/test_db.rs
+++ b/crates/ra_hir_ty/src/test_db.rs
@@ -8,8 +8,10 @@ use std::{
 use hir_def::{db::DefDatabase, AssocItemId, ModuleDefId, ModuleId};
 use hir_expand::{db::AstDatabase, diagnostics::DiagnosticSink};
 use ra_db::{salsa, CrateId, FileId, FileLoader, FileLoaderDelegate, SourceDatabase, Upcast};
-use rustc_hash::FxHashSet;
+use ra_syntax::TextRange;
+use rustc_hash::{FxHashMap, FxHashSet};
 use stdx::format_to;
+use test_utils::extract_annotations;
 
 use crate::{
     db::HirDatabase, diagnostics::Diagnostic, expr::ExprValidator,
@@ -155,17 +157,27 @@ impl TestDB {
         (buf, count)
     }
 
-    pub fn all_files(&self) -> Vec<FileId> {
-        let mut res = Vec::new();
+    pub fn extract_annotations(&self) -> FxHashMap<FileId, Vec<(TextRange, String)>> {
+        let mut files = Vec::new();
         let crate_graph = self.crate_graph();
         for krate in crate_graph.iter() {
             let crate_def_map = self.crate_def_map(krate);
             for (module_id, _) in crate_def_map.modules.iter() {
                 let file_id = crate_def_map[module_id].origin.file_id();
-                res.extend(file_id)
+                files.extend(file_id)
             }
         }
-        res
+        files
+            .into_iter()
+            .filter_map(|file_id| {
+                let text = self.file_text(file_id);
+                let annotations = extract_annotations(&text);
+                if annotations.is_empty() {
+                    return None;
+                }
+                Some((file_id, annotations))
+            })
+            .collect()
     }
 }