From 03efe44d9affc989eee8623651fbdf1bcc7240dc Mon Sep 17 00:00:00 2001
From: Liam Miller-Cushon <cushon@google.com>
Date: Fri, 26 Jan 2024 10:20:00 -0800
Subject: [PATCH] Keep type-use annotation with type when wrapping

PiperOrigin-RevId: 601800310
---
 .../java/JavaInputAstVisitor.java              | 18 ++++++++++++++----
 .../java/testdata/B322210626.input             | 11 +++++++++++
 .../java/testdata/B322210626.output            | 11 +++++++++++
 3 files changed, 36 insertions(+), 4 deletions(-)
 create mode 100644 core/src/test/resources/com/google/googlejavaformat/java/testdata/B322210626.input
 create mode 100644 core/src/test/resources/com/google/googlejavaformat/java/testdata/B322210626.output

diff --git a/core/src/main/java/com/google/googlejavaformat/java/JavaInputAstVisitor.java b/core/src/main/java/com/google/googlejavaformat/java/JavaInputAstVisitor.java
index 8914c8ea2..1f51a7632 100644
--- a/core/src/main/java/com/google/googlejavaformat/java/JavaInputAstVisitor.java
+++ b/core/src/main/java/com/google/googlejavaformat/java/JavaInputAstVisitor.java
@@ -1407,6 +1407,14 @@ public Void visitMethod(MethodTree node, Void unused) {
             annotations,
             Direction.VERTICAL,
             /* declarationAnnotationBreak= */ Optional.empty());
+    if (node.getTypeParameters().isEmpty() && node.getReturnType() != null) {
+      // If there are type parameters, we use a heuristic above to format annotations after the
+      // type parameter declarations as type-use annotations. If there are no type parameters,
+      // use the heuristics in visitModifiers for recognizing well known type-use annotations and
+      // formatting them as annotations on the return type.
+      returnTypeAnnotations = typeAnnotations;
+      typeAnnotations = ImmutableList.of();
+    }
 
     Tree baseReturnType = null;
     Deque<List<? extends AnnotationTree>> dims = null;
@@ -1436,10 +1444,6 @@ public Void visitMethod(MethodTree node, Void unused) {
         }
         token("<");
         typeParametersRest(node.getTypeParameters(), plusFour);
-        if (!returnTypeAnnotations.isEmpty()) {
-          builder.breakToFill(" ");
-          visitAnnotations(returnTypeAnnotations, BreakOrNot.NO, BreakOrNot.NO);
-        }
         first = false;
       }
 
@@ -1456,8 +1460,14 @@ public Void visitMethod(MethodTree node, Void unused) {
           builder.open(make(breakBeforeType, plusFour, ZERO));
           openedNameAndTypeScope = true;
         }
+        builder.open(ZERO);
+        if (!returnTypeAnnotations.isEmpty()) {
+          visitAnnotations(returnTypeAnnotations, BreakOrNot.NO, BreakOrNot.NO);
+          builder.breakOp(" ");
+        }
         scan(baseReturnType, null);
         maybeAddDims(dims);
+        builder.close();
       }
       if (!first) {
         builder.breakOp(Doc.FillMode.INDEPENDENT, " ", ZERO, Optional.of(breakBeforeName));
diff --git a/core/src/test/resources/com/google/googlejavaformat/java/testdata/B322210626.input b/core/src/test/resources/com/google/googlejavaformat/java/testdata/B322210626.input
new file mode 100644
index 000000000..931a7d726
--- /dev/null
+++ b/core/src/test/resources/com/google/googlejavaformat/java/testdata/B322210626.input
@@ -0,0 +1,11 @@
+class WrapTypeUse {
+  public static final <A extends @Nullable Object, B extends @Nullable Object> @Nullable
+      WrapTypeUse get() {
+    return null;
+  }
+
+  public @Nullable
+      Xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx getXxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx() {
+    return null;
+  }
+}
diff --git a/core/src/test/resources/com/google/googlejavaformat/java/testdata/B322210626.output b/core/src/test/resources/com/google/googlejavaformat/java/testdata/B322210626.output
new file mode 100644
index 000000000..b34bdb38c
--- /dev/null
+++ b/core/src/test/resources/com/google/googlejavaformat/java/testdata/B322210626.output
@@ -0,0 +1,11 @@
+class WrapTypeUse {
+  public static final <A extends @Nullable Object, B extends @Nullable Object>
+      @Nullable WrapTypeUse get() {
+    return null;
+  }
+
+  public @Nullable Xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+      getXxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx() {
+    return null;
+  }
+}