Rust: Make trait a base type mention of the self type parameter

Author: paldepind

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

@@ -16,7 +16,7 @@ newtype TType =
   TRefType() or // todo: add mut?
   TTypeParamTypeParameter(TypeParam t) or
   TRefTypeParameter() or
-  TSelfTypeParameter()
+  TSelfTypeParameter(Trait t)
 
 /**
  * A type without type arguments.
@@ -144,9 +144,6 @@ class TraitType extends Type, TTrait {
 
   override TypeParameter getTypeParameter(int i) {
     result = TTypeParamTypeParameter(trait.getGenericParamList().getTypeParam(i))
-    or
-    result = TSelfTypeParameter() and
-    i = -1
   }
 
   pragma[nomagic]
@@ -226,11 +223,9 @@ class ImplType extends Type, TImpl {
 
   override TypeParameter getTypeParameter(int i) {
     result = TTypeParamTypeParameter(impl.getGenericParamList().getTypeParam(i))
-    or
-    result = TSelfTypeParameter() and
-    i = -1
   }
 
+  /** Get the trait implemented by this `impl` block, if any. */
   override TypeMention getABaseTypeMention() { result = impl.getTrait() }
 
   override string toString() { result = impl.toString() }
@@ -334,11 +329,29 @@ class RefTypeParameter extends TypeParameter, TRefTypeParameter {
   override Location getLocation() { result instanceof EmptyLocation }
 }
 
-/** An implicit `Self` type parameter. */
+/**
+ * The implicit `Self` type parameter of a trait, that refers to the
+ * implementing type of the trait.
+ *
+ * The Rust Reference on the implicit `Self` parameter:
+ * https://doc.rust-lang.org/reference/items/traits.html#r-items.traits.self-param
+ */
 class SelfTypeParameter extends TypeParameter, TSelfTypeParameter {
-  override Function getMethod(string name) { none() }
+  private Trait trait;
 
-  override string toString() { result = "(Self)" }
+  SelfTypeParameter() { this = TSelfTypeParameter(trait) }
 
-  override Location getLocation() { result instanceof EmptyLocation }
+  Trait getTrait() { result = trait }
+
+  override TypeMention getABaseTypeMention() { result = trait }
+
+  override Function getMethod(string name) {
+    // The `Self` type parameter is an implementation of the trait, so it has
+    // all the trait's methods.
+    result = trait.(ItemNode).getASuccessor(name)
+  }
+
+  override string toString() { result = "Self [" + trait.toString() + "]" }
+
+  override Location getLocation() { result = trait.getLocation() }
 }

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

@@ -38,21 +38,42 @@ private module Input1 implements InputSig1<Location> {
     }
   }
 
-  class TypeParameterPosition = TypeParam;
+  private newtype TTypeParameterPosition =
+    TTypeParamTypeParameterPosition(TypeParam tp) or
+    TSelfTypeParameterPosition()
+
+  class TypeParameterPosition extends TTypeParameterPosition {
+    TypeParam asTypeParam() { this = TTypeParamTypeParameterPosition(result) }
+
+    predicate isSelf() { this = TSelfTypeParameterPosition() }
+
+    string toString() {
+      result = this.asTypeParam().toString()
+      or
+      result = "Self" and this.isSelf()
+    }
+  }
+
+  /** Holds if `typeParam`, `param` and `ppos` all concern the same `TypeParam`. */
+  additional predicate typeParamMatchPosition(
+    TypeParam typeParam, TypeParamTypeParameter param, TypeParameterPosition ppos
+  ) {
+    typeParam = param.getTypeParam() and typeParam = ppos.asTypeParam()
+  }
 
   bindingset[apos]
   bindingset[ppos]
   predicate typeArgumentParameterPositionMatch(TypeArgumentPosition apos, TypeParameterPosition ppos) {
-    apos.asTypeParam() = ppos
+    apos.asTypeParam() = ppos.asTypeParam()
     or
-    apos.asMethodTypeArgumentPosition() = ppos.getPosition()
+    apos.asMethodTypeArgumentPosition() = ppos.asTypeParam().getPosition()
   }
 
-  private predicate id(Raw::TypeParam x, Raw::TypeParam y) { x = y }
+  private predicate id(Raw::AstNode x, Raw::AstNode y) { x = y }
 
-  private predicate idOfRaw(Raw::TypeParam x, int y) = equivalenceRelation(id/2)(x, y)
+  private predicate idOfRaw(Raw::AstNode x, int y) = equivalenceRelation(id/2)(x, y)
 
