wip11

Author: hvitved

rust/ql/lib/codeql/rust/internal/typeinference/TypeInference.qll

@@ -457,19 +457,6 @@ private module Input3 implements InputSig3 {
     )
   }
 
-  Type inferCallReturnType(AstNode n, TypePath path) {
-    exists(Call call, TypePath path0 |
-      result = M3::inferCallReturnType(call, _, path0) and
-      n = call and
-      if
-        // index expression `x[i]` desugars to `*x.index(i)`, so we must account for
-        // the implicit deref
-        call instanceof IndexExpr
-      then path0.isCons(getRefTypeParameter(_), path)
-      else path = path0
-    )
-  }
-
   Type inferCallArgumentType(AstNode n, TypePath path) {
     exists(FunctionCallMatchingInput::Access call, FunctionPosition pos |
       result = inferCallArgumentType(call, pos.asPosition(), n, _, _, path) and
@@ -752,7 +739,7 @@ private class FunctionDeclaration extends Function {
   TypeParameter getTypeParameter(ImplOrTraitItemNodeOption i, TypeParameterPosition ppos) {
     result = ppos.asTypeParameter() and
     (
-      ppos.asTypeParam() = this.getTypeParam(i)
+      result = TTypeParamTypeParameter(this.getTypeParam(i))
       or
       // For every `TypeParam` of this function, any associated types accessed on
       // the type parameter are also type parameters.
@@ -2604,13 +2591,29 @@ private module FunctionCallMatchingInput implements MatchingWithEnvironmentInput
   import FunctionPositionMatchingInput
 
   private newtype TDeclaration =
-    TFunctionDeclaration(ImplOrTraitItemNodeOption i, FunctionDeclaration f) { f.isFor(i) }
+    TFunctionDeclaration(ImplOrTraitItemNodeOption i, FunctionDeclaration f, boolean forIndexExpr) {
+      f.isFor(i) and
+      (
+        forIndexExpr = false
+        or
+        forIndexExpr = true and
+        exists(Function traitFunction |
+          traitFunction =
+            [any(IndexTrait t).getIndexFunction(), any(IndexMutTrait t).getIndexMutFunction()]
+        |
+          f = traitFunction
+          or
+          f.implements(traitFunction)
+        )
+      )
+    }
 
   final class Declaration extends TFunctionDeclaration {
     ImplOrTraitItemNodeOption i;
     FunctionDeclaration f;
+    boolean forIndexExpr;
 
-    Declaration() { this = TFunctionDeclaration(i, f) }
+    Declaration() { this = TFunctionDeclaration(i, f, forIndexExpr) }
 
     FunctionDeclaration getFunction() { result = f }
 
@@ -2619,9 +2622,10 @@ private module FunctionCallMatchingInput implements MatchingWithEnvironmentInput
       f_ = f
     }
 
-    predicate isAssocFunction(ImplOrTraitItemNode i_, Function f_) {
+    predicate isAssocFunction(ImplOrTraitItemNode i_, Function f_, boolean forIndexExpr_) {
       i_ = i.asSome() and
-      f_ = f
+      f_ = f and
+      forIndexExpr_ = forIndexExpr
     }
 
     TypeParameter getTypeParameter(TypeParameterPosition ppos) {
@@ -2631,8 +2635,16 @@ private module FunctionCallMatchingInput implements MatchingWithEnvironmentInput
     Type getDeclaredType(FunctionPosition pos, TypePath path) {
       result = f.getParameterType(i, pos, path)
       or
+      // index expression `x[i]` desugars to `*x.index(i)`, so we must account for
+      // the implicit deref
       pos.isReturn() and
-      result = f.getReturnType(i, path)
+      exists(TypePath path0 | result = f.getReturnType(i, path0) |
+        forIndexExpr = true and
+        path0.isCons(getRefTypeParameter(_), path)
+        or
+        forIndexExpr = false and
+        path = path0
+      )
     }
 
     string toString() {
@@ -2738,23 +2750,8 @@ private module FunctionCallMatchingInput implements MatchingWithEnvironmentInput
 
     pragma[nomagic]
     private Type getInferredNonSelfType(FunctionPosition pos, TypePath path) {
-      if
-        // index expression `x[i]` desugars to `*x.index(i)`, so we must account for
-        // the implicit deref
-        pos.isReturn() and
-        this instanceof IndexExpr
-      then
-        path.isEmpty() and
-        result instanceof RefType
-        or
-        exists(TypePath suffix |
-          result = super.getTypeAt(pos, suffix) and
-          path = TypePath::cons(getRefTypeParameter(_), suffix)
-        )
-      else (
-        not super.hasReceiverAtPos(pos) and
-        result = super.getTypeAt(pos, path)
-      )
+      not super.hasReceiverAtPos(pos) and
+      result = super.getTypeAt(pos, path)
     }
 
     bindingset[derefChainBorrow]
@@ -2772,7 +2769,10 @@ private module FunctionCallMatchingInput implements MatchingWithEnvironmentInput
     }
 
     override Declaration getTarget(string derefChainBorrow) {
-      exists(ImplOrTraitItemNode i | result.isAssocFunction(i, this.getTarget(i, derefChainBorrow)))
+      exists(ImplOrTraitItemNode i, boolean forIndexExpr |
+        result.isAssocFunction(i, this.getTarget(i, derefChainBorrow), forIndexExpr) and
+        if this instanceof IndexExpr then forIndexExpr = true else forIndexExpr = false
+      )
     }
 
     pragma[nomagic]
@@ -2825,7 +2825,7 @@ private module FunctionCallMatchingInput implements MatchingWithEnvironmentInput
     private Declaration getTarget() {
       result =
         TFunctionDeclaration(ImplOrTraitItemNodeOption::none_(),
-          super.resolveCallTargetViaPathResolution())
+          super.resolveCallTargetViaPathResolution(), false)
     }
 
     override Declaration getTarget(string derefChainBorrow) {
@@ -3151,7 +3151,7 @@ private module OperationMatchingInput implements MatchingInputSig {
 
     Declaration getTarget() {
       exists(ImplOrTraitItemNode i |
-        result.isAssocFunction(i, this.resolveCallTarget(i, _, _, _)) // mutual recursion
+        result.isAssocFunction(i, this.resolveCallTarget(i, _, _, _), false) // mutual recursion
       )
     }
   }

shared/typeinference/codeql/typeinference/internal/TypeInference.qll

@@ -2194,7 +2194,12 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
 
       /** A callable. */
       class Callable {
-        /** Gets the type parameter at position `ppos` of this callable, if any. */
+        /**
+         * Gets the type parameter at position `ppos` of this callable, if any.
+         *
+         * This should include type parameters declared on the callable itself,
+         * as well as type parameters declared on the enclosing declaration(s).
+         */
         TypeParameter getTypeParameter(TypeParameterPosition ppos);
 
         /**
@@ -2213,7 +2218,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
         /**
          * Gets the declared type of the `i`th parameter of this callable at `path`.
          *
-         * This should also include (possibly implicit) `this`/`self` parameters, perhaps
+         * This should also include (possibly implicit) `this`/`self` parameters,
          * using index `-1`.
          */
         Type getParameterType(int i, TypePath path);
@@ -2236,8 +2241,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
         /**
          * Gets the AST node corresponding to the `i`th argument of this call.
          *
-         * This should include the receiver argument for method calls, perhaps
-         * using index `-1`.
+         * This should include the receiver argument for method calls, using index `-1`.
          */
         AstNode getArgument(int i);
 
@@ -2262,14 +2266,6 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
         exists(ctx)
       }
 
-      /**
-       * Gets the inferred return type of `call` at `path`.
-       *
-       * When no post-processing is needed, simply implement this predicate as
-       * `result = inferCallReturnType(n, _, path)`.
-       */
-      Type inferCallReturnType(AstNode n, TypePath path);
-
       /**
        * Gets the top-down inferred type of `call` at `path` and argument position
        * `pos`.
@@ -2520,7 +2516,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
         or
         result = inferTypeFromStep(n, path)
         or
-        result = inferCallReturnType(n, path)
+        result = inferCallReturnType(n, _, path)
         or
         // top-down inference
         exists(TypePath prefix |
@@ -2625,7 +2621,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
 
       private module CallMatching = MatchingWithEnvironment<CallMatchingInput>;
 
-      Type inferCallReturnType(Call call, CallResolutionContext ctx, TypePath path) {
+      private Type inferCallReturnType(Call call, CallResolutionContext ctx, TypePath path) {
         result = CallMatching::inferAccessType(call, ctx, -2, path)
       }