[J2KT] Propagate nullability in method calls to lambda arguments.

It fixes problems in Javac frontend, where nullability in lambdas in not inferred from enclosing method calls. It also fixes some outstanding problems on the boundary between null-marked and non-null-marked scope.

PiperOrigin-RevId: 863626947

Author: Googler

transpiler/java/com/google/j2cl/transpiler/passes/PropagateNullability.java

@@ -22,6 +22,7 @@
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Ordering;
 import com.google.common.collect.Streams;
+import com.google.j2cl.common.visitor.Processor;
 import com.google.j2cl.transpiler.ast.AbstractRewriter;
 import com.google.j2cl.transpiler.ast.AbstractVisitor;
 import com.google.j2cl.transpiler.ast.ArrayLiteral;
@@ -31,7 +32,6 @@
 import com.google.j2cl.transpiler.ast.Expression;
 import com.google.j2cl.transpiler.ast.FunctionExpression;
 import com.google.j2cl.transpiler.ast.IntersectionTypeDescriptor;
-import com.google.j2cl.transpiler.ast.Invocation;
 import com.google.j2cl.transpiler.ast.MethodCall;
 import com.google.j2cl.transpiler.ast.MethodDescriptor;
 import com.google.j2cl.transpiler.ast.MethodLike;
@@ -128,7 +128,7 @@ public class PropagateNullability extends AbstractJ2ktNormalizationPass {
 
   @Override
   public void applyTo(CompilationUnit compilationUnit) {
-    compilationUnit.accept(
+    Processor processor =
         new AbstractRewriter() {
           @Override
           public Node rewriteArrayLiteral(ArrayLiteral arrayLiteral) {
@@ -193,7 +193,17 @@ public Node rewriteMethodCall(MethodCall methodCall) {
                       methodCall.getArguments());
             }
 
-            return Invocation.Builder.from(methodCall).setTarget(rewrittenMethodDescriptor).build();
+            // Propagate nullability from parameters to function expression arguments.
+            ImmutableList<Expression> rewrittenArguments =
+                zip(
+                    methodCall.getArguments(),
+                    rewrittenMethodDescriptor.getParameterTypeDescriptors(),
+                    PropagateNullability.this::propagateNullabilityToFunctionExpression);
+
+            return MethodCall.Builder.from(methodCall)
+                .setTarget(rewrittenMethodDescriptor)
+                .setArguments(rewrittenArguments)
+                .build();
           }
 
           // TODO(b/406815802): See if rewriteInvocation can be refactored to seamlessly
@@ -260,7 +270,14 @@ public Node rewriteFunctionExpression(FunctionExpression functionExpression) {
                 .setTypeDescriptor(inferredFunctionalInterface)
                 .build();
           }
-        });
+        };
+
+    // Propagate nullability twice to ensure that we apply the effect of outward nullability
+    // propagation from the first pass is used in the inward nullability propagation in the
+    // second pass.
+    // TODO(b/406815802): See whether this can be refactored to avoid running twice if not needed.
+    compilationUnit.accept(processor);
+    compilationUnit.accept(processor);
   }
 
   /**
@@ -447,6 +464,19 @@ private MethodDescriptor propagateNullabilityFromArguments(
         methodDescriptor, typeParameterDescriptors, inferredTypeArgumentDescriptors);
   }
 
+  private Expression propagateNullabilityToFunctionExpression(
+      Expression toExpression, TypeDescriptor fromTypeDescriptor) {
+    return switch (toExpression) {
+      case FunctionExpression functionExpression ->
+          FunctionExpression.Builder.from(functionExpression)
+              .setTypeDescriptor(
+                  propagateNullabilityTo(
+                      toExpression.getTypeDescriptor(), fromTypeDescriptor, ImmutableSet.of()))
+              .build();
+      default -> toExpression;
+    };
+  }
+
   // TODO(b/406815802): Add JavaDoc
   private MethodDescriptor propagateNullabilityFromQualifier(
       MethodDescriptor methodDescriptor,

transpiler/javatests/com/google/j2cl/integration/java/nullability/Main.java

@@ -17,7 +17,6 @@
 
 import static com.google.j2cl.integration.testing.Asserts.assertNull;
 import static com.google.j2cl.integration.testing.Asserts.assertTrue;
-import static com.google.j2cl.integration.testing.TestUtils.isJ2Kt;
 
 import java.util.function.Function;
 import org.jspecify.annotations.NullMarked;
@@ -311,12 +310,7 @@ private static void testNullWildcardInLambda() {
   }
 
   private static void testLambdaParameterTypeInference() {
-    if (!isJ2Kt()) {
-      // TODO(b/454662844): The function here gets inferred to Function<Integer, Integer> then
-      // cast to Function<in Integer?, Integer>, which at runtime has a bridge/adaptor
-      // that does an implicit unboxing and ends in NPE.
-      assertTrue(1 == Main.<Integer>applyToNull(value -> (value == null) ? 1 : value + 1));
-    }
+    assertTrue(1 == Main.<Integer>applyToNull(value -> (value == null) ? 1 : value + 1));
   }
 
   private static <T> int applyToNull(Function<? super @Nullable T, Integer> function) {

transpiler/javatests/com/google/j2cl/readable/java/j2kt/NullabilityInferenceProblem10.java

@@ -43,8 +43,7 @@ interface Observer<V extends @Nullable Object> {
     }
 
     static void test(Lazy<Observable<Foo>> lazyObservable) {
-      // TODO(b/...): Uncomment when fixed
-      // lazyObservable.get().addObserver(InNullMarked::observeNullable);
+      lazyObservable.get().addObserver(InNullMarked::observeNullable);
     }
   }
 }

transpiler/javatests/com/google/j2cl/readable/java/j2kt/output_kt/NullabilityInferenceProblem10.kt.txt

@@ -38,7 +38,13 @@ open class NullabilityInferenceProblem10 {
    }
 
    @JvmStatic
-   internal fun test_pp_j2kt(lazyObservable: Lazy<Observable<Foo>>) {}
+   internal fun test_pp_j2kt(lazyObservable: Lazy<Observable<Foo>>) {
+    lazyObservable.get()!!.addObserver(
+     Observer { arg0: Foo? ->
+      return@Observer InNullMarked.observeNullable_pp_j2kt(arg0)
+     },
+    )
+   }
   }
 
   interface Foo

transpiler/javatests/com/google/j2cl/readable/java/j2ktjavac/output_kt/LambdaParameterTypeInference.kt.txt

@@ -57,14 +57,14 @@ open class LambdaParameterTypeInference {
   fun testInference(map: MutableMap<String, Int>) {
    map.compute(
     "",
-    BiFunction { key: String, value: Int ->
-     return@BiFunction if (value == null) 1 else value.toInt() + 1
-    } as BiFunction<in String, in Int?, out Int?>,
+    BiFunction { key: String, value: Int? ->
+     return@BiFunction if (value == null) 1 else value!!.toInt() + 1
+    },
    )
    LambdaParameterTypeInference.applyToNull<Int>(
-    Function { value_1: Int ->
-     return@Function if (value_1 == null) 1 else value_1.toInt() + 1
-    } as Function<in Int?, Int>,
+    Function { value_1: Int? ->
+     return@Function if (value_1 == null) 1 else value_1!!.toInt() + 1
+    },
    )
   }
 

transpiler/javatests/com/google/j2cl/readable/java/j2ktjavac/output_kt/Nullability.kt.txt

@@ -37,8 +37,8 @@ open class Nullability {
 
   internal open fun testBoxedHolder_pp_j2ktjavac(holder: NullMarkedClass.Holder<Int>) {
    holder.acceptNullableConsumer_pp_j2ktjavac(
-    Consumer { s: Int ->
-    } as Consumer<Int?>,
+    Consumer { s: Int? ->
+    },
    )
   }
 

transpiler/javatests/com/google/j2cl/readable/java/j2ktjdt/NullabilityInferenceProblem10.java

@@ -43,8 +43,7 @@ interface Observer<V extends @Nullable Object> {
     }
 
     static void test(Lazy<Observable<Foo>> lazyObservable) {
-      // TODO(b/...): Uncomment when fixed
-      // lazyObservable.get().addObserver(InNullMarked::observeNullable);
+      lazyObservable.get().addObserver(InNullMarked::observeNullable);
     }
   }
 }

transpiler/javatests/com/google/j2cl/readable/java/j2ktjdt/output_kt/NullabilityInferenceProblem10.kt.txt

@@ -38,7 +38,13 @@ open class NullabilityInferenceProblem10 {
    }
 
    @JvmStatic
-   internal fun test_pp_j2ktjdt(lazyObservable: Lazy<Observable<Foo>>) {}
+   internal fun test_pp_j2ktjdt(lazyObservable: Lazy<Observable<Foo>>) {
+    lazyObservable.get()!!.addObserver(
+     Observer { arg0: Foo? ->
+      return@Observer InNullMarked.observeNullable_pp_j2ktjdt(arg0)
+     },
+    )
+   }
   }
 
   interface Foo