Allow more precise span origin tracing

crates/common2/src/diagnostics.rs

@@ -52,11 +52,33 @@ pub struct SubDiagnostic {
 pub struct Span {
     pub file: InputFile,
     pub range: TextRange,
+    pub kind: SpanKind,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum SpanKind {
+    /// A node corresponding is originally written in the source code.
+    Original,
+
+    /// A node corresponding to the span is generated by macro expansion.
+    Expanded,
+
+    /// No span information was found.
+    /// This happens if analysis code tries to get a span for a node that is
+    /// generated in lowering phase.
+    ///
+    /// If span has this kind, it means there is a bug in the analysis code.
+    /// The reason not to panic is that LSP should continue working even if
+    /// there are bugs in the span generation(This also makes easier to identify
+    /// the cause of the bug)
+    ///
+    /// Range is always the first character of the file in this case.
+    NotFound,
 }
 
 impl Span {
-    pub fn new(file: InputFile, range: TextRange) -> Self {
-        Self { file, range }
+    pub fn new(file: InputFile, range: TextRange, kind: SpanKind) -> Self {
+        Self { file, range, kind }
     }
 }
 

crates/hir/src/hir_def/expr.rs

@@ -1,5 +1,7 @@
 use cranelift_entity::entity_impl;
 
+use crate::span::expr::LazyExprSpan;
+
 use super::{Body, GenericArgListId, IdentId, IntegerId, LitKind, Partial, PatId, PathId, StmtId};
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
@@ -44,6 +46,12 @@ pub enum Expr {
 pub struct ExprId(u32);
 entity_impl!(ExprId);
 
+impl ExprId {
+    pub fn lazy_span(self, body: Body) -> LazyExprSpan {
+        LazyExprSpan::new(self, body)
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum FieldIndex {
     /// The field is indexed by its name.

crates/hir/src/hir_def/pat.rs

@@ -1,6 +1,8 @@
 use cranelift_entity::entity_impl;
 
-use super::{IdentId, LitKind, Partial, PathId};
+use crate::span::pat::LazyPatSpan;
+
+use super::{Body, IdentId, LitKind, Partial, PathId};
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum Pat {
@@ -18,6 +20,12 @@ pub enum Pat {
 pub struct PatId(u32);
 entity_impl!(PatId);
 
+impl PatId {
+    pub fn lazy_span(self, body: Body) -> LazyPatSpan {
+        LazyPatSpan::new(self, body)
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct RecordPatField {
     pub label: Partial<IdentId>,

crates/hir/src/hir_def/stmt.rs

@@ -1,6 +1,8 @@
 use cranelift_entity::entity_impl;
 
-use super::{ExprId, PatId, TypeId};
+use crate::span::stmt::LazyStmtSpan;
+
+use super::{Body, ExprId, PatId, TypeId};
 
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum Stmt {
@@ -31,3 +33,9 @@ pub enum Stmt {
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct StmtId(u32);
 entity_impl!(StmtId);
+
+impl StmtId {
+    pub fn lazy_span(self, body: Body) -> LazyStmtSpan {
+        LazyStmtSpan::new(self, body)
+    }
+}

crates/hir/src/lib.rs

@@ -133,16 +133,12 @@ mod test_db {
 
         /// Parses the given source text and returns the first inner item in the
         /// file.
-        pub fn parse_source_to_first_item<T>(&mut self, text: &str) -> T
+        pub fn expect_item<T>(&mut self, text: &str) -> T
         where
             ItemKind: TryInto<T, Error = &'static str>,
         {
             let tree = self.parse_source(text);
-            tree.children(tree.top_mod)
-                .next()
-                .unwrap()
-                .try_into()
-                .unwrap()
+            tree.dfs().find_map(|it| it.try_into().ok()).unwrap()
         }
 
         pub fn text_at(&self, top_mod: TopLevelMod, span: &impl LazySpan) -> &str {

crates/hir/src/lower/parse.rs

@@ -1,5 +1,5 @@
 use common::{
-    diagnostics::{AnalysisPass, CompleteDiagnostic, GlobalErrorCode, Severity, Span},
+    diagnostics::{AnalysisPass, CompleteDiagnostic, GlobalErrorCode, Severity, Span, SpanKind},
     InputFile,
 };
 use parser::GreenNode;
@@ -38,7 +38,7 @@ impl DiagnosticVoucher for ParseDiagnostic {
 
     fn to_complete(self, _db: &dyn SpannedHirDb) -> CompleteDiagnostic {
         let error_code = self.error_code();
-        let span = Span::new(self.file, self.error.range);
+        let span = Span::new(self.file, self.error.range, SpanKind::Original);
         CompleteDiagnostic::new(Severity::Error, self.error.msg, span, vec![], error_code)
     }
 }

crates/hir/src/span/attr.rs

@@ -1,5 +1,7 @@
 use parser::ast::{self, prelude::*};
 
+use crate::span::transition::ResolvedOrigin;
+
 use super::define_lazy_span_node;
 
 define_lazy_span_node!(
@@ -11,17 +13,16 @@ define_lazy_span_node!(
 );
 impl LazyAttrListSpan {
     pub fn normal_attr(&self, idx: usize) -> LazyNormalAttrSpan {
-        fn f(
-            node: parser::SyntaxNode,
-            arg: crate::span::transition::LazyArg,
-        ) -> Option<parser::NodeOrToken> {
+        fn f(origin: ResolvedOrigin, arg: crate::span::transition::LazyArg) -> ResolvedOrigin {
             let idx = match arg {
                 crate::span::transition::LazyArg::Idx(idx) => idx,
                 _ => unreachable!(),
             };
-            ast::AttrList::cast(node)
-                .and_then(|f| f.normal_attrs().nth(idx))
-                .map(|n| n.syntax().clone().into())
+            origin.map(|node| {
+                ast::AttrList::cast(node)
+                    .and_then(|f| f.normal_attrs().nth(idx))
+                    .map(|n| n.syntax().clone().into())
+            })
         }
 
         let lazy_transition = crate::span::transition::LazyTransitionFn {

crates/hir/src/span/expr.rs

@@ -1,5 +1,4 @@
-use common::InputFile;
-use parser::{ast, SyntaxNode};
+use parser::ast;
 
 use crate::{
     hir_def::{Body, ExprId},
@@ -8,8 +7,8 @@ use crate::{
 };
 
 use super::{
-    body_ast, body_source_map, define_lazy_span_node,
-    transition::{ChainInitiator, SpanTransitionChain},
+    body_source_map, define_lazy_span_node,
+    transition::{ChainInitiator, ResolvedOrigin, SpanTransitionChain},
 };
 
 define_lazy_span_node!(LazyExprSpan, ast::Expr,);
@@ -169,15 +168,10 @@ pub(crate) struct ExprRoot {
 }
 
 impl ChainInitiator for ExprRoot {
-    fn init(&self, db: &dyn SpannedHirDb) -> (InputFile, SyntaxNode) {
+    fn init(&self, db: &dyn SpannedHirDb) -> ResolvedOrigin {
         let source_map = body_source_map(db, self.body);
-        let expr_source = source_map.expr_map.node_to_source(self.expr);
-        let ptr = expr_source
-            .syntax_ptr()
-            .unwrap_or_else(|| body_ast(db, self.body).syntax_ptr().unwrap());
-
-        let (file, root_node) = self.body.top_mod(db.upcast()).init(db);
-        let node = ptr.to_node(&root_node);
-        (file, node)
+        let origin = source_map.expr_map.node_to_source(self.expr);
+        let top_mod = self.body.top_mod(db.upcast());
+        ResolvedOrigin::resolve(db, top_mod, origin)
     }
 }

crates/hir/src/span/item.rs

@@ -275,7 +275,7 @@ mod tests {
             }
         "#;
 
-        let mod_ = db.parse_source_to_first_item::<Mod>(text);
+        let mod_ = db.expect_item::<Mod>(text);
         let top_mod = mod_.top_mod(db.upcast());
         let mod_span = mod_.lazy_span();
         assert_eq!(
@@ -296,7 +296,7 @@ mod tests {
                 where U: Add
         "#;
 
-        let fn_ = db.parse_source_to_first_item::<Func>(text);
+        let fn_ = db.expect_item::<Func>(text);
         let top_mod = fn_.top_mod(db.upcast());
         let fn_span = fn_.lazy_span();
         assert_eq!("my_func", db.text_at(top_mod, &fn_span.name()));
@@ -347,7 +347,7 @@ mod tests {
                 pub y: foo::Bar<2>
             }"#;
 
-        let struct_ = db.parse_source_to_first_item::<Struct>(text);
+        let struct_ = db.expect_item::<Struct>(text);
         let top_mod = struct_.top_mod(db.upcast());
         let struct_span = struct_.lazy_span();
         assert_eq!("Foo", db.text_at(top_mod, &struct_span.name()));
@@ -374,7 +374,7 @@ mod tests {
                 Baz(u32, i32)
             }"#;
 
-        let enum_ = db.parse_source_to_first_item::<Enum>(text);
+        let enum_ = db.expect_item::<Enum>(text);
         let top_mod = enum_.top_mod(db.upcast());
         let enum_span = enum_.lazy_span();
         assert_eq!("Foo", db.text_at(top_mod, &enum_span.name()));
@@ -396,7 +396,7 @@ mod tests {
             pub type Foo = u32
         "#;
 
-        let type_alias = db.parse_source_to_first_item::<TypeAlias>(text);
+        let type_alias = db.expect_item::<TypeAlias>(text);
         let top_mod = type_alias.top_mod(db.upcast());
         let type_alias_span = type_alias.lazy_span();
         assert_eq!("Foo", db.text_at(top_mod, &type_alias_span.alias()));
@@ -412,7 +412,7 @@ mod tests {
             use foo::bar::{baz::*, qux as Alias}
         "#;
 
-        let use_ = db.parse_source_to_first_item::<Use>(text);
+        let use_ = db.expect_item::<Use>(text);
         let top_mod = use_.top_mod(db.upcast());
         let use_tree = use_.lazy_span().use_tree();
 

crates/hir/src/span/mod.rs

@@ -128,14 +128,6 @@ where
         Self::Raw(AstPtr::new(ast))
     }
 
-    fn syntax_ptr(&self) -> Option<SyntaxNodePtr> {
-        match self {
-            HirOrigin::Raw(ptr) => Some(ptr.syntax_node_ptr()),
-            HirOrigin::Expanded(ptr) => Some(ptr.clone()),
-            _ => None,
-        }
-    }
-
     pub(crate) fn desugared(origin: impl Into<DesugaredOrigin>) -> Self {
         Self::Desugared(origin.into())
     }

crates/hir/src/span/pat.rs

@@ -1,5 +1,4 @@
-use common::InputFile;
-use parser::{ast, SyntaxNode};
+use parser::ast;
 
 use crate::{
     hir_def::{Body, PatId},
@@ -8,8 +7,8 @@ use crate::{
 };
 
 use super::{
-    body_ast, body_source_map, define_lazy_span_node,
-    transition::{ChainInitiator, SpanTransitionChain},
+    body_source_map, define_lazy_span_node,
+    transition::{ChainInitiator, ResolvedOrigin, SpanTransitionChain},
 };
 
 define_lazy_span_node!(LazyPatSpan, ast::Pat,);
@@ -80,15 +79,10 @@ pub(crate) struct PatRoot {
 }
 
 impl ChainInitiator for PatRoot {
-    fn init(&self, db: &dyn SpannedHirDb) -> (InputFile, SyntaxNode) {
+    fn init(&self, db: &dyn SpannedHirDb) -> ResolvedOrigin {
         let source_map = body_source_map(db, self.body);
-        let pat_source = source_map.pat_map.node_to_source(self.pat);
-        let ptr = pat_source
-            .syntax_ptr()
-            .unwrap_or_else(|| body_ast(db, self.body).syntax_ptr().unwrap());
-
-        let (file, root_node) = self.body.top_mod(db.upcast()).init(db);
-        let node = ptr.to_node(&root_node);
-        (file, node)
+        let origin = source_map.pat_map.node_to_source(self.pat);
+        let top_mod = self.body.top_mod(db.upcast());
+        ResolvedOrigin::resolve(db, top_mod, origin)
     }
 }

crates/hir/src/span/stmt.rs

@@ -1,5 +1,4 @@
-use common::InputFile;
-use parser::{ast, SyntaxNode};
+use parser::ast;
 
 use crate::{
     hir_def::{Body, StmtId},
@@ -8,8 +7,8 @@ use crate::{
 };
 
 use super::{
-    body_ast, body_source_map, define_lazy_span_node,
-    transition::{ChainInitiator, SpanTransitionChain},
+    body_source_map, define_lazy_span_node,
+    transition::{ChainInitiator, ResolvedOrigin, SpanTransitionChain},
 };
 
 define_lazy_span_node!(LazyStmtSpan, ast::Stmt,);
@@ -39,15 +38,44 @@ pub(crate) struct StmtRoot {
 }
 
 impl ChainInitiator for StmtRoot {
-    fn init(&self, db: &dyn SpannedHirDb) -> (InputFile, SyntaxNode) {
+    fn init(&self, db: &dyn SpannedHirDb) -> ResolvedOrigin {
         let source_map = body_source_map(db, self.body);
-        let stmt_source = source_map.stmt_map.node_to_source(self.stmt);
-        let ptr = stmt_source
-            .syntax_ptr()
-            .unwrap_or_else(|| body_ast(db, self.body).syntax_ptr().unwrap());
-
-        let (file, root_node) = self.body.top_mod(db.upcast()).init(db);
-        let node = ptr.to_node(&root_node);
-        (file, node)
+        let origin = source_map.stmt_map.node_to_source(self.stmt);
+        let top_mod = self.body.top_mod(db.upcast());
+        ResolvedOrigin::resolve(db, top_mod, origin)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{hir_def::Body, test_db::TestDb};
+    use common::Upcast;
+
+    #[test]
+    fn aug_assign() {
+        let mut db = TestDb::default();
+
+        let text = r#" {
+            fn foo() {
+                let mut x = 0
+                x += 1
+            }
+        }"#;
+
+        let body: Body = db.expect_item::<Body>(text);
+        let top_mod = body.top_mod(db.upcast());
+        for (i, stmt) in body.stmts(db.upcast()).keys().enumerate() {
+            match i {
+                0 => {
+                    let span = stmt.lazy_span(body);
+                    assert_eq!("let mut x = 0", db.text_at(top_mod, &span));
+                }
+                1 => {
+                    let span = stmt.lazy_span(body);
+                    assert_eq!("x += 1", db.text_at(top_mod, &span));
+                }
+                _ => unreachable!(),
+            }
+        }
     }
 }