Rust: Associated types

rust/ql/lib/codeql/rust/elements/internal/TraitImpl.qll

@@ -26,5 +26,15 @@ module Impl {
    */
   class Trait extends Generated::Trait {
     override string toStringImpl() { result = "trait " + this.getName().getText() }
+
+    /**
+     * Gets the number of generic parameters of this trait.
+     */
+    int getNumberOfGenericParams() {
+      result = this.getGenericParamList().getNumberOfGenericParams()
+      or
+      not this.hasGenericParamList() and
+      result = 0
+    }
   }
 }

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

@@ -2,9 +2,10 @@
 
 private import rust
 private import PathResolution
-private import TypeInference
 private import TypeMention
 private import codeql.rust.internal.CachedStages
+private import codeql.rust.elements.internal.generated.Raw
+private import codeql.rust.elements.internal.generated.Synth
 
 cached
 newtype TType =
@@ -15,6 +16,7 @@ newtype TType =
   TArrayType() or // todo: add size?
   TRefType() or // todo: add mut?
   TTypeParamTypeParameter(TypeParam t) or
+  TAssociatedTypeTypeParameter(TypeAlias t) { any(TraitItemNode trait).getAnAssocItem() = t } or
   TRefTypeParameter() or
   TSelfTypeParameter(Trait t)
 
