From 6e9fd496a6126814ffe113c0343625ff327d7fb6 Mon Sep 17 00:00:00 2001
From: Simon Friis Vindum <simonfv@gmail.com>
Date: Fri, 20 Jun 2025 13:53:29 +0200
Subject: [PATCH] Rust: Implement `await` expression using
 `SatisfiesConstraint` module

---
 .../codeql/rust/internal/PathResolution.qll   |  4 +-
 .../codeql/rust/internal/TypeInference.qll    | 93 ++++---------------
 rust/tools/builtins/await.rs                  |  7 --
 3 files changed, 18 insertions(+), 86 deletions(-)
 delete mode 100644 rust/tools/builtins/await.rs

diff --git a/rust/ql/lib/codeql/rust/internal/PathResolution.qll b/rust/ql/lib/codeql/rust/internal/PathResolution.qll
index f2c487598de..2f8c051d770 100644
--- a/rust/ql/lib/codeql/rust/internal/PathResolution.qll
+++ b/rust/ql/lib/codeql/rust/internal/PathResolution.qll
@@ -1132,7 +1132,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"
@@ -1501,7 +1501,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
diff --git a/rust/ql/lib/codeql/rust/internal/TypeInference.qll b/rust/ql/lib/codeql/rust/internal/TypeInference.qll
index b560ac5ec8c..a3cacc2f0ed 100644
--- a/rust/ql/lib/codeql/rust/internal/TypeInference.qll
+++ b/rust/ql/lib/codeql/rust/internal/TypeInference.qll
@@ -997,79 +997,6 @@ 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
@@ -1077,13 +1004,25 @@ private TraitType inferAsyncBlockExprRootType(AsyncBlockExpr abe) {
   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
diff --git a/rust/tools/builtins/await.rs b/rust/tools/builtins/await.rs
deleted file mode 100644
index c15af9dc529..00000000000
--- a/rust/tools/builtins/await.rs
+++ /dev/null
@@ -1,7 +0,0 @@
-use std::future::Future;
-
-fn await_type_matching<T1, T2: Future<Output = T1>>(x: T2) -> T1 {
-    panic!(
-        "This function exists only in order to implement type inference for `.await` expressions."
-    );
-}