Rust: Associated types are inherited as type parameters by traits and dyn traits

Author: paldepind

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

@@ -5,6 +5,7 @@
  */
 
 private import codeql.rust.elements.internal.generated.Trait
+private import codeql.rust.internal.PathResolution as PathResolution
 
 /**
  * INTERNAL: This module contains the customizable definition of `Trait` and should not
@@ -67,5 +68,11 @@ module Impl {
      * `where` clauses for `Self`.
      */
     TypeBound getATypeBound() { result = this.getTypeBound(_) }
+
+    /** Gets a direct supertrait of this trait, if any. */
+    Trait getSupertrait() {
+      result =
+        PathResolution::resolvePath(this.getATypeBound().getTypeRepr().(PathTypeRepr).getPath())
+    }
   }
 }

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

@@ -9,10 +9,17 @@ private import codeql.rust.elements.internal.generated.Synth
 private import codeql.rust.frameworks.stdlib.Stdlib
 private import codeql.rust.frameworks.stdlib.Builtins as Builtins
 
+/** Gets a type alias of `trait` or of a supertrait of `trait`. */
+private TypeAlias getTraitTypeAlias(Trait trait) {
+  result = trait.getSupertrait*().getAssocItemList().getAnAssocItem()
+}
+
 /**
- * Holds if a dyn trait type should have a type parameter associated with `n`. A
- * dyn trait type inherits the type parameters of the trait it implements. That
- * includes the type parameters corresponding to associated types.
+ * Holds if a dyn trait type for the trait `trait` should have a type parameter
+ * associated with `n`.
+ *
+ * A dyn trait type inherits the type parameters of the trait it implements.
+ * That includes the type parameters corresponding to associated types.
  *
  * For instance in
  * ```rust
@@ -24,10 +31,7 @@ private import codeql.rust.frameworks.stdlib.Builtins as Builtins
  */
 private predicate dynTraitTypeParameter(Trait trait, AstNode n) {
   trait = any(DynTraitTypeRepr dt).getTrait() and
-  (
-    n = trait.getGenericParamList().getATypeParam() or
-    n = trait.(TraitItemNode).getAnAssocItem().(TypeAlias)
-  )
+  n = [trait.getGenericParamList().getATypeParam().(AstNode), getTraitTypeAlias(trait)]
 }
 
 cached
@@ -39,8 +43,10 @@ newtype TType =
   TNeverType() or
   TUnknownType() or
   TTypeParamTypeParameter(TypeParam t) or
-  TAssociatedTypeTypeParameter(TypeAlias t) { any(TraitItemNode trait).getAnAssocItem() = t } or
-  TDynTraitTypeParameter(AstNode n) { dynTraitTypeParameter(_, n) } or
+  TAssociatedTypeTypeParameter(Trait trait, TypeAlias typeAlias) {
+    getTraitTypeAlias(trait) = typeAlias
+  } or
+  TDynTraitTypeParameter(Trait trait, AstNode n) { dynTraitTypeParameter(trait, n) } or
   TImplTraitTypeParameter(ImplTraitTypeRepr implTrait, TypeParam tp) {
     implTraitTypeParam(implTrait, _, tp)
   } or
