Fix Parameter.getACallArgument for parameters of extension methods

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

@@ -558,13 +558,13 @@ private module ControlFlowGraphImpl {
       or
       exists(ExtensionMethodAccess e | e = this |
         // the actual qualifier of the expression method access
-        index = -1 and result.(Expr).isNthChildOf(this, index)
+        index = -1 and result.(Expr).isNthChildOf(this, index) and not result instanceof TypeAccess
         or
         // the extension receiver
         index = 0 and result = e.getQualifier()
         or
         // the arguments
-        result = e.getArgument(index)
+        result = e.getArgument(index - 1)
       )
       or
       exists(StringTemplateExpr e | e = this | result = e.getComponent(index))

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

@@ -1994,9 +1994,9 @@ class ExtensionMethodAccess extends MethodAccess {
   // whereas the actual arguments begin at index 1.
   override Expr getQualifier() { result.isNthChildOf(this, 0) }
 
-  override Expr getAnArgument() { result.getIndex() >= 1 and result.getParent() = this }
+  override Expr getAnArgument() { result = getArgument(_) }
 
-  override Expr getArgument(int index) { result = super.getArgument(index + 1) }
+  override Expr getArgument(int index) { result = super.getArgument(index + 1) and index >= 0 }
 }
 
 /** A type access is a (possibly qualified) reference to a type. */

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

@@ -89,6 +89,11 @@ class Parameter extends Element, @param, LocalScopeVariable {
   /** Holds if this formal parameter is a variable arity parameter. */
   predicate isVarargs() { isVarargsParam(this) }
 
+  /** Holds if this formal parameter is a parameter representing the dispatch receiver in an extension method. */
+  predicate isExtensionParameter() {
+    getPosition() = 0 and this.getCallable() instanceof ExtensionMethod
+  }
+
   /**
    * Gets an argument for this parameter in any call to the callable that declares this formal
    * parameter.
@@ -103,7 +108,20 @@ class Parameter extends Element, @param, LocalScopeVariable {
   pragma[noinline]
   private Expr getACallArgument(int i) {
     exists(Call call |
-      result = call.getArgument(i) and
+      (
+        exists(int idx |
+          (
+            idx = i and not call instanceof ExtensionMethodAccess
+            or
+            idx = i - 1 and call instanceof ExtensionMethodAccess
+          ) and
+          result = call.getArgument(idx)
+        )
+        or
+        i = 0 and
+        call instanceof ExtensionMethodAccess and
+        result = call.getQualifier()
+      ) and
       call.getCallee().getSourceDeclaration().getAParameter() = this
     )
   }

java/ql/test/kotlin/library-tests/extensions/parameters.expected

@@ -0,0 +1,47 @@
+parametersWithArgs
+| extensions.kt:3:25:3:34 | p1 | 0 | extensions.kt:21:34:21:36 | foo |
+| extensions.kt:6:28:6:37 | p1 | 0 | extensions.kt:25:40:25:42 | foo |
+| extensions.kt:9:5:9:13 | <this> | 0 | extensions.kt:22:5:22:15 | new SomeClass(...) |
+| extensions.kt:9:23:9:32 | p1 | 1 | extensions.kt:22:26:22:28 | foo |
+| extensions.kt:10:5:10:16 | <this> | 0 | extensions.kt:26:5:26:18 | new AnotherClass(...) |
+| extensions.kt:10:29:10:38 | p1 | 1 | extensions.kt:26:32:26:34 | foo |
+| extensions.kt:12:5:12:13 | <this> | 0 | extensions.kt:23:5:23:15 | new SomeClass(...) |
+| extensions.kt:12:23:12:32 | p1 | 1 | extensions.kt:23:26:23:28 | foo |
+| extensions.kt:13:5:13:16 | <this> | 0 | extensions.kt:27:5:27:18 | new AnotherClass(...) |
+| extensions.kt:13:26:13:35 | p1 | 1 | extensions.kt:27:29:27:31 | foo |
+| extensions.kt:15:5:15:13 | <this> | 0 | extensions.kt:24:5:24:15 | new SomeClass(...) |
+| extensions.kt:15:32:15:38 | p1 | 1 | extensions.kt:24:34:24:34 | 1 |
+| extensions.kt:16:5:16:16 | <this> | 0 | extensions.kt:28:5:28:18 | new AnotherClass(...) |
+| extensions.kt:16:35:16:44 | p1 | 1 | extensions.kt:28:38:28:40 | foo |
+| extensions.kt:18:5:18:10 | <this> | 0 | extensions.kt:29:6:29:15 | someString |
+| extensions.kt:18:16:18:25 | p1 | 1 | extensions.kt:29:23:29:25 | foo |
+| extensions.kt:30:9:30:14 | <this> | 0 | extensions.kt:31:6:31:15 | someString |
+| extensions.kt:30:20:30:29 | p1 | 1 | extensions.kt:31:23:31:30 | bazParam |
+extensionParameter
+| extensions.kt:9:5:9:13 | <this> |
+| extensions.kt:10:5:10:16 | <this> |
+| extensions.kt:12:5:12:13 | <this> |
+| extensions.kt:13:5:13:16 | <this> |
+| extensions.kt:15:5:15:13 | <this> |
+| extensions.kt:16:5:16:16 | <this> |
+| extensions.kt:18:5:18:10 | <this> |
+| extensions.kt:30:9:30:14 | <this> |
+#select
+| extensions.kt:3:5:3:38 | someClassMethod | extensions.kt:3:25:3:34 | p1 | 0 |
+| extensions.kt:6:5:6:41 | anotherClassMethod | extensions.kt:6:28:6:37 | p1 | 0 |
+| extensions.kt:9:1:9:36 | someFun | extensions.kt:9:5:9:13 | <this> | 0 |
+| extensions.kt:9:1:9:36 | someFun | extensions.kt:9:23:9:32 | p1 | 1 |
+| extensions.kt:10:1:10:42 | anotherFun | extensions.kt:10:5:10:16 | <this> | 0 |
+| extensions.kt:10:1:10:42 | anotherFun | extensions.kt:10:29:10:38 | p1 | 1 |
+| extensions.kt:12:1:12:36 | bothFun | extensions.kt:12:5:12:13 | <this> | 0 |
+| extensions.kt:12:1:12:36 | bothFun | extensions.kt:12:23:12:32 | p1 | 1 |
+| extensions.kt:13:1:13:39 | bothFun | extensions.kt:13:5:13:16 | <this> | 0 |
+| extensions.kt:13:1:13:39 | bothFun | extensions.kt:13:26:13:35 | p1 | 1 |
+| extensions.kt:15:1:15:57 | bothFunDiffTypes | extensions.kt:15:5:15:13 | <this> | 0 |
+| extensions.kt:15:1:15:57 | bothFunDiffTypes | extensions.kt:15:32:15:38 | p1 | 1 |
+| extensions.kt:16:1:16:70 | bothFunDiffTypes | extensions.kt:16:5:16:16 | <this> | 0 |
+| extensions.kt:16:1:16:70 | bothFunDiffTypes | extensions.kt:16:35:16:44 | p1 | 1 |
+| extensions.kt:18:1:18:51 | bar | extensions.kt:18:5:18:10 | <this> | 0 |
+| extensions.kt:18:1:18:51 | bar | extensions.kt:18:16:18:25 | p1 | 1 |
+| extensions.kt:30:5:30:55 | baz | extensions.kt:30:9:30:14 | <this> | 0 |
+| extensions.kt:30:5:30:55 | baz | extensions.kt:30:20:30:29 | p1 | 1 |

java/ql/test/kotlin/library-tests/extensions/parameters.ql

@@ -0,0 +1,11 @@
+import java
+
+from Method m, int n
+where m.fromSource()
+select m, m.getParameter(n), n
+
+query predicate parametersWithArgs(Parameter p, int idx, Expr arg) {
+  p.fromSource() and p.getPosition() = idx and p.getAnArgument() = arg
+}
+
+query predicate extensionParameter(Parameter p) { p.isExtensionParameter() }