Rust: Reimplement type inference for impl Traits and await expressions

Author: hvitved

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
@@ -165,6 +166,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)
   }
 
   /**
@@ -618,6 +621,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)" }
 
@@ -1062,14 +1087,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 `m` depends on crate `dep` named `name`.
  */
+pragma[nomagic]
 private predicate crateDependencyEdge(ModuleLikeNode m, string name, CrateItemNode dep) {
-  exists(CrateItemNode c |
-    dep = c.(Crate).getDependency(name) and
-    m = c.getASourceFile()
-  )
+  exists(CrateItemNode c | dep = c.(Crate).getDependency(name) | m = c.getASourceFile())
+  or
+  // Give builtin files, such as `await.rs`, access to `std`
+  m instanceof BuiltinSourceFile and
+  dep.getName() = name and
+  name = "std"
 }
 
 private predicate useTreeDeclares(UseTree tree, string name) {
@@ -1414,9 +1446,14 @@ private predicate useImportEdge(Use use, string name, ItemNode item) {
  * [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
@@ -1426,12 +1463,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,14 +15,11 @@ newtype TType =
   TTrait(Trait t) or
   TArrayType() or // todo: add size?
   TRefType() or // todo: add mut?
-  TImplTraitType(int bounds) {
-    bounds = any(ImplTraitTypeRepr impl).getTypeBoundList().getNumberOfBounds()
-  } or
+  TImplTraitType(ImplTraitTypeRepr impl) or
   TTypeParamTypeParameter(TypeParam t) or
   TAssociatedTypeTypeParameter(TypeAlias t) { any(TraitItemNode trait).getAnAssocItem() = t } or
   TRefTypeParameter() or
-  TSelfTypeParameter(Trait t) or
-  TImplTraitTypeParameter(ImplTraitType t, int i) { i in [0 .. t.getNumberOfBounds() - 1] }
+  TSelfTypeParameter(Trait t)
 
 /**
  * A type without type arguments.
@@ -184,30 +181,50 @@ class RefType extends Type, TRefType {
 }
 
 /**
- * An [`impl Trait`][1] type.
+ * An [impl Trait][1] type.
  *
- * We represent `impl Trait` types as generic types with as many type parameters
- * as there are bounds.
+ * 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/book/ch10-02-traits.html#traits-as-parameters
+ * [1]: https://doc.rust-lang.org/reference/types/impl-trait.html
  */
 class ImplTraitType extends Type, TImplTraitType {
-  private int bounds;
+  ImplTraitTypeRepr impl;
 
-  ImplTraitType() { this = TImplTraitType(bounds) }
+  ImplTraitType() { this = TImplTraitType(impl) }
 
-  /** Gets the number of bounds of this `impl Trait` type. */
-  int getNumberOfBounds() { result = bounds }
+  /** 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) { result = TImplTraitTypeParameter(this, i) }
+  override TypeParameter getTypeParameter(int i) { none() }
 
-  override string toString() { result = "impl Trait ..." }
+  override string toString() { result = impl.toString() }
 
-  override Location getLocation() { result instanceof EmptyLocation }
+  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. */
@@ -219,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 }
 
@@ -316,23 +333,34 @@ class SelfTypeParameter extends TypeParameter, TSelfTypeParameter {
 }
 
 /**
- *  An `impl Trait` type parameter.
+ * 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 ImplTraitTypeParameter extends TypeParameter, TImplTraitTypeParameter {
-  private ImplTraitType implTraitType;
-  private int i;
+class ImplTraitTypeTypeParameter extends ImplTraitType, TypeParameter {
+  private Function function;
 
-  ImplTraitTypeParameter() { this = TImplTraitTypeParameter(implTraitType, i) }
+  ImplTraitTypeTypeParameter() { impl = function.getParamList().getAParam().getTypeRepr() }
 
-  /** Gets the `impl Trait` type that this parameter belongs to. */
-  ImplTraitType getImplTraitType() { result = implTraitType }
+  override Function getFunction() { result = function }
 
-  /** Gets the index of this type parameter. */
-  int getIndex() { result = i }
+  override StructField getStructField(string name) { none() }
 
-  override string toString() { result = "impl Trait<" + i.toString() + ">" }
+  override TupleField getTupleField(int i) { none() }
 
-  override Location getLocation() { result instanceof EmptyLocation }
+  override TypeParameter getTypeParameter(int i) { none() }
 }
 
 /**
@@ -370,3 +398,7 @@ final class SelfTypeBoundTypeAbstraction extends TypeAbstraction, Name {
 
   override TypeParamTypeParameter getATypeParameter() { none() }
 }
+
+final class ImplTraitTypeReprAbstraction extends TypeAbstraction, ImplTraitTypeRepr {
+  override TypeParamTypeParameter getATypeParameter() { none() }
+}