Merge pull request #19584 from hvitved/rust/type-inference-await

Rust: Type inference for `.await` expressions
2026-04-18 21:44:02 +02:00 · 2025-06-10 14:56:28 +02:00
parent 3af10d2c6c 3d395dde47
commit 791369dd0e
11 changed files with 588 additions and 24 deletions
--- a/rust/ql/lib/codeql/rust/elements/internal/GenericArgListImpl.qll
+++ b/rust/ql/lib/codeql/rust/elements/internal/GenericArgListImpl.qll
@@ -37,5 +37,15 @@ module Impl {

    /** Gets a type argument of this list. */
    TypeRepr getATypeArg() { result = this.getTypeArg(_) }
+
+    /** Gets the associated type argument with the given `name`, if any. */
+    pragma[nomagic]
+    TypeRepr getAssocTypeArg(string name) {
+      exists(AssocTypeArg arg |
+        arg = this.getAGenericArg() and
+        result = arg.getTypeRepr() and
+        name = arg.getIdentifier().getText()
+      )
+    }
  }
 }
--- a/rust/ql/lib/codeql/rust/frameworks/stdlib/Stdlib.qll
+++ b/rust/ql/lib/codeql/rust/frameworks/stdlib/Stdlib.qll
@@ -49,3 +49,19 @@ class ResultEnum extends Enum {
  /** Gets the `Err` variant. */
  Variant getErr() { result = this.getVariant("Err") }
 }
+
+/**
+ * The [`Future` trait][1].
+ *
+ * [1]: https://doc.rust-lang.org/std/future/trait.Future.html
+ */
+class FutureTrait extends Trait {
+  FutureTrait() { this.getCanonicalPath() = "core::future::future::Future" }
+
+  /** Gets the `Output` associated type. */
+  pragma[nomagic]
+  TypeAlias getOutputType() {
+    result = this.getAssocItemList().getAnAssocItem() and
+    result.getName().getText() = "Output"
+  }
+}
--- a/rust/ql/lib/codeql/rust/internal/PathResolution.qll
+++ b/rust/ql/lib/codeql/rust/internal/PathResolution.qll
@@ -5,6 +5,7 @@
 private import rust
 private import codeql.rust.elements.internal.generated.ParentChild
 private import codeql.rust.internal.CachedStages
+private import codeql.rust.frameworks.stdlib.Bultins as Builtins

 private newtype TNamespace =
  TTypeNamespace() or
@@ -178,6 +179,8 @@ abstract class ItemNode extends Locatable {
    or
    // type parameters have access to the associated items of its bounds
    result = this.(TypeParamItemNode).resolveABound().getASuccessorRec(name).(AssocItemNode)
+    or
+    result = this.(ImplTraitTypeReprItemNode).resolveABound().getASuccessorRec(name).(AssocItemNode)
  }

  /**
@@ -645,6 +648,28 @@ class ImplItemNode extends ImplOrTraitItemNode instanceof Impl {
  }
 }

+private class ImplTraitTypeReprItemNode extends ItemNode instanceof ImplTraitTypeRepr {
+  pragma[nomagic]
+  Path getABoundPath() {
+    result = super.getTypeBoundList().getABound().getTypeRepr().(PathTypeRepr).getPath()
+  }
+
+  pragma[nomagic]
+  ItemNode resolveABound() { result = resolvePathFull(this.getABoundPath()) }
+
+  override string getName() { result = "(impl trait)" }
+
+  override Namespace getNamespace() { result.isType() }
+
+  override Visibility getVisibility() { none() }
+
+  override TypeParam getTypeParam(int i) { none() }
+
+  override predicate hasCanonicalPath(Crate c) { none() }
+
+  override string getCanonicalPath(Crate c) { none() }
+}
+
 private class MacroCallItemNode extends AssocItemNode instanceof MacroCall {
  override string getName() { result = "(macro call)" }

@@ -1093,14 +1118,21 @@ private predicate crateDefEdge(CrateItemNode c, string name, ItemNode i) {
  not i instanceof Crate
 }

+private class BuiltinSourceFile extends SourceFileItemNode {
+  BuiltinSourceFile() { this.getFile().getParentContainer() instanceof Builtins::BuiltinsFolder }
+}
+
 /**
 * Holds if `file` depends on crate `dep` named `name`.
 */
