Unified: Type inference for nested types

Author: hvitved

unified/ql/lib/codeql/unified/internal/AstExtra.qll

@@ -138,6 +138,9 @@ module Public {
   final class EnumCaseClass extends ClassLikeDeclaration {
     EnumCaseClass() { this.hasModifier("enum_case") }
 
+    /** Gets the enum declaration that this case class belongs to. */
+    ClassLikeDeclaration getEnumClass() { this = result.getAMember() }
+
     /** Gets the constructor of this algebraic data type. */
     EnumCaseConstructor getConstructor() { result = this.getAMember() }
   }
@@ -168,5 +171,31 @@ module Public {
 
     /** Gets the class containing this constructor. */
     EnumCaseClass getEnumCaseClass() { result.getAMember() = this }
+
+    /** Gets the enum declaration that this constructor belongs to. */
+    ClassLikeDeclaration getEnumClass() { result = this.getEnumCaseClass().getEnumClass() }
+  }
+
+  /**
+   * A field of an enum case in an algebraic data type. For example `a` and
+   * `b` in
+   *
+   * ```swift
+   * enum E {
+   *   case a
+   *   case b
+   * }
+   * ```
+   */
+  final class EnumField extends VariableDeclaration {
+    private ClassLikeDeclaration enum;
+
+    EnumField() {
+      this = enum.getAMember() and
+      this.hasModifier("enum_case")
+    }
+
+    /** Gets the enum declaration that this enum field belongs to. */
+    ClassLikeDeclaration getEnumClass() { result = enum }
   }
 }

unified/ql/lib/codeql/unified/internal/ExprPositions.qll

@@ -3,9 +3,10 @@ private import NameBinding as NameBinding
 private import ExprPositionsPlugin
 
 /**
- * Holds if `expr` appears in a context where it must refer to a type.
+ * Holds if `expr` appears in a context where it must refer to a type,
+ * based entirely on the syntactic context of the expression.
  */
