[Metadata] Add support for rewriting of kotlin module files

Bug: b/242289529
Change-Id: I165ef93f72e828ebecdfc9e48f2c72b0df155830
diff --git a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
index 4932691..a9f393e 100644
--- a/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
+++ b/src/main/java/com/android/tools/r8/dex/ApplicationWriter.java
@@ -47,6 +47,7 @@
 import com.android.tools.r8.graph.InnerClassAttribute;
 import com.android.tools.r8.graph.ObjectToOffsetMapping;
 import com.android.tools.r8.graph.ParameterAnnotationsList;
+import com.android.tools.r8.naming.KotlinModuleSynthesizer;
 import com.android.tools.r8.naming.NamingLens;
 import com.android.tools.r8.naming.ProguardMapSupplier.ProguardMapId;
 import com.android.tools.r8.origin.Origin;
@@ -605,13 +606,19 @@
       ExceptionUtils.withFinishedResourceHandler(options.reporter, options.mainDexListConsumer);
     }
 
+    KotlinModuleSynthesizer kotlinModuleSynthesizer = new KotlinModuleSynthesizer(appView);
+
     DataResourceConsumer dataResourceConsumer = options.dataResourceConsumer;
     if (dataResourceConsumer != null) {
       ImmutableList<DataResourceProvider> dataResourceProviders =
           appView.app().dataResourceProviders;
       ResourceAdapter resourceAdapter = new ResourceAdapter(appView);
       adaptAndPassDataResources(
-          options, dataResourceConsumer, dataResourceProviders, resourceAdapter);
+          options,
+          dataResourceConsumer,
+          dataResourceProviders,
+          resourceAdapter,
+          kotlinModuleSynthesizer);
 
       // Write the META-INF/services resources. Sort on service names and keep the order from
       // the input for the implementation lines for deterministic output.
@@ -638,6 +645,10 @@
                       options.reporter);
                 });
       }
+      // Rewrite/synthesize kotlin_module files
+      kotlinModuleSynthesizer
+          .synthesizeKotlinModuleFiles()
+          .forEach(file -> dataResourceConsumer.accept(file, options.reporter));
     }
 
     if (options.featureSplitConfiguration != null) {
@@ -645,7 +656,11 @@
           options.featureSplitConfiguration.getDataResourceProvidersAndConsumers()) {
         ResourceAdapter resourceAdapter = new ResourceAdapter(appView);
         adaptAndPassDataResources(
-            options, entry.getConsumer(), entry.getProviders(), resourceAdapter);
+            options,
+            entry.getConsumer(),
+            entry.getProviders(),
+            resourceAdapter,
+            kotlinModuleSynthesizer);
       }
     }
   }
@@ -654,7 +669,8 @@
       InternalOptions options,
       DataResourceConsumer dataResourceConsumer,
       Collection<DataResourceProvider> dataResourceProviders,