@@ -270,17 +276,10 @@ class DynTraitType extends Type, TDynTraitType {
   DynTraitType() { this = TDynTraitType(trait) }
 
   override DynTraitTypeParameter getPositionalTypeParameter(int i) {
-    result = TDynTraitTypeParameter(trait.getGenericParamList().getTypeParam(i))
+    result.getTypeParam() = trait.getGenericParamList().getTypeParam(i)
   }
 
-  override TypeParameter getATypeParameter() {
-    result = super.getATypeParameter()
-    or
-    exists(AstNode n |
-      dynTraitTypeParameter(trait, n) and
-      result = TDynTraitTypeParameter(n)
-    )
-  }
+  override DynTraitTypeParameter getATypeParameter() { result.getTrait() = trait }
 
   Trait getTrait() { result = trait }
 
@@ -427,30 +426,54 @@ class TypeParamTypeParameter extends TypeParameter, TTypeParamTypeParameter {
  *   // ...
  * }
  * ```
+ * Furthermore, associated types of a supertrait induce a corresponding type
+ * parameter in any subtraits. E.g., if we have a trait `SubTrait: ATrait` then
+ * `SubTrait` also has a type parameter for the associated type
+ * `AssociatedType`.
  */
 class AssociatedTypeTypeParameter extends TypeParameter, TAssociatedTypeTypeParameter {
+  private Trait trait;
   private TypeAlias typeAlias;
 
-  AssociatedTypeTypeParameter() { this = TAssociatedTypeTypeParameter(typeAlias) }
+  AssociatedTypeTypeParameter() { this = TAssociatedTypeTypeParameter(trait, typeAlias) }
 
   TypeAlias getTypeAlias() { result = typeAlias }
 
   /** Gets the trait that contains this associated type declaration. */
-  TraitItemNode getTrait() { result.getAnAssocItem() = typeAlias }
+  TraitItemNode getTrait() { result = trait }
 
-  override ItemNode getDeclaringItem() { result = this.getTrait() }
+  /**
+   * Holds if this associated type type parameter corresponds directly its
+   * trait, that is, it is not inherited from a supertrait.
+   */
+  predicate isDirect() { trait.(TraitItemNode).getAnAssocItem() = typeAlias }
+
+  override ItemNode getDeclaringItem() { result = trait }
 
-  override string toString() { result = typeAlias.getName().getText() }
+  override string toString() {
+    result = typeAlias.getName().getText() + "[" + trait.getName().toString() + "]"
+  }
 
   override Location getLocation() { result = typeAlias.getLocation() }
 }
 
+/** Gets the associated type type parameter corresponding directly to `typeAlias`. */
+AssociatedTypeTypeParameter getAssociatedTypeTypeParameter(TypeAlias typeAlias) {
+  result.isDirect() and result.getTypeAlias() = typeAlias
+}
+
+/** Gets the dyn type type parameter corresponding directly to `typeAlias`. */
+DynTraitTypeParameter getDynTraitTypeParameter(TypeAlias typeAlias) {
+  result.getTraitTypeParameter() = getAssociatedTypeTypeParameter(typeAlias)
+}
+
 class DynTraitTypeParameter extends TypeParameter, TDynTraitTypeParameter {
+  private Trait trait;
   private AstNode n;
 
-  DynTraitTypeParameter() { this = TDynTraitTypeParameter(n) }
+  DynTraitTypeParameter() { this = TDynTraitTypeParameter(trait, n) }
 
-  Trait getTrait() { dynTraitTypeParameter(result, n) }
+  Trait getTrait() { result = trait }
 
   /** Gets the dyn trait type that this type parameter belongs to. */
   DynTraitType getDynTraitType() { result.getTrait() = this.getTrait() }
@@ -465,7 +488,7 @@ class DynTraitTypeParameter extends TypeParameter, TDynTraitTypeParameter {
   TypeParameter getTraitTypeParameter() {
     result.(TypeParamTypeParameter).getTypeParam() = n
     or
-    result.(AssociatedTypeTypeParameter).getTypeAlias() = n
+    result = TAssociatedTypeTypeParameter(trait, n)
   }
 
   private string toStringInner() {

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

@@ -90,7 +90,7 @@ private module Input1 implements InputSig1<Location> {
     tp =
       rank[result](TypeParameter tp0, int kind, int id1, int id2 |
         kind = 1 and
-        id1 = 0 and
+        id1 = idOfTypeParameterAstNode(tp0.(DynTraitTypeParameter).getTrait()) and
         id2 =
           idOfTypeParameterAstNode([
               tp0.(DynTraitTypeParameter).getTypeParam().(AstNode),
@@ -102,10 +102,13 @@ private module Input1 implements InputSig1<Location> {
         id2 = idOfTypeParameterAstNode(tp0.(ImplTraitTypeParameter).getTypeParam())
         or
         kind = 3 and
+        id1 = idOfTypeParameterAstNode(tp0.(AssociatedTypeTypeParameter).getTrait()) and
+        id2 = idOfTypeParameterAstNode(tp0.(AssociatedTypeTypeParameter).getTypeAlias())
+        or
+        kind = 4 and
         id1 = 0 and
         exists(AstNode node | id2 = idOfTypeParameterAstNode(node) |
           node = tp0.(TypeParamTypeParameter).getTypeParam() or
-          node = tp0.(AssociatedTypeTypeParameter).getTypeAlias() or
           node = tp0.(SelfTypeParameter).getTrait() or
           node = tp0.(ImplTraitTypeTypeParameter).getImplTraitTypeRepr()
         )
@@ -3507,12 +3510,12 @@ private DynTraitType getFutureTraitType() { result.getTrait() instanceof FutureT
 
 pragma[nomagic]
 private AssociatedTypeTypeParameter getFutureOutputTypeParameter() {
-  result.getTypeAlias() = any(FutureTrait ft).getOutputType()
+  result = getAssociatedTypeTypeParameter(any(FutureTrait ft).getOutputType())
 }
 
 pragma[nomagic]
 private DynTraitTypeParameter getDynFutureOutputTypeParameter() {
-  result = TDynTraitTypeParameter(any(FutureTrait ft).getOutputType())
+  result.getTraitTypeParameter() = getFutureOutputTypeParameter()
 }
 
 pragma[nomagic]
@@ -3824,20 +3827,20 @@ private Type invokedClosureFnTypeAt(InvokedClosureExpr ce, TypePath path) {
 
 /** Gets the path to a closure's return type. */
 private TypePath closureReturnPath() {
-  result = TypePath::singleton(TDynTraitTypeParameter(any(FnOnceTrait t).getOutputType()))
+  result = TypePath::singleton(getDynTraitTypeParameter(any(FnOnceTrait t).getOutputType()))
 }
 
 /** Gets the path to a closure with arity `arity`s `index`th parameter type. */
 pragma[nomagic]
 private TypePath closureParameterPath(int arity, int index) {
   result =
-    TypePath::cons(TDynTraitTypeParameter(any(FnOnceTrait t).getTypeParam()),
+    TypePath::cons(TDynTraitTypeParameter(_, any(FnOnceTrait t).getTypeParam()),
       TypePath::singleton(getTupleTypeParameter(arity, index)))
 }
 
 /** Gets the path to the return type of the `FnOnce` trait. */
 private TypePath fnReturnPath() {
-  result = TypePath::singleton(TAssociatedTypeTypeParameter(any(FnOnceTrait t).getOutputType()))
+  result = TypePath::singleton(getAssociatedTypeTypeParameter(any(FnOnceTrait t).getOutputType()))
 }
 
 /**
@@ -3898,7 +3901,7 @@ private Type inferClosureExprType(AstNode n, TypePath path) {
     result = TDynTraitType(any(FnOnceTrait t)) // always exists because of the mention in `builtins/mentions.rs`
     or
     n = ce and
-    path = TypePath::singleton(TDynTraitTypeParameter(any(FnOnceTrait t).getTypeParam())) and
+    path = TypePath::singleton(TDynTraitTypeParameter(_, any(FnOnceTrait t).getTypeParam())) and
     result.(TupleType).getArity() = ce.getNumberOfParams()
     or
     // Propagate return type annotation to body

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

@@ -148,30 +148,11 @@ class NonAliasPathTypeMention extends PathTypeMention {
 
   TypeItemNode getResolved() { result = resolved }
 
-  /**
-   * Gets a type alias with the name `name` of the trait that this path resolves
-   * to, if any.
-   */
-  pragma[nomagic]
-  private TypeAlias getResolvedTraitAlias(string name) {
-    result = resolved.(TraitItemNode).getAnAssocItem() and
-    name = result.getName().getText()
-  }
-
   pragma[nomagic]
   private TypeRepr getAssocTypeArg(string name) {
     result = this.getSegment().getGenericArgList().getAssocTypeArg(name)
   }
 
-  /** 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)
-    )
-  }
-
   /**
    * Gets the type mention that instantiates the implicit `Self` type parameter
    * for this path, if it occurs in the position of a trait bound.
@@ -239,7 +220,7 @@ class NonAliasPathTypeMention extends PathTypeMention {
       tp = TTypeParamTypeParameter(t.getTypeParam()) and
       result = s.getParenthesizedArgList().(TypeMention).resolveTypeAt(path)
       or
-      tp = TAssociatedTypeTypeParameter(t.getOutputType()) and
+      tp = TAssociatedTypeTypeParameter(t, t.getOutputType()) and
       (
         result = s.getRetType().getTypeRepr().(TypeMention).resolveTypeAt(path)
         or
@@ -249,6 +230,28 @@ class NonAliasPathTypeMention extends PathTypeMention {
         path.isEmpty()
       )
     )
+    or
+    // If `path` is the supertrait of a trait block then any associated types
+    // of the supertrait should be instantiated with the subtrait's
+    // corresponding copies.
+    //
+    // As an example, for
+    // ```rust
+    // trait Sub: Super {
+    // //         ^^^^^ this
+    // ```
+    // we do something to the effect of:
+    // ```rust
+    // trait Sub: Super<Assoc=Assoc[Sub]>
+    // ```
+    // Where `Assoc` is an associated type of `Super` and `Assoc[Sub]` denotes
+    // the copy of the type parameter inherited into `Sub`.
+    exists(Trait subtrait, TypeAlias alias |
+      subtrait.getATypeBound().getTypeRepr().(PathTypeRepr).getPath() = this and
+      result = TAssociatedTypeTypeParameter(subtrait, alias) and
+      tp = TAssociatedTypeTypeParameter(resolved, alias) and
+      path.isEmpty()
+    )
   }
 
   pragma[nomagic]
