Proper handling $crate Take 2

2 files changed, 186 insertions, 36 deletions

diff --git a/crates/hir_expand/src/db.rs b/crates/hir_expand/src/db.rs
index 0a0d021e0..ab2637b8c 100644
--- a/crates/hir_expand/src/db.rs
+++ b/crates/hir_expand/src/db.rs
@@ -8,9 +8,9 @@ use parser::FragmentKind;
 use syntax::{algo::diff, ast::NameOwner, AstNode, GreenNode, Parse, SyntaxKind::*, SyntaxNode};
 
 use crate::{
-    ast_id_map::AstIdMap, BuiltinDeriveExpander, BuiltinFnLikeExpander, EagerCallLoc, EagerMacroId,
-    HirFileId, HirFileIdRepr, LazyMacroId, MacroCallId, MacroCallLoc, MacroDefId, MacroDefKind,
-    MacroFile, ProcMacroExpander,
+    ast_id_map::AstIdMap, hygiene::HygieneFrame, BuiltinDeriveExpander, BuiltinFnLikeExpander,
+    EagerCallLoc, EagerMacroId, HirFileId, HirFileIdRepr, LazyMacroId, MacroCallId, MacroCallLoc,
+    MacroDefId, MacroDefKind, MacroFile, ProcMacroExpander,
 };
 
 /// Total limit on the number of tokens produced by any macro invocation.
@@ -94,6 +94,8 @@ pub trait AstDatabase: SourceDatabase {
     fn intern_eager_expansion(&self, eager: EagerCallLoc) -> EagerMacroId;
 
     fn expand_proc_macro(&self, call: MacroCallId) -> Result<tt::Subtree, mbe::ExpandError>;
+
+    fn hygiene_frame(&self, file_id: HirFileId) -> Arc<HygieneFrame>;
 }
 
 /// This expands the given macro call, but with different arguments. This is
@@ -369,6 +371,10 @@ fn parse_macro_with_arg(
     }
 }
 
