Rust: Type inference performance improvements

Author: hvitved

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

@@ -101,7 +101,7 @@ predicate traitTypeParameterOccurrence(
  * 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 +133,30 @@ 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` 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(
+  ImplItemNode impl, Function f, FunctionPosition pos, TypePath path, Type type
+) {
+  functionResolutionDependsOnArgument0(impl, f, pos, path, type) and
+  (
     pos.isPosition()
+    or
+    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

@@ -1159,9 +1159,6 @@ private module ContextTyping {
     inferType(n, path) = TUnknownType()
   }
 
-  pragma[nomagic]
-  private predicate hasUnknownType(AstNode n) { hasUnknownTypeAt(n, _) }
-
   signature Type inferCallTypeSig(AstNode n, boolean isReturn, TypePath path);
 
   /**
@@ -1172,24 +1169,21 @@ private module ContextTyping {
    */
   module CheckContextTyping<inferCallTypeSig/3 inferCallType> {
     pragma[nomagic]
-    private Type inferCallTypeFromContextCand(AstNode n, TypePath path, TypePath prefix) {
-      result = inferCallType(n, false, path) and
-      hasUnknownType(n) and
-      prefix = path
-      or
-      exists(TypePath mid |
-        result = inferCallTypeFromContextCand(n, path, mid) and
-        mid.isSnoc(prefix, _)
-      )
-    }
+    Type inferCallNonReturnType(AstNode n, TypePath path) { result = inferCallType(n, false, path) }
 
     pragma[nomagic]
     Type check(AstNode n, TypePath path) {
       result = inferCallType(n, true, path)
       or
-      exists(TypePath prefix |
-        result = inferCallTypeFromContextCand(n, path, prefix) and
-        hasUnknownTypeAt(n, prefix)
+      exists(Type t, TypePath path0 |
+        t = inferCallNonReturnType(n, path0) and
+        hasUnknownTypeAt(n, path0)
+      |
+        result = t and
+        path = path0
+        or
+        result = TUnknownType() and
+        path = path0.append(TypePath::singleton(t.getATypeParameter()))
       )
     }
   }
@@ -2669,9 +2663,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 +2830,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)
       }