Rust: Type inference for `.await` expressions #19584

hvitved · 2025-05-26T19:31:28Z

Adds type inference support for .await expressions. Since functions can be either implicitly asynchronous, using the async keyword, or explicitly using an impl Future<Output = ...> return type, this PR also add support for associated type argument (Output = ...) and impl Trait types.

Copilot

Pull Request Overview

This PR enhances the Rust type inference engine to support .await expressions, associated type arguments (Output = ...), and impl Trait types.

Add new async_ and impl_trait tests exercising .await and impl Trait inference.
Update expected outputs for type‐inference tests to include await and impl Trait cases.
Extend QL libraries (TypeMention.qll, TypeInference.qll, Type.qll, Stdlib.qll) to resolve associated types, impl Trait parameters, and infer await expressions.

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
rust/ql/test/library-tests/type-inference/main.rs	Add `async_` and `impl_trait` modules and invoke new test functions
rust/ql/test/library-tests/type-inference/type-inference.expected	Update expected inference output to cover `.await` and `impl Trait` cases
rust/ql/lib/codeql/rust/internal/TypeMention.qll	Support associated type arguments and `impl Trait` mentions in types
rust/ql/lib/codeql/rust/internal/TypeInference.qll	Implement `inferAwaitExprType`, extend method resolution for `impl Trait`
rust/ql/lib/codeql/rust/internal/Type.qll	Define `TImplTraitType`, `TImplTraitTypeParameter`, and related classes
rust/ql/lib/codeql/rust/frameworks/stdlib/Stdlib.qll	Introduce `FutureTrait` with its `Output` associated type

Comments suppressed due to low confidence (2)

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

[nitpick] This only resolves the first impl Trait bound. For multi-bound impl Trait types, consider iterating over all TImplTraitTypeParameter(_, _) parameters so methods on any bound are discovered.

result = getTraitMethod(mc.getTypeAt(TypePath::singleton(TImplTraitTypeParameter(_, _))), mc.getMethodName())

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

In the typeEquality predicate, matching n2 = any(BlockExpr be | ...) could match unrelated block expressions. It should constrain be to the actual n2 (e.g. be = n2.(BlockExpr)) to avoid spurious matches.

n2 =

Copilot · 2025-05-28T06:34:18Z

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

@@ -77,6 +77,16 @@ private module Input1 implements InputSig1<Location> {
    apos.asMethodTypeArgumentPosition() = ppos.asTypeParam().getPosition()
  }

+  private int getImplTraitTypeParameterId(ImplTraitTypeParameter tp) {


The getImplTraitTypeParameterId function uses rank[result] but never assigns result. To return the intended integer ID, add result = id after the rank expression or switch to rank[id].

paldepind

Since impl types implement traits we should have a case for them in conditionSatisfiesConstraint. It would also be great to have a test covering that. Something like the following should pass with impl added to conditionSatisfiesConstraint:

trait MyTrait<A> {
    fn get_a(&self) -> A;
}

impl MyTrait<S2> for S1 {
    fn get_a(&self) -> S2 {
        S2
    }
}

fn get_a_my_trait() -> impl MyTrait<S2> {
    S1
}

fn uses_my_trait<A, B: MyTrait<A>>(t: B) -> A {
    t.get_a() // $ method=MyTrait::get_a
}

fn test() {
    let a = get_a_my_trait();
    let b = uses_my_trait(a); // $ type=b:S2
}

paldepind · 2025-05-28T09:01:40Z

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

+ * We represent `impl Trait` types as generic types with as many type parameters
+ * as there are bounds.
+ *
+ * [1] https://doc.rust-lang.org/book/ch10-02-traits.html#traits-as-parameters


I think it's better to link to the reference than to the book.

Suggested change

* [1] https://doc.rust-lang.org/book/ch10-02-traits.html#traits-as-parameters

* [1] https://doc.rust-lang.org/reference/types/impl-trait.html

paldepind · 2025-05-28T09:22:53Z

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

+  TImplTraitType(int bounds) {
+    bounds = any(ImplTraitTypeRepr impl).getTypeBoundList().getNumberOfBounds()
+  } or


If we have two functions like

fn f() -> impl MyTrait { .. } fn g() -> impl MyTrait { .. }

then they return two different types (in the general case) even though they are syntactically identical.

Given that wouldn't it be more accurate to represent impl types like:

Suggested change

TImplTraitType(int bounds) {

bounds = any(ImplTraitTypeRepr impl).getTypeBoundList().getNumberOfBounds()

} or

TImplTraitType(ImplTraitTypeRepr impl) or

I.e., every impl is it own types, similar to how every struct or enum is its own type. I think this would also simplify things a bit as the impl types wouldn't need to have type parameters.

Thanks for suggesting this, which I have implemented. I also learned that impl Trait types in parameter type positions are very different from impl Trait types in return type positions, so I have reworked the implementation to respect that (the former corresponds to a type parameter, the latter to an abstract type).

paldepind · 2025-05-28T09:33:07Z

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

+      not be.isAsync() and
+      n1 = be.getStmtList().getTailExpr() and
+      path1 = path2
+    )


