Clean up computation of macro expansion span and correctly handle spans open inside expansion spans

Author: GuillaumeGomez

src/librustdoc/html/highlight.rs

@@ -162,6 +162,17 @@ struct TokenHandler<'a, 'tcx, F: Write> {
     write_line_number: fn(&mut F, u32, &'static str),
 }
 
+impl<F: Write> std::fmt::Debug for TokenHandler<'_, '_, F> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("TokenHandler")
+            .field("closing_tags", &self.closing_tags)
+            .field("pending_exit_span", &self.pending_exit_span)
+            .field("current_class", &self.current_class)
+            .field("pending_elems", &self.pending_elems)
+            .finish()
+    }
+}
+
 impl<F: Write> TokenHandler<'_, '_, F> {
     fn handle_exit_span(&mut self) {
         // We can't get the last `closing_tags` element using `pop()` because `closing_tags` is
@@ -271,7 +282,7 @@ fn get_next_expansion<'a>(
     span: Span,
 ) -> Option<&'a ExpandedCode> {
     if let Some(expanded_codes) = expanded_codes {
-        expanded_codes.iter().find(|code| code.start_line == line && code.span.lo() >= span.lo())
+        expanded_codes.iter().find(|code| code.start_line == line && code.span.lo() > span.lo())
     } else {
         None
     }
@@ -321,7 +332,7 @@ fn start_expansion(out: &mut Vec<(Cow<'_, str>, Option<Class>)>, expanded_code:
 fn end_expansion<'a, W: Write>(
     token_handler: &mut TokenHandler<'_, '_, W>,
     expanded_codes: Option<&'a Vec<ExpandedCode>>,
-    level: usize,
+    expansion_start_tags: &[(&'static str, Class)],
     line: u32,
     span: Span,
 ) -> Option<&'a ExpandedCode> {
@@ -330,15 +341,27 @@ fn end_expansion<'a, W: Write>(
         token_handler.pending_elems.push((Cow::Borrowed("</span>"), Some(Class::Expansion)));
         return Some(expanded_code);
     }
-    if level == 0 {
+    if expansion_start_tags.is_empty() && token_handler.closing_tags.is_empty() {
+        // No need tag opened so we can just close expansion.
         token_handler.pending_elems.push((Cow::Borrowed("</span></span>"), Some(Class::Expansion)));
         return None;
     }
+
+    // If tags were opened inside the expansion, we need to close them and re-open them outside
+    // of the expansion span.
     let mut out = String::new();
     let mut end = String::new();
-    for (tag, class) in
-        token_handler.closing_tags.iter().skip(token_handler.closing_tags.len() - level)
+
+    let mut closing_tags = token_handler.closing_tags.iter().peekable();
+    let mut start_closing_tags = expansion_start_tags.iter().peekable();
+
+    while let (Some(tag), Some(start_tag)) = (closing_tags.peek(), start_closing_tags.peek())
+        && tag == start_tag
     {
+        closing_tags.next();
+        start_closing_tags.next();
+    }
+    for (tag, class) in start_closing_tags.chain(closing_tags) {
         out.push_str(tag);
         end.push_str(&format!("<span class=\"{}\">", class.as_html()));
     }
@@ -424,7 +447,7 @@ pub(super) fn write_code(
     };
     let mut current_expansion = get_expansion(&mut token_handler, expanded_codes, line, file_span);
     token_handler.write_pending_elems(None);
-    let mut level = 0;
+    let mut expansion_start_tags = Vec::new();
 
     Classifier::new(
         &src,
@@ -464,6 +487,12 @@ pub(super) fn write_code(
                     if current_expansion.is_none() {
                         current_expansion =
                             get_expansion(&mut token_handler, expanded_codes, line, span);
+                        expansion_start_tags = token_handler.closing_tags.clone();
+                    }
+                    if let Some(ref current_expansion) = current_expansion
+                        && current_expansion.span.lo() == span.hi()
+                    {
+                        start_expansion(&mut token_handler.pending_elems, current_expansion);
                     }
                 } else {
                     token_handler.pending_elems.push((Cow::Borrowed(text), class));
@@ -479,11 +508,13 @@ pub(super) fn write_code(
                         }
                     }
                     if need_end {
-                        current_expansion =
-                            end_expansion(&mut token_handler, expanded_codes, level, line, span);
-                        if current_expansion.is_none() {
-                            level = 0;
-                        }
+                        current_expansion = end_expansion(
+                            &mut token_handler,
+                            expanded_codes,
+                            &expansion_start_tags,
+                            line,
+                            span,
+                        );
                     }
                 }
             }
@@ -504,9 +535,6 @@ pub(super) fn write_code(
                 if should_add {
                     let closing_tag =
                         enter_span(token_handler.out, class, &token_handler.href_context);
-                    if current_expansion.is_some() {
-                        level += 1;
-                    }
                     token_handler.closing_tags.push((closing_tag, class));
                 }
 
@@ -515,9 +543,6 @@ pub(super) fn write_code(
             }
             Highlight::ExitSpan => {
                 token_handler.current_class = None;
-                if current_expansion.is_some() {
-                    level -= 1;
-                }
                 token_handler.pending_exit_span = Some(
                     token_handler
                         .closing_tags

src/librustdoc/html/render/span_map.rs

@@ -325,6 +325,8 @@ struct ExpandedCodeInfo {
     span: Span,
     /// Expanded macro source code.
     code: String,
+    /// Expanded span
+    expanded_span: Span,
 }
 
 /// HIR visitor which retrieves expanded macro.
@@ -341,27 +343,40 @@ impl<'tcx> ExpandedCodeVisitor<'tcx> {
         if new_span.is_dummy() || !new_span.from_expansion() {
             return;
         }
-        let new_span = new_span.source_callsite();
+        let callsite_span = new_span.source_callsite();
         if let Some(index) =
-            self.expanded_codes.iter().position(|info| info.span.overlaps(new_span))
+            self.expanded_codes.iter().position(|info| info.span.overlaps(callsite_span))
         {
-            if !self.expanded_codes[index].span.contains(new_span) {
+            let info = &mut self.expanded_codes[index];
+            if new_span.contains(info.expanded_span) {
                 // We replace the item.
-                let info = &mut self.expanded_codes[index];
-                info.span = new_span;
+                info.span = callsite_span;
+                info.expanded_span = new_span;
                 info.code = f(self.tcx);
+            } else {
+                // We push the new item after the existing one.
+                let expanded_code = &mut self.expanded_codes[index];
+                expanded_code.code.push('\n');
+                expanded_code.code.push_str(&f(self.tcx));
+                let lo = BytePos(expanded_code.expanded_span.lo().0.min(new_span.lo().0));
+                let hi = BytePos(expanded_code.expanded_span.hi().0.min(new_span.hi().0));
+                expanded_code.expanded_span = expanded_code.expanded_span.with_lo(lo).with_hi(hi);
             }
         } else {
             // We add a new item.
-            self.expanded_codes.push(ExpandedCodeInfo { span: new_span, code: f(self.tcx) });
+            self.expanded_codes.push(ExpandedCodeInfo {
+                span: callsite_span,
+                code: f(self.tcx),
+                expanded_span: new_span,
+            });
         }
     }
 
     fn compute_expanded(mut self) -> FxHashMap<BytePos, Vec<ExpandedCode>> {
         self.expanded_codes.sort_unstable_by(|item1, item2| item1.span.cmp(&item2.span));
         let source_map = self.tcx.sess.source_map();
         let mut expanded: FxHashMap<BytePos, Vec<ExpandedCode>> = FxHashMap::default();
-        for ExpandedCodeInfo { span, code } in self.expanded_codes {
+        for ExpandedCodeInfo { span, code, .. } in self.expanded_codes {
             if let Ok(lines) = source_map.span_to_lines(span)
                 && !lines.lines.is_empty()
             {
@@ -389,12 +404,18 @@ impl<'tcx> Visitor<'tcx> for ExpandedCodeVisitor<'tcx> {
     }
 
     fn visit_expr(&mut self, expr: &'tcx rustc_hir::Expr<'tcx>) {
-        self.handle_new_span(expr.span, |tcx| rustc_hir_pretty::expr_to_string(&tcx, expr));
-        intravisit::walk_expr(self, expr);
+        if expr.span.from_expansion() {
+            self.handle_new_span(expr.span, |tcx| rustc_hir_pretty::expr_to_string(&tcx, expr));
+        } else {
+            intravisit::walk_expr(self, expr);
+        }
     }
 
     fn visit_item(&mut self, item: &'tcx rustc_hir::Item<'tcx>) {
-        self.handle_new_span(item.span, |tcx| rustc_hir_pretty::item_to_string(&tcx, item));
-        intravisit::walk_item(self, item);
+        if item.span.from_expansion() {
+            self.handle_new_span(item.span, |tcx| rustc_hir_pretty::item_to_string(&tcx, item));
+        } else {
+            intravisit::walk_item(self, item);
+        }
     }
 }