Rust: Disambiguate calls to associated functions

Author: hvitved

rust/ql/lib/codeql/rust/elements/internal/CallExprBaseImpl.qll

@@ -13,6 +13,7 @@ private import codeql.rust.elements.Resolvable
  */
 module Impl {
   private import rust
+  private import codeql.rust.internal.TypeInference as TypeInference
 
   pragma[nomagic]
   Resolvable getCallResolvable(CallExprBase call) {
@@ -27,7 +28,11 @@ module Impl {
    */
   class CallExprBase extends Generated::CallExprBase {
     /** Gets the static target of this call, if any. */
-    Callable getStaticTarget() { none() } // overridden by subclasses, but cannot be made abstract
+    final Callable getStaticTarget() {
+      result = TypeInference::resolveMethodCallTarget(this)
+      or
+      result = TypeInference::resolveCallTarget(this)
+    }
 
     override Expr getArg(int index) { result = this.getArgList().getArg(index) }
   }

rust/ql/lib/codeql/rust/elements/internal/CallExprImpl.qll

@@ -14,7 +14,6 @@ private import codeql.rust.elements.PathExpr
 module Impl {
   private import rust
   private import codeql.rust.internal.PathResolution as PathResolution
-  private import codeql.rust.internal.TypeInference as TypeInference
 
   pragma[nomagic]
   Path getFunctionPath(CallExpr ce) { result = ce.getFunction().(PathExpr).getPath() }
@@ -37,15 +36,6 @@ module Impl {
   class CallExpr extends Generated::CallExpr {
     override string toStringImpl() { result = this.getFunction().toAbbreviatedString() + "(...)" }
 
-    override Callable getStaticTarget() {
-      // If this call is to a trait method, e.g., `Trait::foo(bar)`, then check
-      // if type inference can resolve it to the correct trait implementation.
-      result = TypeInference::resolveMethodCallTarget(this)
-      or
-      not exists(TypeInference::resolveMethodCallTarget(this)) and
-      result = getResolvedFunction(this)
-    }
-
     /** Gets the struct that this call resolves to, if any. */
     Struct getStruct() { result = getResolvedFunction(this) }
 

rust/ql/lib/codeql/rust/elements/internal/CallImpl.qll

@@ -62,7 +62,6 @@ module Impl {
     Function getStaticTarget() {
       result = TypeInference::resolveMethodCallTarget(this)
       or
-      not exists(TypeInference::resolveMethodCallTarget(this)) and
       result = this.(CallExpr).getStaticTarget()
     }
 

rust/ql/lib/codeql/rust/elements/internal/MethodCallExprImpl.qll

@@ -6,8 +6,6 @@
 
 private import rust
 private import codeql.rust.elements.internal.generated.MethodCallExpr
-private import codeql.rust.internal.PathResolution
-private import codeql.rust.internal.TypeInference
 
 /**
  * INTERNAL: This module contains the customizable definition of `MethodCallExpr` and should not
@@ -23,8 +21,6 @@ module Impl {
    * ```
    */
   class MethodCallExpr extends Generated::MethodCallExpr {
-    override Function getStaticTarget() { result = resolveMethodCallTarget(this) }
-
     private string toStringPart(int index) {
       index = 0 and
       result = this.getReceiver().toAbbreviatedString()

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

@@ -11,6 +11,7 @@ private import codeql.rust.frameworks.stdlib.Stdlib
 private import codeql.rust.frameworks.stdlib.Builtins as Builtins
 private import codeql.rust.elements.Call
 private import codeql.rust.elements.internal.CallImpl::Impl as CallImpl
+private import codeql.rust.elements.internal.CallExprImpl::Impl as CallExprImpl
 
 class Type = T::Type;
 
@@ -724,8 +725,6 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
     }
   }
 
-  private import codeql.rust.elements.internal.CallExprImpl::Impl as CallExprImpl
-
   final class Access extends Call {
     pragma[nomagic]
     Type getTypeArgument(TypeArgumentPosition apos, TypePath path) {
@@ -771,7 +770,9 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
     Declaration getTarget() {
       result = resolveMethodCallTarget(this) // mutual recursion; resolving method calls requires resolving types and vice versa
       or
-      result = CallExprImpl::getResolvedFunction(this)
+      result = resolveCallTargetSimple(this)
+      or
+      result = resolveCallTargetComplex(this) // mutual recursion
     }
   }
 
@@ -1350,7 +1351,7 @@ private predicate implSiblingCandidate(
   // siblings).
   not exists(impl.getAttributeMacroExpansion()) and
   // We use this for resolving methods, so exclude traits that do not have methods.
-  exists(Function f | f = trait.getASuccessor(_) and f.getParamList().hasSelfParam()) and
+  exists(Function f | f = trait.getASuccessor(_)) and
   selfTy = impl.getSelfTy() and
   rootType = selfTy.resolveType()
 }
@@ -1499,6 +1500,58 @@ private Function getTraitMethod(ImplTraitReturnType trait, string name) {
   result = getMethodSuccessor(trait.getImplTraitTypeRepr(), name)
 }
 
+pragma[nomagic]
+private predicate assocFuncResolutionDependsOnArgument(Function f, Impl impl, int pos) {
+  methodResolutionDependsOnArgument(impl, _, f, pos, _, _)
+}
+
+private class AssocFuncCallExpr extends CallExpr {
+  private int pos;
+
+  AssocFuncCallExpr() {
+    assocFuncResolutionDependsOnArgument(CallExprImpl::getResolvedFunction(this), _, pos)
+  }
+
+  Function getACandidate(Impl impl) {
+    result = CallExprImpl::getResolvedFunction(this) and
+    assocFuncResolutionDependsOnArgument(result, impl, pos)
+  }
+
+  int getPosition() { result = pos }
+
+  /** Gets the type of the receiver of the associated function call at `path`. */
+  Type getTypeAt(TypePath path) { result = inferType(this.getArg(pos), path) }
+}
+
+private module AssocFuncIsInstantiationOfInput implements
+  IsInstantiationOfInputSig<AssocFuncCallExpr>
+{
+  pragma[nomagic]
+  predicate potentialInstantiationOf(
+    AssocFuncCallExpr ce, TypeAbstraction impl, TypeMention constraint
+  ) {
+    exists(Function cand |
+      cand = ce.getACandidate(impl) and
+      constraint = cand.getParam(ce.getPosition()).getTypeRepr()
+    )
+  }
+}
+
+pragma[nomagic]
+ItemNode resolveCallTargetSimple(CallExpr ce) {
+  result = CallExprImpl::getResolvedFunction(ce) and
+  not assocFuncResolutionDependsOnArgument(result, _, _)
+}
+
+pragma[nomagic]
+Function resolveCallTargetComplex(AssocFuncCallExpr ce) {
+  exists(Impl impl |
+    IsInstantiationOf<AssocFuncCallExpr, AssocFuncIsInstantiationOfInput>::isInstantiationOf(ce,
+      impl, _) and
+    result = getMethodSuccessor(impl, ce.getACandidate(_).getName().getText())
+  )
+}
+
 cached
 private module Cached {
   private import codeql.rust.internal.CachedStages
@@ -1541,6 +1594,14 @@ private module Cached {
     result = getTraitMethod(mc.getTypeAt(TypePath::nil()), mc.getMethodName())
   }
 
+  /** Gets a method that the method call `mc` resolves to, if any. */
+  cached
+  Function resolveCallTarget(CallExpr ce) {
+    result = resolveCallTargetSimple(ce)
+    or
+    result = resolveCallTargetComplex(ce)
+  }
+
   pragma[inline]
   private Type inferRootTypeDeref(AstNode n) {
     result = inferType(n) and
@@ -1685,9 +1746,9 @@ private module Debug {
     result = inferType(n, path)
   }
 
-  Function debugResolveMethodCallTarget(Call mce) {
-    mce = getRelevantLocatable() and
-    result = resolveMethodCallTarget(mce)
+  Function debugResolveCallTarget(Call c) {
+    c = getRelevantLocatable() and
+    result = [resolveMethodCallTarget(c), resolveCallTarget(c)]
   }
 
   predicate debugInferImplicitSelfType(SelfParam self, TypePath path, Type t) {
@@ -1705,6 +1766,11 @@ private module Debug {
     tm.resolveTypeAt(path) = type
   }
 
+  Type debugInferAnnotatedType(AstNode n, TypePath path) {
+    n = getRelevantLocatable() and
+    result = inferAnnotatedType(n, path)
+  }
+
   pragma[nomagic]
   private int countTypesAtPath(AstNode n, TypePath path, Type t) {
     t = inferType(n, path) and