-predicate isInTypeContext(Expr expr) {
+predicate isInSyntacticTypeContext(Expr expr) {
   expr = any(TypeCastExpr n).getType()
   or
   expr = any(TypeTestExpr n).getType()
@@ -36,11 +37,24 @@ predicate isInTypeContext(Expr expr) {
   or
   expr.getParent() instanceof TypeConstraint
   or
+  any(ExprPositionsPlugin p).isInSyntacticTypeContext(expr)
+  or
+  isInSyntacticTypeContext(expr.getEnclosingExpr())
+}
+
+private predicate isInTypeContext0(Expr expr) {
   exists(NameBinding::getTypeFromRef(expr))
   or
-  any(ExprPositionsPlugin p).isInTypeContext(expr)
+  isInTypeContext0(expr.getEnclosingExpr())
+}
+
+/**
+ * Holds if `expr` appears in a context where it must refer to a type.
+ */
+predicate isInTypeContext(Expr expr) {
+  isInSyntacticTypeContext(expr)
   or
-  isInTypeContext(expr.getEnclosingExpr())
+  isInTypeContext0(expr)
 }
 
 /** Holds if `e` appears in a name-binding position inside `declaration` */

unified/ql/lib/codeql/unified/internal/ExprPositionsPlugin.qll

@@ -6,7 +6,7 @@ private module Plugins {
 }
 
 class ExprPositionsPlugin extends Unit {
-  predicate isInTypeContext(Expr e) { none() }
+  predicate isInSyntacticTypeContext(Expr e) { none() }
 
   predicate isReturnedExpr(Callable callable, Expr e) { none() }
 }

unified/ql/lib/codeql/unified/internal/ExprPositionsPluginSwift.qll

@@ -2,7 +2,7 @@ private import unified
 private import ExprPositionsPlugin
 
 private class ExprPositionsPluginSwift extends ExprPositionsPlugin {
-  override predicate isInTypeContext(Expr e) { e = any(GenericTypeExpr g).getBase() }
+  override predicate isInSyntacticTypeContext(Expr e) { e = any(GenericTypeExpr g).getBase() }
 
   override predicate isReturnedExpr(Callable callable, Expr e) {
     // Short-hand return from a function body, e.g. `func foo() -> String { "42" }`

unified/ql/lib/codeql/unified/internal/FacadeAst.qll

@@ -8,6 +8,7 @@ private import unified as U
 private import AstPlugin
 private import codeql.files.FileSystem
 private import codeql.unified.internal.NameBinding as NameBinding
+private import codeql.unified.internal.typeinference.TypeMention as TypeMention
 
 module Unified {
   private import Ast::Unified as G
@@ -50,6 +51,11 @@ module Unified {
       )
     }
 
+    /**
+     * Gets the nearest resolved enclosing class declaration, if any.
+     */
+    ClassLikeDeclaration getEnclosingClassResolved() { result = this.getEnclosingClass().resolve() }
+
     private AstNode overrideEnclosingCallableParent() {
       exists(FunctionExpr func |
         // Capture declarations are evaluated as part of the outer context, and
@@ -239,15 +245,16 @@ module Unified {
     /** Gets the name of this declaration. */
     string getName() { result = this.getNameNode().getValue() }
 
-    /** Gets a direct base class of this class. */
+    /** Gets a direct base class of this class, where type aliases are resolved. */
+    pragma[nomagic]
     ClassLikeDeclaration getABaseClass() {
-      result.getNameNode() =
-        NameBinding::getStaticBindingTargetFromRef(this.getABaseType().getType())
+      result = this.getABaseType().getType().(TypeMention::TypeMention).getClassLikeDeclaration()
     }
 
     /**
      * Gets the target of this extension class, or this class itself if it is not an extension.
      */
+    pragma[nomagic]
     ClassLikeDeclaration resolve() {
       not exists(this.getExtensionTarget()) and
       result = this

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

@@ -4,6 +4,7 @@ module;
 
 import unified // reexport in order to shadow e.g. the `TypeParameter` class
 private import unified as Unified
+private import TypeMention
 private import TypeInference
 private import codeql.unified.internal.StaticNameBinding
 
@@ -38,12 +39,24 @@ predicate associatedTypeParameterInherited(
   name = a.getName()
 }
 
+private predicate isNestedType(Member m, Identifier nested, ClassLikeDeclaration parent) {
+  exists(parent.getNameNode()) and
+  nested = m.getNameNode() and
+  m = parent.getAMemberIncludingExtensions() and
+  (
+    exists(m.(ClassLikeDeclaration).getNameNode())
+    or
+    m instanceof TypeAliasDeclaration
+  )
+}
+
 cached
 newtype TType =
   TClassLikeDeclarationType(ClassLikeDeclaration c) {
     CachedStage::ref() and
     exists(c.getNameNode())
   } or
+  TAliasType(TypeAliasDeclaration alias) or
   TClosureParameterPseudoType(Parameter p) {
     exists(FunctionExpr fe |
       p = fe.getAParameter() and
@@ -56,6 +69,7 @@ newtype TType =
   ) {
     associatedTypeParameter(c, a, inherited)
   } or
+  TParentTypeParameterType(Member m) { isNestedType(m, _, _) } or
   TUnknownType() or
   TUnknownTypeTypeParameter(int i) { i in [0 .. 20] }
 
@@ -78,7 +92,14 @@ abstract private class TypeImpl extends TType {
    * This includes both positional type parameters and synthetic type parameters,
    * such as associated types.
    */
-  TypeParameter getATypeParameter() { result = this.getPositionalTypeParameter(_) }
+  TypeParameter getATypeParameter() {
+    result = this.getPositionalTypeParameter(_)
+    or
+    exists(Member m |
+      result = TParentTypeParameterType(m) and
+      isNestedType(m, this.getNameNode(), _)
+    )
+  }
 
   /** Gets the name node of this type, if any. */
   abstract Identifier getNameNode();
@@ -117,11 +138,35 @@ class ClassLikeDeclarationType extends TypeImpl, TClassLikeDeclarationType {
     result = TAssociatedTypeParameterType(c, _, _)
   }
 
-  override string toString() { result = c.getName() }
+  override string toString() { result = c.getQualifiedName() }
 
   override Location getLocation() { result = c.getLocation() }
 }
 
+/**
+ * A type representing a type alias. For type inference, type aliases are expanded,
+ * but we need them for resolving type mentions.
+ */
+class AliasType extends TypeImpl, TAliasType {
+  TypeAliasDeclaration alias;
+
+  AliasType() { this = TAliasType(alias) }
+
+  TypeAliasDeclaration getTypeAliasDeclaration() { result = alias }
+
+  TypeMention getAlias() { result = alias.getType() }
+
+  override Identifier getNameNode() { result = alias.getNameNode() }
+
+  override TypeParameter getPositionalTypeParameter(int i) {
+    result = TTypeParameterType(alias.getTypeParameter(i))
+  }
+
+  override string toString() { result = alias.getQualifiedName() }
+
+  override Location getLocation() { result = alias.getLocation() }
+}
+
 /**
  * A pseudo type that does not correspond to any concrete type in the source code.
  */
@@ -165,7 +210,8 @@ abstract class TypeParameter extends TypeImpl {
 }
 
 private class IdAstNode =
-  @unified_type_parameter or @unified_class_like_declaration or @unified_associated_type_declaration;
+  @unified_type_parameter or @unified_class_like_declaration or
+      @unified_associated_type_declaration or @unified_member;
 
 private predicate id(IdAstNode x, IdAstNode y) { x = y }
 
@@ -217,6 +263,45 @@ class AssociatedTypeParameterType extends TypeParameter, TAssociatedTypeParamete
   }
 }
 
+/**
+ * A special type parameter used to represent the parent type of a nested class.
+ *
+ * For example, for
+ *
+ * ```swift
+ * class Outer<TOuter> {
+ *   class Inner<TInner> {
+ *   }
+ * }
+ * ```
+ *
+ * the type `Outer<Int>.Inner<String>` is represented as
+ *
+ * `TypePath`              | `Type`
+ * ----------------------- | --------
+ * `""`                    | `Inner`
+ * `"TInner"`              | `String`
+ * `"Inner_Parent"`        | `Outer`
+ * `"Inner_Parent.TOuter"` | `Int`
+ *
+ * where `Inner_Parent` is one of these type parameters.
+ */
+class ParentTypeParameterType extends TypeParameter, TParentTypeParameterType {
+  private Member m;
+
+  ParentTypeParameterType() { this = TParentTypeParameterType(m) }
+
+  ClassLikeDeclaration getParent() { isNestedType(m, _, result) }
+
+  override Member getDeclaringItem() { result = m }
+
+  override Identifier getNameNode() { none() }
+
+  override string toString() { result = "(parent:" + m.getNameOrFallback() + ")" }
+
+  override Location getLocation() { result = m.getLocation() }
+}
+
 /**
  * A type parameter of the special `UnknownType`.
  */

unified/ql/lib/codeql/unified/internal/typeinference/TypeAbstraction.qll

@@ -15,5 +15,10 @@ abstract class TypeAbstraction extends AstNode {
  */
 private class ClassLikeDeclarationTypeAbstraction extends TypeAbstraction instanceof Unified::ClassLikeDeclaration
 {
-  override TypeParameter getATypeParameter() { this = result.getDeclaringItem() }
+  override TypeParameter getATypeParameter() {
+    this = result.getDeclaringItem()
+    or
+    // type parameters from outer classes are also in scope
+    super.resolve().getEnclosingClassResolved*() = result.getDeclaringItem()
+  }
 }