Use the arity of k-style lambda group's interface, not main method.

The signature of the main method (lambda method) can be changed, e.g.,
by unused argument removal. Then, kotlin...Lambda arity would be set
incorrectly if we use the main method's signature, instead of lambda
group's interface arity.

Bug: 133379765
Change-Id: Ia0f83a8595f70fb0e61eaba8e588f801ca9c371c
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/lambda/kotlin/KStyleLambdaGroup.java b/src/main/java/com/android/tools/r8/ir/optimize/lambda/kotlin/KStyleLambdaGroup.java
index 587e172..1595e8e 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/lambda/kotlin/KStyleLambdaGroup.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/lambda/kotlin/KStyleLambdaGroup.java
@@ -41,7 +41,7 @@
 //      -----------------------------------------------------------------------
 //
 // Regular stateless k-style lambda class structure looks like below:
-// NOTE: stateless j-style lambdas do not always have INSTANCE field.
+// NOTE: stateless k-style lambdas do not always have INSTANCE field.
 //
 // -----------------------------------------------------------------------------------------------
 // // signature Lkotlin/jvm/internal/Lambda;Lkotlin/jvm/functions/Function2<
@@ -205,7 +205,7 @@
           group.getLambdaIdField(factory),
           id -> group.getCaptureField(factory, id),
           initializerMethod,
-          id.mainMethodProto.parameters.size(),
+          factory.kotlin.functional.getArity(id.iface),
           callerPosition);
     }
   }
diff --git a/src/main/java/com/android/tools/r8/kotlin/Kotlin.java b/src/main/java/com/android/tools/r8/kotlin/Kotlin.java
index 0a85f89..73ab96a 100644
--- a/src/main/java/com/android/tools/r8/kotlin/Kotlin.java
+++ b/src/main/java/com/android/tools/r8/kotlin/Kotlin.java
@@ -11,8 +11,11 @@
 import com.android.tools.r8.graph.DexString;
 import com.android.tools.r8.graph.DexType;
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.Sets;
-import java.util.Set;
+import it.unimi.dsi.fastutil.objects.Object2IntArrayMap;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 /** Class provides basic information about symbols related to Kotlin support. */
 public final class Kotlin {
@@ -43,22 +46,24 @@
   }
 
   public final class Functional {
-    private final Set<DexType> functions = Sets.newIdentityHashSet();
+    // NOTE: Kotlin stdlib defines interface Function0 till Function22 explicitly, see:
+    //    https://github.com/JetBrains/kotlin/blob/master/libraries/
+    //                    stdlib/jvm/runtime/kotlin/jvm/functions/Functions.kt
+    //
+    // For functions with arity bigger that 22 it is supposed to use FunctionN as described
+    // in https://github.com/JetBrains/kotlin/blob/master/spec-docs/function-types.md,
+    // but in current implementation (v1.2.21) defining such a lambda results in:
+    //   > "Error: A JNI error has occurred, please check your installation and try again"
+    //
+    // This implementation just ignores lambdas with arity > 22.
+    private final Object2IntMap<DexType> functions = new Object2IntArrayMap<>(
+        IntStream.rangeClosed(0, 22).boxed().collect(
+            Collectors.toMap(
+                i -> factory.createType(addKotlinPrefix("jvm/functions/Function") + i + ";"),
+                Function.identity()))
+    );
 
     private Functional() {
-      // NOTE: Kotlin stdlib defines interface Function0 till Function22 explicitly, see:
-      //    https://github.com/JetBrains/kotlin/blob/master/libraries/
-      //                    stdlib/jvm/runtime/kotlin/jvm/functions/Functions.kt
-      //
-      // For functions with arity bigger that 22 it is supposed to use FunctionN as described
-      // in https://github.com/JetBrains/kotlin/blob/master/spec-docs/function-types.md,
-      // but in current implementation (v1.2.21) defining such a lambda results in:
-      //   > "Error: A JNI error has occurred, please check your installation and try again"
-      //
-      // This implementation just ignores lambdas with arity > 22.
-      for (int i = 0; i <= 22; i++) {
-        functions.add(factory.createType(addKotlinPrefix("jvm/functions/Function") + i + ";"));
-      }
     }
 
     public final DexString kotlinStyleLambdaInstanceName = factory.createString("INSTANCE");
@@ -73,7 +78,13 @@
         factory.constructorMethodName);
 
     public boolean isFunctionInterface(DexType type) {
-      return functions.contains(type);
+      return functions.containsKey(type);
+    }
+
+    public int getArity(DexType type) {
+      assert isFunctionInterface(type)
+          : "Request to retrieve the arity from non-Kotlin-Function: " + type.toSourceString();
+      return functions.getInt(type);
     }
   }
 
diff --git a/src/test/java/com/android/tools/r8/kotlin/KotlinUnusedArgumentsInLambdasTest.java b/src/test/java/com/android/tools/r8/kotlin/KotlinUnusedArgumentsInLambdasTest.java
index df7c84b..95ec356 100644
--- a/src/test/java/com/android/tools/r8/kotlin/KotlinUnusedArgumentsInLambdasTest.java
+++ b/src/test/java/com/android/tools/r8/kotlin/KotlinUnusedArgumentsInLambdasTest.java
@@ -5,6 +5,7 @@
 
 import static com.android.tools.r8.utils.codeinspector.Matchers.isPresent;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
 import com.android.tools.r8.ToolHelper.KotlinTargetVersion;
