Merge pull request #20094 from paldepind/rust/type-inference-path-mention

Rust: Refactor `PathTypeMention`

Author: paldepind

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

@@ -50,24 +50,70 @@ class SliceTypeReprMention extends TypeMention instanceof SliceTypeRepr {
   }
 }
 
-class PathTypeMention extends TypeMention, Path {
-  TypeItemNode resolved;
+/** Holds if `path` is used as a type mention during type inference. */
+predicate relevantPathTypeMention(Path path) {
+  path =
+    [
+      any(PathTypeRepr r).getPath(),
+      any(StructExpr s).getPath().getQualifier*(),
+      any(CallExpr ce).getFunction().(PathExpr).getPath().getQualifier*(),
+      any(StructPat p).getPath(),
+      any(TupleStructPat p).getPath()
+    ]
+}
+
+abstract class PathTypeMention extends TypeMention, Path {
+  PathTypeMention() { relevantPathTypeMention(this) }
+}
+
+class AliasPathTypeMention extends PathTypeMention {
+  TypeAlias resolved;
+  TypeMention rhs;
+
+  AliasPathTypeMention() {
+    resolved = resolvePath(this) and
+    rhs = resolved.getTypeRepr()
+  }
+
+  TypeItemNode getResolved() { result = resolved }
 
-  PathTypeMention() {
-    resolved = resolvePath(this)
+  /**
+   * Holds if this path resolved to a type alias with a rhs. that has the
+   * resulting type at `typePath`.
+   */
+  pragma[nomagic]
+  override Type resolveTypeAt(TypePath typePath) {
+    result = rhs.resolveTypeAt(typePath) and
+    not result = pathGetTypeParameter(resolved, _)
     or
-    resolved = resolvePath(this).(Variant).getEnum()
+    exists(TypeParameter tp, TypeMention arg, TypePath prefix, TypePath suffix, int i |
+      tp = rhs.resolveTypeAt(prefix) and
+      tp = pathGetTypeParameter(resolved, pragma[only_bind_into](i)) and
+      arg = this.getSegment().getGenericArgList().getTypeArg(pragma[only_bind_into](i)) and
+      result = arg.resolveTypeAt(suffix) and
+      typePath = prefix.append(suffix)
+    )
+  }
+}
+
+class NonAliasPathTypeMention extends PathTypeMention {
+  TypeItemNode resolved;
+
+  NonAliasPathTypeMention() {
+    resolved = [resolvePath(this), resolvePath(this).(Variant).getEnum().(TypeItemNode)] and
+    not exists(resolved.(TypeAlias).getTypeRepr())
   }
 
   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) {
-    exists(TraitItemNode trait |
-      trait = resolved and
-      result = trait.getAnAssocItem() and
-      name = result.getName().getText()
-    )
+    result = resolved.(TraitItemNode).getAnAssocItem() and
+    name = result.getName().getText()
   }
 
   pragma[nomagic]
@@ -115,92 +161,75 @@ class PathTypeMention extends TypeMention, 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.resolveType().getTypeParameterDefault(i) and
+    result = this.resolveRootType().getTypeParameterDefault(i) and
     // Defaults only apply to type mentions in type annotations
     this = any(PathTypeRepr ptp).getPath().getQualifier*()
   }
 
-  /**
-   * Holds if this path resolved to a type alias with a rhs. that has the
-   * resulting type at `typePath`.
-   */
+  /** Gets the type mention in this path for the type parameter `tp`, if any. */
   pragma[nomagic]
-  private Type aliasResolveTypeAt(TypePath typePath) {
-    exists(TypeAlias alias, TypeMention rhs | alias = resolved and rhs = alias.getTypeRepr() |
-      result = rhs.resolveTypeAt(typePath) and
-      not result = pathGetTypeParameter(alias, _)
-      or
-      exists(TypeParameter tp, TypeMention arg, TypePath prefix, TypePath suffix, int i |
-        tp = rhs.resolveTypeAt(prefix) and
-        tp = pathGetTypeParameter(alias, pragma[only_bind_into](i)) and
-        arg = this.getSegment().getGenericArgList().getTypeArg(pragma[only_bind_into](i)) and
-        result = arg.resolveTypeAt(suffix) and
-        typePath = prefix.append(suffix)
-      )
+  private TypeMention getTypeMentionForTypeParameter(TypeParameter tp) {
+    exists(int i |
+      result = this.getPositionalTypeArgument(pragma[only_bind_into](i)) and
+      tp = this.resolveRootType().getTypeParameter(pragma[only_bind_into](i))
+    )
+    or
+    exists(TypeAlias alias |
+      result = this.getAnAssocTypeArgument(alias) and
+      tp = TAssociatedTypeTypeParameter(alias)
+    )
+    or
+    // If `path` 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 {
+    //      ^^^^^^^ path
+    //   type AssociatedType = i64
+    //                         ^^^ result
+    //   // ...
+    // }
+    // ```
+    // the rhs. of the type alias is a type argument to the trait.
+    exists(ImplItemNode impl, AssociatedTypeTypeParameter param, TypeAlias alias, string name |
+      this = impl.getTraitPath() and
+      param.getTrait() = resolved and
+      name = param.getTypeAlias().getName().getText() and
+      alias = impl.getASuccessor(pragma[only_bind_into](name)) and
+      result = alias.getTypeRepr() and
+      tp =
+        TAssociatedTypeTypeParameter(resolved
+              .(TraitItemNode)
+              .getAssocItem(pragma[only_bind_into](name)))
     )
   }
 
-  override Type resolveTypeAt(TypePath typePath) {
-    result = this.aliasResolveTypeAt(typePath)
+  Type resolveRootType() {
+    result = TStruct(resolved)
     or
-    typePath.isEmpty() and
-    (
-      result = TStruct(resolved)
-      or
-      result = TEnum(resolved)
-      or
-      exists(TraitItemNode trait | trait = resolved |
-        // 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(resolved)
-      or
-      result = TAssociatedTypeTypeParameter(resolved)
+    result = TEnum(resolved)
+    or
+    exists(TraitItemNode trait | trait = resolved |
+      // 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
-    not exists(resolved.(TypeAlias).getTypeRepr()) and
-    exists(TypeParameter tp, TypeMention arg, TypePath suffix |
-      result = arg.resolveTypeAt(suffix) and
+    result = TTypeParamTypeParameter(resolved)
+    or
+    result = TAssociatedTypeTypeParameter(resolved)
+  }
+
+  override Type resolveTypeAt(TypePath typePath) {
+    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(int i |
-        arg = this.getPositionalTypeArgument(pragma[only_bind_into](i)) and
-        tp = this.resolveType().getTypeParameter(pragma[only_bind_into](i))
-      )
-      or
-      exists(TypeAlias alias |
-        arg = this.getAnAssocTypeArgument(alias) and
-        tp = TAssociatedTypeTypeParameter(alias)
-      )
-      or
-      // If `path` 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 {
-      //      ^^^^^^^ path
-      //   type AssociatedType = i64
-      //                         ^^^ result
-      //   // ...
-      // }
-      // ```
-      // the rhs. of the type alias is a type argument to the trait.
-      exists(ImplItemNode impl, AssociatedTypeTypeParameter param, TypeAlias alias, string name |
-        this = impl.getTraitPath() and
-        param.getTrait() = resolved and
-        name = param.getTypeAlias().getName().getText() and
-        alias = impl.getASuccessor(pragma[only_bind_into](name)) and
-        arg = alias.getTypeRepr() and
-        tp =
-          TAssociatedTypeTypeParameter(resolved
-                .(TraitItemNode)
-                .getAssocItem(pragma[only_bind_into](name)))
-      )
     )
   }
 }