google · copybara-service · May 17, 2021
diff --git a/translator/src/main/java/com/google/devtools/j2objc/translate/Rewriter.java b/translator/src/main/java/com/google/devtools/j2objc/translate/Rewriter.java
@@ -20,24 +20,29 @@
 import com.google.devtools.j2objc.ast.AnnotationTypeDeclaration;
 import com.google.devtools.j2objc.ast.Assignment;
 import com.google.devtools.j2objc.ast.Block;
+import com.google.devtools.j2objc.ast.BreakStatement;
 import com.google.devtools.j2objc.ast.CastExpression;
 import com.google.devtools.j2objc.ast.CatchClause;
 import com.google.devtools.j2objc.ast.CompilationUnit;
+import com.google.devtools.j2objc.ast.DoStatement;
 import com.google.devtools.j2objc.ast.EnumDeclaration;
 import com.google.devtools.j2objc.ast.Expression;
 import com.google.devtools.j2objc.ast.ExpressionStatement;
 import com.google.devtools.j2objc.ast.FieldAccess;
 import com.google.devtools.j2objc.ast.FieldDeclaration;
 import com.google.devtools.j2objc.ast.ForStatement;
+import com.google.devtools.j2objc.ast.FunctionInvocation;
 import com.google.devtools.j2objc.ast.IfStatement;
 import com.google.devtools.j2objc.ast.InfixExpression;
+import com.google.devtools.j2objc.ast.LabeledStatement;
 import com.google.devtools.j2objc.ast.MethodDeclaration;
 import com.google.devtools.j2objc.ast.MethodInvocation;
 import com.google.devtools.j2objc.ast.NullLiteral;
 import com.google.devtools.j2objc.ast.PackageDeclaration;
 import com.google.devtools.j2objc.ast.ParenthesizedExpression;
 import com.google.devtools.j2objc.ast.PropertyAnnotation;
 import com.google.devtools.j2objc.ast.QualifiedName;
+import com.google.devtools.j2objc.ast.ReturnStatement;
 import com.google.devtools.j2objc.ast.SimpleName;
 import com.google.devtools.j2objc.ast.SingleVariableDeclaration;
 import com.google.devtools.j2objc.ast.Statement;
@@ -46,6 +51,7 @@
 import com.google.devtools.j2objc.ast.TreeNode;
 import com.google.devtools.j2objc.ast.TreeNode.Kind;
 import com.google.devtools.j2objc.ast.TreeUtil;
+import com.google.devtools.j2objc.ast.TreeVisitor;
 import com.google.devtools.j2objc.ast.TryStatement;
 import com.google.devtools.j2objc.ast.Type;
 import com.google.devtools.j2objc.ast.TypeDeclaration;
@@ -54,6 +60,7 @@
 import com.google.devtools.j2objc.ast.VariableDeclarationFragment;
 import com.google.devtools.j2objc.ast.VariableDeclarationStatement;
 import com.google.devtools.j2objc.types.ExecutablePair;
+import com.google.devtools.j2objc.types.FunctionElement;
 import com.google.devtools.j2objc.types.GeneratedVariableElement;
 import com.google.devtools.j2objc.util.ElementUtil;
 import com.google.devtools.j2objc.util.ErrorUtil;
