From 557df6ff3f90196417df7c7a9903c5dfd09a4f47 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Fri, 20 Aug 2021 13:49:28 +0200
Subject: [PATCH] Use correct HirFileId in find_related_test

---
 crates/hir/src/semantics.rs  |  4 ++++
 crates/hir_expand/src/lib.rs | 13 ++++---------
 crates/ide/src/runnables.rs  | 13 ++++++-------
 3 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 84563fd892..68b52ddbc4 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -173,6 +173,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.descend_node_at_offset(node, offset).find_map(N::cast)
     }
 
+    pub fn hir_file_for(&self, syntax_node: &SyntaxNode) -> HirFileId {
+        self.imp.find_file(syntax_node.clone()).file_id
+    }
+
     pub fn original_range(&self, node: &SyntaxNode) -> FileRange {
         self.imp.original_range(node)
     }
diff --git a/crates/hir_expand/src/lib.rs b/crates/hir_expand/src/lib.rs
index c9165a0cf6..4c09b2ab19 100644
--- a/crates/hir_expand/src/lib.rs
+++ b/crates/hir_expand/src/lib.rs
@@ -22,7 +22,7 @@ use either::Either;
 pub use mbe::{ExpandError, ExpandResult};
 pub use parser::FragmentKind;
 
-use std::{hash::Hash, sync::Arc};
+use std::{hash::Hash, iter, sync::Arc};
 
 use base_db::{impl_intern_key, salsa, CrateId, FileId, FileRange};
 use syntax::{
@@ -454,7 +454,7 @@ impl InFile<SyntaxNode> {
         self,
         db: &dyn db::AstDatabase,
     ) -> impl Iterator<Item = InFile<SyntaxNode>> + '_ {
-        std::iter::successors(Some(self), move |node| match node.value.parent() {
+        iter::successors(Some(self), move |node| match node.value.parent() {
             Some(parent) => Some(node.with_value(parent)),
             None => {
                 let parent_node = node.file_id.call_node(db)?;
@@ -570,19 +570,14 @@ impl<N: AstNode> InFile<N> {
     where
         N: 'db,
     {
-        std::iter::successors(Some(self), move |node| {
+        iter::successors(Some(self), move |node| {
             let InFile { file_id, value } = node.file_id.call_node(db)?;
             N::cast(value).map(|n| InFile::new(file_id, n))
         })
     }
 
     pub fn node_with_attributes(self, db: &dyn db::AstDatabase) -> InFile<N> {
-        std::iter::successors(Some(self), move |node| {
-            let InFile { file_id, value } = node.file_id.call_node(db)?;
-            N::cast(value).map(|n| InFile::new(file_id, n))
-        })
-        .last()
-        .unwrap()
+        self.nodes_with_attributes(db).last().unwrap()
     }
 }
 
diff --git a/crates/ide/src/runnables.rs b/crates/ide/src/runnables.rs
index ed220c9080..d2350db037 100644
--- a/crates/ide/src/runnables.rs
+++ b/crates/ide/src/runnables.rs
@@ -3,7 +3,7 @@ use std::fmt;
 use ast::NameOwner;
 use cfg::CfgExpr;
 use either::Either;
-use hir::{AsAssocItem, HasAttrs, HasSource, HirDisplay, Semantics};
+use hir::{AsAssocItem, HasAttrs, HasSource, HirDisplay, InFile, Semantics};
 use ide_assists::utils::test_related_attribute;
 use ide_db::{
     base_db::{FilePosition, FileRange},
@@ -232,22 +232,21 @@ fn find_related_tests(
             let functions = refs.iter().filter_map(|(range, _)| {
                 let token = file.token_at_offset(range.start()).next()?;
                 let token = sema.descend_into_macros(token);
-                // FIXME: This is the wrong file_id
                 token
                     .ancestors()
                     .find_map(ast::Fn::cast)
-                    .map(|f| hir::InFile::new(file_id.into(), f))
+                    .map(|f| hir::InFile::new(sema.hir_file_for(f.syntax()), f))
             });
 
             for fn_def in functions {
                 // #[test/bench] expands to just the item causing us to lose the attribute, so recover them by going out of the attribute
-                let fn_def = fn_def.node_with_attributes(sema.db);
-                if let Some(runnable) = as_test_runnable(sema, &fn_def.value) {
+                let InFile { value: fn_def, .. } = &fn_def.node_with_attributes(sema.db);
+                if let Some(runnable) = as_test_runnable(sema, fn_def) {
                     // direct test
                     tests.insert(runnable);
-                } else if let Some(module) = parent_test_module(sema, &fn_def.value) {
+                } else if let Some(module) = parent_test_module(sema, fn_def) {
                     // indirect test
-                    find_related_tests_in_module(sema, &fn_def.value, &module, tests);
+                    find_related_tests_in_module(sema, fn_def, &module, tests);
                 }
             }
         }