From 0745f991b5bfd8f37b092f021bc35dd70574c364 Mon Sep 17 00:00:00 2001
From: Googler <noreply@google.com>
Date: Wed, 4 Dec 2024 15:08:58 -0800
Subject: [PATCH] [J2KT] Fix `FixJavaToKotlinCollectionMismatch` pass to handle
 return type conversion explicitly, instead of relying on other passes.

PiperOrigin-RevId: 702876141
---
 ...ixJavaKotlinCollectionMethodsMismatch.java | 41 +++++++++++++++++--
 .../java/j2kt/output_kt/Collections.kt.txt    |  2 +-
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/transpiler/java/com/google/j2cl/transpiler/passes/FixJavaKotlinCollectionMethodsMismatch.java b/transpiler/java/com/google/j2cl/transpiler/passes/FixJavaKotlinCollectionMethodsMismatch.java
index afe2ebfcb0..7be490a04e 100644
--- a/transpiler/java/com/google/j2cl/transpiler/passes/FixJavaKotlinCollectionMethodsMismatch.java
+++ b/transpiler/java/com/google/j2cl/transpiler/passes/FixJavaKotlinCollectionMethodsMismatch.java
@@ -31,7 +31,9 @@
 import com.google.j2cl.transpiler.ast.Method;
 import com.google.j2cl.transpiler.ast.MethodCall;
 import com.google.j2cl.transpiler.ast.MethodDescriptor;
+import com.google.j2cl.transpiler.ast.MethodLike;
 import com.google.j2cl.transpiler.ast.Node;
+import com.google.j2cl.transpiler.ast.ReturnStatement;
 import com.google.j2cl.transpiler.ast.Statement;
 import com.google.j2cl.transpiler.ast.SuperReference;
 import com.google.j2cl.transpiler.ast.TypeDescriptor;
@@ -171,7 +173,7 @@ public Node rewriteMethod(Method method) {
             if (methodMapping == null) {
               return method;
             }
-            return methodMapping.fixMethodParameters(method);
+            return methodMapping.fixMethodParametersAndReturnStatements(method);
           }
 
           @Override
@@ -259,10 +261,10 @@ private boolean isOrOverrides(MethodDescriptor methodDescriptor) {
     }
 
     /**
-     * Fixes parameters in {@code method} to match Kotlin, and rewrites the body of the method to
-     * introduce local variables with original types.
+     * Fixes parameters and return statements in {@code method} to match Kotlin, and rewrites the
+     * body of the method to introduce local variables with original types.
      */
-    private Method fixMethodParameters(Method method) {
+    private Method fixMethodParametersAndReturnStatements(Method method) {
       MethodDescriptor javaMethodDescriptor = getJavaMethodDescriptor();
       MethodDescriptor kotlinMethodDescriptor = getKotlinMethodDescriptor();
       MethodDescriptor methodDescriptor = method.getDescriptor();
@@ -320,6 +322,16 @@ private Method fixMethodParameters(Method method) {
         }
       }
 
+      TypeDescriptor javaDeclarationReturnTypeDescriptor =
+          javaMethodDescriptor.getReturnTypeDescriptor();
+      TypeDescriptor kotlinDeclarationReturnTypeDescriptor =
+          kotlinMethodDescriptor.getReturnTypeDescriptor();
+      if (!javaDeclarationReturnTypeDescriptor.equals(kotlinDeclarationReturnTypeDescriptor)) {
+        TypeDescriptor kotlinReturnTypeDescriptor =
+            kotlinDeclarationReturnTypeDescriptor.specializeTypeVariables(parameterization);
+        body = (Block) insertReturnStatementCasts(body, kotlinReturnTypeDescriptor);
+      }
+
       TypeDescriptor kotlinReturnTypeDescriptor =
           kotlinMethodDescriptor
               .getReturnTypeDescriptor()
@@ -338,6 +350,27 @@ private Method fixMethodParameters(Method method) {
           .build();
     }
 
+    private static Node insertReturnStatementCasts(Node node, TypeDescriptor castTypeDescriptor) {
+      return node.rewrite(
+          new AbstractRewriter() {
+            @Override
+            public Node rewriteReturnStatement(ReturnStatement returnStatement) {
+              if (getParent(MethodLike.class::isInstance) != null) {
+                // This is a return in other method-like node, nothing to do.
+                return returnStatement;
+              }
+
+              return ReturnStatement.Builder.from(returnStatement)
+                  .setExpression(
+                      CastExpression.newBuilder()
+                          .setExpression(returnStatement.getExpression())
+                          .setCastTypeDescriptor(castTypeDescriptor)
+                          .build())
+                  .build();
+            }
+          });
+    }
+
     /**
      * Fixes method call to match Kotlin signature. Ordinary method calls are delegated to "java_"
      * bridge extension functions, and explicit casts are inserted for arguments to super calls.
diff --git a/transpiler/javatests/com/google/j2cl/readable/java/j2kt/output_kt/Collections.kt.txt b/transpiler/javatests/com/google/j2cl/readable/java/j2kt/output_kt/Collections.kt.txt
index 18758467e3..3edf45ad87 100644
--- a/transpiler/javatests/com/google/j2cl/readable/java/j2kt/output_kt/Collections.kt.txt
+++ b/transpiler/javatests/com/google/j2cl/readable/java/j2kt/output_kt/Collections.kt.txt
@@ -487,7 +487,7 @@ open class Collections {
    val defaultValue_1: String? = defaultValue
    var key_1: Any? = key
    key_1 = Collections.convert<Any?>(key_1)
-   return super<Collections.CustomMap>.getOrDefault(key_1 as String, defaultValue_1 as String)!!
+   return super<Collections.CustomMap>.getOrDefault(key_1 as String, defaultValue_1 as String) as String
   }
  }
 }