@@ -89,14 +96,8 @@ public Rewriter(CompilationUnit unit) {
   public boolean visit(MethodDeclaration node) {
     ExecutableElement element = node.getExecutableElement();
     if (ElementUtil.hasAnnotation(element, AutoreleasePool.class)) {
-      if (TypeUtil.isReferenceType(element.getReturnType())) {
-        ErrorUtil.warning(
-            "Ignoring AutoreleasePool annotation on method with retainable return type");
-      } else if (node.getBody() != null) {
-        node.getBody().setHasAutoreleasePool(true);
-      }
+      rewriteAutoreleasePoolMethod(node, element);
     }
-
     if (ElementUtil.hasNullableAnnotation(element) || ElementUtil.hasNonnullAnnotation(element)) {
       unit.setHasNullabilityAnnotations();
     }
@@ -424,4 +425,68 @@ private void maybeAddGenericCastExpression(Expression node, ExecutableElement me
       node.replaceWith(cast);
     }
   }
+
+  private void rewriteAutoreleasePoolMethod(MethodDeclaration node, ExecutableElement element) {
+    if (node.getBody() == null) {
+      ErrorUtil.warning("Ignoring AutoreleasePool annotation on abstract or native method");
+      return;
+    }
+    if (TypeUtil.isReferenceType(element.getReturnType())) {
+      // Reference types need to be saved outside of the autoreleasepool block.
+      Block methodBody = new Block();
+      final VariableElement resultVar =
+          GeneratedVariableElement.newLocalVar("$result$", element.getReturnType(), element);
+
+      methodBody.addStatement(
+          new VariableDeclarationStatement(resultVar, new NullLiteral(typeUtil.getNull())));
+
+      Block innerBody = node.getBody();
+      innerBody.setHasAutoreleasePool(true);
+      innerBody.accept(
+          new TreeVisitor() {
+            @Override
+            public boolean preVisit(TreeNode node) {
+              if (node.getKind() == Kind.RETURN_STATEMENT) {
+                // Convert return statement into assignment to result variable.
+                FunctionInvocation retainCall =
+                    createMemoryMacro("RETAIN_", ((ReturnStatement) node).getExpression().copy());
+                Block replacement = new Block();
+                replacement.addStatement(
+                    new ExpressionStatement(new Assignment(new SimpleName(resultVar), retainCall)));
+                replacement.addStatement(new BreakStatement().setLabel(new SimpleName("$outer$")));
+                node.replaceWith(replacement);
+                return false;
+              }
+              if (node.getKind() == Kind.TYPE_DECLARATION
+                  || node.getKind() == Kind.TYPE_DECLARATION_STATEMENT) {
+                // Don't traverse into other classes.
+                return false;
+              }
+              return super.preVisit(node);
+            }
+          });
+
+      methodBody.addStatement(
+          new LabeledStatement("$outer$")
+              .setBody(
+                  new DoStatement()
+                      .setExpression(TreeUtil.newLiteral(Boolean.FALSE, typeUtil))
+                      .setBody(innerBody.copy())));
+
+      FunctionInvocation autoreleaseCall =
+          createMemoryMacro("AUTORELEASE", new SimpleName(resultVar));
+      methodBody.addStatement(new ReturnStatement(autoreleaseCall));
+      node.getBody().replaceWith(methodBody);
+    } else {
+      // Primitive types just need to set the hasAutoreleasePool boolean.
+      node.getBody().setHasAutoreleasePool(true);
+    }
+  }
+
+  private FunctionInvocation createMemoryMacro(String name, Expression argument) {
+    FunctionElement element =
+        new FunctionElement(name, TypeUtil.ID_TYPE, TypeUtil.NS_OBJECT)
+            .addParameters(TypeUtil.ID_TYPE);
+    return new FunctionInvocation(element, TypeUtil.ID_TYPE).addArgument(argument);
+  }
 }
diff --git a/translator/src/test/java/com/google/devtools/j2objc/translate/RewriterTest.java b/translator/src/test/java/com/google/devtools/j2objc/translate/RewriterTest.java
@@ -521,4 +521,58 @@ public void testTryWithResourceOnEffectivelyFinalVariable() throws IOException {
         "  }",
         "}");
   }
+
+  // Verify that a local variable for the return value is declared, all return
+  // statements are converted to retained assignments, and that outside the
+  // autoreleasepool the result is autoreleased before being returned.
+  // One exception: methods in an inner class aren't converted.
+  public void testAutoreleasePoolWithRetainedReturnType() throws IOException {
+    String translation = translateSourceFile(
+        "import com.google.j2objc.annotations.AutoreleasePool;"
+            + "class Test {"
+            + "  private int state = 42;"
+            + "  @AutoreleasePool"
+            + "  Object test() {"
+            + "    if (state == -1) {"
+            + "      return new Integer(666);"
+            + "    }"
+            + "    outer: for (int i = 0; i < 1; i++) {"
+            + "if (state == 31) break outer;"
+            + "      return new Float(0.1);"
+            + "    }"
+            + "    return Test.class;"
+            + "  }"
+            + "}", "Test", "Test.m");
+
+    assertTranslatedLines(
+        translation,
+        "- (id)test {",
+            "  id $result$ = nil;",
+            "  do {",
+            "    @autoreleasepool {",
+            "      if (state_ == -1) {",
+            "        {",
+            "          $result$ = RETAIN_(create_JavaLangInteger_initWithInt_(666));",
+            "          goto break_$outer$;",
+            "        }",
+            "      }",
+            "      for (jint i = 0; i < 1; i++) {",
+            "        if (state_ == 31) goto break_outer;",
+            "        {",
+            "          $result$ = RETAIN_(create_JavaLangFloat_initWithDouble_(0.1));",
+            "          goto break_$outer$;",
+            "        }",
+            "      }",
+            "      break_outer: ;",
+            "      {",
+            "        $result$ = RETAIN_(Test_class_());",
+            "        goto break_$outer$;",
+            "      }",
+            "    }",
+            "  }",
+            "  while (false);",
+            "  break_$outer$: ;",
+            "  return AUTORELEASE($result$);",
+            "}");
+  }
 }