Merge b8cd9ac into 154b758

uber · Aug 25, 2023 · e2b5ff6 · e2b5ff6
2 parents 154b758 + b8cd9ac
commit e2b5ff6
Show file tree

Hide file tree

Showing 2 changed files with 70 additions and 4 deletions.
diff --git a/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java b/nullaway/src/main/java/com/uber/nullaway/ErrorBuilder.java
@@ -201,7 +201,7 @@ private Description.Builder addSuggestedSuppression(
       case ASSIGN_FIELD_NULLABLE:
       case SWITCH_EXPRESSION_NULLABLE:
         if (config.getCastToNonNullMethod() != null && canBeCastToNonNull(suggestTree)) {
-          builder = addCastToNonNullFix(suggestTree, builder);
+          builder = addCastToNonNullFix(suggestTree, builder, state);
         } else {
           // When there is a castToNonNull method, suggestTree is set to the expression to be
           // casted, which is not suppressible. For simplicity, we just always recompute the
@@ -356,15 +356,16 @@ private boolean canBeCastToNonNull(Tree tree) {
     }
   }
 
-  private Description.Builder addCastToNonNullFix(Tree suggestTree, Description.Builder builder) {
+  private Description.Builder addCastToNonNullFix(
+      Tree suggestTree, Description.Builder builder, VisitorState state) {
     final String fullMethodName = config.getCastToNonNullMethod();
     if (fullMethodName == null) {
       throw new IllegalStateException("cast-to-non-null method not set");
     }
     // Add a call to castToNonNull around suggestTree:
     final String[] parts = fullMethodName.split("\\.");
     final String shortMethodName = parts[parts.length - 1];
-    final String replacement = shortMethodName + "(" + suggestTree.toString() + ")";
+    final String replacement = shortMethodName + "(" + state.getSourceForNode(suggestTree) + ")";
     final SuggestedFix fix =
         SuggestedFix.builder()
             .replace(suggestTree, replacement)
@@ -390,7 +391,7 @@ private Description.Builder removeCastToNonNullFix(
             invTree, suggestTree));
     // Remove the call to castToNonNull:
     final SuggestedFix fix =
-        SuggestedFix.builder().replace(invTree, suggestTree.toString()).build();
+        SuggestedFix.builder().replace(invTree, state.getSourceForNode(suggestTree)).build();
     return builder.addFix(fix);
   }
 

diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayAutoSuggestTest.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayAutoSuggestTest.java
@@ -22,6 +22,8 @@
 
 package com.uber.nullaway;
 
+import static com.google.errorprone.BugCheckerRefactoringTestHelper.TestMode.TEXT_MATCH;
+
 import com.google.errorprone.BugCheckerRefactoringTestHelper;
 import com.google.errorprone.ErrorProneFlags;
 import com.sun.source.tree.Tree;
@@ -215,6 +217,38 @@ public void removeUnnecessaryCastToNonNullFromLibraryModel() throws IOException
         .doTest();
   }
 
+  @Test
+  public void removeUnnecessaryCastToNonNullMultiLine() throws IOException {
+    makeTestHelper()
+        .addInputLines(
+            "Test.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import static com.uber.nullaway.testdata.Util.castToNonNull;",
+            "class Test {",
+            "  static class Foo { Object getObj() { return new Object(); } }",
+            "  Object test1(Foo f) {",
+            "    return castToNonNull(f",
+            "                         // comment that should not be deleted",
+            "                         .getObj());",
+            "  }",
+            "}")
+        .addOutputLines(
+            "out/Test.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "import static com.uber.nullaway.testdata.Util.castToNonNull;",
+            "class Test {",
+            "  static class Foo { Object getObj() { return new Object(); } }",
+            "  Object test1(Foo f) {",
+            "    return f",
+            "        // comment that should not be deleted",
+            "        .getObj();",
+            "  }",
+            "}")
+        .doTest(TEXT_MATCH);
+  }
+
   @Test
   public void suggestSuppressionOnMethodRef() throws IOException {
     makeTestHelper()
@@ -258,6 +292,37 @@ public void suggestSuppressionOnMethodRef() throws IOException {
   }
 
   @Test
+  public void suggestCastToNonNullMultiLine() throws IOException {
+    makeTestHelper()
+        .addInputLines(
+            "Test.java",
+            "package com.uber;",
+            "import javax.annotation.Nullable;",
+            "class Test {",
+            "  static class Foo { @Nullable Object getObj() { return null; } }",
+            "  Object test1(Foo f) {",
+            "    return f",
+            "           // comment that should not be deleted",
+            "           .getObj();",
+            "  }",
+            "}")
+        .addOutputLines(
+            "out/Test.java",
+            "package com.uber;",
+            "import static com.uber.nullaway.testdata.Util.castToNonNull;",
+            "",
+            "import javax.annotation.Nullable;",
+            "class Test {",
+            "  static class Foo { @Nullable Object getObj() { return null; } }",
+            "  Object test1(Foo f) {",
+            "    return castToNonNull(f",
+            "           // comment that should not be deleted",
+            "           .getObj());",
+            "  }",
+            "}")
+        .doTest(TEXT_MATCH);
+  }
+
   public void suggestInitSuppressionOnConstructor() throws IOException {
     makeTestHelper()
         .addInputLines(