Initial work on formatting headers

4 files changed, 130 insertions, 7 deletions

diff --git a/src/header.rs b/src/header.rs
new file mode 100644
index 00000000000..fbf4d127be5
--- /dev/null
+++ b/src/header.rs
@@ -0,0 +1,102 @@
+//! headers are sets of consecutive keywords and tokens, such as
+//! `pub const unsafe fn foo` and `pub(crate) unsafe trait Bar`.
+//!
+//! This module contains general logic for formatting such headers,
+//! where they are always placed on a single line except when there
+//! are comments between parts of the header.
+
+use std::borrow::Cow;
+
+use rustc_ast as ast;
+use rustc_span::symbol::Ident;
+use rustc_span::Span;
+
+use crate::comment::combine_strs_with_missing_comments;
+use crate::rewrite::RewriteContext;
+use crate::shape::Shape;
+use crate::utils::rewrite_ident;
+
+pub(crate) fn format_header(
+    context: &RewriteContext<'_>,
+    shape: Shape,
+    parts: Vec<HeaderPart>,
+) -> String {
+    debug!(?parts, "format_header");
+    let shape = shape.infinite_width();
+
+    // Empty `HeaderPart`s are ignored.
+    let mut parts = parts.into_iter().filter(|x| !x.snippet.is_empty());
+    let Some(part) = parts.next() else {
+        return String::new();
+    };
+
+    let mut result = part.snippet.into_owned();
+    let mut span = part.span;
+
+    for part in parts {
+        debug!(?result, "before combine");
+        result = combine_strs_with_missing_comments(
+            context,
+            &result,
+            &part.snippet,
+            span.between(part.span),
+            shape,
+            true,
+        )
+        .unwrap_or_else(|| format!("{} {}", &result, part.snippet));
+        debug!(?result);
+        span = part.span;
+    }
+
+    result
+}
+
+#[derive(Debug)]
+pub(crate) struct HeaderPart {
+    /// snippet of this part without surrounding space
+    snippet: Cow<'static, str>,
+    span: Span,
+}
+
+impl HeaderPart {
+    pub(crate) fn new(snippet: impl Into<Cow<'static, str>>, span: Span) -> Self {
+        Self {
+            snippet: snippet.into(),
+            span,
+        }
+    }
+
+    pub(crate) fn ident(context: &RewriteContext<'_>, ident: Ident) -> Self {
+        Self {
+            snippet: rewrite_ident(context, ident).to_owned().into(),
+            span: ident.span,
+        }
+    }
+
+    pub(crate) fn visibility(context: &RewriteContext<'_>, vis: &ast::Visibility) -> Self {
+        let snippet = match vis.kind {
+            ast::VisibilityKind::Public => Cow::from("pub"),
+            ast::VisibilityKind::Inherited => Cow::from(""),
+            ast::VisibilityKind::Restricted { ref path, .. } => {
+                let ast::Path { ref segments, .. } = **path;
+                let mut segments_iter =
+                    segments.iter().map(|seg| rewrite_ident(context, seg.ident));
+                if path.is_global() {
+                    segments_iter
+                        .next()
+                        .expect("Non-global path in pub(restricted)?");
+                }
+                let is_keyword = |s: &str| s == "crate" || s == "self" || s == "super";
+                let path = segments_iter.collect::<Vec<_>>().join("::");
+                let in_str = if is_keyword(&path) { "" } else { "in " };
+
+                Cow::from(format!("pub({}{})", in_str, path))
+            }
+        };
+
+        HeaderPart {
+            snippet,
+            span: vis.span,
+        }
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index a67adb1478f..0d1ac7ad958 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -71,6 +71,7 @@ mod emitter;
 mod expr;
 mod format_report_formatter;
 pub(crate) mod formatting;
+pub(crate) mod header;
 mod ignore_path;
 mod imports;
 mod items;
diff --git a/src/macros.rs b/src/macros.rs
index 6e114c76f26..2cc5c692b22 100644
--- a/src/macros.rs
+++ b/src/macros.rs
@@ -27,6 +27,7 @@ use crate::comment::{
 use crate::config::lists::*;
 use crate::config::Version;
 use crate::expr::{rewrite_array, rewrite_assign_rhs, RhsAssignKind};
+use crate::header::{format_header, HeaderPart};
 use crate::lists::{itemize_list, write_list, ListFormatting};
 use crate::overflow;
 use crate::parse::macros::lazy_static::parse_lazy_static;
@@ -36,8 +37,8 @@ use crate::shape::{Indent, Shape};
 use crate::source_map::SpanUtils;
 use crate::spanned::Spanned;
 use crate::utils::{
-    filtered_str_fits, format_visibility, indent_next_line, is_empty_line, mk_sp,
-    remove_trailing_white_spaces, rewrite_ident, trim_left_preserve_layout, NodeIdExt,
+    filtered_str_fits, indent_next_line, is_empty_line, mk_sp, remove_trailing_white_spaces,
+    rewrite_ident, trim_left_preserve_layout, NodeIdExt,
 };
 use crate::visitor::FmtVisitor;
 
@@ -418,14 +419,21 @@ pub(crate) fn rewrite_macro_def(
         None => return snippet,
     };
 
-    let mut result = if def.macro_rules {
-        String::from("macro_rules!")
+    let mut header = if def.macro_rules {
+        let pos = context.snippet_provider.span_after(span, "macro_rules!");
+        vec![HeaderPart::new("macro_rules!", span.with_hi(pos))]
     } else {
-        format!("{}macro", format_visibility(context, vis))
+        let macro_lo = context.snippet_provider.span_before(span, "macro");
+        let macro_hi = macro_lo + BytePos("macro".len() as u32);
+        vec![
+            HeaderPart::visibility(context, vis),
+            HeaderPart::new("macro", mk_sp(macro_lo, macro_hi)),
+        ]
     };
 
-    result += " ";
-    result += rewrite_ident(context, ident);
+    header.push(HeaderPart::ident(context, ident));
+
+    let mut result = format_header(context, shape, header);
 
     let multi_branch_style = def.macro_rules || parsed_def.branches.len() != 1;
 
diff --git a/tests/target/keywords.rs b/tests/target/keywords.rs
new file mode 100644
index 00000000000..22bf7741936
--- /dev/null
+++ b/tests/target/keywords.rs
@@ -0,0 +1,12 @@
+pub // a
+macro // b
+hi(
+    // c
+) {
+    // d
+}
+
+macro_rules! // a
+my_macro {
+    () => {};
+}