Rust: Fix context-based type inference explosion

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

@@ -77,32 +77,23 @@
 private predicate implHasSibling(ImplItemNode impl, Trait trait) { implSiblings(trait, impl, _) }
 
 /**
- * Holds if type parameter `tp` of `trait` occurs in the function `f` with the name
- * `functionName` at position `pos` and path `path`.
- *
- * Note that `pos` can also be the special `return` position, which is sometimes
- * needed to disambiguate associated function calls like `Default::default()`
- * (in this case, `tp` is the special `Self` type parameter).
+ * Holds if `f` is a function declared inside `trait`, and the type of `f` at
+ * `pos` and `path` is `traitTp`, which is a type parameter of `trait`.
  */
-bindingset[trait]
-pragma[inline_late]
+pragma[nomagic]
 predicate traitTypeParameterOccurrence(
   TraitItemNode trait, Function f, string functionName, FunctionPosition pos, TypePath path,
-  TypeParameter tp
+  TypeParameter traitTp
 ) {
-  f = trait.getASuccessor(functionName) and
-  tp = getAssocFunctionTypeAt(f, trait, pos, path) and
-  tp = trait.(TraitTypeAbstraction).getATypeParameter()
+  f = trait.getAssocItem(functionName) and
+  traitTp = getAssocFunctionTypeAt(f, trait, pos, path) and
+  traitTp = 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(
-  ImplItemNode impl, Function f, FunctionPosition pos, TypePath path, Type type
+private predicate functionResolutionDependsOnArgumentCand(
+  ImplItemNode impl, Function f, string functionName, TypeParameter traitTp, FunctionPosition pos,
+  TypePath path
 ) {
   /*
    * As seen in the example below, when an implementation has a sibling for a
@@ -129,11 +120,113 @@
    * method. In that case we will still resolve several methods.
    */
 
-  exists(TraitItemNode trait, string functionName |
+  exists(TraitItemNode trait |
     implHasSibling(impl, trait) and
-    traitTypeParameterOccurrence(trait, _, functionName, pos, path, _) and
-    type = getAssocFunctionTypeAt(f, impl, pos, path) and
+    traitTypeParameterOccurrence(trait, _, functionName, pos, path, traitTp) and
     f = impl.getASuccessor(functionName) and
+    not pos.isSelf()
+  )
+}
+
+private predicate functionResolutionDependsOnPositionalArgument(
+  ImplItemNode impl, Function f, string functionName, TypeParameter traitTp
+) {
+  exists(FunctionPosition pos |
+    functionResolutionDependsOnArgumentCand(impl, f, functionName, traitTp, pos, _) and
+    pos.isPosition()
+  )
+}
+
+pragma[nomagic]
+private Type getAssocFunctionNonTypeParameterTypeAt(
+  ImplItemNode impl, Function f, FunctionPosition pos, TypePath path
+) {
+  result = getAssocFunctionTypeAt(f, impl, pos, path) and
+  not result instanceof TypeParameter
+}
+
+/**
+ * Holds if `f` inside `impl` has a sibling implementation inside `sibling`, where
+ * those two implementations agree on the instantiation of `traitTp`, which occurs
+ * in a positional position inside `f`.
+ */
+pragma[nomagic]
+private predicate hasEquivalentPositionalSibling(
+  ImplItemNode impl, ImplItemNode sibling, Function f, TypeParameter traitTp
+) {
+  exists(string functionName, FunctionPosition pos, TypePath path |
+    functionResolutionDependsOnArgumentCand(impl, f, functionName, traitTp, pos, path) and
     pos.isPosition()
+  |
+    exists(Function f1 |
+      implSiblings(_, impl, sibling) and
+      f1 = sibling.getASuccessor(functionName)
+    |
+      forall(TypePath path0, Type t |
+        t = getAssocFunctionNonTypeParameterTypeAt(impl, f, pos, path0) and
+        (path = path0.getAPrefix() or path = path0)
+      |
+        t = getAssocFunctionNonTypeParameterTypeAt(sibling, f1, pos, path0)
+      ) and
+      forall(TypePath path0, Type t |
+        t = getAssocFunctionNonTypeParameterTypeAt(sibling, f1, pos, path0) and
+        (path = path0.getAPrefix() or path = path0)
+      |
+        t = getAssocFunctionNonTypeParameterTypeAt(impl, f, pos, path0)
+      )
+    )
+  )
+}
+
+/**
+ * Holds if resolving the function `f` in `impl` requires inspecting the type
+ * of applied _arguments_ or possibly knowing the return type.
+ *
+ * `traitTp` is a type parameter of the trait being implemented by `impl`, and
+ * we need to check that the type of `f` corresponding to `traitTp` is satisfied
+ * at any one of the positions `pos` in which that type occurs in `f`.
+ *
+ * Type parameters that only occur in return positions are only included when
+ * all other type parameters that occur in a positional position are insufficient
+ * to disambiguate.
+ *
+ * Example:
+ *
+ * ```rust
+ * trait Trait1<T1> {
+ *   fn f(self, x: T1) -> T1;
+ * }
+ *
+ * impl Trait1<i32> for i32 {
+ *   fn f(self, x: i32) -> i32 { 0 } // f1
+ * }
+ *
+ * impl Trait1<i64> for i32 {
+ *   fn f(self, x: i64) -> i64 { 0 } // f2
+ * }
+ * ```
+ *
+ * The type for `T1` above occurs in both a positional position and a return position
+ * in `f`, so both may be used to disambiguate between `f1` and `f2`. That is, `f(0i32)`
+ * is sufficient to resolve to `f1`, and so is `let y: i64 = f(Default::default())`.
+ */
+pragma[nomagic]
+predicate functionResolutionDependsOnArgument(
+  ImplItemNode impl, Function f, TypeParameter traitTp, FunctionPosition pos
+) {
+  exists(string functionName, TypePath path |
+    functionResolutionDependsOnArgumentCand(impl, f, functionName, traitTp, pos, path)
+  |
+    if functionResolutionDependsOnPositionalArgument(impl, f, functionName, traitTp)
+    then any()
+    else
+      exists(ImplItemNode sibling |
+        implSiblings(_, impl, sibling) and
+        forall(TypeParameter otherTraitTp |
+          functionResolutionDependsOnPositionalArgument(impl, f, functionName, otherTraitTp)
+        |
+          hasEquivalentPositionalSibling(impl, sibling, f, otherTraitTp)
+        )
+      )
   )
 }

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

@@ -294,10 +294,14 @@ module ArgIsInstantiationOf<
  */
 signature module ArgsAreInstantiationsOfInputSig {
   /**
-   * Holds if types need to be matched against the type `t` at position `pos` of
-   * `f` inside `i`.
+   * Holds if `f` implements a trait function with type parameter `traitTp`, where
+   * we need to check that the type of `f` for `traitTp` is satisfied.
+   *
+   * `pos` is one of the positions in `f` in which the relevant type occours.
+   *
+   * For example, i
    */
-  predicate toCheck(ImplOrTraitItemNode i, Function f, FunctionPosition pos, AssocFunctionType t);
+  predicate toCheck(ImplOrTraitItemNode i, Function f, TypeParameter traitTp, FunctionPosition pos);
 
   /** A call whose argument types are to be checked. */
   class Call {
@@ -318,23 +322,28 @@ signature module ArgsAreInstantiationsOfInputSig {
  */
 module ArgsAreInstantiationsOf<ArgsAreInstantiationsOfInputSig Input> {
   pragma[nomagic]
-  private predicate toCheckRanked(ImplOrTraitItemNode i, Function f, FunctionPosition pos, int rnk) {
-    Input::toCheck(i, f, pos, _) and
-    pos =
-      rank[rnk + 1](FunctionPosition pos0, int j |
-        Input::toCheck(i, f, pos0, _) and
-        (
-          j = pos0.asPosition()
-          or
-          pos0.isSelf() and j = -1
-          or
-          pos0.isReturn() and j = -2
-        )
+  private predicate toCheckRanked(
+    ImplOrTraitItemNode i, Function f, TypeParameter traitTp, FunctionPosition pos, int rnk
+  ) {
+    Input::toCheck(i, f, traitTp, pos) and
+    traitTp =
+      rank[rnk + 1](TypeParameter traitTp0, int j |
+        Input::toCheck(i, f, traitTp0, _) and
+        j = getTypeParameterId(traitTp0)
       |
-        pos0 order by j
+        traitTp0 order by j
       )
   }
 
+  pragma[nomagic]
+  private predicate toCheck(
+    ImplOrTraitItemNode i, Function f, TypeParameter traitTp, FunctionPosition pos,
+    AssocFunctionType t
+  ) {
+    Input::toCheck(i, f, traitTp, pos) and
+    t.appliesTo(f, i, pos)
+  }
+
   private newtype TCallAndPos =
     MkCallAndPos(Input::Call call, FunctionPosition pos) { exists(call.getArgType(pos, _)) }
 
@@ -356,36 +365,34 @@ module ArgsAreInstantiationsOf<ArgsAreInstantiationsOfInputSig Input> {
     string toString() { result = call.toString() + " [arg " + pos + "]" }
   }
 
+  pragma[nomagic]
+  private predicate potentialInstantiationOf0(
+    CallAndPos cp, Input::Call call, TypeParameter traitTp, FunctionPosition pos, Function f,
+    TypeAbstraction abs, AssocFunctionType constraint
+  ) {
+    cp = MkCallAndPos(call, pragma[only_bind_into](pos)) and
+    call.hasTargetCand(abs, f) and
+    toCheck(abs, f, traitTp, pragma[only_bind_into](pos), constraint)
+  }
+
   private module ArgIsInstantiationOfToIndexInput implements
     IsInstantiationOfInputSig<CallAndPos, AssocFunctionType>
   {
-    pragma[nomagic]
-    private predicate potentialInstantiationOf0(
-      CallAndPos cp, Input::Call call, FunctionPosition pos, int rnk, Function f,
-      TypeAbstraction abs, AssocFunctionType constraint
-    ) {
-      cp = MkCallAndPos(call, pragma[only_bind_into](pos)) and
-      call.hasTargetCand(abs, f) and
-      toCheckRanked(abs, f, pragma[only_bind_into](pos), rnk) and
-      Input::toCheck(abs, f, pragma[only_bind_into](pos), constraint)
-    }
-
     pragma[nomagic]
     predicate potentialInstantiationOf(
       CallAndPos cp, TypeAbstraction abs, AssocFunctionType constraint
     ) {
-      exists(Input::Call call, int rnk, Function f |
-        potentialInstantiationOf0(cp, call, _, rnk, f, abs, constraint)
+      exists(Input::Call call, TypeParameter traitTp, FunctionPosition pos, int rnk, Function f |
+        potentialInstantiationOf0(cp, call, traitTp, pragma[only_bind_into](pos), f, abs, constraint) and
+        toCheckRanked(abs, f, traitTp, pragma[only_bind_into](pos), rnk)
       |
         rnk = 0
         or
         argsAreInstantiationsOfToIndex(call, abs, f, rnk - 1)
       )
     }
 
-    predicate relevantConstraint(AssocFunctionType constraint) {
-      Input::toCheck(_, _, _, constraint)
-    }
+    predicate relevantConstraint(AssocFunctionType constraint) { toCheck(_, _, _, _, constraint) }
   }
 
   private module ArgIsInstantiationOfToIndex =
@@ -398,39 +405,63 @@ module ArgsAreInstantiationsOf<ArgsAreInstantiationsOfInputSig Input> {
     exists(FunctionPosition pos |
       ArgIsInstantiationOfToIndex::argIsInstantiationOf(MkCallAndPos(call, pos), i, _) and
       call.hasTargetCand(i, f) and
-      toCheckRanked(i, f, pos, rnk)
+      toCheckRanked(i, f, _, pos, rnk)
+    |
+      rnk = 0
+      or
+      argsAreInstantiationsOfToIndex(call, i, f, rnk - 1)
     )
   }
 
   /**
    * Holds if all arguments of `call` have types that are instantiations of the
    * types of the corresponding parameters of `f` inside `i`.
+   *
+   * TODO: Check type parameter constraints as well.
    */
   pragma[nomagic]
   predicate argsAreInstantiationsOf(Input::Call call, ImplOrTraitItemNode i, Function f) {
     exists(int rnk |
       argsAreInstantiationsOfToIndex(call, i, f, rnk) and
-      rnk = max(int r | toCheckRanked(i, f, _, r))
+      rnk = max(int r | toCheckRanked(i, f, _, _, r))
     )
   }
 
+  private module ArgsAreNotInstantiationOfInput implements
+    IsInstantiationOfInputSig<CallAndPos, AssocFunctionType>
+  {
+    pragma[nomagic]
+    predicate potentialInstantiationOf(
+      CallAndPos cp, TypeAbstraction abs, AssocFunctionType constraint
+    ) {
+      potentialInstantiationOf0(cp, _, _, _, _, abs, constraint)
+    }
+
+    predicate relevantConstraint(AssocFunctionType constraint) { toCheck(_, _, _, _, constraint) }
+  }
+
+  private module ArgsAreNotInstantiationOf =
+    ArgIsInstantiationOf<CallAndPos, ArgsAreNotInstantiationOfInput>;
+
   pragma[nomagic]
   private predicate argsAreNotInstantiationsOf0(
     Input::Call call, FunctionPosition pos, ImplOrTraitItemNode i
   ) {
-    ArgIsInstantiationOfToIndex::argIsNotInstantiationOf(MkCallAndPos(call, pos), i, _, _)
+    ArgsAreNotInstantiationOf::argIsNotInstantiationOf(MkCallAndPos(call, pos), i, _, _)
   }
 
   /**
    * Holds if _some_ argument of `call` has a type that is not an instantiation of the
    * type of the corresponding parameter of `f` inside `i`.
+   *
+   * TODO: Check type parameter constraints as well.
    */
   pragma[nomagic]
   predicate argsAreNotInstantiationsOf(Input::Call call, ImplOrTraitItemNode i, Function f) {
     exists(FunctionPosition pos |
       argsAreNotInstantiationsOf0(call, pos, i) and
       call.hasTargetCand(i, f) and
-      Input::toCheck(i, f, pos, _)
+      Input::toCheck(i, f, _, pos)
     )
   }
 }