From 2cd9cf45edb546ac691e37781d8fbe8f2dbea581 Mon Sep 17 00:00:00 2001
From: Summer Ji <summerji@google.com>
Date: Tue, 6 Oct 2020 23:12:43 -0700
Subject: [PATCH] integrate return into method definition (#369)

---
 .../generator/engine/ast/MethodDefinition.java | 18 +++++++++---------
 .../engine/writer/JavaWriterVisitor.java       |  6 +-----
 .../engine/ast/MethodDefinitionTest.java       | 15 +++++++++++++++
 3 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/src/main/java/com/google/api/generator/engine/ast/MethodDefinition.java b/src/main/java/com/google/api/generator/engine/ast/MethodDefinition.java
index 21feb7a1db..687e752f3f 100644
--- a/src/main/java/com/google/api/generator/engine/ast/MethodDefinition.java
+++ b/src/main/java/com/google/api/generator/engine/ast/MethodDefinition.java
@@ -66,7 +66,7 @@ public abstract class MethodDefinition implements AstNode {
   abstract ImmutableList<String> returnTemplateNames();
 
   @Nullable
-  public abstract Expr returnExpr();
+  public abstract ReturnExpr returnExpr();
 
   abstract boolean isOverride();
 
@@ -147,7 +147,11 @@ public Builder setArguments(VariableExpr... arguments) {
 
     public abstract Builder setBody(List<Statement> body);
 
-    public abstract Builder setReturnExpr(Expr returnExpr);
+    public Builder setReturnExpr(Expr expr) {
+      return setReturnExpr(ReturnExpr.withExpr(expr));
+    }
+
+    public abstract Builder setReturnExpr(ReturnExpr returnExpr);
 
     public abstract Builder setIsOverride(boolean isOverride);
 
@@ -292,18 +296,14 @@ public MethodDefinition build() {
         }
 
         if (method.returnExpr() != null && !isLastStatementThrowExpr) {
-          // TODO(miraleung): Refactor this to use ReturnExpr under the covers instead.
-          Preconditions.checkState(
-              !(method.returnExpr() instanceof ReturnExpr),
-              "A method's return expression can only consist of non-ReturnExpr expressions");
           if (method.returnType().isPrimitiveType()
-              || method.returnExpr().type().isPrimitiveType()) {
+              || method.returnExpr().expr().type().isPrimitiveType()) {
             Preconditions.checkState(
-                method.returnType().equals((method.returnExpr().type())),
+                method.returnType().equals((method.returnExpr().expr().type())),
                 "Method return type does not match the return expression type");
           } else {
             Preconditions.checkState(
-                method.returnType().isSupertypeOrEquals(method.returnExpr().type()),
+                method.returnType().isSupertypeOrEquals(method.returnExpr().expr().type()),
                 "Method reference return type is not a supertype of the return expression type");
           }
         }
diff --git a/src/main/java/com/google/api/generator/engine/writer/JavaWriterVisitor.java b/src/main/java/com/google/api/generator/engine/writer/JavaWriterVisitor.java
index 7a1e4d3b6d..538d4a62b0 100644
--- a/src/main/java/com/google/api/generator/engine/writer/JavaWriterVisitor.java
+++ b/src/main/java/com/google/api/generator/engine/writer/JavaWriterVisitor.java
@@ -767,11 +767,7 @@ public void visit(MethodDefinition methodDefinition) {
     newline();
     statements(methodDefinition.body());
     if (methodDefinition.returnExpr() != null) {
-      buffer.append(RETURN);
-      space();
-      methodDefinition.returnExpr().accept(this);
-      semicolon();
-      newline();
+      ExprStatement.withExpr(methodDefinition.returnExpr()).accept(this);
     }
 
     rightBrace();
diff --git a/src/test/java/com/google/api/generator/engine/ast/MethodDefinitionTest.java b/src/test/java/com/google/api/generator/engine/ast/MethodDefinitionTest.java
index 4fe7a5c05b..5cde3a0859 100644
--- a/src/test/java/com/google/api/generator/engine/ast/MethodDefinitionTest.java
+++ b/src/test/java/com/google/api/generator/engine/ast/MethodDefinitionTest.java
@@ -279,6 +279,21 @@ public void validMethodDefinition_boxedReturnType() {
         .build();
   }
 
+  @Test
+  public void validMethodDefinition_setReturnExprUsingReturnExpr() {
+    ReturnExpr returnExpr =
+        ReturnExpr.withExpr(
+            ValueExpr.withValue(
+                PrimitiveValue.builder().setType(TypeNode.INT).setValue("3").build()));
+    MethodDefinition.builder()
+        .setName("foobar")
+        .setScope(ScopeNode.PUBLIC)
+        .setReturnType(TypeNode.INT_OBJECT)
+        .setBody(Arrays.asList(ExprStatement.withExpr(createAssignmentExpr())))
+        .setReturnExpr(returnExpr)
+        .build();
+  }
+
   @Test
   public void invalidMethodDefinition_badTemplateName() {
     assertThrows(