Unify @Metadata rewriting of functions and properties in KmDeclarationContainer.

Bug: 70169921
Change-Id: Idef0f5ad4029e14aabc3c25958c5206a3d724b6b
diff --git a/src/main/java/com/android/tools/r8/graph/DexClass.java b/src/main/java/com/android/tools/r8/graph/DexClass.java
index 815e721..06f579c 100644
--- a/src/main/java/com/android/tools/r8/graph/DexClass.java
+++ b/src/main/java/com/android/tools/r8/graph/DexClass.java
@@ -233,6 +233,9 @@
       List<KmFunction> extensions, AppView<?> appView) {
     ImmutableMap.Builder<DexEncodedMethod, KmFunction> builder = ImmutableMap.builder();
     for (DexEncodedMethod method : directMethods) {
+      if (method.isInitializer()) {
+        continue;
+      }
       KmFunction extension = method.findCompatibleKotlinExtension(extensions, appView);
       if (extension != null) {
         // Found a compatible extension that is likely asked to keep.
@@ -245,7 +248,10 @@
   public List<DexEncodedMethod> kotlinFunctions(
       List<KmFunction> functions, List<KmProperty> properties, AppView<?> appView) {
     ImmutableList.Builder<DexEncodedMethod> builder = ImmutableList.builder();
-    for (DexEncodedMethod method : virtualMethods) {
+    for (DexEncodedMethod method : methods()) {
+      if (method.isInitializer()) {
+        continue;
+      }
       KmFunction function = method.findCompatibleKotlinFunction(functions, appView);
       if (function != null) {
         // Found a compatible function that is likely asked to keep.
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinClass.java b/src/main/java/com/android/tools/r8/kotlin/KotlinClass.java
index f4503ce..ff97c60 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinClass.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinClass.java
@@ -5,11 +5,8 @@
 package com.android.tools.r8.kotlin;
 
 import static com.android.tools.r8.kotlin.Kotlin.addKotlinPrefix;
-import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.isExtension;
 import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toKmType;
 import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toRenamedKmConstructor;
-import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toRenamedKmFunction;
-import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toRenamedKmFunctionAsExtension;
 import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toRenamedKmType;
 
 import com.android.tools.r8.graph.AppView;
@@ -21,11 +18,8 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
-import java.util.stream.Collectors;
 import kotlinx.metadata.KmClass;
 import kotlinx.metadata.KmConstructor;
-import kotlinx.metadata.KmFunction;
-import kotlinx.metadata.KmProperty;
 import kotlinx.metadata.KmType;
 import kotlinx.metadata.jvm.KotlinClassHeader;
 import kotlinx.metadata.jvm.KotlinClassMetadata;
@@ -87,32 +81,7 @@
       }
     }
 
-    List<KmFunction> functions = kmClass.getFunctions();
-    List<KmFunction> originalFunctions =
-        functions.stream()
-            .filter(kmFunction -> !isExtension(kmFunction))
-            .collect(Collectors.toList());
-    List<KmFunction> originalExtensions =
-        functions.stream()
-            .filter(KotlinMetadataSynthesizer::isExtension)
-            .collect(Collectors.toList());
-    functions.clear();
-
-    List<KmProperty> properties = kmClass.getProperties();
-    for (DexEncodedMethod method : clazz.kotlinFunctions(originalFunctions, properties, appView)) {
-      KmFunction function = toRenamedKmFunction(method, null, appView, lens);
-      if (function != null) {
-        functions.add(function);
-      }
-    }
-    for (Map.Entry<DexEncodedMethod, KmFunction> entry :
-        clazz.kotlinExtensions(originalExtensions, appView).entrySet()) {
-      KmFunction extension =
-          toRenamedKmFunctionAsExtension(entry.getKey(), entry.getValue(), appView, lens);
-      if (extension != null) {
-        functions.add(extension);
-      }
-    }
+    rewriteDeclarationContainer(kmClass, appView, lens);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinClassPart.java b/src/main/java/com/android/tools/r8/kotlin/KotlinClassPart.java
index 4306645..a9a4b9c 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinClassPart.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinClassPart.java
@@ -4,17 +4,10 @@
 
 package com.android.tools.r8.kotlin;
 
-import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toRenamedKmFunctionAsExtension;
-
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexClass;
-import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.naming.NamingLens;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import java.util.List;
-import java.util.Map;
-import java.util.stream.Collectors;
-import kotlinx.metadata.KmFunction;
 import kotlinx.metadata.KmPackage;
 import kotlinx.metadata.jvm.KotlinClassHeader;
 import kotlinx.metadata.jvm.KotlinClassMetadata;
@@ -47,21 +40,7 @@
     if (!appView.options().enableKotlinMetadataRewriting) {
       return;
     }
-    List<KmFunction> functions = kmPackage.getFunctions();
-    List<KmFunction> originalExtensions =
-        functions.stream()
-            .filter(KotlinMetadataSynthesizer::isExtension)
-            .collect(Collectors.toList());
-    functions.clear();
-
-    for (Map.Entry<DexEncodedMethod, KmFunction> entry :
-        clazz.kotlinExtensions(originalExtensions, appView).entrySet()) {
-      KmFunction extension =
-          toRenamedKmFunctionAsExtension(entry.getKey(), entry.getValue(), appView, lens);
-      if (extension != null) {
-        functions.add(extension);
-      }
-    }
+    rewriteDeclarationContainer(kmPackage, appView, lens);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinFile.java b/src/main/java/com/android/tools/r8/kotlin/KotlinFile.java
index 71538fc..855a428 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinFile.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinFile.java
@@ -4,17 +4,10 @@
 
 package com.android.tools.r8.kotlin;
 
-import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toRenamedKmFunctionAsExtension;
-
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexClass;
-import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.naming.NamingLens;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
-import java.util.List;
-import java.util.Map;
-import java.util.stream.Collectors;
-import kotlinx.metadata.KmFunction;
 import kotlinx.metadata.KmPackage;
 import kotlinx.metadata.jvm.KotlinClassHeader;
 import kotlinx.metadata.jvm.KotlinClassMetadata;
@@ -47,21 +40,7 @@
     if (!appView.options().enableKotlinMetadataRewriting) {
       return;
     }
-    List<KmFunction> functions = kmPackage.getFunctions();
-    List<KmFunction> originalExtensions =
-        functions.stream()
-            .filter(KotlinMetadataSynthesizer::isExtension)
-            .collect(Collectors.toList());
-    functions.clear();
-
-    for (Map.Entry<DexEncodedMethod, KmFunction> entry :
-        clazz.kotlinExtensions(originalExtensions, appView).entrySet()) {
-      KmFunction extension =
-          toRenamedKmFunctionAsExtension(entry.getKey(), entry.getValue(), appView, lens);
-      if (extension != null) {
-        functions.add(extension);
-      }
-    }
+    rewriteDeclarationContainer(kmPackage, appView, lens);
   }
 
   @Override
diff --git a/src/main/java/com/android/tools/r8/kotlin/KotlinInfo.java b/src/main/java/com/android/tools/r8/kotlin/KotlinInfo.java
index bd8d4c0..9cd71c3 100644
--- a/src/main/java/com/android/tools/r8/kotlin/KotlinInfo.java
+++ b/src/main/java/com/android/tools/r8/kotlin/KotlinInfo.java
@@ -4,10 +4,21 @@
 
 package com.android.tools.r8.kotlin;
 
+import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.isExtension;
+import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toRenamedKmFunction;
+import static com.android.tools.r8.kotlin.KotlinMetadataSynthesizer.toRenamedKmFunctionAsExtension;
+
 import com.android.tools.r8.graph.AppView;
 import com.android.tools.r8.graph.DexClass;
+import com.android.tools.r8.graph.DexEncodedMethod;
 import com.android.tools.r8.naming.NamingLens;
 import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import kotlinx.metadata.KmDeclarationContainer;
+import kotlinx.metadata.KmFunction;
+import kotlinx.metadata.KmProperty;
 import kotlinx.metadata.jvm.KotlinClassHeader;
 import kotlinx.metadata.jvm.KotlinClassMetadata;
 
@@ -28,6 +39,8 @@
   abstract void processMetadata();
 
   // Subtypes will define how to rewrite metadata after shrinking and minification.
+  // Subtypes that represent subtypes of {@link KmDeclarationContainer} can use
+  // {@link #rewriteDeclarationContainer} below.
   abstract void rewrite(AppView<AppInfoWithLiveness> appView, NamingLens lens);
 
   abstract KotlinClassHeader createHeader();
@@ -83,4 +96,40 @@
     return (clazz != null ? clazz.toSourceString() : "<null class?!>")
         + ": " + metadata.toString();
   }
+
+  // {@link KmClass} and {@link KmPackage} are inherited from {@link KmDeclarationContainer} that
+  // abstract functions and properties. Rewriting of those portions can be unified here.
+  void rewriteDeclarationContainer(
+      KmDeclarationContainer kmDeclarationContainer,
+      AppView<AppInfoWithLiveness> appView,
+      NamingLens lens) {
+    assert clazz != null;
+
+    List<KmFunction> functions = kmDeclarationContainer.getFunctions();
+    List<KmFunction> originalFunctions =
+        functions.stream()
+            .filter(kmFunction -> !isExtension(kmFunction))
+            .collect(Collectors.toList());
+    List<KmFunction> originalExtensions =
+        functions.stream()
+            .filter(KotlinMetadataSynthesizer::isExtension)
+            .collect(Collectors.toList());
+    functions.clear();
+
+    List<KmProperty> properties = kmDeclarationContainer.getProperties();
+    for (DexEncodedMethod method : clazz.kotlinFunctions(originalFunctions, properties, appView)) {
+      KmFunction function = toRenamedKmFunction(method, null, appView, lens);
+      if (function != null) {
+        functions.add(function);
+      }
+    }
+    for (Map.Entry<DexEncodedMethod, KmFunction> entry :
+        clazz.kotlinExtensions(originalExtensions, appView).entrySet()) {
+      KmFunction extension =
+          toRenamedKmFunctionAsExtension(entry.getKey(), entry.getValue(), appView, lens);
+      if (extension != null) {
+        functions.add(extension);
+      }
+    }
+  }
 }
diff --git a/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRenameInFunctionTest.java b/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRenameInFunctionTest.java
index 4e117ef..b11d4ee 100644
--- a/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRenameInFunctionTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/metadata/MetadataRenameInFunctionTest.java
@@ -6,16 +6,14 @@
 import static com.android.tools.r8.KotlinCompilerTool.KOTLINC;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
 import static com.android.tools.r8.utils.codeinspector.Matchers.isRenamed;
-import static org.hamcrest.CoreMatchers.containsString;
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.Assert.assertNotEquals;
 import static org.junit.Assert.assertTrue;
 
 import com.android.tools.r8.R8TestCompileResult;
 import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.ToolHelper;
 import com.android.tools.r8.ToolHelper.KotlinTargetVersion;
-import com.android.tools.r8.ToolHelper.ProcessResult;
 import com.android.tools.r8.shaking.ProguardKeepAttributes;
 import com.android.tools.r8.utils.InternalOptions;
 import com.android.tools.r8.utils.codeinspector.ClassSubject;
@@ -90,16 +88,18 @@
     Path libJar = compileResult.writeToZip();
 
     String appFolder = PKG_PREFIX + "/function_app";
-    ProcessResult kotlinTestCompileResult =
+    Path output =
         kotlinc(parameters.getRuntime().asCf(), KOTLINC, KotlinTargetVersion.JAVA_8)
             .addClasspathFiles(libJar)
             .addSourceFiles(getKotlinFileInTest(appFolder, "main"))
             .setOutputPath(temp.newFolder().toPath())
-            // TODO(b/70169921): update to just .compile() once fixed.
-            .compileRaw();
-    // TODO(b/70169921): should be able to compile!
-    assertNotEquals(0, kotlinTestCompileResult.exitCode);
-    assertThat(kotlinTestCompileResult.stderr, containsString("unresolved reference: `fun`"));
+            .compile();
+
+    testForJvm()
+        .addRunClasspathFiles(ToolHelper.getKotlinStdlibJar(), libJar)
+        .addClasspath(output)
+        .run(parameters.getRuntime(), pkg + ".function_app.MainKt")
+        .assertSuccessWithOutputLines("do stuff", "do stuff");
   }
 
   @Test
@@ -141,15 +141,17 @@
     Path libJar = compileResult.writeToZip();
 
     String appFolder = PKG_PREFIX + "/function_app";
-    ProcessResult kotlinTestCompileResult =
+    Path output =
         kotlinc(parameters.getRuntime().asCf(), KOTLINC, KotlinTargetVersion.JAVA_8)
             .addClasspathFiles(libJar)
             .addSourceFiles(getKotlinFileInTest(appFolder, "main"))
             .setOutputPath(temp.newFolder().toPath())
-            // TODO(b/70169921): update to just .compile() once fixed.
-            .compileRaw();
-    // TODO(b/70169921): should be able to compile!
-    assertNotEquals(0, kotlinTestCompileResult.exitCode);
-    assertThat(kotlinTestCompileResult.stderr, containsString("unresolved reference: `fun`"));
+            .compile();
+
+    testForJvm()
+        .addRunClasspathFiles(ToolHelper.getKotlinStdlibJar(), libJar)
+        .addClasspath(output)
+        .run(parameters.getRuntime(), pkg + ".function_app.MainKt")
+        .assertSuccessWithOutputLines("do stuff", "do stuff");
   }
-}
\ No newline at end of file
+}