Rust: Add SatisfiesConstraintInput module in shared type inference

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

@@ -1128,7 +1128,7 @@ pragma[nomagic]
 private predicate crateDependencyEdge(SourceFileItemNode file, string name, CrateItemNode dep) {
   exists(CrateItemNode c | dep = c.(Crate).getDependency(name) | file = c.getASourceFile())
   or
-  // Give builtin files, such as `await.rs`, access to `std`
+  // Give builtin files access to `std`
   file instanceof BuiltinSourceFile and
   dep.getName() = name and
   name = "std"
@@ -1497,7 +1497,7 @@ private predicate preludeEdge(SourceFile f, string name, ItemNode i) {
   exists(Crate stdOrCore, ModuleLikeNode mod, ModuleItemNode prelude, ModuleItemNode rust |
     f = any(Crate c0 | stdOrCore = c0.getDependency(_) or stdOrCore = c0).getASourceFile()
     or
-    // Give builtin files, such as `await.rs`, access to the prelude
+    // Give builtin files access to the prelude
     f instanceof BuiltinSourceFile
   |
     stdOrCore.getName() = ["std", "core"] and

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

@@ -981,93 +981,32 @@ private AssociatedTypeTypeParameter getFutureOutputTypeParameter() {
   result.getTypeAlias() = any(FutureTrait ft).getOutputType()
 }
 
-/**
- * A matching configuration for resolving types of `.await` expressions.
- */
-private module AwaitExprMatchingInput implements MatchingInputSig {
-  private newtype TDeclarationPosition =
-    TSelfDeclarationPosition() or
-    TOutputPos()
-
-  class DeclarationPosition extends TDeclarationPosition {
-    predicate isSelf() { this = TSelfDeclarationPosition() }
-
-    predicate isOutput() { this = TOutputPos() }
-
-    string toString() {
-      this.isSelf() and
-      result = "self"
-      or
-      this.isOutput() and
-      result = "(output)"
-    }
-  }
-
-  private class BuiltinsAwaitFile extends File {
-    BuiltinsAwaitFile() {
-      this.getBaseName() = "await.rs" and
-      this.getParentContainer() instanceof Builtins::BuiltinsFolder
-    }
-  }
-
-  class Declaration extends Function {
-    Declaration() {
-      this.getFile() instanceof BuiltinsAwaitFile and
-      this.getName().getText() = "await_type_matching"
-    }
-
-    TypeParameter getTypeParameter(TypeParameterPosition ppos) {
-      typeParamMatchPosition(this.getGenericParamList().getATypeParam(), result, ppos)
-    }
-
-    Type getDeclaredType(DeclarationPosition dpos, TypePath path) {
-      dpos.isSelf() and
-      result = this.getParam(0).getTypeRepr().(TypeMention).resolveTypeAt(path)
-      or
-      dpos.isOutput() and
-      result = this.getRetType().getTypeRepr().(TypeMention).resolveTypeAt(path)
-    }
-  }
-
-  class AccessPosition = DeclarationPosition;
-
-  class Access extends AwaitExpr {
-    Type getTypeArgument(TypeArgumentPosition apos, TypePath path) { none() }
-
-    AstNode getNodeAt(AccessPosition apos) {
-      result = this.getExpr() and
-      apos.isSelf()
-      or
-      result = this and
-      apos.isOutput()
-    }
-
-    Type getInferredType(AccessPosition apos, TypePath path) {
-      result = inferType(this.getNodeAt(apos), path)
-    }
-
-    Declaration getTarget() { exists(this) and exists(result) }
-  }
-
-  predicate accessDeclarationPositionMatch(AccessPosition apos, DeclarationPosition dpos) {
-    apos = dpos
-  }
-}
-
 pragma[nomagic]
 private TraitType inferAsyncBlockExprRootType(AsyncBlockExpr abe) {
   // `typeEquality` handles the non-root case
   exists(abe) and
   result = getFutureTraitType()
 }
 
-private module AwaitExprMatching = Matching<AwaitExprMatchingInput>;
+final class AwaitTarget extends Expr {
+  AwaitTarget() { this = any(AwaitExpr ae).getExpr() }
+
+  Type getTypeAt(TypePath path) { result = inferType(this, path) }
+}
+
+private module AwaitSatisfiesConstraintInput implements SatisfiesConstraintSig<AwaitTarget> {
+  predicate relevantConstraint(AwaitTarget term, Type constraint) {
+    exists(term) and
+    constraint.(TraitType).getTrait() instanceof FutureTrait
+  }
+}
 
 pragma[nomagic]
 private Type inferAwaitExprType(AstNode n, TypePath path) {
-  exists(AwaitExprMatchingInput::Access a, AwaitExprMatchingInput::AccessPosition apos |
-    n = a.getNodeAt(apos) and
-    result = AwaitExprMatching::inferAccessType(a, apos, path)
+  exists(TypePath exprPath |
+    SatisfiesConstraint<AwaitTarget, AwaitSatisfiesConstraintInput>::satisfiesConstraintTypeMention(n.(AwaitExpr)
+          .getExpr(), _, exprPath, result) and
+    exprPath.isCons(getFutureOutputTypeParameter(), path)
   )
   or
   // This case is needed for `async` functions and blocks, where we assign

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

@@ -866,6 +866,101 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
 
     private import BaseTypes
 
+    signature module SatisfiesConstraintSig<HasTypeTreeSig TypeTree> {
+      /** Holds if it is relevant to know if `term` satisfies `constraint`. */
+      predicate relevantConstraint(TypeTree term, Type constraint);
+    }
+
+    module SatisfiesConstraint<HasTypeTreeSig TypeTree, SatisfiesConstraintSig<TypeTree> Input> {
+      import Input
+
+      private module IsInstantiationOfInput implements IsInstantiationOfInputSig<TypeTree> {
+        predicate potentialInstantiationOf(TypeTree tt, TypeAbstraction abs, TypeMention cond) {
+          exists(Type constraint, Type type |
+            type = tt.getTypeAt(TypePath::nil()) and
+            relevantConstraint(tt, constraint) and
+            rootTypesSatisfaction(type, constraint, abs, cond, _) and
+            // We only need to check instantiations where there are multiple candidates.
+            countConstraintImplementations(type, constraint) > 1
+          )
+        }
+
+        predicate relevantTypeMention(TypeMention constraint) {
+          rootTypesSatisfaction(_, _, _, constraint, _)
+        }
+      }
+
+      /** Holds if the type tree has the type `type` and should satisfy `constraint`. */
+      private predicate hasTypeConstraint(TypeTree term, Type type, Type constraint) {
+        type = term.getTypeAt(TypePath::nil()) and
+        relevantConstraint(term, constraint)
+      }
+
+      /**
+       * Holds if `tt` satisfies `constraint` through `abs`, `sub`, and `constraintMention`.
+       */
+      pragma[nomagic]
+      private predicate hasConstraintMention(
+        TypeTree tt, TypeAbstraction abs, TypeMention sub, Type constraint,
+        TypeMention constraintMention
+      ) {
+        exists(Type type | hasTypeConstraint(tt, type, constraint) |
+          not exists(countConstraintImplementations(type, constraint)) and
+          conditionSatisfiesConstraintTypeAt(abs, sub, constraintMention, _, _) and
+          resolveTypeMentionRoot(sub) = abs.getATypeParameter() and
+          constraint = resolveTypeMentionRoot(constraintMention)
+          or
+          countConstraintImplementations(type, constraint) > 0 and
+          rootTypesSatisfaction(type, constraint, abs, sub, constraintMention) and
+          // When there are multiple ways the type could implement the
+          // constraint we need to find the right implementation, which is the
+          // one where the type instantiates the precondition.
+          if countConstraintImplementations(type, constraint) > 1
+          then IsInstantiationOf<TypeTree, IsInstantiationOfInput>::isInstantiationOf(tt, abs, sub)
+          else any()
+        )
+      }
+
+      pragma[nomagic]
+      private predicate satisfiesConstraintTypeMention0(
+        TypeTree tt, Type constraint, TypeAbstraction abs, TypeMention sub, TypePath path, Type t
+      ) {
+        exists(TypeMention constraintMention |
+          hasConstraintMention(tt, abs, sub, constraint, constraintMention) and
+          conditionSatisfiesConstraintTypeAt(abs, sub, constraintMention, path, t)
+        )
+      }
+
+      pragma[nomagic]
+      private predicate satisfiesConstraintTypeMention1(
+        TypeTree tt, Type constraint, TypePath path, TypePath pathToTypeParamInSub
+      ) {
+        exists(TypeAbstraction abs, TypeMention sub, TypeParameter tp |
+          satisfiesConstraintTypeMention0(tt, constraint, abs, sub, path, tp) and
+          tp = abs.getATypeParameter() and
+          sub.resolveTypeAt(pathToTypeParamInSub) = tp
+        )
+      }
+
+      /**
+       * Holds if the type tree at `tt` satisfies the constraint `constraint`
+       * with the type `t` at `path`.
+       */
+      pragma[nomagic]
+      predicate satisfiesConstraintTypeMention(TypeTree tt, Type constraint, TypePath path, Type t) {
+        exists(TypeAbstraction abs |
+          satisfiesConstraintTypeMention0(tt, constraint, abs, _, path, t) and
+          not t = abs.getATypeParameter()
+        )
+        or
+        exists(TypePath prefix0, TypePath pathToTypeParamInSub, TypePath suffix |
+          satisfiesConstraintTypeMention1(tt, constraint, prefix0, pathToTypeParamInSub) and
+          tt.getTypeAt(pathToTypeParamInSub.appendInverse(suffix)) = t and
+          path = prefix0.append(suffix)
+        )
+      }
+    }
+
     /** Provides the input to `Matching`. */
     signature module MatchingInputSig {
       /**
@@ -1129,11 +1224,8 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
             adjustedAccessType(a, apos, target, path.appendInverse(suffix), result)
           }
 
-          /** Holds if this relevant access has the type `type` and should satisfy `constraint`. */
-          predicate hasTypeConstraint(Type type, Type constraint) {
-            adjustedAccessType(a, apos, target, path, type) and
-            relevantAccessConstraint(a, target, apos, path, constraint)
-          }
+          /** Holds if this relevant access should satisfy `constraint`. */
+          Type getConstraint() { relevantAccessConstraint(a, target, apos, path, result) }
 
           string toString() {
             result = a.toString() + ", " + apos.toString() + ", " + path.toString()
@@ -1142,94 +1234,18 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
           Location getLocation() { result = a.getLocation() }
         }
 
-        private module IsInstantiationOfInput implements IsInstantiationOfInputSig<RelevantAccess> {
-          predicate potentialInstantiationOf(
-            RelevantAccess at, TypeAbstraction abs, TypeMention cond
-          ) {
-            exists(Type constraint, Type type |
-              at.hasTypeConstraint(type, constraint) and
-              rootTypesSatisfaction(type, constraint, abs, cond, _) and
-              // We only need to check instantiations where there are multiple candidates.
-              countConstraintImplementations(type, constraint) > 1
-            )
-          }
-
-          predicate relevantTypeMention(TypeMention constraint) {
-            rootTypesSatisfaction(_, _, _, constraint, _)
+        private module SatisfiesConstraintInput implements SatisfiesConstraintSig<RelevantAccess> {
+          predicate relevantConstraint(RelevantAccess at, Type constraint) {
+            constraint = at.getConstraint()
           }
         }
 
-        /**
-         * Holds if `at` satisfies `constraint` through `abs`, `sub`, and `constraintMention`.
-         */
-        pragma[nomagic]
-        private predicate hasConstraintMention(
-          RelevantAccess at, TypeAbstraction abs, TypeMention sub, Type constraint,
-          TypeMention constraintMention
-        ) {
-          exists(Type type | at.hasTypeConstraint(type, constraint) |
-            not exists(countConstraintImplementations(type, constraint)) and
-            conditionSatisfiesConstraintTypeAt(abs, sub, constraintMention, _, _) and
-            resolveTypeMentionRoot(sub) = abs.getATypeParameter() and
-            constraint = resolveTypeMentionRoot(constraintMention)
-            or
-            countConstraintImplementations(type, constraint) > 0 and
-            rootTypesSatisfaction(type, constraint, abs, sub, constraintMention) and
-            // When there are multiple ways the type could implement the
-            // constraint we need to find the right implementation, which is the
-            // one where the type instantiates the precondition.
-            if countConstraintImplementations(type, constraint) > 1
-            then
-              IsInstantiationOf<RelevantAccess, IsInstantiationOfInput>::isInstantiationOf(at, abs,
-                sub)
-            else any()
-          )
-        }
-
-        pragma[nomagic]
-        predicate satisfiesConstraintTypeMention0(
-          RelevantAccess at, Access a, AccessPosition apos, TypePath prefix, Type constraint,
-          TypeAbstraction abs, TypeMention sub, TypePath path, Type t
-        ) {
-          exists(TypeMention constraintMention |
-            at = MkRelevantAccess(a, _, apos, prefix) and
-            hasConstraintMention(at, abs, sub, constraint, constraintMention) and
-            conditionSatisfiesConstraintTypeAt(abs, sub, constraintMention, path, t)
-          )
-        }
-
-        pragma[nomagic]
-        predicate satisfiesConstraintTypeMention1(
-          RelevantAccess at, Access a, AccessPosition apos, TypePath prefix, Type constraint,
-          TypePath path, TypePath pathToTypeParamInSub
-        ) {
-          exists(TypeAbstraction abs, TypeMention sub, TypeParameter tp |
-            satisfiesConstraintTypeMention0(at, a, apos, prefix, constraint, abs, sub, path, tp) and
-            tp = abs.getATypeParameter() and
-            sub.resolveTypeAt(pathToTypeParamInSub) = tp
-          )
-        }
-
-        /**
-         * Holds if the type at `a`, `apos`, and `path` satisfies the constraint
-         * `constraint` with the type `t` at `path`.
-         */
-        pragma[nomagic]
         predicate satisfiesConstraintTypeMention(
           Access a, AccessPosition apos, TypePath prefix, Type constraint, TypePath path, Type t
         ) {
-          exists(TypeAbstraction abs |
-            satisfiesConstraintTypeMention0(_, a, apos, prefix, constraint, abs, _, path, t) and
-            not t = abs.getATypeParameter()
-          )
-          or
-          exists(
-            RelevantAccess at, TypePath prefix0, TypePath pathToTypeParamInSub, TypePath suffix
-          |
-            satisfiesConstraintTypeMention1(at, a, apos, prefix, constraint, prefix0,
-              pathToTypeParamInSub) and
-            at.getTypeAt(pathToTypeParamInSub.appendInverse(suffix)) = t and
-            path = prefix0.append(suffix)
+          exists(RelevantAccess at | at = MkRelevantAccess(a, _, apos, prefix) |
+            SatisfiesConstraint<RelevantAccess, SatisfiesConstraintInput>::satisfiesConstraintTypeMention(at,
+              constraint, path, t)
           )
         }
       }