Rust: Implement await expression usingSatisfiesConstraint module

Author: paldepind

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,36 @@ 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
+  }
+}
+
+predicate awaitSatisfiesConstraintTypeMention(AwaitTarget at, Type constraint, TypePath path, Type t) {
+  SatisfiesConstraint<AwaitTarget, AwaitSatisfiesConstraintInput>::satisfiesConstraintTypeMention(at,
+    constraint, path, t)
+}
 
 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 |
+    awaitSatisfiesConstraintTypeMention(n.(AwaitExpr).getExpr(), _, exprPath, result) and
+    exprPath.isCons(getFutureOutputTypeParameter(), path)
   )
   or
   // This case is needed for `async` functions and blocks, where we assign