-      ResourceAdapter resourceAdapter) {
+      ResourceAdapter resourceAdapter,
+      KotlinModuleSynthesizer kotlinModuleSynthesizer) {
     Set<String> generatedResourceNames = new HashSet<>();
 
     for (DataResourceProvider dataResourceProvider : dataResourceProviders) {
@@ -676,7 +692,10 @@
                   // META-INF/services resources are handled below.
                   return;
                 }
-
+                if (kotlinModuleSynthesizer.isKotlinModuleFile(file)) {
+                  // .kotlin_module files are synthesized.
+                  return;
+                }
                 DataEntryResource adapted = resourceAdapter.adaptIfNeeded(file);
                 if (generatedResourceNames.add(adapted.getName())) {
                   dataResourceConsumer.accept(adapted, options.reporter);
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinFileFacadeInfo.java b/src/main/java/com/android/tools/r8/kotlin/KotlinFileFacadeInfo.java
index f112b4f..3daa33d 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinFileFacadeInfo.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinFileFacadeInfo.java
@@ -70,6 +70,10 @@
     return packageName;
   }
 
+  public String getModuleName() {
+    return packageInfo.getModuleName();
+  }
+
   @Override
   public int[] getMetadataVersion() {
     return metadataVersion;
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 63e7c95..f5b40c2 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinFunctionInfo.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinFunctionInfo.java
@@ -121,7 +121,9 @@
       AppView<?> appView) {
     // TODO(b/154348683): Check method for flags to pass in.
     boolean rewritten = false;
-    String finalName = this.name;
+    String finalName = name;
+    // Only rewrite the kotlin method name if it was equal to the method name when reading the
+    // metadata.
     if (method != null) {
       String methodName = method.getReference().name.toString();
       String rewrittenName = appView.getNamingLens().lookupName(method.getReference()).toString();
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinMultiFileClassPartInfo.java b/src/main/java/com/android/tools/r8/kotlin/KotlinMultiFileClassPartInfo.java
index 6c96b8d..62460ae 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinMultiFileClassPartInfo.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinMultiFileClassPartInfo.java
@@ -18,7 +18,6 @@
 // Holds information about Metadata.MultiFileClassPartInfo
 public class KotlinMultiFileClassPartInfo implements KotlinClassLevelInfo {
 
-  // TODO(b/157630779): Maybe model facadeClassName.
   private final String facadeClassName;
   private final KotlinPackageInfo packageInfo;
   private final String packageName;
@@ -78,6 +77,10 @@
     return packageName;
   }
 
+  public String getModuleName() {
+    return packageInfo.getModuleName();
+  }
+
   @Override
   public int[] getMetadataVersion() {
     return metadataVersion;
@@ -87,4 +90,8 @@
   public void trace(DexDefinitionSupplier definitionSupplier) {
     packageInfo.trace(definitionSupplier);
   }
+
+  public String getFacadeClassName() {
+    return facadeClassName;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinPackageInfo.java b/src/main/java/com/android/tools/r8/kotlin/KotlinPackageInfo.java
index f1e2573..a0f0e81 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinPackageInfo.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinPackageInfo.java
@@ -97,4 +97,8 @@
     containerInfo.trace(definitionSupplier);
     localDelegatedProperties.trace(definitionSupplier);
   }
+
+  public String getModuleName() {
+    return moduleName;
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/naming/KotlinModuleSynthesizer.java b/src/main/java/com/android/tools/r8/naming/KotlinModuleSynthesizer.java
new file mode 100644
index 0000000..7686a1e
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/naming/KotlinModuleSynthesizer.java
@@ -0,0 +1,188 @@
+// Copyright (c) 2022, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.naming;
+
+import static com.android.tools.r8.utils.FunctionUtils.ignoreArgument;
+
+import com.android.tools.r8.DataEntryResource;
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexItemFactory;
+import com.android.tools.r8.graph.DexProgramClass;
+import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.GraphLens;
+import com.android.tools.r8.kotlin.KotlinClassLevelInfo;
+import com.android.tools.r8.kotlin.KotlinMultiFileClassPartInfo;
+import com.android.tools.r8.origin.Origin;
+import com.android.tools.r8.utils.Box;
+import com.android.tools.r8.utils.Pair;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+import kotlinx.metadata.jvm.KotlinModuleMetadata.Writer;
+
+/**
+ * The kotlin module synthesizer will scan through all file facades and multiclass files to figure
+ * out the residual package destination of these and bucket them into their original module names.
+ */
+public class KotlinModuleSynthesizer {
+
+  private final AppView<?> appView;
+
+  public KotlinModuleSynthesizer(AppView<?> appView) {
+    this.appView = appView;
+  }
+
+  public boolean isKotlinModuleFile(DataEntryResource file) {
+    return file.getName().endsWith(".kotlin_module");
+  }
+
+  public List<DataEntryResource> synthesizeKotlinModuleFiles() {
+    Map<String, KotlinModuleInfoBuilder> kotlinModuleBuilders = new HashMap<>();
+    // We cannot obtain the module name for multi class file facades. But, we can for a multi class
+    // part obtain both the module name and the multi class facade. We therefore iterate over all
+    // classes to find a multi class facade -> module name mapping, and then iterate over all
+    // classes to assign multi class facades to modules.
+    Map<String, String> moduleNamesForParts = new HashMap<>();
+    for (DexProgramClass clazz : appView.app().classesWithDeterministicOrder()) {
+      KotlinClassLevelInfo kotlinInfo = clazz.getKotlinInfo();
+      if (kotlinInfo.isFileFacade()) {
+        kotlinModuleBuilders
+            .computeIfAbsent(
+                kotlinInfo.asFileFacade().getModuleName(),
+                moduleName -> new KotlinModuleInfoBuilder(moduleName, appView))
+            .add(clazz);
+      } else if (kotlinInfo.isMultiFileClassPart()) {
+        KotlinMultiFileClassPartInfo kotlinMultiFileClassPartInfo =
+            kotlinInfo.asMultiFileClassPart();
+        moduleNamesForParts.computeIfAbsent(
+            kotlinMultiFileClassPartInfo.getFacadeClassName(),
+            ignored -> kotlinMultiFileClassPartInfo.getModuleName());
+        kotlinModuleBuilders
+            .computeIfAbsent(
+                kotlinMultiFileClassPartInfo.getModuleName(),
+                moduleName -> new KotlinModuleInfoBuilder(moduleName, appView))
+            .add(clazz);
+      }
+    }
+    for (DexProgramClass clazz : appView.app().classesWithDeterministicOrder()) {
+      KotlinClassLevelInfo kotlinInfo = clazz.getKotlinInfo();
+      if (kotlinInfo.isMultiFileFacade()) {
+        DexType originalType = appView.graphLens().getOriginalType(clazz.getType());
+        if (originalType != null) {
+          String moduleNameForPart = moduleNamesForParts.get(originalType.toBinaryName());
+          // If module name is null then we did not find any multi class file parts and therefore
+          // do not have to do anything for the facade.
+          if (moduleNameForPart != null) {
+            KotlinModuleInfoBuilder kotlinModuleInfoBuilder =
+                kotlinModuleBuilders.get(moduleNameForPart);
+            assert kotlinModuleInfoBuilder != null;
+            kotlinModuleInfoBuilder.add(clazz);
+          }
+        }
+      }
+    }
+    if (kotlinModuleBuilders.isEmpty()) {
+      return Collections.emptyList();
+    }
+    List<DataEntryResource> newResources = new ArrayList<>();
+    kotlinModuleBuilders.values().forEach(builder -> builder.build().ifPresent(newResources::add));
+    return newResources;
+  }
+
+  private static class KotlinModuleInfoBuilder {
+
+    private final String moduleName;
+    private final GraphLens graphLens;
+    private final NamingLens namingLens;
+    private final DexItemFactory factory;
+
+    private final Map<String, List<String>> newFacades = new HashMap<>();
+    private final Map<String, List<Pair<String, String>>> multiClassFacadeOriginalToRenamed =
+        new LinkedHashMap<>();
+    private final Map<String, List<String>> multiClassPartToOriginal = new HashMap<>();
+    private final Box<int[]> metadataVersion = new Box<>();
+
+    private KotlinModuleInfoBuilder(String moduleName, AppView<?> appView) {
+      this.moduleName = moduleName;
+      this.graphLens = appView.graphLens();
+      this.namingLens = appView.getNamingLens();
+      this.factory = appView.dexItemFactory();
+    }
+
+    private void add(DexProgramClass clazz) {
+      DexType classType = clazz.getType();
+      KotlinClassLevelInfo classKotlinInfo = clazz.getKotlinInfo();
+      DexType renamedType = namingLens.lookupType(classType, factory);
+      if (classKotlinInfo.isFileFacade()) {
+        metadataVersion.computeIfAbsent(classKotlinInfo::getMetadataVersion);
+        newFacades
+            .computeIfAbsent(renamedType.getPackageName(), ignoreArgument(ArrayList::new))
+            .add(renamedType.toBinaryName());
+      } else if (classKotlinInfo.isMultiFileFacade()) {
+        metadataVersion.computeIfAbsent(classKotlinInfo::getMetadataVersion);
+        DexType originalType = graphLens.getOriginalType(classType);
+        multiClassFacadeOriginalToRenamed
+            .computeIfAbsent(renamedType.getPackageName(), ignoreArgument(ArrayList::new))
+            .add(Pair.create(originalType.toBinaryName(), renamedType.toBinaryName()));
+      } else {
+        assert classKotlinInfo.isMultiFileClassPart();
+        metadataVersion.computeIfAbsent(classKotlinInfo::getMetadataVersion);
+        KotlinMultiFileClassPartInfo classPart = classKotlinInfo.asMultiFileClassPart();
+        multiClassPartToOriginal
+            .computeIfAbsent(classPart.getFacadeClassName(), ignoreArgument(ArrayList::new))
+            .add(renamedType.toBinaryName());
+      }
+    }
+
+    private Optional<DataEntryResource> build() {
+      // If multiClassParts are non empty but multiFileFacade is, then we have no place to put
+      // the parts anyway, so we can just return empty.
+      if (newFacades.isEmpty() && multiClassFacadeOriginalToRenamed.isEmpty()) {
+        return Optional.empty();
+      }
+      assert metadataVersion.isSet();
+      List<String> packagesSorted = new ArrayList<>(newFacades.keySet());
+      for (String newPackage : multiClassFacadeOriginalToRenamed.keySet()) {
+        if (!newFacades.containsKey(newPackage)) {
+          packagesSorted.add(newPackage);
+        }
+      }
+      Collections.sort(packagesSorted);
+      Writer writer = new Writer();
+      for (String newPackage : packagesSorted) {
+        // Calling other visitors than visitPackageParts are currently not supported.
+        // https://github.com/JetBrains/kotlin/blob/master/libraries/kotlinx-metadata/
+        //  jvm/src/kotlinx/metadata/jvm/KotlinModuleMetadata.kt#L70
+        Map<String, String> newMultiFiles = new LinkedHashMap<>();
+        multiClassFacadeOriginalToRenamed
+            .getOrDefault(newPackage, Collections.emptyList())
+            .forEach(
+                pair -> {
+                  String originalName = pair.getFirst();
+                  String rewrittenName = pair.getSecond();
+                  multiClassPartToOriginal
+                      .getOrDefault(originalName, Collections.emptyList())
+                      .forEach(
+                          classPart -> {
+                            newMultiFiles.put(classPart, rewrittenName);
+                          });
+                });
+        writer.visitPackageParts(
+            newPackage,
+            newFacades.getOrDefault(newPackage, Collections.emptyList()),
+            newMultiFiles);
+      }
+      return Optional.of(
+          DataEntryResource.fromBytes(
+              writer.write(metadataVersion.get()).getBytes(),
+              "META-INF/" + moduleName + ".kotlin_module",
+              Origin.unknown()));
+    }
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/utils/DescriptorUtils.java b/src/main/java/com/android/tools/r8/utils/DescriptorUtils.java
index 7cce06b..99645fa 100644
--- a/src/main/java/com/android/tools/r8/utils/DescriptorUtils.java
+++ b/src/main/java/com/android/tools/r8/utils/DescriptorUtils.java
@@ -404,6 +404,17 @@
   }
 
   /**
+   * Get package java name from a class type name.
+   *
+   * @param typeName a class descriptor i.e. "java.lang.Object"
+   * @return java package name i.e. "java.lang"
+   */
+  public static String getPackageNameFromTypeName(String typeName) {
+    return getPackageNameFromBinaryName(
+        getClassBinaryNameFromDescriptor(javaTypeToDescriptor(typeName)));
+  }
+
+  /**
    * Convert package name to a binary name.
    *
    * @param packageName a package name i.e., "java.lang"
diff --git a/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInExtensionFunctionTest.java b/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInExtensionFunctionTest.java
index 619f4a1..99cdd63 100644
--- a/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInExtensionFunctionTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInExtensionFunctionTest.java
@@ -15,10 +15,13 @@
 import static org.junit.Assume.assumeTrue;
 
 import com.android.tools.r8.KotlinTestParameters;
+import com.android.tools.r8.R8TestCompileResult;
 import com.android.tools.r8.TestParameters;
 import com.android.tools.r8.shaking.ProguardKeepAttributes;
-import com.android.tools.r8.utils.Box;
+import com.android.tools.r8.utils.DescriptorUtils;
+import com.android.tools.r8.utils.FileUtils;
 import com.android.tools.r8.utils.StringUtils;
+import com.android.tools.r8.utils.ZipUtils;
 import com.android.tools.r8.utils.codeinspector.ClassSubject;
 import com.android.tools.r8.utils.codeinspector.CodeInspector;
 import com.android.tools.r8.utils.codeinspector.KmClassSubject;
@@ -28,8 +31,13 @@
 import com.android.tools.r8.utils.codeinspector.KmTypeSubject;
 import com.android.tools.r8.utils.codeinspector.KmValueParameterSubject;
 import com.android.tools.r8.utils.codeinspector.Matchers;
+import com.android.tools.r8.utils.codeinspector.MethodSubject;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -184,8 +192,7 @@
   @Test
   public void testMetadataInExtensionFunction_renamedKotlinSources() throws Exception {
     assumeTrue(kotlinc.getCompilerVersion().isGreaterThanOrEqualTo(KOTLINC_1_4_20));
-    Box<String> renamedKtHolder = new Box<>();
-    Path libJar =
+    R8TestCompileResult r8LibraryResult =
         testForR8(parameters.getBackend())
             .addClasspathFiles(kotlinc.getKotlinStdlibJar(), kotlinc.getKotlinAnnotationJar())
             .addProgramFiles(extLibJarMap.getForConfiguration(kotlinc, targetVersion))
@@ -201,21 +208,65 @@
             .addKeepAttributes(ProguardKeepAttributes.SIGNATURE)
             .addKeepAttributes(ProguardKeepAttributes.INNER_CLASSES)
             .addKeepAttributes(ProguardKeepAttributes.ENCLOSING_METHOD)
-            .compile()
-            .inspect(
-                inspector -> {
-                  ClassSubject clazz = inspector.clazz(PKG + ".extension_function_lib.BKt");
-                  assertThat(clazz, isPresentAndRenamed());
-                  renamedKtHolder.set(clazz.getFinalName());
-                })
-            .writeToZip();
+            .compile();
+    Path kotlinSourcePath = getKotlinFileInTest(PKG_PREFIX + "/extension_function_app", "main");
 
-    kotlinc(parameters.getRuntime().asCf(), kotlinc, targetVersion)
-        .addClasspathFiles(libJar)
-        .addSourceFiles(getKotlinFileInTest(PKG_PREFIX + "/extension_function_app", "main"))
-        .setOutputPath(temp.newFolder().toPath())
-        // TODO(b/242289529): Expect that we can compile without errors.
-        .compile(true);
+    String kotlinSource = FileUtils.readTextFile(kotlinSourcePath, StandardCharsets.UTF_8);
+
+    CodeInspector inspector = r8LibraryResult.inspector();
+
+    ClassSubject clazz = inspector.clazz(PKG + ".extension_function_lib.BKt");
+    assertThat(clazz, isPresentAndRenamed());
+
+    // Rewrite the source kotlin files that reference the four extension methods into their renamed
+    // name by changing the import statement and the actual call.
+    String[] methodNames = new String[] {"extension", "csHash", "longArrayHash", "myApply"};
+    for (String methodName : methodNames) {
+      MethodSubject method = clazz.uniqueMethodWithName(methodName);
+      assertThat(method, isPresentAndRenamed());
+      String finalMethodName = method.getFinalName();
+      kotlinSource =
+          kotlinSource.replace(
+              "import com.android.tools.r8.kotlin.metadata.extension_function_lib." + methodName,
+              "import "
+                  + DescriptorUtils.getPackageNameFromTypeName(clazz.getFinalName())
+                  + "."
+                  + finalMethodName);
+      kotlinSource = kotlinSource.replace(")." + methodName, ")." + finalMethodName);
+    }
+
+    Path newSource = temp.newFolder().toPath().resolve("main.kt");
+    Files.write(newSource, kotlinSource.getBytes(StandardCharsets.UTF_8));
+
+    Path libJar = r8LibraryResult.writeToZip();
+    Path tempUnzipPath = temp.newFolder().toPath();
+    List<String> kotlinModuleFiles = new ArrayList<>();
+    ZipUtils.unzip(
+        libJar,
+        tempUnzipPath,
+        f -> {
+          if (f.getName().endsWith(".kotlin_module")) {
+            kotlinModuleFiles.add(f.getName());
+          }
+          return false;
+        });
+    assertEquals(Collections.singletonList("META-INF/main.kotlin_module"), kotlinModuleFiles);
+    Path output =
+        kotlinc(parameters.getRuntime().asCf(), kotlinc, targetVersion)
+            .addClasspathFiles(libJar)
+            .addSourceFiles(newSource)
+            .setOutputPath(temp.newFolder().toPath())
+            .compile();
+
+    if (kotlinParameters.isOlderThan(KOTLINC_1_4_20)) {
+      return;
+    }
+
+    testForJvm()
+        .addRunClasspathFiles(kotlinc.getKotlinStdlibJar(), libJar)
+        .addClasspath(output)
+        .run(parameters.getRuntime(), PKG + ".extension_function_app.MainKt")
+        .assertSuccessWithOutput(EXPECTED);
   }
 
   private void inspectRenamed(CodeInspector inspector) {
diff --git a/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInMultifileClassTest.java b/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInMultifileClassTest.java
index 5a3960e..ec80ba8 100644
--- a/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInMultifileClassTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRewriteInMultifileClassTest.java
@@ -36,7 +36,7 @@
 
 @RunWith(Parameterized.class)
 public class MetadataRewriteInMultifileClassTest extends KotlinMetadataTestBase {
-  private static final String EXPECTED = StringUtils.lines(", 1, 2, 3");
+  private static final String EXPECTED = StringUtils.lines(", 1, 2, 3", ", 1, 2, 3");
 
   private final TestParameters parameters;
 
@@ -95,9 +95,7 @@
             .addClasspathFiles(libJar)
             .addSourceFiles(getKotlinFileInTest(PKG_PREFIX + "/multifileclass_app", "main"))
             .setOutputPath(temp.newFolder().toPath())
-            // TODO(b/151193860): update to just .compile() once fixed.
             .compileRaw();
-    // TODO(b/151193860): should be able to compile!
     assertNotEquals(0, kotlinTestCompileResult.exitCode);
     assertThat(kotlinTestCompileResult.stderr, containsString("unresolved reference: join"));
   }
@@ -113,8 +111,6 @@
     assertThat(joinOfInt, not(isPresent()));
 
     inspectMetadataForFacade(inspector, util);
-    // TODO(b/156290332): Seems like this test is incorrect and should never work.
-    // inspectSignedKt(inspector);
   }
 
   @Test
@@ -126,6 +122,7 @@
             // Keep UtilKt#comma*Join*().
             .addKeepRules("-keep class **.UtilKt")
             .addKeepRules("-keep,allowobfuscation class **.UtilKt__SignedKt")
+            .addKeepRules("-keep,allowobfuscation class **.UtilKt__UnsignedKt")
             .addKeepRules("-keepclassmembers class * { ** comma*Join*(...); }")
             // Keep yet rename joinOf*(String).
             .addKeepRules("-keepclassmembers,allowobfuscation class * { ** joinOf*(...); }")
@@ -134,16 +131,18 @@
             .inspect(this::inspectRenamed)
             .writeToZip();
 
-    ProcessResult kotlinTestCompileResult =
+    Path output =
         kotlinc(parameters.getRuntime().asCf(), kotlinc, targetVersion)
             .addClasspathFiles(libJar)
             .addSourceFiles(getKotlinFileInTest(PKG_PREFIX + "/multifileclass_app", "main"))
             .setOutputPath(temp.newFolder().toPath())
-            // TODO(b/151193860): update to just .compile() once fixed.
-            .compileRaw();
-    // TODO(b/151193860): should be able to compile!
-    assertNotEquals(0, kotlinTestCompileResult.exitCode);
-    assertThat(kotlinTestCompileResult.stderr, containsString("unresolved reference: join"));
+            .compile();
+
+    testForJvm()
+        .addRunClasspathFiles(kotlinc.getKotlinStdlibJar(), libJar)
+        .addClasspath(output)
+        .run(parameters.getRuntime(), PKG + ".multifileclass_app.MainKt")
+        .assertSuccessWithOutput(EXPECTED);
   }
 
   private void inspectRenamed(CodeInspector inspector) {
diff --git a/src/test/java/com/android/tools/r8/kotlin/metadata/extension_function_app/main.kt b/src/test/java/com/android/tools/r8/kotlin/metadata/extension_function_app/main.kt
index 72e86b4..dc7d16e 100644
--- a/src/test/java/com/android/tools/r8/kotlin/metadata/extension_function_app/main.kt
+++ b/src/test/java/com/android/tools/r8/kotlin/metadata/extension_function_app/main.kt
@@ -13,7 +13,7 @@
   B().doStuff()
   B().extension()
 
-  "R8".csHash()
+  ("R8").csHash()
   longArrayOf(42L).longArrayHash()
   B().myApply { this.doStuff() }
 }
diff --git a/src/test/java/com/android/tools/r8/kotlin/metadata/multifileclass_app/main.kt b/src/test/java/com/android/tools/r8/kotlin/metadata/multifileclass_app/main.kt
index ce098bf..6d81d8c 100644
--- a/src/test/java/com/android/tools/r8/kotlin/metadata/multifileclass_app/main.kt
+++ b/src/test/java/com/android/tools/r8/kotlin/metadata/multifileclass_app/main.kt
@@ -5,7 +5,18 @@
 
 import com.android.tools.r8.kotlin.metadata.multifileclass_lib.join
 
-fun main() {
+fun signed() {
   val s = sequenceOf(1, 2, 3)
   println(s.join())
 }
+
+@OptIn(ExperimentalUnsignedTypes::class)
+fun unsigned() {
+  val s = sequenceOf(1u, 2u, 3u)
+  println(s.join())
+}
+
+fun main() {
+  signed()
+  unsigned()
+}