Handle macro-generated expressions slightly less wrong

2 files changed, 20 insertions, 32 deletions

diff --git a/crates/ra_hir_def/src/body.rs b/crates/ra_hir_def/src/body.rs
index 65fefd912..75bba31c2 100644
--- a/crates/ra_hir_def/src/body.rs
+++ b/crates/ra_hir_def/src/body.rs
@@ -21,7 +21,6 @@ use crate::{
 
 pub struct Expander {
     crate_def_map: Arc<CrateDefMap>,
-    original_file_id: HirFileId,
     current_file_id: HirFileId,
     hygiene: Hygiene,
     module: ModuleId,
@@ -31,13 +30,7 @@ impl Expander {
     pub fn new(db: &impl DefDatabase2, current_file_id: HirFileId, module: ModuleId) -> Expander {
         let crate_def_map = db.crate_def_map(module.krate);
         let hygiene = Hygiene::new(db, current_file_id);
-        Expander {
-            crate_def_map,
-            original_file_id: current_file_id,
-            current_file_id,
-            hygiene,
-            module,
-        }
+        Expander { crate_def_map, current_file_id, hygiene, module }
     }
 
     fn expand(
@@ -82,11 +75,6 @@ impl Expander {
         std::mem::forget(mark);
     }
 
-    // FIXME: remove this.
-    fn is_in_expansion(&self) -> bool {
-        self.original_file_id != self.current_file_id
-    }
-
     fn to_source<T>(&self, ast: T) -> Source<T> {
         Source { file_id: self.current_file_id, ast }
     }
@@ -147,9 +135,9 @@ pub type PatSource = Source<PatPtr>;
 /// this properly for macros.
 #[derive(Default, Debug, Eq, PartialEq)]
 pub struct BodySourceMap {
-    expr_map: FxHashMap<ExprPtr, ExprId>,
+    expr_map: FxHashMap<ExprSource, ExprId>,
     expr_map_back: ArenaMap<ExprId, ExprSource>,
-    pat_map: FxHashMap<PatPtr, PatId>,
+    pat_map: FxHashMap<PatSource, PatId>,
     pat_map_back: ArenaMap<PatId, PatSource>,
     field_map: FxHashMap<(ExprId, usize), AstPtr<ast::RecordField>>,
 }
@@ -202,16 +190,18 @@ impl BodySourceMap {
         self.expr_map_back.get(expr).copied()
     }
 
-    pub fn node_expr(&self, node: &ast::Expr) -> Option<ExprId> {
-        self.expr_map.get(&Either::A(AstPtr::new(node))).cloned()
+    pub fn node_expr(&self, node: Source<&ast::Expr>) -> Option<ExprId> {
+        let src = node.map(|it| Either::A(AstPtr::new(it)));
+        self.expr_map.get(&src).cloned()
     }
 
     pub fn pat_syntax(&self, pat: PatId) -> Option<PatSource> {
         self.pat_map_back.get(pat).copied()
     }
 
-    pub fn node_pat(&self, node: &ast::Pat) -> Option<PatId> {
-        self.pat_map.get(&Either::A(AstPtr::new(node))).cloned()
+    pub fn node_pat(&self, node: Source<&ast::Pat>) -> Option<PatId> {
+        let src = node.map(|it| Either::A(AstPtr::new(it)));
+        self.pat_map.get(&src).cloned()
     }
 
     pub fn field_syntax(&self, expr: ExprId, field: usize) -> AstPtr<ast::RecordField> {
diff --git a/crates/ra_hir_def/src/body/lower.rs b/crates/ra_hir_def/src/body/lower.rs
index c45500195..f6d79ddf0 100644
--- a/crates/ra_hir_def/src/body/lower.rs
+++ b/crates/ra_hir_def/src/body/lower.rs
@@ -94,10 +94,9 @@ where
     fn alloc_expr(&mut self, expr: Expr, ptr: AstPtr<ast::Expr>) -> ExprId {
         let ptr = Either::A(ptr);
         let id = self.body.exprs.alloc(expr);
-        if !self.expander.is_in_expansion() {
-            self.source_map.expr_map.insert(ptr, id);
-        }
-        self.source_map.expr_map_back.insert(id, self.expander.to_source(ptr));
+        let src = self.expander.to_source(ptr);
+        self.source_map.expr_map.insert(src, id);
+        self.source_map.expr_map_back.insert(id, src);
         id
     }
     // desugared exprs don't have ptr, that's wrong and should be fixed
@@ -108,18 +107,16 @@ where
     fn alloc_expr_field_shorthand(&mut self, expr: Expr, ptr: AstPtr<ast::RecordField>) -> ExprId {
         let ptr = Either::B(ptr);
         let id = self.body.exprs.alloc(expr);
-        if !self.expander.is_in_expansion() {
-            self.source_map.expr_map.insert(ptr, id);
-        }
-        self.source_map.expr_map_back.insert(id, self.expander.to_source(ptr));
+        let src = self.expander.to_source(ptr);
+        self.source_map.expr_map.insert(src, id);
+        self.source_map.expr_map_back.insert(id, src);
         id
     }
     fn alloc_pat(&mut self, pat: Pat, ptr: PatPtr) -> PatId {
         let id = self.body.pats.alloc(pat);
-        if !self.expander.is_in_expansion() {
-            self.source_map.pat_map.insert(ptr, id);
-        }
-        self.source_map.pat_map_back.insert(id, self.expander.to_source(ptr));
+        let src = self.expander.to_source(ptr);
+        self.source_map.pat_map.insert(src, id);
+        self.source_map.pat_map_back.insert(id, src);
         id
     }
 
@@ -277,7 +274,8 @@ where
             ast::Expr::ParenExpr(e) => {
                 let inner = self.collect_expr_opt(e.expr());
                 // make the paren expr point to the inner expression as well
-                self.source_map.expr_map.insert(Either::A(syntax_ptr), inner);
+                let src = self.expander.to_source(Either::A(syntax_ptr));
+                self.source_map.expr_map.insert(src, inner);
                 inner
             }
             ast::Expr::ReturnExpr(e) => {