Rust: Improve handling of Self type parameter

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

@@ -143,33 +143,73 @@ class NonAliasPathTypeMention extends PathTypeMention {
     )
   }
 
-  private TypeMention getPositionalTypeArgument0(int i) {
+  /**
+   * Gets the positional type argument at index `i` that occurs in this path, if
+   * any.
+   */
+  private TypeMention getPathPositionalTypeArgument(int i) {
     result = this.getSegment().getGenericArgList().getTypeArg(i)
     or
     // `Option::<i32>::Some` is valid in addition to `Option::Some::<i32>`
     resolvePath(this) instanceof Variant and
     result = this.getQualifier().getSegment().getGenericArgList().getTypeArg(i)
   }
 
-  private TypeMention getPositionalTypeArgument(int i) {
-    result = this.getPositionalTypeArgument0(i)
+  /**
+   * Gets the type mention that instantiates the implicit `Self` type parameter
+   * for this path, if it occurs in the position of a trait bound.
+   */
+  private TypeMention getSelfTraitBoundArg() {
+    exists(ImplItemNode impl | this = impl.getTraitPath() and result = impl.(Impl).getSelfTy())
     or
+    exists(Trait subTrait |
+      this = subTrait.getATypeBound().getTypeRepr().(PathTypeRepr).getPath() and
+      result.(SelfTypeParameterMention).getTrait() = subTrait
+    )
+    or
+    exists(TypeParamItemNode tp | this = tp.getABoundPath() and result = tp)
+  }
+
+  private Type getDefaultPositionalTypeArgument(int i, TypePath path) {
     // If a type argument is not given in the path, then we use the default for
     // the type parameter if one exists for the type.
-    not exists(this.getPositionalTypeArgument0(i)) and
-    result = this.resolveRootType().getTypeParameterDefault(i) and
+    not exists(this.getPathPositionalTypeArgument(i)) and
     // Defaults only apply to type mentions in type annotations
-    this = any(PathTypeRepr ptp).getPath().getQualifier*()
+    this = any(PathTypeRepr ptp).getPath().getQualifier*() and
+    exists(Type ty, TypePath prefix |
+      ty = this.resolveRootType().getTypeParameterDefault(i).resolveTypeAt(prefix) and
+      if not ty = TSelfTypeParameter(resolved)
+      then result = ty and path = prefix
+      else
+        // When a default contains an implicit `Self` type parameter, it should
+        // be substituted for the type that implements the trait.
+        exists(TypePath suffix |
+          path = prefix.append(suffix) and
+          result = this.getSelfTraitBoundArg().resolveTypeAt(suffix)
+        )
+    )
   }
 
-  /** Gets the type mention in this path for the type parameter `tp`, if any. */
-  pragma[nomagic]
-  private TypeMention getTypeMentionForTypeParameter(TypeParameter tp) {
+  private Type getPositionalTypeArgument(int i, TypePath path) {
+    result = this.getPathPositionalTypeArgument(i).resolveTypeAt(path)
+    or
+    result = this.getDefaultPositionalTypeArgument(i, path)
+  }
+
+  /**
+   * Gets the type for this path for the type parameter `tp` at `path`, when the
+   * type parameter does not correspond directly to a type mention.
+   */
+  private Type getTypeForTypeParameterAt(TypeParameter tp, TypePath path) {
     exists(int i |
-      result = this.getPositionalTypeArgument(pragma[only_bind_into](i)) and
+      result = this.getPositionalTypeArgument(pragma[only_bind_into](i), path) and
       tp = this.resolveRootType().getPositionalTypeParameter(pragma[only_bind_into](i))
     )
-    or
+  }
+
+  /** 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)
@@ -237,9 +277,17 @@ class NonAliasPathTypeMention extends PathTypeMention {
     typePath.isEmpty() and
     result = this.resolveRootType()
     or
-    exists(TypeParameter tp, TypePath suffix |
-      result = this.getTypeMentionForTypeParameter(tp).resolveTypeAt(suffix) and
-      typePath = TypePath::cons(tp, suffix)
+    exists(TypeParameter tp, TypePath suffix | typePath = TypePath::cons(tp, suffix) |
+      result = this.getTypeForTypeParameterAt(tp, suffix)
+      or
+      result = this.getTypeMentionForTypeParameter(tp).resolveTypeAt(suffix)
+    )
+    or
+    // When the path refers to a trait, then the implicit `Self` type parameter
+    // should be instantiated from the context.
+    exists(TypePath suffix |
+      result = this.getSelfTraitBoundArg().resolveTypeAt(suffix) and
+      typePath = TypePath::cons(TSelfTypeParameter(resolved), suffix)
     )
   }
 }
@@ -296,6 +344,11 @@ class TraitMention extends TypeMention instanceof TraitItemNode {
     typePath.isEmpty() and
     result = TTrait(this)
     or
+    // The implicit `Self` type parameter occurs at the `Self` type parameter
+    // position.
+    typePath = TypePath::singleton(TSelfTypeParameter(this)) and
+    result = TSelfTypeParameter(this)
+    or
     exists(TypeAlias alias |
       alias = super.getAnAssocItem() and
       typePath = TypePath::singleton(result) and

rust/ql/test/library-tests/type-inference/CONSISTENCY/PathResolutionConsistency.expected

@@ -5,9 +5,9 @@ multipleCallTargets
 | dereference.rs:184:17:184:30 | ... .foo() |
 | dereference.rs:186:17:186:25 | S.bar(...) |
 | dereference.rs:187:17:187:29 | S.bar(...) |
-| main.rs:2437:13:2437:31 | ...::from(...) |
-| main.rs:2438:13:2438:31 | ...::from(...) |
-| main.rs:2439:13:2439:31 | ...::from(...) |
-| main.rs:2445:13:2445:31 | ...::from(...) |
-| main.rs:2446:13:2446:31 | ...::from(...) |
-| main.rs:2447:13:2447:31 | ...::from(...) |
+| main.rs:2481:13:2481:31 | ...::from(...) |
+| main.rs:2482:13:2482:31 | ...::from(...) |
+| main.rs:2483:13:2483:31 | ...::from(...) |
+| main.rs:2489:13:2489:31 | ...::from(...) |
+| main.rs:2490:13:2490:31 | ...::from(...) |
+| main.rs:2491:13:2491:31 | ...::from(...) |

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

@@ -652,6 +652,50 @@ mod type_parameter_bounds {
     }
 }
 
+mod trait_default_self_type_parameter {
+    // A trait with a type parameter that defaults to `Self`.
+    trait TraitWithSelfTp<A = Option<Self>> {
+        // TraitWithSelfTp::get_a
+        fn get_a(&self) -> A;
+    }
+
+    fn get_a<A, T: TraitWithSelfTp<A>>(thing: &T) -> A {
+        thing.get_a() // $ target=TraitWithSelfTp::get_a
+    }
+
+    // The trait bound on `T` uses the default for `A` which contains `Self`
+    fn tp_uses_default<S: TraitWithSelfTp>(thing: S) -> i64 {
+        let _ms = thing.get_a(); // $ target=TraitWithSelfTp::get_a type=_ms:T.S
+        0
+    }
+
+    // The supertrait uses the default for `A` which contains `Self`
+    trait SubTraitOfTraitWithSelfTp: TraitWithSelfTp + Sized {}
+
+    fn get_a_through_tp<S: SubTraitOfTraitWithSelfTp>(thing: &S) {
+        // `thing` is a `TraitWithSelfTp` through the trait hierarchy
+        let _ms = get_a(thing); // $ target=get_a type=_ms:T.S
+    }
+
+    struct MyStruct {
+        value: i32,
+    }
+
+    // The implementing trait uses the default for `A` which contains `Self`
+    impl TraitWithSelfTp for MyStruct {
+        fn get_a(&self) -> Option<Self> {
+            Some(MyStruct { value: self.value }) // $ fieldof=MyStruct
+        }
+    }
+
+    impl SubTraitOfTraitWithSelfTp for MyStruct {}
+
+    pub fn test() {
+        let s = MyStruct { value: 0 };
+        let _ms = get_a(&s); // $ target=get_a type=_ms:T.MyStruct
+    }
+}
+
 mod function_trait_bounds {
     #[derive(Debug, Clone, Copy)]
     struct MyThing<T> {
@@ -2753,6 +2797,7 @@ fn main() {
     method_impl::g(method_impl::Foo {}, method_impl::Foo {}); // $ target=g
     method_non_parametric_impl::f(); // $ target=f
     method_non_parametric_trait_impl::f(); // $ target=f
+    trait_default_self_type_parameter::test(); // $ target=test
     function_trait_bounds::f(); // $ target=f
     associated_type_in_trait::f(); // $ target=f
     generic_enum::f(); // $ target=f