Rust: Refactor using shared constraint satisfaction

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

@@ -896,7 +896,7 @@ private predicate assocFunctionInfo(
 ) {
   f = i.getASuccessor(name) and
   arity = f.getParamList().getNumberOfParams() and
-  t.appliesTo(f, pos, i)
+  t.appliesTo(f, i, pos)
 }
 
 /**
@@ -1791,7 +1791,7 @@ private module MethodResolution {
     predicate toCheck(ImplOrTraitItemNode i, Function f, FunctionPosition pos, AssocFunctionType t) {
       exists(TypePath path, Type t0 |
         FunctionOverloading::functionResolutionDependsOnArgument(i, f, pos, path, t0) and
-        t.appliesTo(f, pos, i) and
+        t.appliesTo(f, i, pos) and
         // for now, we do not handle ambiguous targets when one of the types it iself
         // a type parameter; we should be checking the constraints on that type parameter
         // in this case
@@ -2024,7 +2024,7 @@ private module NonMethodResolution {
     ImplItemNode impl, NonMethodFunction implFunction
   ) {
     exists(TypePath path |
-      assocFunctionTypeAt(implFunction, impl, pos, path, type) and
+      type = assocFunctionTypeAt(implFunction, impl, pos, path) and
       implFunction.implements(traitFunction) and
       FunctionOverloading::traitTypeParameterOccurrence(trait, traitFunction, _, pos, path, _)
     |
@@ -2238,7 +2238,7 @@ private module NonMethodResolution {
 
   private module NonMethodArgsAreInstantiationsOfInput implements ArgsAreInstantiationsOfInputSig {
     predicate toCheck(ImplOrTraitItemNode i, Function f, FunctionPosition pos, AssocFunctionType t) {
-      t.appliesTo(f, pos, i) and
+      t.appliesTo(f, i, pos) and
       (
         exists(Type t0 |
           // for now, we do not handle ambiguous targets when one of the types it iself

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

@@ -79,7 +79,7 @@ predicate traitTypeParameterOccurrence(
   TypeParameter tp
 ) {
   f = trait.getASuccessor(functionName) and
-  assocFunctionTypeAt(f, trait, pos, path, tp) and
+  tp = assocFunctionTypeAt(f, trait, pos, path) and
   tp = trait.(TraitTypeAbstraction).getATypeParameter()
 }
 
@@ -120,7 +120,7 @@ predicate functionResolutionDependsOnArgument(
   exists(TraitItemNode trait, string functionName |
     implHasSibling(impl, trait) and
     traitTypeParameterOccurrence(trait, _, functionName, pos, path, _) and
-    assocFunctionTypeAt(f, impl, pos, path, type) and
+    type = assocFunctionTypeAt(f, impl, pos, path) and
     f = impl.getASuccessor(functionName) and
     pos.isPosition()
   )

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

@@ -72,28 +72,54 @@ module FunctionPositionMatchingInput {
 }
 
 private newtype TAssocFunctionType =
-  MkAssocFunctionType(Function f, FunctionPosition pos, ImplOrTraitItemNode i) {
-    f = i.getAnAssocItem() and
-    exists(pos.getTypeMention(f))
-  } or
-  MkInheritedAssocFunctionType(
-    Function f, FunctionPosition pos, TypeMention parentMention, TraitItemNode parent,
-    ImplOrTraitItemNode i
-  ) {
-    exists(AssocFunctionType inherited |
-      inherited.appliesTo(f, pos, parent) and
-      f = i.getASuccessor(_)
-    |
-      parent = i.(ImplItemNode).resolveTraitTy() and
-      parentMention = i.(ImplItemNode).getTraitPath()
-      or
-      parent = i.(TraitItemNode).resolveBound(parentMention)
-    )
+  /** An associated function `f` that should be specialized for `i` at `pos`. */
+  MkAssocFunctionType(Function f, ImplOrTraitItemNode i, FunctionPosition pos) {
+    f = i.getASuccessor(_) and exists(pos.getTypeMention(f))
   }
 
+bindingset[condition, constraint, tp]
+private Type traitConstraintTypeAt(
+  TypeMention condition, TypeMention constraint, TypeParameter tp, TypePath path
+) {
+  BaseTypes::conditionSatisfiesConstraintTypeAt(_, condition, constraint,
+    TypePath::singleton(tp).appendInverse(path), result)
+}
+
 /**
- * The type of an associated function at a given position, when viewed as a member
- * of a given trait or `impl` block.
+ * Gets if the type of the function `f` when specialized for `i` at position
+ * `pos` at path `path`
+ */
+pragma[nomagic]
+Type assocFunctionTypeAt(Function f, ImplOrTraitItemNode i, FunctionPosition pos, TypePath path) {
+  exists(MkAssocFunctionType(f, i, pos)) and
+  (
+    // No specialization needed when the function is directly in the trait or
+    // impl block or the declared type is not a type parameter
+    (i.getAnAssocItem() = f or not result instanceof TypeParameter) and
+    result = pos.getTypeMention(f).resolveTypeAt(path)
+    or
+    not i.getAnAssocItem() = f and
+    exists(TypePath prefix, TypePath suffix, TypeParameter tp |
+      path = prefix.append(suffix) and
+      tp = pos.getTypeMention(f).resolveTypeAt(prefix)
+    |
+      if tp = TSelfTypeParameter(_)
+      then result = resolveImplOrTraitType(i, suffix)
+      else
+        exists(TraitItemNode trait, TypeMention condition, TypeMention constraint |
+          trait.getAnAssocItem() = f and
+          BaseTypes::rootTypesSatisfaction(_, TTrait(trait), _, condition, constraint) and
+          result = traitConstraintTypeAt(condition, constraint, tp, suffix)
+        |
+          condition = i.(Trait) or condition = i.(Impl).getSelfTy()
+        )
+    )
+  )
+}
+
+/**
+ * The type of an associated function at a given position, when its implicit
+ * `Self` type parameter is specialized to a given trait or `impl` block.
  *
  * Example:
  *
@@ -126,64 +152,13 @@ private newtype TAssocFunctionType =
  * `self4` | `impl T2 for X` | `X`
  * `self5` | `impl T2 for X` | `X`
  */
