1 files changed, 93 insertions, 2 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index 442c5fb5b..ab77d924a 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -1,6 +1,10 @@
 //! A higher level attributes based on TokenTree, with also some shortcuts.
 
-use std::{ops, sync::Arc};
+use std::{
+    convert::{TryFrom, TryInto},
+    ops,
+    sync::Arc,
+};
 
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
@@ -12,7 +16,7 @@ use mbe::ast_to_token_tree;
 use smallvec::{smallvec, SmallVec};
 use syntax::{
     ast::{self, AstNode, AttrsOwner},
-    match_ast, AstToken, SmolStr, SyntaxNode,
+    match_ast, AstToken, SmolStr, SyntaxNode, TextRange, TextSize,
 };
 use tt::Subtree;
 
@@ -452,6 +456,55 @@ impl AttrsWithOwner {
                 .collect(),
         }
     }
+
+    pub fn docs_with_rangemap(
+        &self,
+        db: &dyn DefDatabase,
+    ) -> Option<(Documentation, DocsRangeMap)> {
+        // FIXME: code duplication in `docs` above
+        let docs = self.by_key("doc").attrs().flat_map(|attr| match attr.input.as_ref()? {
+            AttrInput::Literal(s) => Some((s, attr.index)),
+            AttrInput::TokenTree(_) => None,
+        });
+        let indent = docs
+            .clone()
+            .flat_map(|(s, _)| s.lines())
+            .filter(|line| !line.chars().all(|c| c.is_whitespace()))
+            .map(|line| line.chars().take_while(|c| c.is_whitespace()).count())
+            .min()
+            .unwrap_or(0);
+        let mut buf = String::new();
+        let mut mapping = Vec::new();
+        for (doc, idx) in docs {
+            // str::lines doesn't yield anything for the empty string
+            if !doc.is_empty() {
+                for line in doc.split('\n') {
+                    let line = line.trim_end();
+                    let line_len = line.len();
+                    let (offset, line) = match line.char_indices().nth(indent) {
+                        Some((offset, _)) => (offset, &line[offset..]),
+                        None => (0, line),
+                    };
+                    let buf_offset = buf.len();
+                    buf.push_str(line);
+                    mapping.push((
+                        TextRange::new(buf_offset.try_into().ok()?, buf.len().try_into().ok()?),
+                        idx,
+                        TextRange::new(offset.try_into().ok()?, line_len.try_into().ok()?),
+                    ));
+                    buf.push('\n');
+                }
+            } else {
+                buf.push('\n');
+            }
+        }
+        buf.pop();
+        if buf.is_empty() {
+            None
+        } else {
+            Some((Documentation(buf), DocsRangeMap { mapping, source: self.source_map(db).attrs }))
+        }
+    }
 }
 
 fn inner_attributes(
@@ -508,6 +561,44 @@ impl AttrSourceMap {
     }
 }
 
+/// 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>>>,
+    // (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
+    mapping: Vec<(TextRange, u32, TextRange)>,
+}
+
+impl DocsRangeMap {
+    pub fn map(&self, range: TextRange) -> Option<InFile<TextRange>> {
+        let found = self.mapping.binary_search_by(|(probe, ..)| probe.ordering(range)).ok()?;
+        let (line_docs_range, idx, original_line_src_range) = self.mapping[found].clone();
+        if !line_docs_range.contains_range(range) {
+            return None;
+        }
+
+        let relative_range = range - line_docs_range.start();
+
+        let &InFile { file_id, value: ref source } = &self.source[idx as usize];
+        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
+            Either::Right(comment) => {
+                let text_range = comment.syntax().text_range();
+                let range = TextRange::at(
+                    text_range.start()
+                        + TextSize::try_from(comment.prefix().len()).ok()?
+                        + original_line_src_range.start()
+                        + relative_range.start(),
+                    text_range.len().min(range.len()),
+                );
+                Some(InFile { file_id, value: range })
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Attr {
     index: u32,