@@ -13,7 +14,6 @@
 import com.android.tools.r8.utils.codeinspector.MethodSubject;
 import com.google.common.collect.ImmutableList;
 import java.util.function.Consumer;
-import org.junit.Ignore;
 import org.junit.Test;
 
 public class KotlinUnusedArgumentsInLambdasTest extends AbstractR8KotlinTestBase {
@@ -30,9 +30,8 @@
     super(targetVersion, allowAccessModification);
   }
 
-  @Ignore("b/133379765")
   @Test
-  public void testMergingKStyleLambdasAndReprocessing() throws Exception {
+  public void testMergingKStyleLambdasAfterUnusedArgumentRemoval() throws Exception {
     final String mainClassName = "unused_arg_in_lambdas_kstyle.MainKt";
     runTest("unused_arg_in_lambdas_kstyle", mainClassName, optionsModifier, app -> {
       CodeInspector inspector = new CodeInspector(app);
@@ -40,8 +39,27 @@
         if (classSubject.getOriginalDescriptor().contains("$ks")) {
           MethodSubject init = classSubject.init(ImmutableList.of("int"));
           assertThat(init, isPresent());
-          // Arity 1 should appear.
-          assertTrue(init.iterateInstructions(i -> i.isConstNumber(1)).hasNext());
+          // Arity 2 should appear.
+          assertTrue(init.iterateInstructions(i -> i.isConstNumber(2)).hasNext());
+
+          MethodSubject invoke = classSubject.uniqueMethodWithName("invoke");
+          assertThat(invoke, isPresent());
+          assertEquals(2, invoke.getMethod().method.proto.parameters.size());
+        }
+      });
+    });
+  }
+
+  @Test
+  public void testMergingJStyleLambdasAfterUnusedArgumentRemoval() throws Exception {
+    final String mainClassName = "unused_arg_in_lambdas_jstyle.MainKt";
+    runTest("unused_arg_in_lambdas_jstyle", mainClassName, optionsModifier, app -> {
+      CodeInspector inspector = new CodeInspector(app);
+      inspector.forAllClasses(classSubject -> {
+        if (classSubject.getOriginalDescriptor().contains("$js")) {
+          MethodSubject get = classSubject.uniqueMethodWithName("get");
+          assertThat(get, isPresent());
+          assertEquals(3, get.getMethod().method.proto.parameters.size());
         }
       });
     });
diff --git a/src/test/kotlinR8TestResources/lambdas_jstyle_trivial/main.kt b/src/test/kotlinR8TestResources/lambdas_jstyle_trivial/main.kt
index 6cdae60..7fd2d0c 100644
--- a/src/test/kotlinR8TestResources/lambdas_jstyle_trivial/main.kt
+++ b/src/test/kotlinR8TestResources/lambdas_jstyle_trivial/main.kt
@@ -4,7 +4,6 @@
 
 package lambdas_jstyle_trivial
 
-import lambdas_jstyle_trivial.Lambdas
 import lambdas_jstyle_trivial.inner.testInner
 
 private var COUNT = 0
diff --git a/src/test/kotlinR8TestResources/unused_arg_in_lambdas_jstyle/Lambdas.java b/src/test/kotlinR8TestResources/unused_arg_in_lambdas_jstyle/Lambdas.java
new file mode 100644
index 0000000..bb33bd5
--- /dev/null
+++ b/src/test/kotlinR8TestResources/unused_arg_in_lambdas_jstyle/Lambdas.java
@@ -0,0 +1,16 @@
+// Copyright (c) 2019, 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 unused_arg_in_lambdas_jstyle;
+
+public class Lambdas {
+
+  public interface MultiFunction<R, P1, P2, P3> {
+    R get(P1 a, P2 b, P3 c);
+  }
+
+  public synchronized static <R, P1, P2, P3>
+  void acceptMultiFunction(MultiFunction<R, P1, P2, P3> s, P1 a, P2 b, P3 c) {
+    System.out.println(s.get(a, b, c));
+  }
+}
diff --git a/src/test/kotlinR8TestResources/unused_arg_in_lambdas_jstyle/main.kt b/src/test/kotlinR8TestResources/unused_arg_in_lambdas_jstyle/main.kt
new file mode 100644
index 0000000..3302903
--- /dev/null
+++ b/src/test/kotlinR8TestResources/unused_arg_in_lambdas_jstyle/main.kt
@@ -0,0 +1,25 @@
+// Copyright (c) 2019, 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 unused_arg_in_lambdas_jstyle
+
+private var COUNT = 0
+
+fun nextInt() = COUNT++
+fun next() = "${nextInt()}".padStart(3, '0')
+
+fun main(args: Array<String>) {
+    multiFunctionLambdaFactory(next(), next(), next())
+}
+
+private data class Local<out T>(val id: T)
+
+private fun <P1, P2, P3> multiFunctionLambdaFactory(a: P1, b: P2, c:P3) {
+    Lambdas.acceptMultiFunction({ x, _, z -> "$x:unused:$z" }, a, b, c)
+    Lambdas.acceptMultiFunction({ x, _, z -> "$x:unused:$z" }, c, a, b)
+    Lambdas.acceptMultiFunction({ x, _, z -> "$x:unused:$z" }, b, c, a)
+
+    Lambdas.acceptMultiFunction({ x, _, z -> "$x:unused:$z" }, Local(a), b, c)
+    Lambdas.acceptMultiFunction({ x, _, z -> "$x:unused:$z" }, Local(b), a, c)
+    Lambdas.acceptMultiFunction({ x, _, z -> "$x:unused:$z" }, Local(c), b, a)
+}