From f29d74f87f71bf56ef6f89ad3aa670961cd0680d Mon Sep 17 00:00:00 2001
From: mohamedsamehsalah <mohamedsamehsalah.mohamed@teampicnic.com>
Date: Sun, 30 Jul 2023 15:11:31 +0200
Subject: [PATCH] Exclude classes annotated with lombok's @Data from
 `DirectReturn` bug checker

While here, add a new Matcher that determines whether a class is annotated with lombok's @Data.
---
 .../errorprone/bugpatterns/DirectReturn.java  |  8 ++++++-
 .../bugpatterns/util/MoreMatchers.java        | 17 ++++++++++++++
 .../bugpatterns/DirectReturnTest.java         | 22 +++++++++++++++++++
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/DirectReturn.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/DirectReturn.java
index 212d0b12455..62900529f19 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/DirectReturn.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/DirectReturn.java
@@ -12,6 +12,7 @@
 import static com.google.errorprone.matchers.Matchers.staticMethod;
 import static com.google.errorprone.matchers.Matchers.toType;
 import static tech.picnic.errorprone.bugpatterns.util.Documentation.BUG_PATTERNS_BASE_URL;
+import static tech.picnic.errorprone.bugpatterns.util.MoreMatchers.isClassAnnotatedWithLombokData;
 
 import com.google.auto.service.AutoService;
 import com.google.common.collect.Streams;
@@ -25,6 +26,7 @@
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.AssignmentTree;
 import com.sun.source.tree.BlockTree;
+import com.sun.source.tree.ClassTree;
 import com.sun.source.tree.ExpressionStatementTree;
 import com.sun.source.tree.ExpressionTree;
 import com.sun.source.tree.IdentifierTree;
@@ -60,6 +62,8 @@ public final class DirectReturn extends BugChecker implements BlockTreeMatcher {
       allOf(
           not(toType(MethodInvocationTree.class, argument(0, isSameType(Class.class.getName())))),
           staticMethod().onClass("org.mockito.Mockito").namedAnyOf("mock", "spy"));
+  private static final Matcher<ClassTree> IS_CLASS_ANNOTATED_WITH_LOMBOK_DATA =
+      isClassAnnotatedWithLombokData();
 
   /** Instantiates a new {@link DirectReturn} instance. */
   public DirectReturn() {}
@@ -67,7 +71,9 @@ public DirectReturn() {}
   @Override
   public Description matchBlock(BlockTree tree, VisitorState state) {
     List<? extends StatementTree> statements = tree.getStatements();
-    if (statements.size() < 2) {
+    ClassTree enclosingNode = ASTHelpers.findEnclosingNode(state.getPath(), ClassTree.class);
+    if (statements.size() < 2
+        || IS_CLASS_ANNOTATED_WITH_LOMBOK_DATA.matches(enclosingNode, state)) {
       return Description.NO_MATCH;
     }
 
diff --git a/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/util/MoreMatchers.java b/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/util/MoreMatchers.java
index b62e761db00..a71ca36b088 100644
--- a/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/util/MoreMatchers.java
+++ b/error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/util/MoreMatchers.java
@@ -5,7 +5,9 @@
 import com.google.errorprone.predicates.TypePredicate;
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.AnnotationTree;
+import com.sun.source.tree.ClassTree;
 import com.sun.tools.javac.code.Symbol;
+import lombok.Data;
 
 /**
  * A collection of general-purpose {@link Matcher}s.
@@ -32,6 +34,21 @@ public static <T extends AnnotationTree> Matcher<T> hasMetaAnnotation(String ann
     };
   }
 
+  /**
+   * Returns a {@link Matcher} that determines whether a given {@link ClassTree} is annotated with
+   * {@link Data}.
+   *
+   * @param <T> The type of tree to match against.
+   * @return A {@link Matcher} that matches trees with {@link Data} annotation.
+   */
+  public static <T extends ClassTree> Matcher<T> isClassAnnotatedWithLombokData() {
+    TypePredicate typePredicate = hasAnnotation("lombok.Data");
+    return (tree, state) -> {
+      Symbol sym = ASTHelpers.getSymbol(tree);
+      return sym != null && typePredicate.apply(sym.type, state);
+    };
+  }
+
   // XXX: Consider moving to a `MoreTypePredicates` utility class.
   private static TypePredicate hasAnnotation(String annotationClassName) {
     return (type, state) -> ASTHelpers.hasAnnotation(type.tsym, annotationClassName, state);
diff --git a/error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/DirectReturnTest.java b/error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/DirectReturnTest.java
index 42e18894148..b51a7da3aa5 100644
--- a/error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/DirectReturnTest.java
+++ b/error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/DirectReturnTest.java
@@ -223,4 +223,26 @@ void replacement() {
             "}")
         .doTest(TestMode.TEXT_MATCH);
   }
+
+  @Test
+  void excludeClassesAnnotatedWithLombokData() {
+    BugCheckerRefactoringTestHelper.newInstance(DirectReturn.class, getClass())
+        .addInputLines(
+            "A.java",
+            "import lombok.Data;",
+            "",
+            "@Data",
+            "public class A {",
+            "  private String field;",
+            "}")
+        .addOutputLines(
+            "A.java",
+            "import lombok.Data;",
+            "",
+            "@Data",
+            "public class A {",
+            "  private String field;",
+            "}")
+        .doTest(TestMode.TEXT_MATCH);
+  }
 }