Rust: Fix context-based type inference explosion

Author: hvitved

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

@@ -95,13 +95,8 @@ predicate traitTypeParameterOccurrence(
   tp = trait.(TraitTypeAbstraction).getATypeParameter()
 }
 
-/**
- * Holds if resolving the function `f` in `impl` with the name `functionName`
- * requires inspecting the type of applied _arguments_ at position `pos` in
- * order to determine whether it is the correct resolution.
- */
 pragma[nomagic]
-predicate functionResolutionDependsOnArgument(
+private predicate functionResolutionDependsOnArgument0(
   ImplItemNode impl, Function f, FunctionPosition pos, TypePath path, Type type
 ) {
   /*
@@ -133,7 +128,32 @@ predicate functionResolutionDependsOnArgument(
     implHasSibling(impl, trait) and
     traitTypeParameterOccurrence(trait, _, functionName, pos, path, _) and
     type = getAssocFunctionTypeAt(f, impl, pos, path) and
-    f = impl.getASuccessor(functionName) and
+    f = impl.getASuccessor(functionName)
+  )
+}
+
+/**
+ * Holds if resolving the function `f` in `impl` requires inspecting the type
+ * of applied _arguments_ at position `pos` (including the return type) in
+ * order to determine whether it is the correct resolution.
+ */
+pragma[nomagic]
+predicate functionResolutionDependsOnArgument(
+  ImplItemNode impl, Function f, FunctionPosition pos, TypePath path, Type type
+) {
+  functionResolutionDependsOnArgument0(impl, f, pos, path, type) and
+  (
     pos.isPosition()
+    or
+    // Only disambiguate based on return type when all other positions are trivially
+    // satisfied for all arguments.
+    pos.isReturn() and
+    forall(FunctionPosition pos0, TypePath path0, Type type0 |
+      pos0.isPosition() and
+      functionResolutionDependsOnArgument0(impl, f, pos0, path0, type0)
+    |
+      path0.isEmpty() and
+      type0.(TypeParamTypeParameter).getTypeParam() = any(TypeParam tp | not tp.hasTypeBound())
+    )
   )
 }

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

@@ -1172,13 +1172,13 @@ private module ContextTyping {
    */
   module CheckContextTyping<inferCallTypeSig/3 inferCallType> {
     pragma[nomagic]
-    private Type inferCallTypeFromContextCand(AstNode n, TypePath path, TypePath prefix) {
+    private Type inferCallTypeFromContextCand(AstNode n, TypePath prefix, TypePath path) {
       result = inferCallType(n, false, path) and
       hasUnknownType(n) and
       prefix = path
       or
       exists(TypePath mid |
-        result = inferCallTypeFromContextCand(n, path, mid) and
+        result = inferCallTypeFromContextCand(n, mid, path) and
         mid.isSnoc(prefix, _)
       )
     }
@@ -1188,7 +1188,7 @@ private module ContextTyping {
       result = inferCallType(n, true, path)
       or
       exists(TypePath prefix |
-        result = inferCallTypeFromContextCand(n, path, prefix) and
+        result = inferCallTypeFromContextCand(n, prefix, path) and
         hasUnknownTypeAt(n, prefix)
       )
     }
@@ -2669,9 +2669,8 @@ private module NonMethodResolution {
       then
         // We only check that the context of the call provides relevant type information
         // when no argument can
-        not exists(FunctionPosition pos0 |
-          FunctionOverloading::traitTypeParameterOccurrence(trait, traitFunction, _, pos0, _, _) and
-          not pos0.isReturn()
+        not exists(FunctionPosition pos0 | not pos0.isReturn() |
+          FunctionOverloading::traitTypeParameterOccurrence(trait, traitFunction, _, pos0, _, _)
           or
           FunctionOverloading::functionResolutionDependsOnArgument(impl, implFunction, pos0, _, _)
         )
@@ -2837,7 +2836,7 @@ private module NonMethodResolution {
     NonMethodFunction resolveTraitFunctionViaPathResolution(TraitItemNode trait) {
       this.hasTrait() and
       result = this.getPathResolutionResolved() and
-      result = trait.getASuccessor(_)
+      result = trait.getAnAssocItem()
     }
   }
 

rust/ql/test/query-tests/security/CWE-312/CONSISTENCY/PathResolutionConsistency.expected

@@ -980,6 +980,11 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
         not t = abs.getATypeParameter()
       }
 
+      pragma[nomagic]
+      private predicate hasTypeConstraint(HasTypeTree term, Type constraint) {
+        hasTypeConstraint(term, constraint, constraint)
+      }
+
       /**
        * Holds if the type tree at `tt` satisfies the constraint `constraint`
        * with the type `t` at `path`.
@@ -994,7 +999,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
           path = prefix0.append(suffix)
         )
         or
-        hasTypeConstraint(tt, constraint, constraint) and
+        hasTypeConstraint(tt, constraint) and
         t = getTypeAt(tt, path)
       }