-class AssocFunctionType extends TAssocFunctionType {
-  private predicate isFunctionType(Function f, FunctionPosition pos, ImplOrTraitItemNode i) {
-    this = MkAssocFunctionType(f, pos, i)
-  }
-
-  private predicate isInheritedFunctionType(
-    Function f, FunctionPosition pos, TypeMention parentMention, TraitItemNode parent,
-    ImplOrTraitItemNode i
-  ) {
-    this = MkInheritedAssocFunctionType(f, pos, parentMention, parent, i)
-  }
-
+class AssocFunctionType extends MkAssocFunctionType {
   /**
    * Holds if this function type applies to the function `f` at position `pos`,
    * when viewed as a member of the `impl` or trait item `i`.
    */
-  predicate appliesTo(Function f, FunctionPosition pos, ImplOrTraitItemNode i) {
-    this.isFunctionType(f, pos, i)
-    or
-    this.isInheritedFunctionType(f, pos, _, _, i)
-  }
-
-  /** Gets the type at the given path. */
-  pragma[nomagic]
-  Type getDeclaredTypeAt(TypePath path) {
-    exists(Function f, FunctionPosition pos |
-      this.isFunctionType(f, pos, _) and
-      result = pos.getTypeMention(f).resolveTypeAt(path)
-    )
-    or
-    exists(
-      Function f, FunctionPosition pos, TypeMention parentMention, TraitItemNode parent,
-      AssocFunctionType parentType, ImplOrTraitItemNode i
-    |
-      this.isInheritedFunctionType(f, pos, parentMention, parent, i) and
-      parentType.appliesTo(f, pos, parent)
-    |
-      result = parentType.getDeclaredTypeAt(path) and
-      not result instanceof TypeParameter
-      or
-      exists(TypePath prefix, TypePath suffix | path = prefix.append(suffix) |
-        parentType.hasSelfTypeParameterAt(prefix) and
-        result = resolveImplOrTraitType(i, suffix)
-        or
-        exists(TypeParameter tp |
-          tp = parentType.getTypeParameterAt(prefix) and
-          result = parentMention.resolveTypeAt(TypePath::singleton(tp).appendInverse(suffix))
-        )
-      )
-    )
-  }
-
-  pragma[nomagic]
-  private TypeParameter getTypeParameterAt(TypePath path) { result = this.getDeclaredTypeAt(path) }
-
-  pragma[nomagic]
-  private predicate hasSelfTypeParameterAt(TypePath path) {
-    this.getTypeParameterAt(path) = TSelfTypeParameter(_)
+  predicate appliesTo(Function f, ImplOrTraitItemNode i, FunctionPosition pos) {
+    this = MkAssocFunctionType(f, i, pos)
   }
 
   /**
@@ -196,7 +171,10 @@ class AssocFunctionType extends TAssocFunctionType {
    * traits when matching.
    */
   Type getTypeAt(TypePath path) {
-    exists(Type t | t = this.getDeclaredTypeAt(path) |
+    exists(Function f, FunctionPosition pos, ImplOrTraitItemNode i, Type t |
+      this.appliesTo(f, i, pos) and
+      t = assocFunctionTypeAt(f, i, pos, path)
+    |
       not t instanceof SelfTypeParameter and
       result = t
       or
@@ -206,7 +184,7 @@ class AssocFunctionType extends TAssocFunctionType {
 
   private TypeMention getTypeMention() {
     exists(Function f, FunctionPosition pos |
-      this.appliesTo(f, pos, _) and
+      this.appliesTo(f, _, pos) and
       result = pos.getTypeMention(f)
     )
   }
@@ -216,20 +194,6 @@ class AssocFunctionType extends TAssocFunctionType {
   Location getLocation() { result = this.getTypeMention().getLocation() }
 }
 
-/**
- * Holds if the type of the function `f` at position `pos` and path `path` inside
- * `i` is `type`.
- */
-pragma[nomagic]
-predicate assocFunctionTypeAt(
-  Function f, ImplOrTraitItemNode i, FunctionPosition pos, TypePath path, Type type
-) {
-  exists(AssocFunctionType aft |
-    aft.appliesTo(f, pos, i) and
-    type = aft.getDeclaredTypeAt(path)
-  )
-}
-
 private Trait getALookupTrait(Type t) {
   result = t.(TypeParamTypeParameter).getTypeParam().(TypeParamItemNode).resolveABound()
   or

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

@@ -729,7 +729,7 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
     }
 
     /** Provides logic related to base types. */
-    private module BaseTypes {
+    module BaseTypes {
       /**
        * Holds if, when `tm1` is considered an instantiation of `tm2`, then at
        * the type parameter `tp` it has the type `t` at `path`.