Correctly build BodySourceMap for macro-expanded expressions

6 files changed, 119 insertions, 66 deletions

diff --git a/crates/ra_hir/src/code_model/src.rs b/crates/ra_hir/src/code_model/src.rs
index 7c9454c0b..b9ffb0c7a 100644
--- a/crates/ra_hir/src/code_model/src.rs
+++ b/crates/ra_hir/src/code_model/src.rs
@@ -1,11 +1,15 @@
-use ra_syntax::ast::{self, AstNode};
+use ra_syntax::{
+    ast::{self, AstNode},
+    SyntaxNode,
+};
 
 use crate::{
-    ids::AstItemDef, AstDatabase, Const, DefDatabase, Enum, EnumVariant, FieldSource, Function,
-    HasBody, HirDatabase, HirFileId, MacroDef, Module, ModuleSource, Static, Struct, StructField,
-    Trait, TypeAlias, Union,
+    ids::AstItemDef, AstDatabase, Const, DefDatabase, Either, Enum, EnumVariant, FieldSource,
+    Function, HasBody, HirDatabase, HirFileId, MacroDef, Module, ModuleSource, Static, Struct,
+    StructField, Trait, TypeAlias, Union,
 };
 
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
 pub struct Source<T> {
     pub file_id: HirFileId,
     pub ast: T,
@@ -16,6 +20,15 @@ pub trait HasSource {
     fn source(self, db: &(impl DefDatabase + AstDatabase)) -> Source<Self::Ast>;
 }
 
+impl<T> Source<T> {
+    pub(crate) fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
+        Source { file_id: self.file_id, ast: f(self.ast) }
+    }
+    pub(crate) fn file_syntax(&self, db: &impl AstDatabase) -> SyntaxNode {
+        db.parse_or_expand(self.file_id).expect("source created from invalid file")
+    }
+}
+
 /// NB: Module is !HasSource, because it has two source nodes at the same time:
 /// definition and declaration.
 impl Module {
@@ -117,12 +130,12 @@ where
         self,
         db: &impl HirDatabase,
         expr_id: crate::expr::ExprId,
-    ) -> Option<Source<ast::Expr>> {
+    ) -> Option<Source<Either<ast::Expr, ast::RecordField>>> {
         let source_map = self.body_source_map(db);
-        let expr_syntax = source_map.expr_syntax(expr_id)?.a()?;
-        let source = self.source(db);
-        let ast = expr_syntax.to_node(&source.ast.syntax());
-        Some(Source { file_id: source.file_id, ast })
+        let source_ptr = source_map.expr_syntax(expr_id)?;
+        let root = source_ptr.file_syntax(db);
+        let source = source_ptr.map(|ast| ast.map(|it| it.to_node(&root), |it| it.to_node(&root)));
+        Some(source)
     }
 }
 
diff --git a/crates/ra_hir/src/expr.rs b/crates/ra_hir/src/expr.rs
index 5f6a4b320..fc21e269f 100644
--- a/crates/ra_hir/src/expr.rs
+++ b/crates/ra_hir/src/expr.rs
@@ -12,7 +12,7 @@ use crate::{
     path::GenericArgs,
     ty::primitive::{UncertainFloatTy, UncertainIntTy},
     type_ref::{Mutability, TypeRef},
-    DefWithBody, Either, HasSource, HirDatabase, Name, Path, Resolver,
+    DefWithBody, Either, HasSource, HirDatabase, Name, Path, Resolver, Source,
 };
 
 pub use self::scope::ExprScopes;
@@ -43,23 +43,32 @@ pub struct Body {
     body_expr: ExprId,
 }
 
+type ExprPtr = Either<AstPtr<ast::Expr>, AstPtr<ast::RecordField>>;
+type ExprSource = Source<ExprPtr>;
+
+type PatPtr = Either<AstPtr<ast::Pat>, AstPtr<ast::SelfParam>>;
+type PatSource = Source<PatPtr>;
+
 /// An item body together with the mapping from syntax nodes to HIR expression
 /// IDs. This is needed to go from e.g. a position in a file to the HIR
 /// expression containing it; but for type inference etc., we want to operate on
 /// a structure that is agnostic to the actual positions of expressions in the
 /// file, so that we don't recompute types whenever some whitespace is typed.