+pragma[nomagic]
 private predicate crateDependencyEdge(SourceFileItemNode file, string name, CrateItemNode dep) {
-  exists(CrateItemNode c |
-    dep = c.(Crate).getDependency(name) and
-    file = c.getASourceFile()
-  )
+  exists(CrateItemNode c | dep = c.(Crate).getDependency(name) | file = c.getASourceFile())
+  or
+  // Give builtin files, such as `await.rs`, access to `std`
+  file instanceof BuiltinSourceFile and
+  dep.getName() = name and
+  name = "std"
 }

 private predicate useTreeDeclares(UseTree tree, string name) {
@@ -1461,9 +1493,14 @@ private predicate externCrateEdge(ExternCrateItemNode ec, string name, CrateItem
 * [1]: https://doc.rust-lang.org/core/prelude/index.html
 * [2]: https://doc.rust-lang.org/std/prelude/index.html
 */
+pragma[nomagic]
 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() and
+    f = any(Crate c0 | stdOrCore = c0.getDependency(_) or stdOrCore = c0).getASourceFile()
+    or
+    // Give builtin files, such as `await.rs`, access to the prelude
+    f instanceof BuiltinSourceFile
+  |
    stdOrCore.getName() = ["std", "core"] and
    mod = stdOrCore.getSourceFile() and
    prelude = mod.getASuccessorRec("prelude") and
@@ -1473,12 +1510,10 @@ private predicate preludeEdge(SourceFile f, string name, ItemNode i) {
  )
 }

-private import codeql.rust.frameworks.stdlib.Bultins as Builtins
-
 pragma[nomagic]
 private predicate builtin(string name, ItemNode i) {
-  exists(SourceFileItemNode builtins |
-    builtins.getFile().getParentContainer() instanceof Builtins::BuiltinsFolder and
+  exists(BuiltinSourceFile builtins |
+    builtins.getFile().getBaseName() = "types.rs" and
    i = builtins.getASuccessorRec(name)
  )
 }
--- a/rust/ql/lib/codeql/rust/internal/Type.qll
+++ b/rust/ql/lib/codeql/rust/internal/Type.qll
@@ -15,6 +15,7 @@ newtype TType =
  TTrait(Trait t) or
  TArrayType() or // todo: add size?
  TRefType() or // todo: add mut?
+  TImplTraitType(ImplTraitTypeRepr impl) or
  TTypeParamTypeParameter(TypeParam t) or
  TAssociatedTypeTypeParameter(TypeAlias t) { any(TraitItemNode trait).getAnAssocItem() = t } or
  TRefTypeParameter() or
@@ -115,6 +116,9 @@ class TraitType extends Type, TTrait {

  TraitType() { this = TTrait(trait) }

+  /** Gets the underlying trait. */
+  Trait getTrait() { result = trait }
+
  override StructField getStructField(string name) { none() }

  override TupleField getTupleField(int i) { none() }
@@ -176,6 +180,53 @@ class RefType extends Type, TRefType {
  override Location getLocation() { result instanceof EmptyLocation }
 }

+/**
+ * An [impl Trait][1] type.
+ *
+ * Each syntactic `impl Trait` type gives rise to its own type, even if
+ * two `impl Trait` types have the same bounds.
+ *
+ * [1]: https://doc.rust-lang.org/reference/types/impl-trait.html
+ */
+class ImplTraitType extends Type, TImplTraitType {
+  ImplTraitTypeRepr impl;
+
+  ImplTraitType() { this = TImplTraitType(impl) }
+
+  /** Gets the underlying AST node. */
+  ImplTraitTypeRepr getImplTraitTypeRepr() { result = impl }
+
+  /** Gets the function that this `impl Trait` belongs to. */
+  abstract Function getFunction();
+
+  override StructField getStructField(string name) { none() }
+
+  override TupleField getTupleField(int i) { none() }
+
+  override TypeParameter getTypeParameter(int i) { none() }
+
+  override string toString() { result = impl.toString() }
+
+  override Location getLocation() { result = impl.getLocation() }
+}
+
+/**
+ * An [impl Trait in return position][1] type, for example:
+ *
+ * ```rust
+ * fn foo() -> impl Trait
+ * ```
+ *
+ * [1]: https://doc.rust-lang.org/reference/types/impl-trait.html#r-type.impl-trait.return
+ */
+class ImplTraitReturnType extends ImplTraitType {
+  private Function function;
+
+  ImplTraitReturnType() { impl = function.getRetType().getTypeRepr() }
+
+  override Function getFunction() { result = function }
+}
+
 /** A type parameter. */
 abstract class TypeParameter extends Type {
  override StructField getStructField(string name) { none() }
@@ -185,7 +236,7 @@ abstract class TypeParameter extends Type {
  override TypeParameter getTypeParameter(int i) { none() }
 }

-private class RawTypeParameter = @type_param or @trait or @type_alias;
+private class RawTypeParameter = @type_param or @trait or @type_alias or @impl_trait_type_repr;

 private predicate id(RawTypeParameter x, RawTypeParameter y) { x = y }

@@ -281,6 +332,37 @@ class SelfTypeParameter extends TypeParameter, TSelfTypeParameter {
  override Location getLocation() { result = trait.getLocation() }
 }

+/**
+ * An [impl Trait in argument position][1] type, for example:
+ *
+ * ```rust
+ * fn foo(arg: impl Trait)
+ * ```
+ *
+ * Such types are syntactic sugar for type parameters, that is
+ *
+ * ```rust
+ * fn foo<T: Trait>(arg: T)
+ * ```
+ *
+ * so we model them as type parameters.
+ *
+ * [1]: https://doc.rust-lang.org/reference/types/impl-trait.html#r-type.impl-trait.param
+ */
+class ImplTraitTypeTypeParameter extends ImplTraitType, TypeParameter {
+  private Function function;
+
+  ImplTraitTypeTypeParameter() { impl = function.getParamList().getAParam().getTypeRepr() }
+
+  override Function getFunction() { result = function }
+
+  override StructField getStructField(string name) { none() }
+
+  override TupleField getTupleField(int i) { none() }
+
+  override TypeParameter getTypeParameter(int i) { none() }
+}
+
 /**
 * A type abstraction. I.e., a place in the program where type variables are
 * introduced.
@@ -316,3 +398,7 @@ final class SelfTypeBoundTypeAbstraction extends TypeAbstraction, Name {

  override TypeParamTypeParameter getATypeParameter() { none() }
 }
+
+final class ImplTraitTypeReprAbstraction extends TypeAbstraction, ImplTraitTypeRepr {
+  override TypeParamTypeParameter getATypeParameter() { none() }
+}
--- a/rust/ql/lib/codeql/rust/internal/TypeInference.qll
+++ b/rust/ql/lib/codeql/rust/internal/TypeInference.qll
@@ -88,7 +88,8 @@ private module Input1 implements InputSig1<Location> {
        exists(AstNode node | id = idOfTypeParameterAstNode(node) |
          node = tp0.(TypeParamTypeParameter).getTypeParam() or
          node = tp0.(AssociatedTypeTypeParameter).getTypeAlias() or
-          node = tp0.(SelfTypeParameter).getTrait()
+          node = tp0.(SelfTypeParameter).getTrait() or
+          node = tp0.(ImplTraitTypeTypeParameter).getImplTraitTypeRepr()
        )
      |
        tp0 order by kind, id
@@ -117,6 +118,13 @@ private module Input2 implements InputSig2 {
    result = tp.(TypeParamTypeParameter).getTypeParam().getTypeBoundList().getABound().getTypeRepr()
    or
    result = tp.(SelfTypeParameter).getTrait()
+    or
+    result =
+      tp.(ImplTraitTypeTypeParameter)
+          .getImplTraitTypeRepr()
+          .getTypeBoundList()
+          .getABound()
+          .getTypeRepr()
  }

  /**
@@ -156,6 +164,12 @@ private module Input2 implements InputSig2 {
      condition = self and
      constraint = self.getTrait()
    )
+    or
+    exists(ImplTraitTypeRepr impl |
+      abs = impl and
+      condition = impl and
+      constraint = impl.getTypeBoundList().getABound().getTypeRepr()
+    )
  }
 }

@@ -228,8 +242,6 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
    or
    n1 = n2.(ParenExpr).getExpr()
    or
-    n1 = n2.(BlockExpr).getStmtList().getTailExpr()
-    or
    n1 = n2.(IfExpr).getABranch()
    or
    n1 = n2.(MatchExpr).getAnArm().getExpr()
@@ -248,6 +260,19 @@ private predicate typeEquality(AstNode n1, TypePath prefix1, AstNode n2, TypePat
  n1 = n2.(DerefExpr).getExpr() and
  prefix1 = TypePath::singleton(TRefTypeParameter()) and
  prefix2.isEmpty()
+  or
+  exists(BlockExpr be |
+    n1 = be and
+    n2 = be.getStmtList().getTailExpr() and
+    if be.isAsync()
+    then
+      prefix1 = TypePath::singleton(getFutureOutputTypeParameter()) and
+      prefix2.isEmpty()
+    else (
+      prefix1.isEmpty() and
+      prefix2.isEmpty()
+    )
+  )
 }

 pragma[nomagic]
@@ -563,6 +588,9 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
        ppos.isImplicit() and
        result.(AssociatedTypeTypeParameter).getTrait() = trait
      )
+      or
+      ppos.isImplicit() and
+      this = result.(ImplTraitTypeTypeParameter).getFunction()
    }

    override Type getParameterType(DeclarationPosition dpos, TypePath path) {
@@ -582,9 +610,22 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
      )
    }

-    override Type getReturnType(TypePath path) {
+    private Type resolveRetType(TypePath path) {
      result = this.getRetType().getTypeRepr().(TypeMention).resolveTypeAt(path)
    }
+
+    override Type getReturnType(TypePath path) {
+      if this.isAsync()
+      then
+        path.isEmpty() and
+        result = getFutureTraitType()
+        or
+        exists(TypePath suffix |
+          result = this.resolveRetType(suffix) and
+          path = TypePath::cons(getFutureOutputTypeParameter(), suffix)
+        )
+      else result = this.resolveRetType(path)
+    }
  }

  private predicate argPos(CallExprBase call, Expr e, int pos, boolean isMethodCall) {
@@ -1010,6 +1051,113 @@ private StructType inferLiteralType(LiteralExpr le) {
  )
 }

+pragma[nomagic]
+private TraitType getFutureTraitType() { result.getTrait() instanceof FutureTrait }
+
+pragma[nomagic]
+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.getParamList().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>;
+
+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)
+  )
+  or
+  // This case is needed for `async` functions and blocks, where we assign
+  // the type `Future<Output = T>` directly instead of `impl Future<Output = T>`
+  //
+  // TODO: It would be better if we could handle this in the shared library
+  exists(TypePath exprPath |
+    result = inferType(n.(AwaitExpr).getExpr(), exprPath) and
+    exprPath.isCons(getFutureOutputTypeParameter(), path)
+  )
+}
+
 private module MethodCall {
  /** An expression that calls a method. */
  abstract private class MethodCallImpl extends Expr {
@@ -1119,12 +1267,17 @@ private predicate methodCandidateTrait(Type type, Trait trait, string name, int
 }

 private module IsInstantiationOfInput implements IsInstantiationOfInputSig<MethodCall> {
+  pragma[nomagic]
+  private predicate isMethodCall(MethodCall mc, Type rootType, string name, int arity) {
+    rootType = mc.getTypeAt(TypePath::nil()) and
+    name = mc.getMethodName() and
+    arity = mc.getArity()
+  }
+
  pragma[nomagic]
  predicate potentialInstantiationOf(MethodCall mc, TypeAbstraction impl, TypeMention constraint) {
    exists(Type rootType, string name, int arity |
-      rootType = mc.getTypeAt(TypePath::nil()) and
-      name = mc.getMethodName() and
-      arity = mc.getArity() and
+      isMethodCall(mc, rootType, name, arity) and
      constraint = impl.(ImplTypeAbstraction).getSelfTy()
    |
      methodCandidateTrait(rootType, mc.getTrait(), name, arity, impl)
@@ -1151,6 +1304,8 @@ private Function getTypeParameterMethod(TypeParameter tp, string name) {
  result = getMethodSuccessor(tp.(TypeParamTypeParameter).getTypeParam(), name)
  or
  result = getMethodSuccessor(tp.(SelfTypeParameter).getTrait(), name)
+  or
+  result = getMethodSuccessor(tp.(ImplTraitTypeTypeParameter).getImplTraitTypeRepr(), name)
 }

 /** Gets a method from an `impl` block that matches the method call `mc`. */
@@ -1161,6 +1316,12 @@ private Function getMethodFromImpl(MethodCall mc) {
  )
 }

+bindingset[trait, name]
+pragma[inline_late]
+private Function getTraitMethod(ImplTraitReturnType trait, string name) {
+  result = getMethodSuccessor(trait.getImplTraitTypeRepr(), name)
+}
+
 /**
 * Gets a method that the method call `mc` resolves to based on type inference,
 * if any.
@@ -1172,6 +1333,9 @@ private Function inferMethodCallTarget(MethodCall mc) {
  // The type of the receiver is a type parameter and the method comes from a
  // trait bound on the type parameter.
  result = getTypeParameterMethod(mc.getTypeAt(TypePath::nil()), mc.getMethodName())
+  or
+  // The type of the receiver is an `impl Trait` type.
+  result = getTraitMethod(mc.getTypeAt(TypePath::nil()), mc.getMethodName())
 }

 cached
@@ -1347,6 +1511,11 @@ private module Cached {
    or
    result = inferLiteralType(n) and
    path.isEmpty()
+    or
+    result = inferAsyncBlockExprRootType(n) and
+    path.isEmpty()
+    or
+    result = inferAwaitExprType(n, path)
  }
 }

@@ -1363,7 +1532,7 @@ private module Debug {
    exists(string filepath, int startline, int startcolumn, int endline, int endcolumn |
      result.getLocation().hasLocationInfo(filepath, startline, startcolumn, endline, endcolumn) and
      filepath.matches("%/main.rs") and
-      startline = 948
+      startline = 1718
    )
  }

--- a/rust/ql/lib/codeql/rust/internal/TypeMention.qll
+++ b/rust/ql/lib/codeql/rust/internal/TypeMention.qll
@@ -15,7 +15,7 @@ abstract class TypeMention extends AstNode {

  /** Gets the sub mention at `path`. */
  pragma[nomagic]
-  private TypeMention getMentionAt(TypePath path) {
+  TypeMention getMentionAt(TypePath path) {
    path.isEmpty() and
    result = this
    or
@@ -56,6 +56,29 @@ class PathTypeReprMention extends TypeMention instanceof PathTypeRepr {

  ItemNode getResolved() { result = resolved }

+  pragma[nomagic]
+  private TypeAlias getResolvedTraitAlias(string name) {
+    exists(TraitItemNode trait |
+      trait = resolvePath(path) and
+      result = trait.getAnAssocItem() and
+      name = result.getName().getText()
+    )
+  }
+
+  pragma[nomagic]
+  private TypeRepr getAssocTypeArg(string name) {
+    result = path.getSegment().getGenericArgList().getAssocTypeArg(name)
+  }
+
+  /** Gets the type argument for the associated type `alias`, if any. */
+  pragma[nomagic]
+  private TypeRepr getAnAssocTypeArgument(TypeAlias alias) {
+    exists(string name |
+      alias = this.getResolvedTraitAlias(name) and
+      result = this.getAssocTypeArg(name)
+    )
+  }
+
  override TypeMention getTypeArgument(int i) {
    result = path.getSegment().getGenericArgList().getTypeArg(i)
    or
@@ -96,6 +119,11 @@ class PathTypeReprMention extends TypeMention instanceof PathTypeRepr {
      result = alias.getTypeRepr() and
      param.getIndex() = i
    )
+    or
+    exists(TypeAlias alias |
+      result = this.getAnAssocTypeArgument(alias) and
+      traitAliasIndex(_, i, alias)
+    )
  }

  /**
@@ -152,6 +180,12 @@ class PathTypeReprMention extends TypeMention instanceof PathTypeRepr {
  }
 }

+class ImplTraitTypeReprMention extends TypeMention instanceof ImplTraitTypeRepr {
+  override TypeMention getTypeArgument(int i) { none() }
+
+  override ImplTraitType resolveType() { result.getImplTraitTypeRepr() = this }
+}
+
 private TypeParameter pathGetTypeParameter(TypeAlias alias, int i) {
  result = TTypeParamTypeParameter(alias.getGenericParamList().getTypeParam(i))
 }
--- a/rust/ql/test/library-tests/type-inference/main.rs
+++ b/rust/ql/test/library-tests/type-inference/main.rs
@@ -1630,6 +1630,110 @@ mod overloadable_operators {
    }
 }

+mod async_ {
+    use std::future::Future;
+
+    struct S1;
+
+    impl S1 {
+        pub fn f(self) {} // S1f
+    }
+
+    async fn f1() -> S1 {
+        S1
+    }
+
+    fn f2() -> impl Future<Output = S1> {
+        async {
+            S1
+        }
+    }
+
+    struct S2;
+
+    impl Future for S2 {
+        type Output = S1;
+
+        fn poll(self: std::pin::Pin<&mut Self>, _cx: &mut std::task::Context<'_>) -> std::task::Poll<Self::Output> {
+            std::task::Poll::Ready(S1)
+        }
+    }
+
+    fn f3() -> impl Future<Output = S1> {
+        S2
+    }
+
+    pub async fn f() {
+        f1().await.f(); // $ method=S1f
+        f2().await.f(); // $ method=S1f
+        f3().await.f(); // $ method=S1f
+        S2.await.f(); // $ method=S1f
+        let b = async {
+            S1
+        };
+        b.await.f(); // $ method=S1f
+    }
+}
+
+
+mod impl_trait {
+    struct S1;
+    struct S2;
+
+    trait Trait1 {
+        fn f1(&self) {} // Trait1f1
+    }
+
+    trait Trait2 {
+        fn f2(&self) {} // Trait2f2
+    }
+
+    impl Trait1 for S1 {
+        fn f1(&self) {} // S1f1
+    }
+
+    impl Trait2 for S1 {
+        fn f2(&self) {} // S1f2
+    }
+
+    fn f1() -> impl Trait1 + Trait2 {
+        S1
+    }
+
+    trait MyTrait<A> {
+        fn get_a(&self) -> A; // MyTrait::get_a
+    }
+
+    impl MyTrait<S2> for S1 {
+        fn get_a(&self) -> S2 {
+            S2
+        }
+    }
+
+    fn get_a_my_trait() -> impl MyTrait<S2> {
+        S1
+    }
+
+    fn uses_my_trait1<A, B: MyTrait<A>>(t: B) -> A {
+        t.get_a() // $ method=MyTrait::get_a
+    }
+
+    fn uses_my_trait2<A>(t: impl MyTrait<A>) -> A {
+        t.get_a() // $ method=MyTrait::get_a
+    }
+
+    pub fn f() {
+        let x = f1();
+        x.f1(); // $ method=Trait1f1
+        x.f2(); // $ method=Trait2f2
+        let a = get_a_my_trait();
+        let b = uses_my_trait1(a); // $ type=b:S2
+        let a = get_a_my_trait();
+        let c = uses_my_trait2(a); // $ type=c:S2
+        let d = uses_my_trait2(S1); // $ type=d:S2
+    }
+}
+
 fn main() {
    field_access::f();
    method_impl::f();
@@ -1649,4 +1753,6 @@ fn main() {
    try_expressions::f();
    builtins::f();
    operators::f();
+    async_::f();
+    impl_trait::f();
 }
--- a/rust/ql/test/library-tests/type-inference/type-inference.expected
+++ b/rust/ql/test/library-tests/type-inference/type-inference.expected
@@ -2374,8 +2374,97 @@ inferType
 | main.rs:1629:13:1629:20 | vec2_not |  | main.rs:1278:5:1283:5 | Vec2 |
 | main.rs:1629:24:1629:26 | ! ... |  | main.rs:1278:5:1283:5 | Vec2 |
 | main.rs:1629:25:1629:26 | v1 |  | main.rs:1278:5:1283:5 | Vec2 |
-| main.rs:1635:5:1635:20 | ...::f(...) |  | main.rs:67:5:67:21 | Foo |
-| main.rs:1636:5:1636:60 | ...::g(...) |  | main.rs:67:5:67:21 | Foo |
-| main.rs:1636:20:1636:38 | ...::Foo {...} |  | main.rs:67:5:67:21 | Foo |
-| main.rs:1636:41:1636:59 | ...::Foo {...} |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1639:18:1639:21 | SelfParam |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1642:25:1644:5 | { ... } |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1643:9:1643:10 | S1 |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1646:41:1650:5 | { ... } |  | {EXTERNAL LOCATION} | trait Future |
+| main.rs:1646:41:1650:5 | { ... } |  | main.rs:1646:16:1646:39 | ImplTraitTypeRepr |
+| main.rs:1646:41:1650:5 | { ... } | Output | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1647:9:1649:9 | { ... } |  | {EXTERNAL LOCATION} | trait Future |
+| main.rs:1647:9:1649:9 | { ... } |  | main.rs:1646:16:1646:39 | ImplTraitTypeRepr |
+| main.rs:1647:9:1649:9 | { ... } | Output | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1648:13:1648:14 | S1 |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1657:17:1657:46 | SelfParam |  | {EXTERNAL LOCATION} | Pin |
+| main.rs:1657:17:1657:46 | SelfParam | Ptr | file://:0:0:0:0 | & |
+| main.rs:1657:17:1657:46 | SelfParam | Ptr.&T | main.rs:1652:5:1652:14 | S2 |
+| main.rs:1657:49:1657:51 | _cx |  | file://:0:0:0:0 | & |
+| main.rs:1657:49:1657:51 | _cx | &T | {EXTERNAL LOCATION} | Context |
+| main.rs:1657:116:1659:9 | { ... } |  | {EXTERNAL LOCATION} | Poll |
+| main.rs:1657:116:1659:9 | { ... } | T | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1658:13:1658:38 | ...::Ready(...) |  | {EXTERNAL LOCATION} | Poll |
+| main.rs:1658:13:1658:38 | ...::Ready(...) | T | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1658:36:1658:37 | S1 |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1662:41:1664:5 | { ... } |  | main.rs:1652:5:1652:14 | S2 |
+| main.rs:1662:41:1664:5 | { ... } |  | main.rs:1662:16:1662:39 | ImplTraitTypeRepr |
+| main.rs:1663:9:1663:10 | S2 |  | main.rs:1652:5:1652:14 | S2 |
+| main.rs:1663:9:1663:10 | S2 |  | main.rs:1662:16:1662:39 | ImplTraitTypeRepr |
+| main.rs:1667:9:1667:12 | f1(...) |  | {EXTERNAL LOCATION} | trait Future |
+| main.rs:1667:9:1667:12 | f1(...) | Output | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1667:9:1667:18 | await ... |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1668:9:1668:12 | f2(...) |  | main.rs:1646:16:1646:39 | ImplTraitTypeRepr |
+| main.rs:1668:9:1668:18 | await ... |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1669:9:1669:12 | f3(...) |  | main.rs:1662:16:1662:39 | ImplTraitTypeRepr |
+| main.rs:1669:9:1669:18 | await ... |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1670:9:1670:10 | S2 |  | main.rs:1652:5:1652:14 | S2 |
+| main.rs:1670:9:1670:16 | await S2 |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1671:13:1671:13 | b |  | {EXTERNAL LOCATION} | trait Future |
+| main.rs:1671:13:1671:13 | b | Output | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1671:17:1673:9 | { ... } |  | {EXTERNAL LOCATION} | trait Future |
+| main.rs:1671:17:1673:9 | { ... } | Output | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1672:13:1672:14 | S1 |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1674:9:1674:9 | b |  | {EXTERNAL LOCATION} | trait Future |
+| main.rs:1674:9:1674:9 | b | Output | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1674:9:1674:15 | await b |  | main.rs:1636:5:1636:14 | S1 |
+| main.rs:1684:15:1684:19 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1684:15:1684:19 | SelfParam | &T | main.rs:1683:5:1685:5 | Self [trait Trait1] |
+| main.rs:1688:15:1688:19 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1688:15:1688:19 | SelfParam | &T | main.rs:1687:5:1689:5 | Self [trait Trait2] |
+| main.rs:1692:15:1692:19 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1692:15:1692:19 | SelfParam | &T | main.rs:1680:5:1680:14 | S1 |
+| main.rs:1696:15:1696:19 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1696:15:1696:19 | SelfParam | &T | main.rs:1680:5:1680:14 | S1 |
+| main.rs:1699:37:1701:5 | { ... } |  | main.rs:1680:5:1680:14 | S1 |
+| main.rs:1699:37:1701:5 | { ... } |  | main.rs:1699:16:1699:35 | ImplTraitTypeRepr |
+| main.rs:1700:9:1700:10 | S1 |  | main.rs:1680:5:1680:14 | S1 |
+| main.rs:1700:9:1700:10 | S1 |  | main.rs:1699:16:1699:35 | ImplTraitTypeRepr |
+| main.rs:1704:18:1704:22 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1704:18:1704:22 | SelfParam | &T | main.rs:1703:5:1705:5 | Self [trait MyTrait] |
+| main.rs:1708:18:1708:22 | SelfParam |  | file://:0:0:0:0 | & |
+| main.rs:1708:18:1708:22 | SelfParam | &T | main.rs:1680:5:1680:14 | S1 |
+| main.rs:1708:31:1710:9 | { ... } |  | main.rs:1681:5:1681:14 | S2 |
+| main.rs:1709:13:1709:14 | S2 |  | main.rs:1681:5:1681:14 | S2 |
+| main.rs:1713:45:1715:5 | { ... } |  | main.rs:1680:5:1680:14 | S1 |
+| main.rs:1713:45:1715:5 | { ... } |  | main.rs:1713:28:1713:43 | ImplTraitTypeRepr |
+| main.rs:1714:9:1714:10 | S1 |  | main.rs:1680:5:1680:14 | S1 |
+| main.rs:1714:9:1714:10 | S1 |  | main.rs:1713:28:1713:43 | ImplTraitTypeRepr |
+| main.rs:1717:41:1717:41 | t |  | main.rs:1717:26:1717:38 | B |
+| main.rs:1717:52:1719:5 | { ... } |  | main.rs:1717:23:1717:23 | A |
+| main.rs:1718:9:1718:9 | t |  | main.rs:1717:26:1717:38 | B |
+| main.rs:1718:9:1718:17 | t.get_a() |  | main.rs:1717:23:1717:23 | A |
+| main.rs:1721:26:1721:26 | t |  | main.rs:1721:29:1721:43 | ImplTraitTypeRepr |
+| main.rs:1721:51:1723:5 | { ... } |  | main.rs:1721:23:1721:23 | A |
+| main.rs:1722:9:1722:9 | t |  | main.rs:1721:29:1721:43 | ImplTraitTypeRepr |
+| main.rs:1722:9:1722:17 | t.get_a() |  | main.rs:1721:23:1721:23 | A |
+| main.rs:1726:13:1726:13 | x |  | main.rs:1699:16:1699:35 | ImplTraitTypeRepr |
+| main.rs:1726:17:1726:20 | f1(...) |  | main.rs:1699:16:1699:35 | ImplTraitTypeRepr |
+| main.rs:1727:9:1727:9 | x |  | main.rs:1699:16:1699:35 | ImplTraitTypeRepr |
+| main.rs:1728:9:1728:9 | x |  | main.rs:1699:16:1699:35 | ImplTraitTypeRepr |
+| main.rs:1729:13:1729:13 | a |  | main.rs:1713:28:1713:43 | ImplTraitTypeRepr |
+| main.rs:1729:17:1729:32 | get_a_my_trait(...) |  | main.rs:1713:28:1713:43 | ImplTraitTypeRepr |
+| main.rs:1730:13:1730:13 | b |  | main.rs:1681:5:1681:14 | S2 |
+| main.rs:1730:17:1730:33 | uses_my_trait1(...) |  | main.rs:1681:5:1681:14 | S2 |
+| main.rs:1730:32:1730:32 | a |  | main.rs:1713:28:1713:43 | ImplTraitTypeRepr |
+| main.rs:1731:13:1731:13 | a |  | main.rs:1713:28:1713:43 | ImplTraitTypeRepr |
+| main.rs:1731:17:1731:32 | get_a_my_trait(...) |  | main.rs:1713:28:1713:43 | ImplTraitTypeRepr |
+| main.rs:1732:13:1732:13 | c |  | main.rs:1681:5:1681:14 | S2 |
+| main.rs:1732:17:1732:33 | uses_my_trait2(...) |  | main.rs:1681:5:1681:14 | S2 |
+| main.rs:1732:32:1732:32 | a |  | main.rs:1713:28:1713:43 | ImplTraitTypeRepr |
+| main.rs:1733:13:1733:13 | d |  | main.rs:1681:5:1681:14 | S2 |
+| main.rs:1733:17:1733:34 | uses_my_trait2(...) |  | main.rs:1681:5:1681:14 | S2 |
+| main.rs:1733:32:1733:33 | S1 |  | main.rs:1680:5:1680:14 | S1 |
+| main.rs:1739:5:1739:20 | ...::f(...) |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1740:5:1740:60 | ...::g(...) |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1740:20:1740:38 | ...::Foo {...} |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1740:41:1740:59 | ...::Foo {...} |  | main.rs:67:5:67:21 | Foo |
+| main.rs:1756:5:1756:15 | ...::f(...) |  | {EXTERNAL LOCATION} | trait Future |
 testFailures
--- a/rust/ql/test/query-tests/security/CWE-022/CONSISTENCY/PathResolutionConsistency.expected
+++ b/rust/ql/test/query-tests/security/CWE-022/CONSISTENCY/PathResolutionConsistency.expected
@@ -4,6 +4,8 @@ multiplePathResolutions
 | src/main.rs:8:21:8:33 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:8:21:8:33 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:8:21:8:33 | ...::from | file://:0:0:0:0 | fn from |
+| src/main.rs:8:21:8:33 | ...::from | file://:0:0:0:0 | fn from |
+| src/main.rs:19:21:19:33 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:19:21:19:33 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:19:21:19:33 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:19:21:19:33 | ...::from | file://:0:0:0:0 | fn from |
@@ -14,6 +16,8 @@ multiplePathResolutions
 | src/main.rs:25:23:25:35 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:25:23:25:35 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:25:23:25:35 | ...::from | file://:0:0:0:0 | fn from |
+| src/main.rs:25:23:25:35 | ...::from | file://:0:0:0:0 | fn from |
+| src/main.rs:26:38:26:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:26:38:26:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:26:38:26:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:26:38:26:50 | ...::from | file://:0:0:0:0 | fn from |
@@ -24,6 +28,8 @@ multiplePathResolutions
 | src/main.rs:39:23:39:35 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:39:23:39:35 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:39:23:39:35 | ...::from | file://:0:0:0:0 | fn from |
+| src/main.rs:39:23:39:35 | ...::from | file://:0:0:0:0 | fn from |
+| src/main.rs:40:38:40:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:40:38:40:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:40:38:40:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:40:38:40:50 | ...::from | file://:0:0:0:0 | fn from |
@@ -34,6 +40,8 @@ multiplePathResolutions
 | src/main.rs:52:23:52:35 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:52:23:52:35 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:52:23:52:35 | ...::from | file://:0:0:0:0 | fn from |
+| src/main.rs:52:23:52:35 | ...::from | file://:0:0:0:0 | fn from |
+| src/main.rs:53:38:53:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:53:38:53:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:53:38:53:50 | ...::from | file://:0:0:0:0 | fn from |
 | src/main.rs:53:38:53:50 | ...::from | file://:0:0:0:0 | fn from |
--- a/rust/tools/builtins/await.rs
+++ b/rust/tools/builtins/await.rs
@@ -0,0 +1,7 @@
+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."
+    );
+}