@@ -259,9 +262,10 @@ class NonAliasPathTypeMention extends PathTypeMention {
   /** Gets the type mention in this path for the type parameter `tp`, if any. */
   pragma[nomagic]
   private TypeMention getTypeMentionForTypeParameter(TypeParameter tp) {
-    exists(TypeAlias alias |
-      result = this.getAnAssocTypeArgument(alias) and
-      tp = TAssociatedTypeTypeParameter(alias)
+    exists(TypeAlias alias, string name |
+      result = this.getAssocTypeArg(name) and
+      tp = TAssociatedTypeTypeParameter(resolved, alias) and
+      alias = resolved.(TraitItemNode).getASuccessor(name)
     )
     or
     // If `path` is the trait of an `impl` block then any associated types
@@ -281,7 +285,8 @@ class NonAliasPathTypeMention extends PathTypeMention {
       this = impl.getTraitPath() and
       alias = impl.getASuccessor(pragma[only_bind_into](name)) and
       result = alias.getTypeRepr() and
-      tp = TAssociatedTypeTypeParameter(this.getResolvedAlias(pragma[only_bind_into](name)))
+      tp =
+        TAssociatedTypeTypeParameter(resolved, this.getResolvedAlias(pragma[only_bind_into](name)))
     )
   }
 
@@ -299,7 +304,7 @@ class NonAliasPathTypeMention extends PathTypeMention {
     or
     result = TTypeParamTypeParameter(resolved)
     or
-    result = TAssociatedTypeTypeParameter(resolved)
+    result = TAssociatedTypeTypeParameter(resolvePath(this.getQualifier()), resolved)
   }
 
   override Type resolvePathTypeAt(TypePath typePath) {
@@ -384,9 +389,8 @@ class TraitMention extends TypeMention instanceof TraitItemNode {
     result = TSelfTypeParameter(this)
     or
     exists(TypeAlias alias |
-      alias = super.getAnAssocItem() and
       typePath = TypePath::singleton(result) and
-      result = TAssociatedTypeTypeParameter(alias)
+      result = TAssociatedTypeTypeParameter(this, alias)
     )
     or
     exists(TypeParam tp |
@@ -540,7 +544,7 @@ class DynTraitTypeReprMention extends TypeMention instanceof DynTraitTypeRepr {
 // impl<A, B, ..> Trait<A, B, ..> for (dyn Trait)<A, B, ..>
 // ```
 // To achieve this:
-// - `DynTypeAbstraction` is an abstraction over type parameters of the trait.
+// - `DynTypeAbstraction` is an abstraction over the type parameters of the trait.
 // - `DynTypeBoundListMention` (this class) is a type mention which has `dyn
 //   Trait` at the root and which for every type parameter of `dyn Trait` has the
 //   corresponding type parameter of the trait.

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

@@ -69,7 +69,7 @@ mod default_method_using_associated_type {
             Self::AssociatedType: Default,
             Self: Sized,
         {
-            self.m1(); // $ target=MyTrait::m1 type=self.m1():AssociatedType
+            self.m1(); // $ target=MyTrait::m1 type=self.m1():AssociatedType[MyTrait]
             let _default = Self::AssociatedType::default(); // $ MISSING: target=default _default:AssociatedType
             Self::AssociatedType::default() // $ MISSING: target=default
         }
@@ -144,8 +144,8 @@ mod equality_on_associated_type {
     where
         T: AnotherGet<Output = i32, AnotherOutput = bool>,
     {
-        let _a1 = x.get(); // $ target=GetSet::get MISSING: type=_a1:i32
-        let _a2 = get(&x); // $ target=get MISSING: type=_a2:i32
+        let _a1 = x.get(); // $ target=GetSet::get type=_a1:i32
+        let _a2 = get(&x); // $ target=get type=_a2:i32
         let _b = x.get_another(); // $ type=_b:bool target=AnotherGet::get_another
     }
 
@@ -346,9 +346,9 @@ mod dyn_trait {
     }
 
     fn _assoc_type_from_supertrait(t: &dyn AnotherGet<Output = i32, AnotherOutput = bool>) {
-        let _a1 = (*t).get(); // $ target=deref target=GetSet::get MISSING: type=_a1:i32
-        let _a2 = t.get(); // $ target=GetSet::get MISSING: type=_a2:i32
-        let _a3 = get(t); // $ target=get MISSING: type=_a3:i32
+        let _a1 = (*t).get(); // $ target=deref target=GetSet::get type=_a1:i32
+        let _a2 = t.get(); // $ target=GetSet::get type=_a2:i32
+        let _a3 = get(t); // $ target=get type=_a3:i32
         let _b1 = (*t).get_another(); // $ target=deref target=AnotherGet::get_another type=_b1:bool
         let _b2 = t.get_another(); // $ target=AnotherGet::get_another type=_b2:bool
     }