6 files changed, 111 insertions, 73 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index a2479016e..aadd4e44a 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -9,7 +9,7 @@ use std::{
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
 use either::Either;
-use hir_expand::{hygiene::Hygiene, name::AsName, AstId, AttrId, InFile};
+use hir_expand::{hygiene::Hygiene, name::AsName, AstId, InFile};
 use itertools::Itertools;
 use la_arena::ArenaMap;
 use mbe::ast_to_token_tree;
@@ -101,17 +101,13 @@ impl RawAttrs {
         hygiene: &Hygiene,
     ) -> Self {
         let entries = collect_attrs(owner)
-            .enumerate()
-            .flat_map(|(i, attr)| {
-                let index = AttrId(i as u32);
-                match attr {
-                    Either::Left(attr) => Attr::from_src(db, attr, hygiene, index),
-                    Either::Right(comment) => comment.doc_comment().map(|doc| Attr {
-                        id: index,
-                        input: Some(AttrInput::Literal(SmolStr::new(doc))),
-                        path: Interned::new(ModPath::from(hir_expand::name!(doc))),
-                    }),
-                }
+            .flat_map(|(id, attr)| match attr {
+                Either::Left(attr) => Attr::from_src(db, attr, hygiene, id),
+                Either::Right(comment) => comment.doc_comment().map(|doc| Attr {
+                    id,
+                    input: Some(AttrInput::Literal(SmolStr::new(doc))),
+                    path: Interned::new(ModPath::from(hir_expand::name!(doc))),
+                }),
             })
             .collect::<Arc<_>>();
 
@@ -124,6 +120,7 @@ impl RawAttrs {
     }
 
     pub(crate) fn merge(&self, other: Self) -> Self {
+        // FIXME: This needs to fixup `AttrId`s
         match (&self.entries, &other.entries) {
             (None, None) => Self::EMPTY,
             (Some(entries), None) | (None, Some(entries)) => {
@@ -375,39 +372,26 @@ impl AttrsWithOwner {
 
                 let def_map = module.def_map(db);
                 let mod_data = &def_map[module.local_id];
-                let attrs = match mod_data.declaration_source(db) {
+                match mod_data.declaration_source(db) {
                     Some(it) => {
-                        let mut attrs: Vec<_> = collect_attrs(&it.value as &dyn ast::AttrsOwner)
-                            .map(|attr| InFile::new(it.file_id, attr))
-                            .collect();
+                        let mut map = AttrSourceMap::new(InFile::new(it.file_id, &it.value));
                         if let InFile { file_id, value: ModuleSource::SourceFile(file) } =
                             mod_data.definition_source(db)
                         {
-                            attrs.extend(
-                                collect_attrs(&file as &dyn ast::AttrsOwner)
-                                    .map(|attr| InFile::new(file_id, attr)),
-                            )
+                            map.merge(AttrSourceMap::new(InFile::new(file_id, &file)));
                         }
-                        attrs
+                        return map;
                     }
                     None => {
                         let InFile { file_id, value } = mod_data.definition_source(db);
-                        match &value {
-                            ModuleSource::SourceFile(file) => {
-                                collect_attrs(file as &dyn ast::AttrsOwner)
-                            }
-                            ModuleSource::Module(module) => {
-                                collect_attrs(module as &dyn ast::AttrsOwner)
-                            }
-                            ModuleSource::BlockExpr(block) => {
-                                collect_attrs(block as &dyn ast::AttrsOwner)
-                            }
-                        }
-                        .map(|attr| InFile::new(file_id, attr))
-                        .collect()
+                        let attrs_owner = match &value {
+                            ModuleSource::SourceFile(file) => file as &dyn ast::AttrsOwner,
+                            ModuleSource::Module(module) => module as &dyn ast::AttrsOwner,
+                            ModuleSource::BlockExpr(block) => block as &dyn ast::AttrsOwner,
+                        };
+                        return AttrSourceMap::new(InFile::new(file_id, attrs_owner));
                     }
-                };
-                return AttrSourceMap { attrs };
+                }
             }
             AttrDefId::FieldId(id) => {
                 let map = db.fields_attrs_source_map(id.parent);
@@ -462,11 +446,7 @@ impl AttrsWithOwner {
             },
         };
 
-        AttrSourceMap {
-            attrs: collect_attrs(&owner.value)
-                .map(|attr| InFile::new(owner.file_id, attr))
-                .collect(),
-        }
+        AttrSourceMap::new(owner.as_ref().map(|node| node as &dyn AttrsOwner))
     }
 
     pub fn docs_with_rangemap(
@@ -518,7 +498,7 @@ impl AttrsWithOwner {
         if buf.is_empty() {
             None
         } else {
-            Some((Documentation(buf), DocsRangeMap { mapping, source: self.source_map(db).attrs }))
+            Some((Documentation(buf), DocsRangeMap { mapping, source_map: self.source_map(db) }))
         }
     }
 }
@@ -559,27 +539,59 @@ fn inner_attributes(
 }
 
 pub struct AttrSourceMap {
-    attrs: Vec<InFile<Either<ast::Attr, ast::Comment>>>,
+    attrs: Vec<InFile<ast::Attr>>,
+    doc_comments: Vec<InFile<ast::Comment>>,
 }
 
 impl AttrSourceMap {
+    fn new(owner: InFile<&dyn ast::AttrsOwner>) -> Self {
+        let mut attrs = Vec::new();
+        let mut doc_comments = Vec::new();
+        for (_, attr) in collect_attrs(owner.value) {
+            match attr {
+                Either::Left(attr) => attrs.push(owner.with_value(attr)),
+                Either::Right(comment) => doc_comments.push(owner.with_value(comment)),
+            }
+        }
+
+        Self { attrs, doc_comments }
+    }
+
+    fn merge(&mut self, other: Self) {
+        self.attrs.extend(other.attrs);
+        self.doc_comments.extend(other.doc_comments);
+    }
+
     /// Maps the lowered `Attr` back to its original syntax node.
     ///
     /// `attr` must come from the `owner` used for AttrSourceMap
     ///
     /// Note that the returned syntax node might be a `#[cfg_attr]`, or a doc comment, instead of
     /// the attribute represented by `Attr`.
-    pub fn source_of(&self, attr: &Attr) -> InFile<&Either<ast::Attr, ast::Comment>> {
-        self.attrs
-            .get(attr.id.0 as usize)
-            .unwrap_or_else(|| panic!("cannot find `Attr` at index {:?}", attr.id))
-            .as_ref()
+    pub fn source_of(&self, attr: &Attr) -> InFile<Either<ast::Attr, ast::Comment>> {
+        self.source_of_id(attr.id)
+    }
+
+    fn source_of_id(&self, id: AttrId) -> InFile<Either<ast::Attr, ast::Comment>> {
+        if id.is_doc_comment {
+            self.doc_comments
+                .get(id.ast_index as usize)
+                .unwrap_or_else(|| panic!("cannot find doc comment at index {:?}", id))
+                .clone()
+                .map(|attr| Either::Right(attr))
+        } else {
+            self.attrs
+                .get(id.ast_index as usize)
+                .unwrap_or_else(|| panic!("cannot find `Attr` at index {:?}", id))
+                .clone()
+                .map(|attr| Either::Left(attr))
+        }
     }
 }
 
 /// A struct to map text ranges from [`Documentation`] back to TextRanges in the syntax tree.
 pub struct DocsRangeMap {
-    source: Vec<InFile<Either<ast::Attr, ast::Comment>>>,
+    source_map: AttrSourceMap,
     // (docstring-line-range, attr_index, attr-string-range)
     // a mapping from the text range of a line of the [`Documentation`] to the attribute index and
     // the original (untrimmed) syntax doc line
@@ -596,7 +608,7 @@ impl DocsRangeMap {
 
         let relative_range = range - line_docs_range.start();
 
-        let &InFile { file_id, value: ref source } = &self.source[idx.0 as usize];
+        let &InFile { file_id, value: ref source } = &self.source_map.source_of_id(idx);
         match source {
             Either::Left(_) => None, // FIXME, figure out a nice way to handle doc attributes here
             // as well as for whats done in syntax highlight doc injection
@@ -615,6 +627,12 @@ impl DocsRangeMap {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub(crate) struct AttrId {
+    is_doc_comment: bool,
+    pub(crate) ast_index: u32,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Attr {
     pub(crate) id: AttrId,
@@ -749,22 +767,32 @@ fn attrs_from_item_tree<N: ItemTreeNode>(id: ItemTreeId<N>, db: &dyn DefDatabase
 
 fn collect_attrs(
     owner: &dyn ast::AttrsOwner,
-) -> impl Iterator<Item = Either<ast::Attr, ast::Comment>> {
+) -> impl Iterator<Item = (AttrId, Either<ast::Attr, ast::Comment>)> {
     let (inner_attrs, inner_docs) = inner_attributes(owner.syntax())
         .map_or((None, None), |(attrs, docs)| (Some(attrs), Some(docs)));
 
     let outer_attrs = owner.attrs().filter(|attr| attr.kind().is_outer());
-    let attrs = outer_attrs
-        .chain(inner_attrs.into_iter().flatten())
-        .map(|attr| (attr.syntax().text_range().start(), Either::Left(attr)));
+    let attrs =
+        outer_attrs.chain(inner_attrs.into_iter().flatten()).enumerate().map(|(idx, attr)| {
+            (
+                AttrId { ast_index: idx as u32, is_doc_comment: false },
+                attr.syntax().text_range().start(),
+                Either::Left(attr),
+            )
+        });
 
     let outer_docs =
         ast::CommentIter::from_syntax_node(owner.syntax()).filter(ast::Comment::is_outer);
-    let docs = outer_docs
-        .chain(inner_docs.into_iter().flatten())
-        .map(|docs_text| (docs_text.syntax().text_range().start(), Either::Right(docs_text)));
+    let docs =
+        outer_docs.chain(inner_docs.into_iter().flatten()).enumerate().map(|(idx, docs_text)| {
+            (
+                AttrId { ast_index: idx as u32, is_doc_comment: true },
+                docs_text.syntax().text_range().start(),
+                Either::Right(docs_text),
+            )
+        });
     // sort here by syntax node offset because the source can have doc attributes and doc strings be interleaved
-    docs.chain(attrs).sorted_by_key(|&(offset, _)| offset).map(|(_, attr)| attr)
+    docs.chain(attrs).sorted_by_key(|&(_, offset, _)| offset).map(|(id, _, attr)| (id, attr))
 }
 
 pub(crate) fn variants_attrs_source_map(
diff --git a/crates/hir_def/src/body/lower.rs b/crates/hir_def/src/body/lower.rs
index 75dc19c11..9f278d35b 100644
--- a/crates/hir_def/src/body/lower.rs
+++ b/crates/hir_def/src/body/lower.rs
@@ -205,7 +205,7 @@ impl ExprCollector<'_> {
         self.maybe_collect_expr(expr).unwrap_or_else(|| self.missing_expr())
     }
 
-    /// Returns `None` if the expression is `#[cfg]`d out.
+    /// Returns `None` if and only if the expression is `#[cfg]`d out.
     fn maybe_collect_expr(&mut self, expr: ast::Expr) -> Option<ExprId> {
         let syntax_ptr = AstPtr::new(&expr);
         self.check_cfg(&expr)?;
@@ -668,7 +668,7 @@ impl ExprCollector<'_> {
                 if self.check_cfg(&stmt).is_none() {
                     return;
                 }
-
+                let has_semi = stmt.semicolon_token().is_some();
                 // Note that macro could be expended to multiple statements
                 if let Some(ast::Expr::MacroCall(m)) = stmt.expr() {
                     let macro_ptr = AstPtr::new(&m);
@@ -685,18 +685,19 @@ impl ExprCollector<'_> {
                                 statements.statements().for_each(|stmt| this.collect_stmt(stmt));
                                 if let Some(expr) = statements.expr() {
                                     let expr = this.collect_expr(expr);
-                                    this.statements_in_scope.push(Statement::Expr(expr));
+                                    this.statements_in_scope
+                                        .push(Statement::Expr { expr, has_semi });
                                 }
                             }
                             None => {
                                 let expr = this.alloc_expr(Expr::Missing, syntax_ptr.clone());
-                                this.statements_in_scope.push(Statement::Expr(expr));
+                                this.statements_in_scope.push(Statement::Expr { expr, has_semi });
                             }
                         },
                     );
                 } else {
                     let expr = self.collect_expr_opt(stmt.expr());
-                    self.statements_in_scope.push(Statement::Expr(expr));
+                    self.statements_in_scope.push(Statement::Expr { expr, has_semi });
                 }
             }
             ast::Stmt::Item(item) => {
@@ -725,8 +726,17 @@ impl ExprCollector<'_> {
         let prev_statements = std::mem::take(&mut self.statements_in_scope);
 
         block.statements().for_each(|s| self.collect_stmt(s));
-
-        let tail = block.tail_expr().map(|e| self.collect_expr(e));
+        block.tail_expr().and_then(|e| {
+            let expr = self.maybe_collect_expr(e)?;
+            Some(self.statements_in_scope.push(Statement::Expr { expr, has_semi: false }))
+        });
+
+        let mut tail = None;
+        if let Some(Statement::Expr { expr, has_semi: false }) = self.statements_in_scope.last() {
+            tail = Some(*expr);
+            self.statements_in_scope.pop();
+        }
+        let tail = tail;
         let statements = std::mem::replace(&mut self.statements_in_scope, prev_statements);
         let syntax_node_ptr = AstPtr::new(&block.into());
         let expr_id = self.alloc_expr(
diff --git a/crates/hir_def/src/body/scope.rs b/crates/hir_def/src/body/scope.rs
index bd7005ca6..6764de3a7 100644
--- a/crates/hir_def/src/body/scope.rs
+++ b/crates/hir_def/src/body/scope.rs
@@ -157,7 +157,7 @@ fn compute_block_scopes(
                 scope = scopes.new_scope(scope);
                 scopes.add_bindings(body, scope, *pat);
             }
-            Statement::Expr(expr) => {
+            Statement::Expr { expr, .. } => {
                 scopes.set_scope(*expr, scope);
                 compute_expr_scopes(*expr, body, scopes, scope);
             }
diff --git a/crates/hir_def/src/expr.rs b/crates/hir_def/src/expr.rs
index b4ad984bd..0c3b41080 100644
--- a/crates/hir_def/src/expr.rs
+++ b/crates/hir_def/src/expr.rs
@@ -242,7 +242,7 @@ pub struct RecordLitField {
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum Statement {
     Let { pat: PatId, type_ref: Option<Interned<TypeRef>>, initializer: Option<ExprId> },
-    Expr(ExprId),
+    Expr { expr: ExprId, has_semi: bool },
 }
 
 impl Expr {
@@ -265,7 +265,7 @@ impl Expr {
                                 f(*expr);
                             }
                         }
-                        Statement::Expr(e) => f(*e),
+                        Statement::Expr { expr: expression, .. } => f(*expression),
                     }
                 }
                 if let Some(expr) = tail {
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
index e96ca953f..a82ea5957 100644
--- a/crates/hir_def/src/lib.rs
+++ b/crates/hir_def/src/lib.rs
@@ -62,14 +62,14 @@ use hir_expand::{
     ast_id_map::FileAstId,
     eager::{expand_eager_macro, ErrorEmitted, ErrorSink},
     hygiene::Hygiene,
-    AstId, AttrId, FragmentKind, HirFileId, InFile, MacroCallId, MacroCallKind, MacroDefId,
-    MacroDefKind,
+    AstId, FragmentKind, HirFileId, InFile, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
 };
 use la_arena::Idx;
 use nameres::DefMap;
 use path::ModPath;
 use syntax::ast;
 
+use crate::attr::AttrId;
 use crate::builtin_type::BuiltinType;
 use item_tree::{
     Const, Enum, Function, Impl, ItemTreeId, ItemTreeNode, ModItem, Static, Struct, Trait,
@@ -753,7 +753,7 @@ fn derive_macro_as_call_id(
             MacroCallKind::Derive {
                 ast_id: item_attr.ast_id,
                 derive_name: last_segment.to_string(),
-                derive_attr,
+                derive_attr_index: derive_attr.ast_index,
             },
         )
         .into();
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index e89136ed1..adfb78c94 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -13,14 +13,14 @@ use hir_expand::{
     builtin_macro::find_builtin_macro,
     name::{AsName, Name},
     proc_macro::ProcMacroExpander,
-    AttrId, FragmentKind, HirFileId, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
+    FragmentKind, HirFileId, MacroCallId, MacroCallKind, MacroDefId, MacroDefKind,
 };
 use hir_expand::{InFile, MacroCallLoc};
 use rustc_hash::{FxHashMap, FxHashSet};
 use syntax::ast;
 
 use crate::{
-    attr::Attrs,
+    attr::{AttrId, Attrs},
     db::DefDatabase,
     derive_macro_as_call_id,
     intern::Interned,