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).getADescendant() = 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,55 @@ class TypeParamTypeParameter extends TypeParameter, TTypeParamTypeParameter {
   }
 }
 
+/** Gets type alias that is the `i`th type parameter of `trait`. */
+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 where
+ * ```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.getADescendant() = typeAlias }
+
+  int getIndex() { traitAliasIndex(this.getTrait(), 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

@@ -75,6 +75,27 @@
       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() {
@@ -93,7 +114,11 @@
       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()
+      )
     )
   }
 }
@@ -106,6 +131,13 @@
   override Type resolveType() { result = TTypeParamTypeParameter(this) }
 }
 
+// Used to represent implicit associated type type arguments in traits.
+class TypeAliasMention extends TypeMention, TypeAlias {
+  override TypeReprMention getTypeArgument(int i) { none() }
+
+  override Type resolveType() { result = TAssociatedTypeTypeParameter(this) }
+}
+
 /**
  * Holds if the `i`th type argument of `selfPath`, belonging to `impl`, resolves
  * to type parameter `tp`.
@@ -157,7 +189,11 @@
 }
 
 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) }
 }

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

@@ -329,38 +329,118 @@ mod function_trait_bounds {
 }
 
 mod trait_associated_type {
+    #[derive(Debug)]
+    struct Wrapper<A> {
+        field: A,
+    }
+
+    impl<A> Wrapper<A> {
+        fn unwrap(self) -> A {
+            self.field // $ fieldof=Wrapper
+        }
+    }
+
     trait MyTrait {
         type AssociatedType;
 
+        // MyTrait::m1
         fn m1(self) -> Self::AssociatedType;
 
         fn m2(self) -> Self::AssociatedType
         where
             Self::AssociatedType: Default,
             Self: Sized,
         {
+            self.m1(); // $ method=MyTrait::m1 type=self.m1():AssociatedType
             Self::AssociatedType::default()
         }
     }
 
+    trait MyTraitAssoc2 {
+        type GenericAssociatedType<AssociatedParam>;
+
+        // MyTrait::put
+        fn put<A>(&self, a: A) -> Self::GenericAssociatedType<A>;
+
+        fn putTwo<A>(&self, a: A, b: A) -> Self::GenericAssociatedType<A> {
+            self.put(a); // $ method=MyTrait::put
+            self.put(b) // $ method=MyTrait::put
+        }
+    }
+
     #[derive(Debug, Default)]
     struct S;
 
+    #[derive(Debug, Default)]
+    struct S2;
+
+    #[derive(Debug, Default)]
+    struct AT;
+
     impl MyTrait for S {
-        type AssociatedType = S;
+        type AssociatedType = AT;
 
         // S::m1
         fn m1(self) -> Self::AssociatedType {
-            S
+            AT
+        }
+    }
+
+    impl MyTraitAssoc2 for S {
+        // Associated type with a type parameter
+        type GenericAssociatedType<AssociatedParam> = Wrapper<AssociatedParam>;
+
+        // S::put
+        fn put<A>(&self, a: A) -> Wrapper<A> {
+            Wrapper { field: a }
+        }
+    }
+
+    impl MyTrait for S2 {
+        // Associated type definition with a type argument
+        type AssociatedType = Wrapper<S2>;
+
+        fn m1(self) -> Self::AssociatedType {
+            Wrapper { field: self }
+        }
+    }
+
+    // NOTE: This implementation is just to make it possible to call `m2` on `S2.`
+    impl Default for Wrapper<S2> {
+        fn default() -> Self {
+            Wrapper { field: S2 }
         }
     }
 
+    // Function that returns an associated type from a trait bound
+    fn g<T: MyTrait>(thing: T) -> <T as MyTrait>::AssociatedType {
+        thing.m1() // $ method=MyTrait::m1
+    }
+
     pub fn f() {
-        let x = S;
-        println!("{:?}", x.m1()); // $ method=S::m1
+        let x1 = S;
+        // Call to method in `impl` block
+        println!("{:?}", x1.m1()); // $ method=S::m1 type=x1.m1():AT
 
-        let x = S;
-        println!("{:?}", x.m2()); // $ method=m2
+        let x2 = S;
+        // Call to default method in `trait` block
+        let y = x2.m2(); // $ method=m2 type=y:AT
+        println!("{:?}", y);
+
+        let x3 = S;
+        // Call to the method in `impl` block
+        println!("{:?}", x3.put(1).unwrap()); // $ method=S::put method=unwrap
+
+        // Call to default implementation in `trait` block
+        println!("{:?}", x3.putTwo(2, 3).unwrap()); // $ method=putTwo MISSING: method=unwrap
+
+        let x4 = g(S); // $ MISSING: type=x4:AT
+        println!("{:?}", x4);
+
+        let x5 = S2;
+        println!("{:?}", x5.m1()); // $ method=m1 MISSING: type=x5.m1():A.S2
+        let x6 = S2;
+        println!("{:?}", x6.m2()); // $ method=m2 type=x6.m2():A.S2
     }
 }