+fn hygiene_frame(db: &dyn AstDatabase, file_id: HirFileId) -> Arc<HygieneFrame> {
+    Arc::new(HygieneFrame::new(db, file_id))
+}
+
 /// Given a `MacroCallId`, return what `FragmentKind` it belongs to.
 /// FIXME: Not completed
 fn to_fragment_kind(db: &dyn AstDatabase, id: MacroCallId) -> FragmentKind {
diff --git a/crates/hir_expand/src/hygiene.rs b/crates/hir_expand/src/hygiene.rs
index 7ab0a5e52..8db581b77 100644
--- a/crates/hir_expand/src/hygiene.rs
+++ b/crates/hir_expand/src/hygiene.rs
@@ -2,65 +2,209 @@
 //!
 //! Specifically, `ast` + `Hygiene` allows you to create a `Name`. Note that, at
 //! this moment, this is horribly incomplete and handles only `$crate`.
+use std::sync::Arc;
+
 use base_db::CrateId;
 use either::Either;
-use syntax::ast;
+use mbe::Origin;
+use parser::SyntaxKind;
+use syntax::{ast, AstNode, SyntaxNode, TextRange, TextSize};
 
 use crate::{
-    db::AstDatabase,
+    db::{self, AstDatabase},
     name::{AsName, Name},
-    HirFileId, HirFileIdRepr, MacroCallId, MacroDefKind,
+    HirFileId, HirFileIdRepr, InFile, MacroCallId, MacroCallLoc, MacroDefKind, MacroFile,
 };
 
 #[derive(Clone, Debug)]
 pub struct Hygiene {
-    // This is what `$crate` expands to
-    def_crate: Option<CrateId>,
-
-    // Indicate this is a local inner macro
-    local_inner: bool,
+    frames: Option<HygieneFrames>,
 }
 
 impl Hygiene {
     pub fn new(db: &dyn AstDatabase, file_id: HirFileId) -> Hygiene {
-        let (def_crate, local_inner) = match file_id.0 {
-            HirFileIdRepr::FileId(_) => (None, false),
-            HirFileIdRepr::MacroFile(macro_file) => match macro_file.macro_call_id {
-                MacroCallId::LazyMacro(id) => {
-                    let loc = db.lookup_intern_macro(id);
-                    match loc.def.kind {
-                        MacroDefKind::Declarative => (Some(loc.def.krate), loc.def.local_inner),
-                        MacroDefKind::BuiltIn(_) => (Some(loc.def.krate), false),
-                        MacroDefKind::BuiltInDerive(_) => (None, false),
-                        MacroDefKind::BuiltInEager(_) => (None, false),
-                        MacroDefKind::ProcMacro(_) => (None, false),
-                    }
-                }
-                MacroCallId::EagerMacro(_id) => (None, false),
-            },
-        };
-        Hygiene { def_crate, local_inner }
+        Hygiene { frames: Some(HygieneFrames::new(db, file_id.clone())) }
     }
 
     pub fn new_unhygienic() -> Hygiene {
-        Hygiene { def_crate: None, local_inner: false }
+        Hygiene { frames: None }
     }
 
     // FIXME: this should just return name
     pub fn name_ref_to_name(&self, name_ref: ast::NameRef) -> Either<Name, CrateId> {
-        if let Some(def_crate) = self.def_crate {
+        if let Some(frames) = &self.frames {
             if name_ref.text() == "$crate" {
-                return Either::Right(def_crate);
+                if let Some(krate) = frames.root_crate(name_ref.syntax()) {
+                    return Either::Right(krate);
+                }
             }
         }
+
         Either::Left(name_ref.as_name())
     }
 
-    pub fn local_inner_macros(&self) -> Option<CrateId> {
-        if self.local_inner {
-            self.def_crate
-        } else {
-            None
+    pub fn local_inner_macros(&self, path: ast::Path) -> Option<CrateId> {
+        let mut token = path.syntax().first_token()?.text_range();
+        let frames = self.frames.as_ref()?;
+        let mut current = frames.0.clone();
+
+        loop {
+            let (mapped, origin) = current.expansion.as_ref()?.map_ident_up(token)?;
+            if origin == Origin::Def {
+                return if current.local_inner { frames.root_crate(path.syntax()) } else { None };
+            }
+            current = current.call_site.as_ref()?.clone();
+            token = mapped.value;
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+struct HygieneFrames(Arc<HygieneFrame>);
+
+#[derive(Clone, Debug, Eq, PartialEq)]
+pub struct HygieneFrame {
+    expansion: Option<HygieneInfo>,
+
+    // Indicate this is a local inner macro
+    local_inner: bool,
+    krate: Option<CrateId>,
+
+    call_site: Option<Arc<HygieneFrame>>,
+    def_site: Option<Arc<HygieneFrame>>,
+}
+
+impl HygieneFrames {
+    fn new(db: &dyn AstDatabase, file_id: HirFileId) -> Self {
+        HygieneFrames(Arc::new(HygieneFrame::new(db, file_id)))
+    }
+
+    fn root_crate(&self, node: &SyntaxNode) -> Option<CrateId> {
+        let mut token = node.first_token()?.text_range();
+        let mut result = self.0.krate;
+        let mut current = self.0.clone();
+
+        while let Some((mapped, origin)) =
+            current.expansion.as_ref().and_then(|it| it.map_ident_up(token))
+        {
+            result = current.krate;
+
+            let site = match origin {
+                Origin::Def => &current.def_site,
+                Origin::Call => &current.call_site,
+            };
+
+            let site = match site {
+                None => break,
+                Some(it) => it,
+            };
+
+            current = site.clone();
+            token = mapped.value;
         }
+
+        result
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct HygieneInfo {
+    arg_start: InFile<TextSize>,
+    /// The `macro_rules!` arguments.
+    def_start: Option<InFile<TextSize>>,
+
+    macro_def: Arc<(db::TokenExpander, mbe::TokenMap)>,
+    macro_arg: Arc<(tt::Subtree, mbe::TokenMap)>,
+    exp_map: Arc<mbe::TokenMap>,
+}
+
+impl HygieneInfo {
+    fn map_ident_up(&self, token: TextRange) -> Option<(InFile<TextRange>, Origin)> {
+        let token_id = self.exp_map.token_by_range(token)?;
+
+        let (token_id, origin) = self.macro_def.0.map_id_up(token_id);
+        let (token_map, tt) = match origin {
+            mbe::Origin::Call => (&self.macro_arg.1, self.arg_start),
+            mbe::Origin::Def => (
+                &self.macro_def.1,
+                self.def_start
+                    .as_ref()
+                    .expect("`Origin::Def` used with non-`macro_rules!` macro")
+                    .clone(),
+            ),
+        };
+
+        let range = token_map.range_by_token(token_id)?.by_kind(SyntaxKind::IDENT)?;
+        Some((tt.with_value(range + tt.value), origin))
+    }
+}
+
+fn make_hygiene_info(
+    db: &dyn AstDatabase,
+    macro_file: MacroFile,
+    loc: &MacroCallLoc,
+) -> Option<HygieneInfo> {
+    let arg_tt = loc.kind.arg(db)?;
+
+    let def_offset = loc.def.ast_id.and_then(|id| {
+        let def_tt = match id.to_node(db) {
+            ast::Macro::MacroRules(mac) => mac.token_tree()?.syntax().text_range().start(),
+            ast::Macro::MacroDef(_) => return None,
+        };
+        Some(InFile::new(id.file_id, def_tt))
+    });
+
+    let macro_def = db.macro_def(loc.def)?;
+    let (_, exp_map) = db.parse_macro_expansion(macro_file).value?;
+    let macro_arg = db.macro_arg(macro_file.macro_call_id)?;
+
+    Some(HygieneInfo {
+        arg_start: InFile::new(loc.kind.file_id(), arg_tt.text_range().start()),
+        def_start: def_offset,
+        macro_arg,
+        macro_def,
+        exp_map,
+    })
+}
+
+impl HygieneFrame {
+    pub(crate) fn new(db: &dyn AstDatabase, file_id: HirFileId) -> HygieneFrame {
+        let (info, krate, local_inner) = match file_id.0 {
+            HirFileIdRepr::FileId(_) => (None, None, false),
+            HirFileIdRepr::MacroFile(macro_file) => match macro_file.macro_call_id {
+                MacroCallId::EagerMacro(_id) => (None, None, false),
+                MacroCallId::LazyMacro(id) => {
+                    let loc = db.lookup_intern_macro(id);
+                    let info = make_hygiene_info(db, macro_file, &loc);
+                    match loc.def.kind {
+                        MacroDefKind::Declarative => {
+                            (info, Some(loc.def.krate), loc.def.local_inner)
+                        }
+                        MacroDefKind::BuiltIn(_) => (info, Some(loc.def.krate), false),
+                        MacroDefKind::BuiltInDerive(_) => (info, None, false),
+                        MacroDefKind::BuiltInEager(_) => (info, None, false),
+                        MacroDefKind::ProcMacro(_) => (info, None, false),
+                    }
+                }
+            },
+        };
+
+        let info = match info {
+            None => {
+                return HygieneFrame {
+                    expansion: None,
+                    local_inner,
+                    krate,
+                    call_site: None,
+                    def_site: None,
+                };
+            }
+            Some(it) => it,
+        };
+
+        let def_site = info.def_start.map(|it| db.hygiene_frame(it.file_id));
+        let call_site = Some(db.hygiene_frame(info.arg_start.file_id));
+
+        HygieneFrame { expansion: Some(info), local_inner, krate, call_site, def_site }
     }
 }