Using exists here would be more consistent with the surrounding code.

paldepind · 2025-05-28T09:36:12Z

rust/ql/test/library-tests/type-inference/main.rs

+    pub async fn f() {
+        f1().await.f(); // $ method=S1f
+        f2().await.f(); // $ method=S1f
+        f3().await.f(); // $ method=S1f


Would be nice to test this case as well:

Suggested change

f3().await.f(); // $ method=S1f

f3().await.f(); // $ method=S1f

S2.await.f(); // $ MISSING: method=S1f

paldepind · 2025-05-28T09:43:13Z

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

+    not exprPath
+        .isCons(TImplTraitTypeParameter(_, _),
+          any(TypePath path0 | path0.isCons(getFutureOutputTypeParameter(), _)))
+  )


Is it correct that this covers most of the special cases that occurs in practice but doesn't handle the general case where the awaited expression is some arbitrary type that implements Future?

It might be nice to add some comments about which cases are covered and about what is not covered.

For full generality I guess that we might have to expose conditionSatisfiesConstraintTypeAt in some form or implement await using the Matching module? (EDIT: Just to be clear I'm not saying we should do that.)

I have reimplemented this.

hvitved · 2025-06-03T08:09:05Z

@paldepind : PR is ready for review again. I have rebased to resolve merge conflicts; only the last two commits need reviewing.

github-actions bot added the Rust Pull requests that update Rust code label May 26, 2025

hvitved force-pushed the rust/type-inference-await branch 2 times, most recently from 92e4659 to 25b928d Compare May 27, 2025 11:59

hvitved marked this pull request as ready for review May 28, 2025 06:32

Copilot AI review requested due to automatic review settings May 28, 2025 06:32

hvitved requested a review from a team as a code owner May 28, 2025 06:32

hvitved requested a review from paldepind May 28, 2025 06:32

Copilot AI reviewed May 28, 2025

View reviewed changes

paldepind requested changes May 28, 2025

View reviewed changes

hvitved added 4 commits June 2, 2025 19:06

Rust: Add async type inference tests

03d6528

Rust: Add type inference tests for impl Trait

15abd47

Rust: Type inference for .await expressions

61bf20c

Rust: Add more type inference tests

7e60cfc

hvitved force-pushed the rust/type-inference-await branch from 25b928d to 80d91e2 Compare June 2, 2025 19:29

hvitved added the no-change-note-required This PR does not need a change note label Jun 2, 2025

hvitved force-pushed the rust/type-inference-await branch from 80d91e2 to 14b8af8 Compare June 3, 2025 07:26

Rust: Reimplement type inference for impl Traits and await expressions

15ca62e

hvitved force-pushed the rust/type-inference-await branch from 14b8af8 to 15ca62e Compare June 3, 2025 07:50

hvitved requested a review from paldepind June 3, 2025 08:08

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Rust: Type inference for `.await` expressions #19584

Rust: Type inference for `.await` expressions #19584

Uh oh!

hvitved commented May 26, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI May 28, 2025

Uh oh!

paldepind left a comment

Uh oh!

paldepind May 28, 2025

Uh oh!

paldepind May 28, 2025 •

edited

Loading

Uh oh!

hvitved Jun 3, 2025

Uh oh!

paldepind May 28, 2025

Uh oh!

paldepind May 28, 2025

Uh oh!

paldepind May 28, 2025 •

edited

Loading

Uh oh!

hvitved Jun 3, 2025

Uh oh!

hvitved commented Jun 3, 2025

Uh oh!

Uh oh!

	* [1] https://doc.rust-lang.org/book/ch10-02-traits.html#traits-as-parameters
	* [1] https://doc.rust-lang.org/reference/types/impl-trait.html

	f3().await.f(); // $ method=S1f
	f3().await.f(); // $ method=S1f
	S2.await.f(); // $ MISSING: method=S1f

Rust: Type inference for .await expressions #19584

Are you sure you want to change the base?

Rust: Type inference for .await expressions #19584

Uh oh!

Conversation

hvitved commented May 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Copilot AI May 28, 2025

Choose a reason for hiding this comment

Uh oh!

paldepind left a comment

Choose a reason for hiding this comment

Uh oh!

paldepind May 28, 2025

Choose a reason for hiding this comment

Uh oh!

paldepind May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hvitved Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

paldepind May 28, 2025

Choose a reason for hiding this comment

Uh oh!

paldepind May 28, 2025

Choose a reason for hiding this comment

Uh oh!

paldepind May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hvitved Jun 3, 2025

Choose a reason for hiding this comment

Uh oh!

hvitved commented Jun 3, 2025

Uh oh!

Uh oh!

Rust: Type inference for `.await` expressions #19584

Rust: Type inference for `.await` expressions #19584

hvitved commented May 26, 2025 •

edited

Loading

paldepind May 28, 2025 •

edited

Loading

paldepind May 28, 2025 •

edited

Loading