@@ -144,6 +146,9 @@ class TraitType extends Type, TTrait {
 
   override TypeParameter getTypeParameter(int i) {
     result = TTypeParamTypeParameter(trait.getGenericParamList().getTypeParam(i))
+    or
+    result =
+      any(AssociatedTypeTypeParameter param | param.getTrait() = trait and param.getIndex() = i)
   }
 
   pragma[nomagic]
@@ -297,6 +302,14 @@ abstract class TypeParameter extends Type {
   override TypeParameter getTypeParameter(int i) { none() }
 }
 
+private class RawTypeParameter = @type_param or @trait or @type_alias;
+
+private predicate id(RawTypeParameter x, RawTypeParameter y) { x = y }
+
+private predicate idOfRaw(RawTypeParameter x, int y) = equivalenceRelation(id/2)(x, y)
+
+int idOfTypeParameterAstNode(AstNode node) { idOfRaw(Synth::convertAstNodeToRaw(node), result) }
+
 /** A type parameter from source code. */
 class TypeParamTypeParameter extends TypeParameter, TTypeParamTypeParameter {
   private TypeParam typeParam;
@@ -320,6 +333,59 @@ class TypeParamTypeParameter extends TypeParameter, TTypeParamTypeParameter {
   }
 }
 
+/**
+ * Gets the type alias that is the `i`th type parameter of `trait`. Type aliases
+ * are numbered consecutively but in arbitrary order, starting from the index
+ * following the last ordinary type parameter.
+ */
+predicate traitAliasIndex(Trait trait, int i, TypeAlias typeAlias) {
+  typeAlias =
+    rank[i + 1 - trait.getNumberOfGenericParams()](TypeAlias alias |
+      trait.(TraitItemNode).getADescendant() = alias
+    |
+      alias order by idOfTypeParameterAstNode(alias)
+    )
+}
+
+/**
+ * A type parameter corresponding to an associated type in a trait.
+ *
+ * We treat associated type declarations in traits as type parameters. E.g., a
+ * trait such as
+ * ```rust
+ * trait ATrait {
+ *   type AssociatedType;
+ *   // ...
+ * }
+ * ```
+ * is treated as if it was
+ * ```rust
+ * trait ATrait<AssociatedType> {
+ *   // ...
+ * }
+ * ```
+ */
+class AssociatedTypeTypeParameter extends TypeParameter, TAssociatedTypeTypeParameter {
+  private TypeAlias typeAlias;
+
+  AssociatedTypeTypeParameter() { this = TAssociatedTypeTypeParameter(typeAlias) }
+
+  TypeAlias getTypeAlias() { result = typeAlias }
+
+  /** Gets the trait that contains this associated type declaration. */
+  TraitItemNode getTrait() { result.getAnAssocItem() = typeAlias }
+
+  int getIndex() { traitAliasIndex(_, result, typeAlias) }
+
+  override Function getMethod(string name) { none() }
+
+  override string toString() { result = typeAlias.getName().getText() }
+
+  override Location getLocation() { result = typeAlias.getLocation() }
+
+  override TypeMention getABaseTypeMention() { none() }
+}
+
 /** An implicit reference type parameter. */
 class RefTypeParameter extends TypeParameter, TRefTypeParameter {
   override Function getMethod(string name) { none() }

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

@@ -40,17 +40,21 @@ private module Input1 implements InputSig1<Location> {
 
   private newtype TTypeParameterPosition =
     TTypeParamTypeParameterPosition(TypeParam tp) or
-    TSelfTypeParameterPosition()
+    TImplicitTypeParameterPosition()
 
   class TypeParameterPosition extends TTypeParameterPosition {
     TypeParam asTypeParam() { this = TTypeParamTypeParameterPosition(result) }
 
-    predicate isSelf() { this = TSelfTypeParameterPosition() }
+    /**
+     * Holds if this is the implicit type parameter position used to represent
+     * parameters that are never passed explicitly as arguments.
+     */
+    predicate isImplicit() { this = TImplicitTypeParameterPosition() }
 
     string toString() {
       result = this.asTypeParam().toString()
       or
-      result = "Self" and this.isSelf()
+      result = "Implicit" and this.isImplicit()
     }
   }
 
@@ -69,15 +73,6 @@ private module Input1 implements InputSig1<Location> {
     apos.asMethodTypeArgumentPosition() = ppos.asTypeParam().getPosition()
   }
 
-  /** A raw AST node that might correspond to a type parameter. */
-  private class RawTypeParameter = @type_param or @trait;
-
-  private predicate id(RawTypeParameter x, RawTypeParameter y) { x = y }
-
-  private predicate idOfRaw(RawTypeParameter x, int y) = equivalenceRelation(id/2)(x, y)
-
-  private int idOf(AstNode node) { idOfRaw(Synth::convertAstNodeToRaw(node), result) }
-
   int getTypeParameterId(TypeParameter tp) {
     tp =
       rank[result](TypeParameter tp0, int kind, int id |
@@ -86,8 +81,9 @@ private module Input1 implements InputSig1<Location> {
         id = 0
         or
         kind = 1 and
-        exists(AstNode node | id = idOf(node) |
+        exists(AstNode node | id = idOfTypeParameterAstNode(node) |
           node = tp0.(TypeParamTypeParameter).getTypeParam() or
+          node = tp0.(AssociatedTypeTypeParameter).getTypeAlias() or
           node = tp0.(SelfTypeParameter).getTrait()
         )
       |
@@ -500,7 +496,10 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
       exists(TraitItemNode trait | this = trait.getAnAssocItem() |
         typeParamMatchPosition(trait.getTypeParam(_), result, ppos)
         or
-        ppos.isSelf() and result = TSelfTypeParameter(trait)
+        ppos.isImplicit() and result = TSelfTypeParameter(trait)
+        or
+        ppos.isImplicit() and
+        result.(AssociatedTypeTypeParameter).getTrait() = trait
       )
     }
 

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

@@ -95,6 +95,27 @@ class NonAliasPathMention extends PathMention {
       this = node.getASelfPath() and
       result = node.(ImplItemNode).getSelfPath().getSegment().getGenericArgList().getTypeArg(i)
     )
+    or
+    // If `this` is the trait of an `impl` block then any associated types
+    // defined in the `impl` block are type arguments to the trait.
+    //
+    // For instance, for a trait implementation like this
+    // ```rust
+    // impl MyTrait for MyType {
+    //      ^^^^^^^ this
+    //   type AssociatedType = i64
+    //                         ^^^ result
+    //   // ...
+    // }
+    // ```
+    // the rhs. of the type alias is a type argument to the trait.
+    exists(ImplItemNode impl, AssociatedTypeTypeParameter param, TypeAlias alias |
+      this = impl.getTraitPath() and
+      param.getTrait() = resolvePath(this) and
+      alias = impl.getASuccessor(param.getTypeAlias().getName().getText()) and
+      result = alias.getTypeRepr() and
+      param.getIndex() = i
+    )
   }
 
   override Type resolveType() {
@@ -113,7 +134,11 @@ class NonAliasPathMention extends PathMention {
       or
       result = TTypeParamTypeParameter(i)
       or
-      result = i.(TypeAlias).getTypeRepr().(TypeReprMention).resolveType()
+      exists(TypeAlias alias | alias = i |
+        result.(AssociatedTypeTypeParameter).getTypeAlias() = alias
+        or
+        result = alias.getTypeRepr().(TypeReprMention).resolveType()
+      )
     )
   }
 }
@@ -153,6 +178,17 @@ class TypeParamMention extends TypeMention, TypeParam {
   override Type resolveType() { result = TTypeParamTypeParameter(this) }
 }
 
+// Used to represent implicit type arguments for associated types in traits.
+class TypeAliasMention extends TypeMention, TypeAlias {
+  private Type t;
+
+  TypeAliasMention() { t = TAssociatedTypeTypeParameter(this) }
+
+  override TypeReprMention getTypeArgument(int i) { none() }
+
+  override Type resolveType() { result = t }
+}
+
 /**
  * Holds if the `i`th type argument of `selfPath`, belonging to `impl`, resolves
  * to type parameter `tp`.
@@ -204,7 +240,11 @@ class ImplMention extends TypeMention, ImplItemNode {
 }
 
 class TraitMention extends TypeMention, TraitItemNode {
-  override TypeMention getTypeArgument(int i) { result = this.getTypeParam(i) }
+  override TypeMention getTypeArgument(int i) {
+    result = this.getTypeParam(i)
+    or
+    traitAliasIndex(this, i, result)
+  }
 
   override Type resolveType() { result = TTrait(this) }
 }