Merge pull request github#11588 from tamasvajk/kotlin-extension-defaults

Kotlin: Fix extraction of `$default` extension functions

Author: tamasvajk

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

@@ -957,7 +957,7 @@ open class KotlinFileExtractor(
         val locId = getLocation(f, null)
         val extReceiver = f.extensionReceiverParameter
         val dispatchReceiver = if (f.shouldExtractAsStatic) null else f.dispatchReceiverParameter
-        val parameterTypes = listOfNotNull(extReceiver?.let { erase(it.type) }) + getDefaultsMethodArgTypes(f)
+        val parameterTypes = getDefaultsMethodArgTypes(f)
         val allParamTypeResults = parameterTypes.mapIndexed { i, paramType ->
             val paramId = tw.getLabelFor<DbParam>(getValueParameterLabel(id, i))
             extractValueParameter(paramId, paramType, "p$i", locId, id, i, paramId, isVararg = false, syntheticParameterNames = true, isCrossinline = false, isNoinline = false).also {
@@ -975,6 +975,11 @@ open class KotlinFileExtractor(
             val methodId = id.cast<DbMethod>()
             extractMethod(methodId, locId, shortName, erase(f.returnType), paramsSignature, parentId, methodId, origin = null, extractTypeAccess = extractMethodAndParameterTypeAccesses)
             addModifiers(id, "static")
+            if (extReceiver != null) {
+                val idx = if (dispatchReceiver != null) 1 else 0
+                val extendedType = allParamTypeResults[idx]
+                tw.writeKtExtensionFunctions(methodId, extendedType.javaResult.id, extendedType.kotlinResult.id)
+            }
         }
         tw.writeHasLocation(id, locId)
         if (f.visibility != DescriptorVisibilities.PRIVATE && f.visibility != DescriptorVisibilities.PRIVATE_TO_THIS) {
@@ -1040,8 +1045,8 @@ open class KotlinFileExtractor(
                                 val realFnIdxOffset = if (f.extensionReceiverParameter != null) 1 else 0
                                 val paramMappings = f.valueParameters.mapIndexed { idx, param -> Triple(param.type, idx + paramIdxOffset, idx + realFnIdxOffset) } +
                                     listOfNotNull(
-                                        dispatchReceiver?.let { Triple(it.type, realFnIdxOffset, -1) },
-                                        extReceiver?.let { Triple(it.type, 0, 0) }
+                                        dispatchReceiver?.let { Triple(it.type, 0, -1) },
+                                        extReceiver?.let { Triple(it.type, if (dispatchReceiver != null) 1 else 0, 0) }
                                     )
                                 paramMappings.forEach { (type, fromIdx, toIdx) ->
                                     extractVariableAccess(tw.getLabelFor<DbParam>(getValueParameterLabel(id, fromIdx)), type, locId, thisCallId, toIdx, id, returnId)
@@ -1184,6 +1189,7 @@ open class KotlinFileExtractor(
                             id
 
                 val extReceiver = f.extensionReceiverParameter
+                // The following parameter order is correct, because member $default methods (where the order would be [dispatchParam], [extensionParam], normalParams) are not extracted here
                 val fParameters = listOfNotNull(extReceiver) + (overriddenAttributes?.valueParameters ?: f.valueParameters)
                 val paramTypes = fParameters.mapIndexed { i, vp ->
                     extractValueParameter(vp, id, i, typeSubstitution, sourceDeclaration, classTypeArgsIncludingOuterClasses, extractTypeAccess = extractMethodAndParameterTypeAccesses, overriddenAttributes?.sourceLoc)
@@ -1793,11 +1799,12 @@ open class KotlinFileExtractor(
         ) ?: pluginContext.irBuiltIns.anyType
 
     private fun getDefaultsMethodArgTypes(f: IrFunction) =
-        // The $default method has type ([extensionReceiver], [dispatchReceiver], paramTypes..., int, Object)
+        // The $default method has type ([dispatchReceiver], [extensionReceiver], paramTypes..., int, Object)
         // All parameter types are erased. The trailing int is a mask indicating which parameter values are real
         // and which should be replaced by defaults. The final Object parameter is apparently always null.
         (
             listOfNotNull(if (f.shouldExtractAsStatic) null else f.dispatchReceiverParameter?.type) +
+            listOfNotNull(f.extensionReceiverParameter?.type) +
             f.valueParameters.map { it.type } +
             listOf(pluginContext.irBuiltIns.intType, getDefaultsMethodLastArgType(f))
         ).map { erase(it) }
@@ -1816,17 +1823,16 @@ open class KotlinFileExtractor(
 
     private fun getDefaultsMethodLabel(f: IrFunction): Label<out DbCallable> {
         val defaultsMethodName = if (f is IrConstructor) "<init>" else getDefaultsMethodName(f)
-        val normalArgTypes = getDefaultsMethodArgTypes(f)
-        val extensionParamType = f.extensionReceiverParameter?.let { erase(it.type) }
+        val argTypes = getDefaultsMethodArgTypes(f)
 
         val defaultMethodLabelStr = getFunctionLabel(
             f.parent,
             maybeParentId = null,
             defaultsMethodName,
-            normalArgTypes,
+            argTypes,
             erase(f.returnType),
-            extensionParamType,
-            listOf(),
+            extensionParamType = null, // if there's any, that's included already in argTypes
+            functionTypeParameters = listOf(),
             classTypeArgsIncludingOuterClasses = null,
             overridesCollectionsMethod = false,
             javaSignature = null,
@@ -1886,13 +1892,14 @@ open class KotlinFileExtractor(
         extensionReceiver: IrExpression?
     ) {
         var nextIdx = 0
-        if (extensionReceiver != null) {
-            extractExpressionExpr(extensionReceiver, enclosingCallable, id, nextIdx++, enclosingStmt)
-        }
         if (dispatchReceiver != null && !callTarget.shouldExtractAsStatic) {
             extractExpressionExpr(dispatchReceiver, enclosingCallable, id, nextIdx++, enclosingStmt)
         }
 
+        if (extensionReceiver != null) {
+            extractExpressionExpr(extensionReceiver, enclosingCallable, id, nextIdx++, enclosingStmt)
+        }
+
         val valueArgsWithDummies = valueArguments.zip(callTarget.valueParameters).map {
                 (expr, param) -> expr ?: IrConstImpl.defaultValueForType(0, 0, param.type)
         }
@@ -4046,8 +4053,7 @@ open class KotlinFileExtractor(
                     // Use of 'this' in a function where the dispatch receiver is passed like an ordinary parameter,
                     // such as a `$default` static function that substitutes in default arguments as needed.
                     val paramDeclarerId = overriddenAttributes.id ?: useDeclarationParent(thisParamParent, false)
-                    val extensionParamOffset = if (thisParamParent.extensionReceiverParameter != null) 1 else 0
-                    val replacementParamId = tw.getLabelFor<DbParam>(getValueParameterLabel(paramDeclarerId, replaceWithParamIdx + extensionParamOffset))
+                    val replacementParamId = tw.getLabelFor<DbParam>(getValueParameterLabel(paramDeclarerId, replaceWithParamIdx))
                     extractVariableAccess(replacementParamId, e.type, locId, exprParent.parent, exprParent.idx, callable, exprParent.enclosingStmt)
                     return
                 }

java/ql/lib/change-notes/2022-12-09-default-extension-methods.md

@@ -0,0 +1,4 @@
+---
+category: minorAnalysis
+---
+* The extraction of Kotlin extension methods has been improved when default parameter values are present. The dispatch and extension receiver parameters are extracted in the correct order. The `ExtensionMethod::getExtensionReceiverParameterIndex` predicate has been introduced to facilitate getting the correct extension parameter index.

java/ql/lib/semmle/code/java/Expr.qll

@@ -1896,7 +1896,8 @@ class VarAccess extends Expr, @varaccess {
 class ExtensionReceiverAccess extends VarAccess {
   ExtensionReceiverAccess() {
     exists(Parameter p |
-      this.getVariable() = p and p.getPosition() = 0 and p.getCallable() instanceof ExtensionMethod
+      this.getVariable() = p and
+      p.isExtensionParameter()
     )
   }
 

java/ql/lib/semmle/code/java/Member.qll

@@ -327,18 +327,8 @@ class Callable extends StmtParent, Member, @callable {
         this instanceof Method and
         result instanceof Method and
         this.getName() + "$default" = result.getName() and
-        extraLeadingParams <= 1 and
-        (
-          if ktExtensionFunctions(this, _, _)
-          then
-            // Both extension receivers are expected to occur at arg0, with any
-            // dispatch receiver inserted afterwards in the $default proxy's parameter list.
-            // Check the extension receiver matches here, and note regular args
-            // are bumped one position to the right.
-            regularParamsStartIdx = extraLeadingParams + 1 and
-            this.getParameterType(0).getErasure() = eraseRaw(result.getParameterType(0))
-          else regularParamsStartIdx = extraLeadingParams
-        ) and
+        extraLeadingParams <= 1 and // 0 for static methods, 1 for instance methods
+        regularParamsStartIdx = extraLeadingParams and
         lastParamType instanceof TypeObject
       )
     |
@@ -824,4 +814,19 @@ class ExtensionMethod extends Method {
   KotlinType getExtendedKotlinType() { result = extendedKotlinType }
 
   override string getAPrimaryQlClass() { result = "ExtensionMethod" }
+
+  /**
+   * Gets the index of the parameter that is the extension receiver. This is typically index 0. In case of `$default`
+   * extension methods that are defined as members, the index is 1. Index 0 is the dispatch receiver of the `$default`
+   * method.
+   */
+  int getExtensionReceiverParameterIndex() {
+    if
+      exists(Method src |
+        this = src.getKotlinParameterDefaultsProxy() and
+        src.getNumberOfParameters() = this.getNumberOfParameters() - 3 // 2 extra parameters + 1 dispatch receiver
+      )
+    then result = 1
+    else result = 0
+  }
 }

java/ql/lib/semmle/code/java/Variable.qll

@@ -91,7 +91,7 @@ class Parameter extends Element, @param, LocalScopeVariable {
 
   /** Holds if this formal parameter is a parameter representing the dispatch receiver in an extension method. */
   predicate isExtensionParameter() {
-    this.getPosition() = 0 and this.getCallable() instanceof ExtensionMethod
+    this.getPosition() = this.getCallable().(ExtensionMethod).getExtensionReceiverParameterIndex()
   }
 
   /**

java/ql/lib/semmle/code/java/dataflow/internal/FlowSummaryImplSpecific.qll

@@ -120,8 +120,11 @@ private predicate correspondingKotlinParameterDefaultsArgSpec(
           exists(int oldArgParsed |
             oldArgParsed = AccessPathSyntax::AccessPath::parseInt(oldArgNumber.splitAt(",").trim())
           |
-            if ktExtensionFunctions(originalCallable, _, _) and oldArgParsed = 0
-            then defaultsArgSpec = "Argument[0]"
+            if
+              ktExtensionFunctions(originalCallable, _, _) and
+              ktExtensionFunctions(defaultsCallable, _, _) and
+              oldArgParsed = 0
+            then defaultsArgSpec = "Argument[" + paramOffset + "]" // 1 if dispatch receiver is present, 0 otherwise.
             else defaultsArgSpec = "Argument[" + (oldArgParsed + paramOffset) + "]" + rest
           )
         )

java/ql/lib/semmle/code/java/security/PathSanitizer.qll

@@ -298,8 +298,9 @@ private class PathNormalizeSanitizer extends MethodAccess {
  * what `getQualifier` actually gets in Java and Kotlin.
  */
 private Expr getVisualQualifier(MethodAccess ma) {
-  if getSourceMethod(ma.getMethod()) instanceof ExtensionMethod
-  then result = ma.getArgument(0)
+  if ma.getMethod() instanceof ExtensionMethod
+  then
+    result = ma.getArgument(ma.getMethod().(ExtensionMethod).getExtensionReceiverParameterIndex())
   else result = ma.getQualifier()
 }
 
@@ -310,8 +311,11 @@ private Expr getVisualQualifier(MethodAccess ma) {
  */
 bindingset[argPos]
 private Argument getVisualArgument(MethodAccess ma, int argPos) {
-  if getSourceMethod(ma.getMethod()) instanceof ExtensionMethod
-  then result = ma.getArgument(argPos + 1)
+  if ma.getMethod() instanceof ExtensionMethod
+  then
+    result =
+      ma.getArgument(argPos + ma.getMethod().(ExtensionMethod).getExtensionReceiverParameterIndex() +
+          1)
   else result = ma.getArgument(argPos)
 }
 

java/ql/test/kotlin/library-tests/jvmoverloads-annotation/PrintAst.expected

@@ -72,7 +72,7 @@ test.kt:
 #   45|       5: [BlockStmt] { ... }
 #   45|         0: [ReturnStmt] return ...
 #   45|           0: [VarAccess] a
-#   45|     5: [Method] testExtensionFunction$default
+#   45|     5: [ExtensionMethod] testExtensionFunction$default
 #   45|       3: [TypeAccess] int
 #-----|       4: (Parameters)
 #   45|         0: [Parameter] p0
@@ -116,7 +116,7 @@ test.kt:
 #   45|         2: [ReturnStmt] return ...
 #   45|           0: [MethodAccess] testExtensionFunction(...)
 #   45|             -1: [TypeAccess] TestKt
-#   45|             0: [VarAccess] p0
+#   45|             0: [ExtensionReceiverAccess] this
 #   45|             1: [VarAccess] p1
 #   45|             2: [VarAccess] p2
 #   45|             3: [VarAccess] p3
@@ -356,9 +356,9 @@ test.kt:
 #   12|         0: [ReturnStmt] return ...
 #   12|           0: [MethodAccess] testMemberExtensionFunction$default(...)
 #   12|             -1: [TypeAccess] Test
-#    0|             0: [ExtensionReceiverAccess] this
-#    0|             1: [ThisAccess] Test.this
+#    0|             0: [ThisAccess] Test.this
 #    0|               0: [TypeAccess] Test
+#    0|             1: [ExtensionReceiverAccess] this
 #    0|             2: [VarAccess] a
 #    1|             3: [NullLiteral] null
 #    0|             4: [VarAccess] c
@@ -383,9 +383,9 @@ test.kt:
 #   12|         0: [ReturnStmt] return ...
 #   12|           0: [MethodAccess] testMemberExtensionFunction$default(...)
 #   12|             -1: [TypeAccess] Test
-#    0|             0: [ExtensionReceiverAccess] this
-#    0|             1: [ThisAccess] Test.this
+#    0|             0: [ThisAccess] Test.this
 #    0|               0: [TypeAccess] Test
+#    0|             1: [ExtensionReceiverAccess] this
 #    0|             2: [VarAccess] a
 #    0|             3: [VarAccess] b
 #    0|             4: [VarAccess] c
@@ -411,13 +411,13 @@ test.kt:
 #   12|       5: [BlockStmt] { ... }
 #   12|         0: [ReturnStmt] return ...
 #   12|           0: [VarAccess] a
-#   12|     13: [Method] testMemberExtensionFunction$default
+#   12|     13: [ExtensionMethod] testMemberExtensionFunction$default
 #   12|       3: [TypeAccess] int
 #-----|       4: (Parameters)
 #   12|         0: [Parameter] p0
-#   12|           0: [TypeAccess] Test2
-#   12|         1: [Parameter] p1
 #   12|           0: [TypeAccess] Test
+#   12|         1: [Parameter] p1
+#   12|           0: [TypeAccess] Test2
 #   12|         2: [Parameter] p2
 #   12|           0: [TypeAccess] int
 #   12|         3: [Parameter] p3
@@ -456,8 +456,8 @@ test.kt:
 #   12|               1: [FloatLiteral] 1.0
 #   12|         2: [ReturnStmt] return ...
 #   12|           0: [MethodAccess] testMemberExtensionFunction(...)
-#   12|             -1: [VarAccess] p1
-#   12|             0: [VarAccess] p0
+#   12|             -1: [VarAccess] p0
+#   12|             0: [ExtensionReceiverAccess] this
 #   12|             1: [VarAccess] p2
 #   12|             2: [VarAccess] p3
 #   12|             3: [VarAccess] p4

java/ql/test/kotlin/library-tests/jvmoverloads-annotation/test.expected

@@ -29,7 +29,7 @@
 | test.kt:3:1:14:1 | Test | test.kt:12:3:12:121 | testMemberExtensionFunction | testMemberExtensionFunction(Test2,int,double,boolean) |
 | test.kt:3:1:14:1 | Test | test.kt:12:3:12:121 | testMemberExtensionFunction | testMemberExtensionFunction(Test2,int,java.lang.String,double,boolean) |
 | test.kt:3:1:14:1 | Test | test.kt:12:3:12:121 | testMemberExtensionFunction | testMemberExtensionFunction(Test2,int,java.lang.String,double,float,boolean) |
-| test.kt:3:1:14:1 | Test | test.kt:12:3:12:121 | testMemberExtensionFunction$default | testMemberExtensionFunction$default(Test2,Test,int,java.lang.String,double,float,boolean,int,java.lang.Object) |
+| test.kt:3:1:14:1 | Test | test.kt:12:3:12:121 | testMemberExtensionFunction$default | testMemberExtensionFunction$default(Test,Test2,int,java.lang.String,double,float,boolean,int,java.lang.Object) |
 | test.kt:16:1:28:1 | Test2 | test.kt:16:34:28:1 | Test2 | Test2(int,double,boolean) |
 | test.kt:16:1:28:1 | Test2 | test.kt:16:34:28:1 | Test2 | Test2(int,java.lang.String,double,boolean) |
 | test.kt:16:1:28:1 | Test2 | test.kt:16:34:28:1 | Test2 | Test2(int,java.lang.String,double,float,boolean) |

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

@@ -293,12 +293,47 @@
 | methods2.kt:8:5:9:5 | Unit | TypeAccess |
 | methods2.kt:8:27:8:32 | int | TypeAccess |
 | methods2.kt:8:35:8:40 | int | TypeAccess |
-| methods3.kt:3:1:3:42 | Unit | TypeAccess |
-| methods3.kt:3:5:3:7 | int | TypeAccess |
-| methods3.kt:3:33:3:38 | int | TypeAccess |
-| methods3.kt:6:5:6:46 | Unit | TypeAccess |
-| methods3.kt:6:9:6:11 | int | TypeAccess |
-| methods3.kt:6:37:6:42 | int | TypeAccess |
+| methods3.kt:3:1:3:49 | 0 | IntegerLiteral |
+| methods3.kt:3:1:3:49 | 1 | IntegerLiteral |
+| methods3.kt:3:1:3:49 | ... & ... | AndBitwiseExpr |
+| methods3.kt:3:1:3:49 | ... == ... | EQExpr |
+| methods3.kt:3:1:3:49 | ...=... | AssignExpr |
+| methods3.kt:3:1:3:49 | Methods3Kt | TypeAccess |
+| methods3.kt:3:1:3:49 | Object | TypeAccess |
+| methods3.kt:3:1:3:49 | String | TypeAccess |
+| methods3.kt:3:1:3:49 | Unit | TypeAccess |
+| methods3.kt:3:1:3:49 | Unit | TypeAccess |
+| methods3.kt:3:1:3:49 | fooBarTopLevelMethodExt(...) | MethodAccess |
+| methods3.kt:3:1:3:49 | int | TypeAccess |
+| methods3.kt:3:1:3:49 | int | TypeAccess |
+| methods3.kt:3:1:3:49 | p1 | VarAccess |
+| methods3.kt:3:1:3:49 | p1 | VarAccess |
+| methods3.kt:3:1:3:49 | p2 | VarAccess |
+| methods3.kt:3:1:3:49 | this | ExtensionReceiverAccess |
+| methods3.kt:3:5:3:10 | String | TypeAccess |
+| methods3.kt:3:36:3:45 | int | TypeAccess |
+| methods3.kt:3:45:3:45 | 1 | IntegerLiteral |
+| methods3.kt:6:5:6:45 | 0 | IntegerLiteral |
+| methods3.kt:6:5:6:45 | 1 | IntegerLiteral |
+| methods3.kt:6:5:6:45 | ... & ... | AndBitwiseExpr |
+| methods3.kt:6:5:6:45 | ... == ... | EQExpr |
+| methods3.kt:6:5:6:45 | ...=... | AssignExpr |
+| methods3.kt:6:5:6:45 | Class3 | TypeAccess |
+| methods3.kt:6:5:6:45 | Object | TypeAccess |
+| methods3.kt:6:5:6:45 | String | TypeAccess |
+| methods3.kt:6:5:6:45 | Unit | TypeAccess |
+| methods3.kt:6:5:6:45 | Unit | TypeAccess |
+| methods3.kt:6:5:6:45 | fooBarMethodExt(...) | MethodAccess |
+| methods3.kt:6:5:6:45 | int | TypeAccess |
+| methods3.kt:6:5:6:45 | int | TypeAccess |
+| methods3.kt:6:5:6:45 | p0 | VarAccess |
+| methods3.kt:6:5:6:45 | p2 | VarAccess |
+| methods3.kt:6:5:6:45 | p2 | VarAccess |
+| methods3.kt:6:5:6:45 | p3 | VarAccess |
+| methods3.kt:6:5:6:45 | this | ExtensionReceiverAccess |
+| methods3.kt:6:9:6:14 | String | TypeAccess |
+| methods3.kt:6:32:6:41 | int | TypeAccess |
+| methods3.kt:6:41:6:41 | 1 | IntegerLiteral |
 | methods4.kt:7:5:7:34 | Unit | TypeAccess |
 | methods4.kt:7:11:7:29 | InsideNestedTest | TypeAccess |
 | methods5.kt:3:1:11:1 | Unit | TypeAccess |