Fix DB inconsistencies with KFunction and KFunction::invoke call extraction

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

@@ -1131,11 +1131,13 @@ open class KotlinFileExtractor(
         extractTypeArguments(typeArguments, locId, id, enclosingCallable, enclosingStmt, -2, true)
 
         val drType = dispatchReceiver?.type
-        val isBigArityFunctionInvoke = drType != null
+
+        val isFunctionInvoke = drType != null
                 && drType is IrSimpleType
                 && drType.isFunctionOrKFunction()
                 && callTarget.name.asString() == OperatorNameConventions.INVOKE.asString()
-                && drType.arguments.size > BuiltInFunctionArity.BIG_ARITY
+        val isBigArityFunctionInvoke = isFunctionInvoke
+                && (drType as IrSimpleType).arguments.size > BuiltInFunctionArity.BIG_ARITY
 
         if (callTarget.isLocalFunction()) {
             val ids = getLocallyVisibleFunctionLabels(callTarget)
@@ -1147,12 +1149,19 @@ open class KotlinFileExtractor(
         } else {
             val methodId =
                 if (drType != null && extractClassTypeArguments && drType is IrSimpleType && !isUnspecialised(drType)) {
-                    if (isBigArityFunctionInvoke) {
-                        val interfaceType = getFunctionalInterfaceTypeWithTypeArgs(drType.arguments)
-                        val invokeMethod = findFunction(interfaceType.classOrNull!!.owner, OperatorNameConventions.INVOKE.asString())!!
-                        useFunction<DbCallable>(invokeMethod, listOf(drType.arguments.last()))
+
+                    val extractionMethod = if (isFunctionInvoke) {
+                        val interfaceType = getFunctionalInterfaceTypeWithTypeArgs(drType.arguments).classOrNull!!.owner
+                        val substituted = getJavaEquivalentClass(interfaceType) ?: interfaceType
+                        findFunction(substituted, OperatorNameConventions.INVOKE.asString())!!
                     } else {
-                        useFunction<DbCallable>(callTarget, getDeclaringTypeArguments(callTarget, drType))
+                        callTarget
+                    }
+
+                    if (isBigArityFunctionInvoke) {
+                        useFunction<DbCallable>(extractionMethod, listOf(drType.arguments.last()))
+                    } else {
+                        useFunction<DbCallable>(extractionMethod, getDeclaringTypeArguments(callTarget, drType))
                     }
                 }
                 else

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

@@ -195,7 +195,14 @@ open class KotlinUsesExtractor(
 
         val extractClass = substituteClass ?: c
 
-        val classTypeResult = addClassLabel(extractClass, typeArgs, inReceiverContext)
+        // `KFunction1<T1,T2>` is substituted by `KFunction<T>`. The last type argument is the return type.
+        val extractedTypeArgs = if (c.symbol.isKFunction() && typeArgs != null && typeArgs.isNotEmpty()) {
+            listOf(typeArgs.last())
+        } else {
+            typeArgs
+        }
+
+        val classTypeResult = addClassLabel(extractClass, extractedTypeArgs, inReceiverContext)
 
         // Extract both the Kotlin and equivalent Java classes, so that we have database entries
         // for both even if all internal references to the Kotlin type are substituted.

java/ql/test/kotlin/library-tests/exprs/PrintAst.expected

@@ -4219,6 +4219,56 @@ funcExprs.kt:
 #   79|         5: [BlockStmt] { ... }
 #   79|           0: [SuperConstructorInvocationStmt] super(...)
 #   79|           1: [BlockStmt] { ... }
+kFunctionInvoke.kt:
+#    0| [CompilationUnit] kFunctionInvoke
+#    0|   1: [Class] KFunctionInvokeKt
+#    7|     1: [Method] useRef
+#-----|       4: (Parameters)
+#    7|         0: [Parameter] a
+#    7|         1: [Parameter] s
+#    7|       5: [BlockStmt] { ... }
+#    8|         0: [LocalVariableDeclStmt] var ...;
+#    8|           1: [LocalVariableDeclExpr] toCall
+#    8|             0: [MemberRefExpr] ...::...
+#    8|               -4: [AnonymousClass] new Function1<String,Unit>(...) { ... }
+#    8|                 1: [Constructor] 
+#-----|                   4: (Parameters)
+#    8|                     0: [Parameter] <dispatchReceiver>
+#    8|                   5: [BlockStmt] { ... }
+#    8|                     0: [SuperConstructorInvocationStmt] super(...)
+#    8|                     1: [ExprStmt] <Expr>;
+#    8|                       0: [AssignExpr] ...=...
+#    8|                         0: [VarAccess] this.<dispatchReceiver>
+#    8|                           -1: [ThisAccess] this
+#    8|                         1: [VarAccess] <dispatchReceiver>
+#    8|                 1: [Method] invoke
+#-----|                   4: (Parameters)
+#    8|                     0: [Parameter] a0
+#    8|                   5: [BlockStmt] { ... }
+#    8|                     0: [ReturnStmt] return ...
+#    8|                       0: [MethodAccess] f(...)
+#    8|                         -1: [VarAccess] this.<dispatchReceiver>
+#    8|                           -1: [ThisAccess] this
+#    8|                         0: [VarAccess] a0
+#    8|                 1: [FieldDeclaration] A <dispatchReceiver>;
+#    8|                   -1: [TypeAccess] A
+#    8|               -3: [TypeAccess] Function1<String,Unit>
+#    8|                 0: [TypeAccess] String
+#    8|                 1: [TypeAccess] Unit
+#    8|               0: [VarAccess] a
+#    9|         1: [ExprStmt] <Expr>;
+#    9|           0: [MethodAccess] invoke(...)
+#    9|             -1: [VarAccess] toCall
+#    9|             0: [VarAccess] s
+#    3|   2: [Class] A
+#    3|     1: [Constructor] A
+#    3|       5: [BlockStmt] { ... }
+#    3|         0: [SuperConstructorInvocationStmt] super(...)
+#    3|         1: [BlockStmt] { ... }
+#    4|     2: [Method] f
+#-----|       4: (Parameters)
+#    4|         0: [Parameter] s
+#    4|       5: [BlockStmt] { ... }
 localFunctionCalls.kt:
 #    0| [CompilationUnit] localFunctionCalls
 #    0|   1: [Class] LocalFunctionCallsKt

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

@@ -2311,6 +2311,24 @@
 | funcExprs.kt:75:12:75:22 | Integer | funcExprs.kt:74:5:76:5 | call | TypeAccess |
 | funcExprs.kt:75:12:75:22 | String | funcExprs.kt:74:5:76:5 | call | TypeAccess |
 | funcExprs.kt:75:20:75:20 | a | funcExprs.kt:75:12:75:22 | invoke | StringLiteral |
+| kFunctionInvoke.kt:8:5:8:47 | toCall | kFunctionInvoke.kt:7:1:10:1 | useRef | LocalVariableDeclExpr |
+| kFunctionInvoke.kt:8:44:8:44 | a | kFunctionInvoke.kt:7:1:10:1 | useRef | VarAccess |
+| kFunctionInvoke.kt:8:44:8:47 | ...::... | kFunctionInvoke.kt:7:1:10:1 | useRef | MemberRefExpr |
+| kFunctionInvoke.kt:8:44:8:47 | ...=... | kFunctionInvoke.kt:8:44:8:47 |  | AssignExpr |
+| kFunctionInvoke.kt:8:44:8:47 | <dispatchReceiver> | kFunctionInvoke.kt:8:44:8:47 |  | VarAccess |
+| kFunctionInvoke.kt:8:44:8:47 | A | file://:0:0:0:0 | <none> | TypeAccess |
+| kFunctionInvoke.kt:8:44:8:47 | Function1<String,Unit> | kFunctionInvoke.kt:7:1:10:1 | useRef | TypeAccess |
+| kFunctionInvoke.kt:8:44:8:47 | String | kFunctionInvoke.kt:7:1:10:1 | useRef | TypeAccess |
+| kFunctionInvoke.kt:8:44:8:47 | Unit | kFunctionInvoke.kt:7:1:10:1 | useRef | TypeAccess |
+| kFunctionInvoke.kt:8:44:8:47 | a0 | kFunctionInvoke.kt:8:44:8:47 | invoke | VarAccess |
+| kFunctionInvoke.kt:8:44:8:47 | f(...) | kFunctionInvoke.kt:8:44:8:47 | invoke | MethodAccess |
+| kFunctionInvoke.kt:8:44:8:47 | this | kFunctionInvoke.kt:8:44:8:47 |  | ThisAccess |
+| kFunctionInvoke.kt:8:44:8:47 | this | kFunctionInvoke.kt:8:44:8:47 | invoke | ThisAccess |
+| kFunctionInvoke.kt:8:44:8:47 | this.<dispatchReceiver> | kFunctionInvoke.kt:8:44:8:47 |  | VarAccess |
+| kFunctionInvoke.kt:8:44:8:47 | this.<dispatchReceiver> | kFunctionInvoke.kt:8:44:8:47 | invoke | VarAccess |
+| kFunctionInvoke.kt:9:5:9:10 | toCall | kFunctionInvoke.kt:7:1:10:1 | useRef | VarAccess |
+| kFunctionInvoke.kt:9:5:9:13 | invoke(...) | kFunctionInvoke.kt:7:1:10:1 | useRef | MethodAccess |
+| kFunctionInvoke.kt:9:12:9:12 | s | kFunctionInvoke.kt:7:1:10:1 | useRef | VarAccess |
 | localFunctionCalls.kt:4:5:4:13 | x | localFunctionCalls.kt:3:1:12:1 | x | LocalVariableDeclExpr |
 | localFunctionCalls.kt:4:13:4:13 | 5 | localFunctionCalls.kt:3:1:12:1 | x | IntegerLiteral |
 | localFunctionCalls.kt:5:25:5:25 | i | localFunctionCalls.kt:5:5:5:29 | a | VarAccess |

java/ql/test/kotlin/library-tests/exprs/funcExprs.expected

@@ -34,5 +34,6 @@ memberRefExprs
 | funcExprs.kt:46:30:46:41 | ...::... | funcExprs.kt:46:30:46:41 | invoke | invoke(java.lang.Object[]) | funcExprs.kt:46:30:46:41 | new FunctionN<String>(...) { ... } |
 | funcExprs.kt:49:26:49:32 | ...::... | funcExprs.kt:49:26:49:32 | invoke | invoke() | funcExprs.kt:49:26:49:32 | new Function0<Integer>(...) { ... } |
 | funcExprs.kt:51:8:51:16 | ...::... | funcExprs.kt:51:8:51:16 | invoke | invoke() | funcExprs.kt:51:8:51:16 | new Function0<FuncRef>(...) { ... } |
+| kFunctionInvoke.kt:8:44:8:47 | ...::... | kFunctionInvoke.kt:8:44:8:47 | invoke | invoke(java.lang.String) | kFunctionInvoke.kt:8:44:8:47 | new Function1<String,Unit>(...) { ... } |
 | samConversion.kt:5:27:5:31 | ...::... | samConversion.kt:5:27:5:31 | invoke | invoke(int,int) | samConversion.kt:5:27:5:31 | new Function2<Integer,Integer,Unit>(...) { ... } |
 | samConversion.kt:41:13:41:16 | ...::... | samConversion.kt:41:13:41:16 | invoke | invoke(java.lang.Object[]) | samConversion.kt:41:13:41:16 | new FunctionN<Boolean>(...) { ... } |

java/ql/test/kotlin/library-tests/exprs/kFunctionInvoke.kt

@@ -0,0 +1,10 @@
+import kotlin.reflect.*
+
+class A {
+    fun f(s: String) { }
+}
+
+fun useRef(a: A, s: String) {
+    val toCall: KFunction1<String, Unit> = a::f
+    toCall(s)
+}