Kotlin: fix array iterator extraction

java/kotlin-extractor/src/main/kotlin/KotlinFileExtractor.kt

@@ -1600,11 +1600,13 @@ open class KotlinFileExtractor(
         return result
     }
 
-    private fun findTopLevelFunctionOrWarn(functionFilter: String, type: String, warnAgainstElement: IrElement): IrFunction? {
+    private fun findTopLevelFunctionOrWarn(functionFilter: String, type: String, parameterTypes: Array<String>, warnAgainstElement: IrElement): IrFunction? {
 
         val fn = pluginContext.referenceFunctions(FqName(functionFilter))
-            .firstOrNull { it.owner.parentClassOrNull?.fqNameWhenAvailable?.asString() == type }
-            ?.owner
+            .firstOrNull { fnSymbol ->
+                fnSymbol.owner.parentClassOrNull?.fqNameWhenAvailable?.asString() == type &&
+                fnSymbol.owner.valueParameters.map { it.type.classFqName?.asString() }.toTypedArray() contentEquals parameterTypes
+            }?.owner
 
         if (fn != null) {
             if (fn.parentClassOrNull != null) {
@@ -1756,6 +1758,12 @@ open class KotlinFileExtractor(
             else -> false
         }
 
+    private fun isGenericArrayType(typeName: String) =
+        when(typeName) {
+            "Array" -> true
+            else -> false
+        }
+
     private fun extractCall(c: IrCall, callable: Label<out DbCallable>, stmtExprParent: StmtExprParent) {
         with("call", c) {
             val target = tryReplaceSyntheticFunction(c.symbol.owner)
@@ -2177,8 +2185,20 @@ open class KotlinFileExtractor(
                         extractRawMethodAccess(getter, c, callable, parent, idx, enclosingStmt, listOf(), null, ext, typeArguments)
                     }
                 }
-                isFunction(target, "kotlin", "(some array type)", { isArrayType(it) }, "iterator") && c.origin == IrStatementOrigin.FOR_LOOP_ITERATOR -> {
-                    findTopLevelFunctionOrWarn("kotlin.jvm.internal.iterator", "kotlin.jvm.internal.ArrayIteratorKt", c)?.let { iteratorFn ->
+                isFunction(target, "kotlin", "(some array type)", { isArrayType(it) }, "iterator") -> {
+                    val parentClass = target.parent
+                    if (parentClass !is IrClass) {
+                        logger.errorElement("Iterator parent is not a class", c)
+                        return
+                    }
+
+                    var typeFilter = if (isGenericArrayType(parentClass.name.asString())) {
+                        "kotlin.jvm.internal.ArrayIteratorKt"
+                    } else {
+                        "kotlin.jvm.internal.ArrayIteratorsKt"
+                    }
+
+                    findTopLevelFunctionOrWarn("kotlin.jvm.internal.iterator", typeFilter, arrayOf(parentClass.kotlinFqName.asString()), c)?.let { iteratorFn ->
                         val dispatchReceiver = c.dispatchReceiver
                         if (dispatchReceiver == null) {
                             logger.errorElement("No dispatch receiver found for array iterator call", c)

java/ql/test/kotlin/library-tests/arrays/arrayIterators.expected

@@ -0,0 +1,6 @@
+| arrayIterators.kt:6:15:6:15 | iterator(...) | iterator(java.lang.Object[]) | kotlin.jvm.internal.ArrayIteratorKt |
+| arrayIterators.kt:7:15:7:15 | iterator(...) | iterator(int[]) | kotlin.jvm.internal.ArrayIteratorsKt |
+| arrayIterators.kt:8:15:8:15 | iterator(...) | iterator(boolean[]) | kotlin.jvm.internal.ArrayIteratorsKt |
+| arrayIterators.kt:10:16:10:25 | iterator(...) | iterator(java.lang.Object[]) | kotlin.jvm.internal.ArrayIteratorKt |
+| arrayIterators.kt:11:16:11:25 | iterator(...) | iterator(int[]) | kotlin.jvm.internal.ArrayIteratorsKt |
+| arrayIterators.kt:12:16:12:25 | iterator(...) | iterator(boolean[]) | kotlin.jvm.internal.ArrayIteratorsKt |

java/ql/test/kotlin/library-tests/arrays/arrayIterators.kt

@@ -0,0 +1,13 @@
+fun test(
+    a: Array<Int>,
+    b: IntArray,
+    c: BooleanArray) {
+
+    for (i in a) { }
+    for (i in b) { }
+    for (i in c) { }
+
+    val i1 = a.iterator()
+    val i2 = b.iterator()
+    val i3 = c.iterator()
+}

java/ql/test/kotlin/library-tests/arrays/arrayIterators.ql

@@ -0,0 +1,10 @@
+import java
+
+query predicate iterator(MethodAccess ma, string mn, string t) {
+  exists(Method m |
+    ma.getMethod() = m and
+    m.getName() = "iterator" and
+    mn = m.getSignature() and
+    t = ma.getMethod().getDeclaringType().getQualifiedName()
+  )
+}

java/ql/test/kotlin/library-tests/arrays/test.expected

@@ -18,6 +18,7 @@ sourceSignatures
 | arrayCreations.kt:27:24:27:38 | invoke | invoke(int) |
 | arrayGetsSets.kt:1:1:22:1 | arrayGetSet | arrayGetSet(int[],short[],byte[],long[],float[],double[],boolean[],char[],java.lang.Object[]) |
 | arrayGetsSets.kt:24:1:41:1 | arrayGetSetInPlace | arrayGetSetInPlace(int[],long[],float[],double[]) |
+| arrayIterators.kt:1:1:13:1 | test | test(java.lang.Integer[],int[],boolean[]) |
 | primitiveArrays.kt:3:1:7:1 | Test | Test() |
 | primitiveArrays.kt:5:3:5:123 | test | test(java.lang.Integer[],java.lang.Integer[],int[],java.lang.Integer[][],java.lang.Integer[][],int[][]) |
 #select