From fe2fe6e6db1a70b8081be6f14787b845a1f6b9cb Mon Sep 17 00:00:00 2001
From: Liam Miller-Cushon <cushon@google.com>
Date: Thu, 1 Aug 2024 09:29:26 -0700
Subject: [PATCH] Handle `var` in MustBeClosedChecker

Fixes https://github.com/google/error-prone/issues/4164

PiperOrigin-RevId: 658429448
---
 .../AbstractMustBeClosedChecker.java          | 23 ++++++++--
 .../bugpatterns/MustBeClosedCheckerTest.java  | 43 +++++++++++++++++++
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/core/src/main/java/com/google/errorprone/bugpatterns/AbstractMustBeClosedChecker.java b/core/src/main/java/com/google/errorprone/bugpatterns/AbstractMustBeClosedChecker.java
index 96306a48d29..d6ed6678c41 100644
--- a/core/src/main/java/com/google/errorprone/bugpatterns/AbstractMustBeClosedChecker.java
+++ b/core/src/main/java/com/google/errorprone/bugpatterns/AbstractMustBeClosedChecker.java
@@ -16,6 +16,7 @@
 
 package com.google.errorprone.bugpatterns;
 
+import static com.google.errorprone.fixes.SuggestedFixes.prettyType;
 import static com.google.errorprone.fixes.SuggestedFixes.qualifyType;
 import static com.google.errorprone.matchers.Description.NO_MATCH;
 import static com.google.errorprone.matchers.Matchers.instanceMethod;
@@ -28,6 +29,7 @@
 import static com.google.errorprone.util.ASTHelpers.getSymbol;
 import static com.google.errorprone.util.ASTHelpers.getType;
 import static com.google.errorprone.util.ASTHelpers.hasAnnotation;
+import static com.google.errorprone.util.ASTHelpers.hasExplicitSource;
 import static com.google.errorprone.util.ASTHelpers.isConsideredFinal;
 import static com.google.errorprone.util.ASTHelpers.isSameType;
 import static com.google.errorprone.util.ASTHelpers.isSubtype;
@@ -467,16 +469,29 @@ private static Optional<Change> extractToResourceInCurrentTry(
   private static Optional<Change> splitVariableDeclarationAroundTry(
       ExpressionTree tree, VariableTree var, VisitorState state, NameSuggester suggester) {
     int initPos = getStartPosition(var.getInitializer());
-    int afterTypePos = state.getEndPosition(var.getType());
+    Tree type = var.getType();
+    String typePrefix;
+    int startPos;
+    if (hasExplicitSource(type, state)) {
+      startPos = state.getEndPosition(type);
+      typePrefix = "";
+    } else {
+      startPos = getStartPosition(var);
+      typePrefix = prettyType(getType(type), state);
+    }
     String name = suggester.suggestName(tree);
     return Change.builder(
             SuggestedFix.builder()
                 .replace(
-                    afterTypePos,
+                    startPos,
                     initPos,
                     String.format(
-                        " %s;\ntry (var %s = %s) {\n%s =",
-                        var.getName(), name, state.getSourceForNode(tree), var.getName()))
+                        "%s %s;\ntry (var %s = %s) {\n%s =",
+                        typePrefix,
+                        var.getName(),
+                        name,
+                        state.getSourceForNode(tree),
+                        var.getName()))
                 .replace(tree, name)
                 .build())
         .closeBraceAfter(var)
diff --git a/core/src/test/java/com/google/errorprone/bugpatterns/MustBeClosedCheckerTest.java b/core/src/test/java/com/google/errorprone/bugpatterns/MustBeClosedCheckerTest.java
index d4b48d6bad0..cb67d7eaa92 100644
--- a/core/src/test/java/com/google/errorprone/bugpatterns/MustBeClosedCheckerTest.java
+++ b/core/src/test/java/com/google/errorprone/bugpatterns/MustBeClosedCheckerTest.java
@@ -155,4 +155,47 @@ public void forLoopUnfixable() {
         .expectUnchanged()
         .doTest();
   }
+
+  @Test
+  public void localVariableTypeInference() {
+    refactoringHelper
+        .addInputLines(
+            "Closeable.java",
+            "class Closeable implements AutoCloseable {",
+            "    @Override",
+            "    public void close() {}",
+            "    public int method() {",
+            "      return 1;",
+            "    }",
+            "  }")
+        .expectUnchanged()
+        .addInputLines(
+            "Foo.java",
+            "import com.google.errorprone.annotations.MustBeClosed;",
+            "class Foo {",
+            "  @MustBeClosed",
+            "  Closeable mustBeClosedMethod() {",
+            "    return null;",
+            "  }",
+            "}")
+        .expectUnchanged()
+        .addInputLines(
+            "Test.java",
+            "class Test {",
+            "  void test(Foo foo) {",
+            "    var bar = foo.mustBeClosedMethod().method();",
+            "  }",
+            "}")
+        .addOutputLines(
+            "Test.java",
+            "class Test {",
+            "  void test(Foo foo) {",
+            "    int bar;",
+            "    try (var closeable = foo.mustBeClosedMethod()) {",
+            "      bar = closeable.method();",
+            "    }",
+            "  }",
+            "}")
+        .doTest();
+  }
 }