-  private int idOf(TypeParam node) { idOfRaw(Synth::convertAstNodeToRaw(node), result) }
+  private int idOf(AstNode node) { idOfRaw(Synth::convertAstNodeToRaw(node), result) }
 
   int getTypeParameterId(TypeParameter tp) {
     tp =
@@ -65,8 +86,11 @@ private module Input1 implements InputSig1<Location> {
         kind = 0 and
         id = 1
         or
-        id = idOf(tp0.(TypeParamTypeParameter).getTypeParam()) and
-        kind = 1
+        kind = 1 and
+        exists(AstNode node | id = idOf(node) |
+          node = tp0.(TypeParamTypeParameter).getTypeParam() or
+          node = tp0.(SelfTypeParameter).getTrait()
+        )
       |
         tp0 order by kind, id
       )
@@ -211,15 +235,6 @@ private Type inferImplSelfType(Impl i, TypePath path) {
   result = i.getSelfTy().(TypeReprMention).resolveTypeAt(path)
 }
 
-pragma[nomagic]
-private Type inferTraitSelfType(Trait t, TypePath path) {
-  result = TTrait(t) and
-  path.isEmpty()
-  or
-  result = TTypeParamTypeParameter(t.getGenericParamList().getATypeParam()) and
-  path = TypePath::singleton(result)
-}
-
 /** Gets the type at `path` of the implicitly typed `self` parameter. */
 pragma[nomagic]
 private Type inferImplicitSelfType(SelfParam self, TypePath path) {
@@ -230,7 +245,7 @@ private Type inferImplicitSelfType(SelfParam self, TypePath path) {
   |
     t = inferImplSelfType(i, suffix)
     or
-    t = inferTraitSelfType(i, suffix)
+    t = TSelfTypeParameter(i) and suffix.isEmpty()
   )
 }
 
@@ -273,8 +288,7 @@ private module StructExprMatchingInput implements MatchingInputSig {
     abstract TypeParam getATypeParam();
 
     final TypeParameter getTypeParameter(TypeParameterPosition ppos) {
-      result.(TypeParamTypeParameter).getTypeParam() = ppos and
-      ppos = this.getATypeParam()
+      typeParamMatchPosition(this.getATypeParam(), result, ppos)
     }
 
     abstract StructField getField(string name);
@@ -417,12 +431,7 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
   }
 
   abstract class Declaration extends AstNode {
-    abstract TypeParam getATypeParam();
-
-    final TypeParameter getTypeParameter(TypeParameterPosition ppos) {
-      result.(TypeParamTypeParameter).getTypeParam() = ppos and
-      ppos = this.getATypeParam()
-    }
+    abstract TypeParameter getTypeParameter(TypeParameterPosition ppos);
 
     pragma[nomagic]
     abstract Type getParameterType(DeclarationPosition dpos, TypePath path);
@@ -440,7 +449,9 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
   private class TupleStructDecl extends Declaration, Struct {
     TupleStructDecl() { this.isTuple() }
 
-    override TypeParam getATypeParam() { result = this.getGenericParamList().getATypeParam() }
+    override TypeParameter getTypeParameter(TypeParameterPosition ppos) {
+      typeParamMatchPosition(this.getGenericParamList().getATypeParam(), result, ppos)
+    }
 
     override Type getParameterType(DeclarationPosition dpos, TypePath path) {
       exists(int pos |
@@ -461,8 +472,8 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
   private class TupleVariantDecl extends Declaration, Variant {
     TupleVariantDecl() { this.isTuple() }
 
-    override TypeParam getATypeParam() {
-      result = this.getEnum().getGenericParamList().getATypeParam()
+    override TypeParameter getTypeParameter(TypeParameterPosition ppos) {
+      typeParamMatchPosition(this.getEnum().getGenericParamList().getATypeParam(), result, ppos)
     }
 
     override Type getParameterType(DeclarationPosition dpos, TypePath path) {
@@ -483,38 +494,36 @@ private module CallExprBaseMatchingInput implements MatchingInputSig {
     }
   }
 
-  pragma[nomagic]
-  private Type inferAnnotatedTypeInclSelf(AstNode n, TypePath path) {
-    result = getTypeAnnotation(n).resolveTypeAtInclSelf(path)
-  }
-
   private class FunctionDecl extends Declaration, Function {
-    override TypeParam getATypeParam() { result = this.getGenericParamList().getATypeParam() }
+    override TypeParameter getTypeParameter(TypeParameterPosition ppos) {
+      typeParamMatchPosition(this.getGenericParamList().getATypeParam(), result, ppos)
+      or
+      exists(TraitItemNode trait | this = trait.getAnAssocItem() |
+        typeParamMatchPosition(trait.getTypeParam(_), result, ppos)
+        or
+        ppos.isSelf() and result = TSelfTypeParameter(trait)
+      )
+    }
 
     override Type getParameterType(DeclarationPosition dpos, TypePath path) {
       exists(Param p, int i, boolean inMethod |
         paramPos(this.getParamList(), p, i, inMethod) and
         dpos = TPositionalDeclarationPosition(i, inMethod) and
-        result = inferAnnotatedTypeInclSelf(p.getPat(), path)
+        result = inferAnnotatedType(p.getPat(), path)
       )
       or
       exists(SelfParam self |
         self = pragma[only_bind_into](this.getParamList().getSelfParam()) and
         dpos.isSelf()
       |
-        // `self` parameter with type annotation
-        result = inferAnnotatedTypeInclSelf(self, path)
-        or
-        // `self` parameter without type annotation
-        result = inferImplicitSelfType(self, path)
+        result = inferAnnotatedType(self, path) // `self` parameter with type annotation
         or
-        // `self` parameter without type annotation should also have the special `Self` type
-        result = getRefAdjustImplicitSelfType(self, TypePath::nil(), TSelfTypeParameter(), path)
+        result = inferImplicitSelfType(self, path) // `self` parameter without type annotation
       )
     }
 
     override Type getReturnType(TypePath path) {
-      result = this.getRetType().getTypeRepr().(TypeReprMention).resolveTypeAtInclSelf(path)
+      result = this.getRetType().getTypeRepr().(TypeReprMention).resolveTypeAt(path)
     }
   }
 

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

@@ -29,27 +29,6 @@ abstract class TypeMention extends AstNode {
 
   /** Gets the type that the sub mention at `path` resolves to, if any. */
   Type resolveTypeAt(TypePath path) { result = this.getMentionAt(path).resolveType() }
-
-  /**
-   * Like `resolveTypeAt`, but also resolves `Self` mentions to the implicit
-   * `Self` type parameter.
-   *
-   * This is only needed when resolving types for calls to methods; inside the
-   * methods themselves, `Self` only resolves to the relevant trait or type
-   * being implemented.
-   */
-  final Type resolveTypeAtInclSelf(TypePath path) {
-    result = this.resolveTypeAt(path)
-    or
-    exists(TypeMention tm, ImplOrTraitItemNode node |
-      tm = this.getMentionAt(path) and
-      result = TSelfTypeParameter()
-    |
-      tm = node.getASelfPath()
-      or
-      tm.(PathTypeRepr).getPath() = node.getASelfPath()
-    )
-  }
 }
 
 class TypeReprMention extends TypeMention, TypeRepr {
@@ -80,22 +59,21 @@ class PathMention extends TypeMention, Path {
   override TypeMention getTypeArgument(int i) {
     result = this.getPart().getGenericArgList().getTypeArg(i)
     or
-    // `Self` paths inside traits and `impl` blocks have implicit type arguments
-    // that are the type parameters of the trait or impl. For example, in
+    // `Self` paths inside `impl` blocks have implicit type arguments that are
+    // the type parameters of the `impl` block. For example, in
     //
     // ```rust
-    // impl Foo<T> {
+    // impl<T> Foo<T> {
     //   fn m(self) -> Self {
     //     self
     //   }
     // }
     // ```
     //
     // the `Self` return type is shorthand for `Foo<T>`.
-    exists(ImplOrTraitItemNode node | this = node.getASelfPath() |
+    exists(ImplItemNode node |
+      this = node.getASelfPath() and
       result = node.(ImplItemNode).getSelfPath().getPart().getGenericArgList().getTypeArg(i)
-      or
-      result = node.(Trait).getGenericParamList().getTypeParam(i)
     )
   }
 
@@ -105,7 +83,13 @@ class PathMention extends TypeMention, Path {
       or
       result = TEnum(i)
       or
-      result = TTrait(i)
+      exists(TraitItemNode trait | trait = i |
+        // If this is a `Self` path, then it resolves to the implicit `Self`
+        // type parameter, otherwise it is a trait bound.
+        if this = trait.getASelfPath()
+        then result = TSelfTypeParameter(trait)
+        else result = TTrait(trait)
+      )
       or
       result = TTypeParamTypeParameter(i)
       or
@@ -171,3 +155,9 @@ class ImplMention extends TypeMention, ImplItemNode {
     )
   }
 }
+
+class TraitMention extends TypeMention, TraitItemNode {
+  override TypeMention getTypeArgument(int i) { result = this.getTypeParam(i) }
+
+  override Type resolveType() { result = TTrait(this) }
+}