Rust: Type inference for .await expressions

Author: hvitved

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/Type.qll

@@ -15,10 +15,14 @@ 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
   TTypeParamTypeParameter(TypeParam t) or
   TAssociatedTypeTypeParameter(TypeAlias t) { any(TraitItemNode trait).getAnAssocItem() = t } or
   TRefTypeParameter() or
-  TSelfTypeParameter(Trait t)
+  TSelfTypeParameter(Trait t) or
+  TImplTraitTypeParameter(ImplTraitType t, int i) { i in [0 .. t.getNumberOfBounds() - 1] }
 
 /**
  * A type without type arguments.
@@ -115,6 +119,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 +183,33 @@ class RefType extends Type, TRefType {
   override Location getLocation() { result instanceof EmptyLocation }
 }
 
+/**
+ * An [`impl Trait`][1] type.
+ *
+ * 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
+ */
+class ImplTraitType extends Type, TImplTraitType {
+  private int bounds;
+
+  ImplTraitType() { this = TImplTraitType(bounds) }
+
+  /** Gets the number of bounds of this `impl Trait` type. */
+  int getNumberOfBounds() { result = bounds }
+
+  override StructField getStructField(string name) { none() }
+
+  override TupleField getTupleField(int i) { none() }
+
+  override TypeParameter getTypeParameter(int i) { result = TImplTraitTypeParameter(this, i) }
+
+  override string toString() { result = "impl Trait ..." }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
 /** A type parameter. */
 abstract class TypeParameter extends Type {
   override StructField getStructField(string name) { none() }
@@ -281,6 +315,26 @@ class SelfTypeParameter extends TypeParameter, TSelfTypeParameter {
   override Location getLocation() { result = trait.getLocation() }
 }
 
+/**
+ *  An `impl Trait` type parameter.
+ */
+class ImplTraitTypeParameter extends TypeParameter, TImplTraitTypeParameter {
+  private ImplTraitType implTraitType;
+  private int i;
+
+  ImplTraitTypeParameter() { this = TImplTraitTypeParameter(implTraitType, i) }
+
+  /** Gets the `impl Trait` type that this parameter belongs to. */
+  ImplTraitType getImplTraitType() { result = implTraitType }
+
+  /** Gets the index of this type parameter. */
+  int getIndex() { result = i }
+
+  override string toString() { result = "impl Trait<" + i.toString() + ">" }
+
+  override Location getLocation() { result instanceof EmptyLocation }
+}
+
 /**
  * A type abstraction. I.e., a place in the program where type variables are
  * introduced.

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) {
+    tp =
+      rank[result](ImplTraitTypeParameter tp0, int bounds, int i |
+        bounds = tp0.getImplTraitType().getNumberOfBounds() and
+        i = tp0.getIndex()
+      |
+        tp0 order by bounds, i
+      )
+  }
+
   int getTypeParameterId(TypeParameter tp) {
     tp =
       rank[result](TypeParameter tp0, int kind, int id |
@@ -90,6 +100,9 @@ private module Input1 implements InputSig1<Location> {
           node = tp0.(AssociatedTypeTypeParameter).getTypeAlias() or
           node = tp0.(SelfTypeParameter).getTrait()
         )
+        or
+        kind = 2 and
+        id = getImplTraitTypeParameterId(tp0)
       |
         tp0 order by kind, id
       )
@@ -232,8 +245,12 @@ private predicate typeEquality(AstNode n1, TypePath path1, AstNode n2, TypePath
   n1 = n2.(ParenExpr).getExpr() and
   path1 = path2
   or
-  n1 = n2.(BlockExpr).getStmtList().getTailExpr() and
-  path1 = path2
+  n2 =
+    any(BlockExpr be |
+      not be.isAsync() and
+      n1 = be.getStmtList().getTailExpr() and
+      path1 = path2
+    )
   or
   n1 = n2.(IfExpr).getABranch() and
   path1 = path2
@@ -1000,6 +1017,29 @@ private StructType inferLiteralType(LiteralExpr le) {
   )
 }
 
+pragma[nomagic]
+private AssociatedTypeTypeParameter getFutureOutputTypeParameter() {
+  result.getTypeAlias() = any(FutureTrait ft).getOutputType()
+}
+
+pragma[nomagic]
+private Type inferAwaitExprType(AwaitExpr ae, TypePath path) {
+  exists(TypePath exprPath | result = inferType(ae.getExpr(), exprPath) |
+    exprPath
+        .isCons(TImplTraitTypeParameter(_, _),
+          any(TypePath path0 | path0.isCons(getFutureOutputTypeParameter(), path)))
+    or
+    path = exprPath and
+    not (
+      exprPath = TypePath::singleton(TImplTraitTypeParameter(_, _)) and
+      result.(TraitType).getTrait() instanceof FutureTrait
+    ) and
+    not exprPath
+        .isCons(TImplTraitTypeParameter(_, _),
+          any(TypePath path0 | path0.isCons(getFutureOutputTypeParameter(), _)))
+  )
+}
+
 private module MethodCall {
   /** An expression that calls a method. */
   abstract private class MethodCallImpl extends Expr {
@@ -1089,12 +1129,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)
@@ -1131,6 +1176,12 @@ private Function getMethodFromImpl(MethodCall mc) {
   )
 }
 
+bindingset[trait, name]
+pragma[inline_late]
+private Function getTraitMethod(TraitType trait, string name) {
+  result = getMethodSuccessor(trait.getTrait(), name)
+}
+
 /**
  * Gets a method that the method call `mc` resolves to based on type inference,
  * if any.
@@ -1142,6 +1193,11 @@ 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::singleton(TImplTraitTypeParameter(_, _))),
+      mc.getMethodName())
 }
 
 cached
@@ -1317,6 +1373,8 @@ private module Cached {
     or
     result = inferLiteralType(n) and
     path.isEmpty()
+    or
+    result = inferAwaitExprType(n, path)
   }
 }
 
@@ -1333,7 +1391,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 = 1334
     )
   }
 

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
@@ -44,6 +44,12 @@ class TypeReprMention extends TypeMention, TypeRepr {
     result = this.(PathTypeRepr).getPath().(PathMention).getTypeArgument(i)
   }
 
+  pragma[nomagic]
+  private TypeBoundList getImplTraitBoundList(ImplTraitType t) {
+    result = this.(ImplTraitTypeRepr).getTypeBoundList() and
+    t.getNumberOfBounds() = result.getNumberOfBounds()
+  }
+
   override Type resolveType() {
     this instanceof ArrayTypeRepr and
     result = TArrayType()
@@ -52,11 +58,20 @@ class TypeReprMention extends TypeMention, TypeRepr {
     result = TRefType()
     or
     result = this.(PathTypeRepr).getPath().(PathMention).resolveType()
+    or
+    exists(this.getImplTraitBoundList(result))
   }
 
   override Type resolveTypeAt(TypePath path) {
     result = this.(PathTypeRepr).getPath().(PathMention).resolveTypeAt(path)
     or
+    exists(TypeBoundList tbl, ImplTraitType t, int i, TypePath suffix |
+      tbl = this.getImplTraitBoundList(t) and
+      result =
+        tbl.getBound(i).getTypeRepr().(PathTypeRepr).getPath().(PathMention).resolveTypeAt(suffix) and
+      path = TypePath::cons(TImplTraitTypeParameter(t, i), suffix)
+    )
+    or
     not exists(this.(PathTypeRepr).getPath()) and
     result = super.resolveTypeAt(path)
   }
@@ -71,6 +86,44 @@ abstract class PathMention extends TypeMention, Path {
   override TypeMention getTypeArgument(int i) {
     result = this.getSegment().getGenericArgList().getTypeArg(i)
   }
+
+  pragma[nomagic]
+  private TypeAlias getResolvedTraitAlias(string name) {
+    exists(TraitItemNode trait |
+      trait = resolvePath(this) and
+      result = trait.getAnAssocItem() and
+      name = result.getName().getText()
+    )
+  }
+
+  pragma[nomagic]
+  private TypeRepr getAssocTypeArg(string name) {
+    exists(AssocTypeArg arg |
+      arg = this.getSegment().getGenericArgList().getAGenericArg() and
+      result = arg.getTypeRepr() and
+      name = arg.getIdentifier().getText()
+    )
+  }
+
+  /** 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 getMentionAt(TypePath path) {
+    result = super.getMentionAt(path)
+    or
+    exists(TypeAlias alias, AssociatedTypeTypeParameter tp, TypeMention arg, TypePath suffix |
+      arg = this.getAnAssocTypeArgument(alias) and
+      result = arg.getMentionAt(suffix) and
+      path = TypePath::cons(tp, suffix) and
+      tp.getTypeAlias() = alias
+    )
+  }
 }
 
 class NonAliasPathMention extends PathMention {

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

@@ -1299,9 +1299,9 @@ mod async_ {
     }
 
     pub async fn f() {
-        f1().await.f(); // $ MISSING: method=S1f
-        f2().await.f(); // $ MISSING: method=S1f
-        f3().await.f(); // $ MISSING: method=S1f
+        f1().await.f(); // $ method=S1f
+        f2().await.f(); // $ method=S1f
+        f3().await.f(); // $ method=S1f
     }
 }
 
@@ -1331,8 +1331,8 @@ mod impl_trait {
 
     pub fn f() {
         let x = f1();
-        x.f1(); // $ MISSING: method=Trait1f1
-        x.f2(); // $ MISSING: method=Trait2f2
+        x.f1(); // $ method=Trait1f1
+        x.f2(); // $ method=Trait2f2
     }
 }