Always parse kotlin metadata lambda structures

Bug: 157448903
Change-Id: I7ffd3bcaffe64ef662a9a1170ba9f3b2fc7fd2ea
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinFunctionInfo.java b/src/main/java/com/android/tools/r8/kotlin/KotlinFunctionInfo.java
index 6756be4..95e55c0 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinFunctionInfo.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinFunctionInfo.java
@@ -132,4 +132,8 @@
   public boolean isExtensionFunction() {
     return receiverParameterType != null;
   }
+
+  public KotlinJvmMethodSignatureInfo getSignature() {
+    return signature;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinJvmMethodSignatureInfo.java b/src/main/java/com/android/tools/r8/kotlin/KotlinJvmMethodSignatureInfo.java
index cb5557e..6b16a51 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinJvmMethodSignatureInfo.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinJvmMethodSignatureInfo.java
@@ -25,12 +25,21 @@
   private final String name;
   private final KotlinTypeReference returnType;
   private final List<KotlinTypeReference> parameters;
+  private final String invalidDescriptor;
 
   private KotlinJvmMethodSignatureInfo(
       String name, KotlinTypeReference returnType, List<KotlinTypeReference> parameters) {
     this.name = name;
     this.returnType = returnType;
     this.parameters = parameters;
+    this.invalidDescriptor = null;
+  }
+
+  private KotlinJvmMethodSignatureInfo(String name, String invalidDescriptor) {
+    this.name = name;
+    this.invalidDescriptor = invalidDescriptor;
+    this.parameters = EMPTY_PARAMETERS_LIST;
+    this.returnType = null;
   }
 
   public static KotlinJvmMethodSignatureInfo create(
@@ -39,6 +48,10 @@
       return null;
     }
     String kotlinDescriptor = methodSignature.getDesc();
+    if (!KotlinMetadataUtils.isValidMethodDescriptor(kotlinDescriptor)) {
+      // If the method descriptor is invalid, keep it as invalid.
+      return new KotlinJvmMethodSignatureInfo(methodSignature.getName(), kotlinDescriptor);
+    }
     String returnTypeDescriptor = DescriptorUtils.getReturnTypeDescriptor(kotlinDescriptor);
     KotlinTypeReference returnType =
         KotlinTypeReference.createFromDescriptor(returnTypeDescriptor, definitionSupplier);
@@ -57,6 +70,10 @@
 
   public JvmMethodSignature rewrite(
       DexEncodedMethod method, AppView<AppInfoWithLiveness> appView, NamingLens namingLens) {
+    if (invalidDescriptor != null) {
+      return new JvmMethodSignature(name, invalidDescriptor);
+    }
+    assert returnType != null;
     String finalName = name;
     if (method != null) {
       String methodName = method.method.name.toString();
@@ -75,4 +92,21 @@
     descBuilder.append(returnType.toRenamedDescriptorOrDefault(appView, namingLens, defValue));
     return new JvmMethodSignature(finalName, descBuilder.toString());
   }
+
+  @Override
+  public String toString() {
+    if (invalidDescriptor != null) {
+      return name + "(" + invalidDescriptor + ")";
+    }
+    assert returnType != null;
+    StringBuilder descBuilder = new StringBuilder();
+    descBuilder.append(name);
+    descBuilder.append("(");
+    for (KotlinTypeReference parameter : parameters) {
+      descBuilder.append(parameter.toString());
+    }
+    descBuilder.append(")");
+    descBuilder.append(returnType.toString());
+    return descBuilder.toString();
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinLambdaInfo.java b/src/main/java/com/android/tools/r8/kotlin/KotlinLambdaInfo.java
index d25f105..d2c3c67 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinLambdaInfo.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinLambdaInfo.java
@@ -14,7 +14,6 @@
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
 import com.android.tools.r8.utils.Reporter;
 import kotlinx.metadata.KmLambda;
-import kotlinx.metadata.KmLambdaVisitor;
 import kotlinx.metadata.jvm.JvmExtensionsKt;
 import kotlinx.metadata.jvm.JvmMethodSignature;
 
@@ -22,9 +21,11 @@
 public class KotlinLambdaInfo {
 
   private final KotlinFunctionInfo function;
+  private final boolean hasBacking;
 
-  private KotlinLambdaInfo(KotlinFunctionInfo function) {
+  private KotlinLambdaInfo(KotlinFunctionInfo function, boolean hasBacking) {
     this.function = function;
+    this.hasBacking = hasBacking;
   }
 
   static KotlinLambdaInfo create(
@@ -36,22 +37,18 @@
       assert false;
       return null;
     }
+    KotlinFunctionInfo kotlinFunctionInfo =
+        KotlinFunctionInfo.create(lambda.function, definitionSupplier, reporter);
     JvmMethodSignature signature = JvmExtensionsKt.getSignature(lambda.function);
-    if (signature == null) {
-      assert false;
-      return null;
-    }
-    for (DexEncodedMethod method : clazz.methods()) {
-      if (toJvmMethodSignature(method.method).asString().equals(signature.asString())) {
-        KotlinFunctionInfo kotlinFunctionInfo =
-            KotlinFunctionInfo.create(lambda.function, definitionSupplier, reporter);
-        method.setKotlinMemberInfo(kotlinFunctionInfo);
-        return new KotlinLambdaInfo(kotlinFunctionInfo);
+    if (signature != null) {
+      for (DexEncodedMethod method : clazz.methods()) {
+        if (toJvmMethodSignature(method.method).asString().equals(signature.asString())) {
+          method.setKotlinMemberInfo(kotlinFunctionInfo);
+          return new KotlinLambdaInfo(kotlinFunctionInfo, true);
+        }
       }
     }
-    // TODO(b/155536535): Resolve this assert for NestTreeShakeJarVerificationTest.
-    // assert false;
-    return null;
+    return new KotlinLambdaInfo(kotlinFunctionInfo, false);
   }
 
   boolean rewrite(
@@ -59,13 +56,26 @@
       DexClass clazz,
       AppView<AppInfoWithLiveness> appView,
       NamingLens namingLens) {
+    if (!hasBacking) {
+      function.rewrite(visitorProvider.get()::visitFunction, null, appView, namingLens);
+      return true;
+    }
+    DexEncodedMethod backing = null;
     for (DexEncodedMethod method : clazz.methods()) {
       if (method.getKotlinMemberInfo() == function) {
-        KmLambdaVisitor kmLambdaVisitor = visitorProvider.get();
-        function.rewrite(kmLambdaVisitor::visitFunction, method, appView, namingLens);
-        return true;
+        backing = method;
+        break;
       }
     }
-    return false;
+    if (backing == null) {
+      appView
+          .options()
+          .reporter
+          .info(
+              KotlinMetadataDiagnostic.lambdaBackingNotFound(clazz.type, function.getSignature()));
+      return false;
+    }
+    function.rewrite(visitorProvider.get()::visitFunction, backing, appView, namingLens);
+    return true;
   }
 }
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinMetadataDiagnostic.java b/src/main/java/com/android/tools/r8/kotlin/KotlinMetadataDiagnostic.java
index 83947d7..3415e13 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinMetadataDiagnostic.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinMetadataDiagnostic.java
@@ -73,4 +73,18 @@
             + StringUtils.LINE_SEPARATOR
             + StringUtils.stacktraceAsString(t));
   }
+
+  static KotlinMetadataDiagnostic lambdaBackingNotFound(
+      DexType type, KotlinJvmMethodSignatureInfo signatureInfo) {
+    return new KotlinMetadataDiagnostic(
+        Origin.unknown(),
+        Position.UNKNOWN,
+        "The lambda function "
+            + signatureInfo.toString()
+            + " could no longer be found in "
+            + type.toSourceString()
+            + " . The method is most likely pruned and would require a specific keep rule to keep"
+            + " alive. As a result, the metadata information regarding the lambda structure has"
+            + " been discarded.");
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinTypeReference.java b/src/main/java/com/android/tools/r8/kotlin/KotlinTypeReference.java
index 79bf0a7..212cecc 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinTypeReference.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinTypeReference.java
@@ -61,6 +61,10 @@
     if (unknown != null) {
       return unknown;
     }
+    assert known != null;
+    if (!known.isClassType()) {
+      return known.descriptor.toString();
+    }
     if (!appView.appInfo().isNonProgramTypeOrLiveProgramType(known)) {
       return defaultValue;
     }
@@ -87,4 +91,9 @@
     }
     return DescriptorUtils.getBinaryNameFromDescriptor(descriptor);
   }
+
+  @Override
+  public String toString() {
+    return known != null ? known.descriptor.toString() : unknown;
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/kotlin/coroutines/KotlinxCoroutinesTestRunner.java b/src/test/java/com/android/tools/r8/kotlin/coroutines/KotlinxCoroutinesTestRunner.java
index ba63861..afe3166 100644
--- a/src/test/java/com/android/tools/r8/kotlin/coroutines/KotlinxCoroutinesTestRunner.java
+++ b/src/test/java/com/android/tools/r8/kotlin/coroutines/KotlinxCoroutinesTestRunner.java
@@ -5,10 +5,10 @@
 package com.android.tools.r8.kotlin.coroutines;
 
 import static com.android.tools.r8.KotlinCompilerTool.KOTLINC;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.MatcherAssert.assertThat;
 import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertThrows;
 
-import com.android.tools.r8.CompilationFailedException;
 import com.android.tools.r8.KotlinTestBase;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.ToolHelper;
@@ -79,24 +79,35 @@
   }
 
   @Test
-  public void runKotlinxCoroutinesNewTests_r8() {
-    // TODO(b/157023682): This should be able to compile.
-    assertThrows(
-        CompilationFailedException.class,
-        () -> {
-          testForR8(parameters.getBackend())
-              .addProgramFiles(BASE_LIBRARY)
-              .addKeepAllClassesRule()
-              .addKeepAllAttributes()
-              // The BASE_LIBRARY contains proguard rules that do not match.
-              .allowUnusedProguardConfigurationRules()
-              .allowDiagnosticInfoMessages()
-              .addKeepRules(
-                  "-dontwarn reactor.blockhound.integration.BlockHoundIntegration",
-                  "-dontwarn org.junit.runners.model.Statement",
-                  "-dontwarn org.junit.rules.TestRule")
-              .compile();
-        });
+  public void runKotlinxCoroutinesNewTests_r8() throws Exception {
+    Path baseLib =
+        testForR8(parameters.getBackend())
+            .addProgramFiles(BASE_LIBRARY)
+            .addKeepAllClassesRule()
+            .addKeepAllAttributes()
+            // The BASE_LIBRARY contains proguard rules that do not match.
+            .allowUnusedProguardConfigurationRules()
+            .addKeepRules(
+                "-dontwarn reactor.blockhound.integration.BlockHoundIntegration",
+                "-dontwarn org.junit.runners.model.Statement",
+                "-dontwarn org.junit.rules.TestRule")
+            .compile()
+            .writeToZip();
+    ProcessResult processResult =
+        kotlinc(KOTLINC, targetVersion)
+            .addArguments(
+                "-Xuse-experimental=kotlinx.coroutines.InternalCoroutinesApi",
+                "-Xuse-experimental=kotlinx.coroutines.ObsoleteCoroutinesApi",
+                "-Xuse-experimental=kotlinx.coroutines.ExperimentalCoroutinesApi")
+            .addClasspathFiles(DEPENDENCIES)
+            .addClasspathFiles(baseLib)
+            .addSourceFiles(TEST_SOURCES)
+            .setOutputPath(temp.newFolder().toPath())
+            .compileRaw();
+    assertEquals(1, processResult.exitCode);
+    assertThat(
+        processResult.stderr,
+        containsString("Couldn't inline method call 'CoroutineExceptionHandler'"));
   }
 
   private void runTestsInJar(Path testJar, Path deps) throws Exception {