From a1795582948d6c721544b920c86a4642df3b39a0 Mon Sep 17 00:00:00 2001
From: Min Zhu <zhumin@google.com>
Date: Mon, 11 Jul 2022 13:46:00 -0400
Subject: [PATCH] feat(ast): add support for LambdaExpr to infer type from
 return expr type (#1011)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This will enable it to be assigned to an appropriate variable

Before this change: cannot assign a lambda expression to a variable because its type is always VOID.

After this change: lambda expression builder will infer type from returnExpr’s expression type. Allowing it to be assigned to an appropriate variable. Because returnExpr is a required field, its type should always be accessible.

Example: a naive sample usage is provided in JavaWriterVisitorTest.java as test.
---
 .../api/generator/engine/ast/LambdaExpr.java    |  4 +---
 .../generator/engine/ast/LambdaExprTest.java    | 10 ++++++++++
 .../engine/writer/JavaWriterVisitorTest.java    | 17 +++++++++++++++++
 3 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/google/api/generator/engine/ast/LambdaExpr.java b/src/main/java/com/google/api/generator/engine/ast/LambdaExpr.java
index bb1bc1a397..dde9989825 100644
--- a/src/main/java/com/google/api/generator/engine/ast/LambdaExpr.java
+++ b/src/main/java/com/google/api/generator/engine/ast/LambdaExpr.java
@@ -27,9 +27,7 @@
 public abstract class LambdaExpr implements Expr {
   @Override
   public TypeNode type() {
-    // TODO(v2): Support set of FunctionalInterface  parameterized on the args and return type,
-    // which would enable assignment to an appropriate variable.
-    return TypeNode.VOID;
+    return returnExpr().expr().type();
   }
 
   public abstract ImmutableList<VariableExpr> arguments();
diff --git a/src/test/java/com/google/api/generator/engine/ast/LambdaExprTest.java b/src/test/java/com/google/api/generator/engine/ast/LambdaExprTest.java
index e4ead3cc87..8011d1a89e 100644
--- a/src/test/java/com/google/api/generator/engine/ast/LambdaExprTest.java
+++ b/src/test/java/com/google/api/generator/engine/ast/LambdaExprTest.java
@@ -14,6 +14,7 @@
 
 package com.google.api.generator.engine.ast;
 
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertThrows;
 
 import java.util.Arrays;
@@ -27,6 +28,15 @@ public void validLambdaExpr_noArguments() {
         .build();
   }
 
+  @Test
+  public void validLambdaExpr_inferTypeFromReturnExpr() {
+    LambdaExpr lambdaExpr =
+        LambdaExpr.builder()
+            .setReturnExpr(ValueExpr.withValue(StringObjectValue.withValue("foo")))
+            .build();
+    assertEquals(TypeNode.STRING, lambdaExpr.type());
+  }
+
   @Test
   public void validLambdaExpr_severalArguments() {
     VariableExpr argOneVarExpr =
diff --git a/src/test/java/com/google/api/generator/engine/writer/JavaWriterVisitorTest.java b/src/test/java/com/google/api/generator/engine/writer/JavaWriterVisitorTest.java
index 39e448a263..75f25c5b94 100644
--- a/src/test/java/com/google/api/generator/engine/writer/JavaWriterVisitorTest.java
+++ b/src/test/java/com/google/api/generator/engine/writer/JavaWriterVisitorTest.java
@@ -2345,6 +2345,23 @@ public void writeLambdaExpr_noParameters() {
     assertEquals("() -> \"foo\"", writerVisitor.write());
   }
 
+  @Test
+  public void writeLambdaExpr_assignToVariable() {
+    LambdaExpr lambdaExpr =
+        LambdaExpr.builder()
+            .setReturnExpr(ValueExpr.withValue(StringObjectValue.withValue("foo")))
+            .build();
+    AssignmentExpr assignmentExpr =
+        AssignmentExpr.builder()
+            .setVariableExpr(
+                VariableExpr.withVariable(
+                    Variable.builder().setName("word").setType(TypeNode.STRING).build()))
+            .setValueExpr(lambdaExpr)
+            .build();
+    assignmentExpr.accept(writerVisitor);
+    assertEquals("word = () -> \"foo\"", writerVisitor.write());
+  }
+
   @Test
   public void writeLambdaExpr_oneParameter() {
     VariableExpr argVarExpr =