+///
+/// One complication here is that, due to macro expansion, a single `Body` might
+/// be spread across several files. So, for each ExprId and PatId, we record
+/// both the HirFileId and the position inside the file. However, we only store
+/// AST -> ExprId mapping for non-macro files, as it is not clear how to handle
+/// this properly for macros.
 #[derive(Default, Debug, Eq, PartialEq)]
 pub struct BodySourceMap {
     expr_map: FxHashMap<ExprPtr, ExprId>,
-    expr_map_back: ArenaMap<ExprId, ExprPtr>,
+    expr_map_back: ArenaMap<ExprId, ExprSource>,
     pat_map: FxHashMap<PatPtr, PatId>,
-    pat_map_back: ArenaMap<PatId, PatPtr>,
+    pat_map_back: ArenaMap<PatId, PatSource>,
     field_map: FxHashMap<(ExprId, usize), AstPtr<ast::RecordField>>,
 }
 
-type ExprPtr = Either<AstPtr<ast::Expr>, AstPtr<ast::RecordField>>;
-type PatPtr = Either<AstPtr<ast::Pat>, AstPtr<ast::SelfParam>>;
-
 impl Body {
     pub fn params(&self) -> &[PatId] {
         &self.params
@@ -123,16 +132,16 @@ impl Index<PatId> for Body {
 }
 
 impl BodySourceMap {
-    pub(crate) fn expr_syntax(&self, expr: ExprId) -> Option<ExprPtr> {
-        self.expr_map_back.get(expr).cloned()
+    pub(crate) fn expr_syntax(&self, expr: ExprId) -> Option<ExprSource> {
+        self.expr_map_back.get(expr).copied()
     }
 
     pub(crate) fn node_expr(&self, node: &ast::Expr) -> Option<ExprId> {
         self.expr_map.get(&Either::A(AstPtr::new(node))).cloned()
     }
 
-    pub(crate) fn pat_syntax(&self, pat: PatId) -> Option<PatPtr> {
-        self.pat_map_back.get(pat).cloned()
+    pub(crate) fn pat_syntax(&self, pat: PatId) -> Option<PatSource> {
+        self.pat_map_back.get(pat).copied()
     }
 
     pub(crate) fn node_pat(&self, node: &ast::Pat) -> Option<PatId> {
diff --git a/crates/ra_hir/src/expr/lower.rs b/crates/ra_hir/src/expr/lower.rs
index 7b3e55b7e..6afd80989 100644
--- a/crates/ra_hir/src/expr/lower.rs
+++ b/crates/ra_hir/src/expr/lower.rs
@@ -14,7 +14,7 @@ use crate::{
     ty::primitive::{FloatTy, IntTy, UncertainFloatTy, UncertainIntTy},
     type_ref::TypeRef,
     DefWithBody, Either, HirDatabase, HirFileId, MacroCallLoc, MacroFileKind, Mutability, Path,
-    Resolver,
+    Resolver, Source,
 };
 
 use super::{
@@ -103,11 +103,13 @@ where
         let id = self.body.exprs.alloc(expr);
         if self.current_file_id == self.original_file_id {
             self.source_map.expr_map.insert(ptr, id);
-            self.source_map.expr_map_back.insert(id, ptr);
         }
+        self.source_map
+            .expr_map_back
+            .insert(id, Source { file_id: self.current_file_id, ast: ptr });
         id
     }
-    // deshugared exprs don't have ptr, that's wrong and should be fixed
+    // desugared exprs don't have ptr, that's wrong and should be fixed
     // somehow.
     fn alloc_expr_desugared(&mut self, expr: Expr) -> ExprId {
         self.body.exprs.alloc(expr)
@@ -117,18 +119,18 @@ where
         let id = self.body.exprs.alloc(expr);
         if self.current_file_id == self.original_file_id {
             self.source_map.expr_map.insert(ptr, id);
-            self.source_map.expr_map_back.insert(id, ptr);
         }
+        self.source_map
+            .expr_map_back
+            .insert(id, Source { file_id: self.current_file_id, ast: ptr });
         id
     }
     fn alloc_pat(&mut self, pat: Pat, ptr: PatPtr) -> PatId {
         let id = self.body.pats.alloc(pat);
-
         if self.current_file_id == self.original_file_id {
             self.source_map.pat_map.insert(ptr, id);
-            self.source_map.pat_map_back.insert(id, ptr);
         }
-
+        self.source_map.pat_map_back.insert(id, Source { file_id: self.current_file_id, ast: ptr });
         id
     }
 
diff --git a/crates/ra_hir/src/expr/validation.rs b/crates/ra_hir/src/expr/validation.rs
index 6fdaf1fce..1202913e2 100644
--- a/crates/ra_hir/src/expr/validation.rs
+++ b/crates/ra_hir/src/expr/validation.rs
@@ -1,9 +1,8 @@
 use std::sync::Arc;
 
-use ra_syntax::ast::{self, AstNode};
+use ra_syntax::ast;
 use rustc_hash::FxHashSet;
 
-use super::{Expr, ExprId, RecordLitField};
 use crate::{
     adt::AdtDef,
     diagnostics::{DiagnosticSink, MissingFields, MissingOkInTailExpr},
@@ -11,9 +10,11 @@ use crate::{
     name,
     path::{PathKind, PathSegment},
     ty::{ApplicationTy, InferenceResult, Ty, TypeCtor},
-    Function, HasSource, HirDatabase, ModuleDef, Name, Path, PerNs, Resolution,
+    Function, HirDatabase, ModuleDef, Name, Path, PerNs, Resolution,
 };
 
+use super::{Expr, ExprId, RecordLitField};
+
 pub(crate) struct ExprValidator<'a, 'b: 'a> {
     func: Function,
     infer: Arc<InferenceResult>,
@@ -78,25 +79,20 @@ impl<'a, 'b> ExprValidator<'a, 'b> {
             return;
         }
         let source_map = self.func.body_source_map(db);
-        let file_id = self.func.source(db).file_id;
-        let parse = db.parse(file_id.original_file(db));
-        let source_file = parse.tree();
-        if let Some(field_list_node) = source_map
-            .expr_syntax(id)
-            .and_then(|ptr| ptr.a())
-            .map(|ptr| ptr.to_node(source_file.syntax()))
-            .and_then(|expr| match expr {
-                ast::Expr::RecordLit(it) => Some(it),
-                _ => None,
-            })
-            .and_then(|lit| lit.record_field_list())
-        {
-            let field_list_ptr = AstPtr::new(&field_list_node);
-            self.sink.push(MissingFields {
-                file: file_id,
-                field_list: field_list_ptr,
-                missed_fields,
-            })
+
+        if let Some(source_ptr) = source_map.expr_syntax(id) {
+            if let Some(expr) = source_ptr.ast.a() {
+                let root = source_ptr.file_syntax(db);
+                if let ast::Expr::RecordLit(record_lit) = expr.to_node(&root) {
+                    if let Some(field_list) = record_lit.record_field_list() {
+                        self.sink.push(MissingFields {
+                            file: source_ptr.file_id,
+                            field_list: AstPtr::new(&field_list),
+                            missed_fields,
+                        })
+                    }
+                }
+            }
         }
     }
 
@@ -136,10 +132,11 @@ impl<'a, 'b> ExprValidator<'a, 'b> {
 
         if params.len() == 2 && &params[0] == &mismatch.actual {
             let source_map = self.func.body_source_map(db);
-            let file_id = self.func.source(db).file_id;
 
-            if let Some(expr) = source_map.expr_syntax(id).and_then(|n| n.a()) {
-                self.sink.push(MissingOkInTailExpr { file: file_id, expr });
+            if let Some(source_ptr) = source_map.expr_syntax(id) {
+                if let Some(expr) = source_ptr.ast.a() {
+                    self.sink.push(MissingOkInTailExpr { file: source_ptr.file_id, expr });
+                }
             }
         }
     }
diff --git a/crates/ra_hir/src/source_binder.rs b/crates/ra_hir/src/source_binder.rs
index e5f4d11a6..fdbe5e8b0 100644
--- a/crates/ra_hir/src/source_binder.rs
+++ b/crates/ra_hir/src/source_binder.rs
@@ -228,7 +228,7 @@ impl SourceAnalyzer {
             let scopes = db.expr_scopes(def);
             let scope = match offset {
                 None => scope_for(&scopes, &source_map, &node),
-                Some(offset) => scope_for_offset(&scopes, &source_map, offset),
+                Some(offset) => scope_for_offset(&scopes, &source_map, file_id.into(), offset),
             };
             let resolver = expr::resolver_for_scope(def.body(db), db, scope);
             SourceAnalyzer {
@@ -330,6 +330,7 @@ impl SourceAnalyzer {
                     .body_source_map
                     .as_ref()?
                     .pat_syntax(it)?
+                    .ast // FIXME: ignoring file_id here is definitelly wrong
                     .map_a(|ptr| ptr.cast::<ast::BindPat>().unwrap());
                 PathResolution::LocalBinding(pat_ptr)
             }
@@ -354,7 +355,7 @@ impl SourceAnalyzer {
         ret.and_then(|entry| {
             Some(ScopeEntryWithSyntax {
                 name: entry.name().clone(),
-                ptr: source_map.pat_syntax(entry.pat())?,
+                ptr: source_map.pat_syntax(entry.pat())?.ast,
             })
         })
     }
@@ -470,20 +471,27 @@ fn scope_for(
 fn scope_for_offset(
     scopes: &ExprScopes,
     source_map: &BodySourceMap,
+    file_id: HirFileId,
     offset: TextUnit,
 ) -> Option<ScopeId> {
     scopes
         .scope_by_expr()
         .iter()
         .filter_map(|(id, scope)| {
-            let ast_ptr = source_map.expr_syntax(*id)?.a()?;
-            Some((ast_ptr.syntax_node_ptr(), scope))
+            let source = source_map.expr_syntax(*id)?;
+            // FIXME: correctly handle macro expansion
+            if source.file_id != file_id {
+                return None;
+            }
+            let syntax_node_ptr =
+                source.ast.either(|it| it.syntax_node_ptr(), |it| it.syntax_node_ptr());
+            Some((syntax_node_ptr, scope))
         })
         // find containing scope
         .min_by_key(|(ptr, _scope)| {
             (!(ptr.range().start() <= offset && offset <= ptr.range().end()), ptr.range().len())
         })
-        .map(|(ptr, scope)| adjust(scopes, source_map, ptr, offset).unwrap_or(*scope))
+        .map(|(ptr, scope)| adjust(scopes, source_map, ptr, file_id, offset).unwrap_or(*scope))
 }
 
 // XXX: during completion, cursor might be outside of any particular
@@ -492,6 +500,7 @@ fn adjust(
     scopes: &ExprScopes,
     source_map: &BodySourceMap,
     ptr: SyntaxNodePtr,
+    file_id: HirFileId,
     offset: TextUnit,
 ) -> Option<ScopeId> {
     let r = ptr.range();
@@ -499,8 +508,14 @@ fn adjust(
         .scope_by_expr()
         .iter()
         .filter_map(|(id, scope)| {
-            let ast_ptr = source_map.expr_syntax(*id)?.a()?;
-            Some((ast_ptr.syntax_node_ptr(), scope))
+            let source = source_map.expr_syntax(*id)?;
+            // FIXME: correctly handle macro expansion
+            if source.file_id != file_id {
+                return None;
+            }
+            let syntax_node_ptr =
+                source.ast.either(|it| it.syntax_node_ptr(), |it| it.syntax_node_ptr());
+            Some((syntax_node_ptr, scope))
         })
         .map(|(ptr, scope)| (ptr.range(), scope))
         .filter(|(range, _)| range.start() <= offset && range.is_subrange(&r) && *range != r);
diff --git a/crates/ra_hir/src/ty/tests.rs b/crates/ra_hir/src/ty/tests.rs
index d344ab12e..cde9801f6 100644
--- a/crates/ra_hir/src/ty/tests.rs
+++ b/crates/ra_hir/src/ty/tests.rs
@@ -2793,6 +2793,10 @@ fn main() {
 }
 "#),
         @r###"
+    ![0; 17) '{Foo(v...,2,])}': Foo
+    ![1; 4) 'Foo': Foo({unknown}) -> Foo
+    ![1; 16) 'Foo(vec![1,2,])': Foo
+    ![5; 15) 'vec![1,2,]': {unknown}
     [156; 182) '{     ...,2); }': ()
     [166; 167) 'x': Foo
     "###
@@ -3566,7 +3570,6 @@ fn infer(content: &str) -> String {
     let source_file = db.parse(file_id).ok().unwrap();
 
     let mut acc = String::new();
-    // acc.push_str("\n");
 
     let mut infer_def = |inference_result: Arc<InferenceResult>,
                          body_source_map: Arc<BodySourceMap>| {
@@ -3574,7 +3577,9 @@ fn infer(content: &str) -> String {
 
         for (pat, ty) in inference_result.type_of_pat.iter() {
             let syntax_ptr = match body_source_map.pat_syntax(pat) {
-                Some(sp) => sp.either(|it| it.syntax_node_ptr(), |it| it.syntax_node_ptr()),
+                Some(sp) => {
+                    sp.map(|ast| ast.either(|it| it.syntax_node_ptr(), |it| it.syntax_node_ptr()))
+                }
                 None => continue,
             };
             types.push((syntax_ptr, ty));
@@ -3582,22 +3587,34 @@ fn infer(content: &str) -> String {
 
         for (expr, ty) in inference_result.type_of_expr.iter() {
             let syntax_ptr = match body_source_map.expr_syntax(expr) {
-                Some(sp) => sp.either(|it| it.syntax_node_ptr(), |it| it.syntax_node_ptr()),
+                Some(sp) => {
+                    sp.map(|ast| ast.either(|it| it.syntax_node_ptr(), |it| it.syntax_node_ptr()))
+                }
                 None => continue,
             };
             types.push((syntax_ptr, ty));
         }
 
         // sort ranges for consistency
-        types.sort_by_key(|(ptr, _)| (ptr.range().start(), ptr.range().end()));
-        for (syntax_ptr, ty) in &types {
-            let node = syntax_ptr.to_node(source_file.syntax());
+        types.sort_by_key(|(src_ptr, _)| (src_ptr.ast.range().start(), src_ptr.ast.range().end()));
+        for (src_ptr, ty) in &types {
+            let node = src_ptr.ast.to_node(&src_ptr.file_syntax(&db));
+
             let (range, text) = if let Some(self_param) = ast::SelfParam::cast(node.clone()) {
                 (self_param.self_kw_token().text_range(), "self".to_string())
             } else {
-                (syntax_ptr.range(), node.text().to_string().replace("\n", " "))
+                (src_ptr.ast.range(), node.text().to_string().replace("\n", " "))
             };
-            write!(acc, "{} '{}': {}\n", range, ellipsize(text, 15), ty.display(&db)).unwrap();
+            let macro_prefix = if src_ptr.file_id != file_id.into() { "!" } else { "" };
+            write!(
+                acc,
+                "{}{} '{}': {}\n",
+                macro_prefix,
+                range,
+                ellipsize(text, 15),
+                ty.display(&db)
+            )
+            .unwrap();
         }
     };