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

Rust: Type inference for `.await` expressions

Author: hvitved

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()
+      )
+    }
   }
 }

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"
+  }
+}

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)
   )
 }

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() }
+}