Unified: Type inference for nested types

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 getEnum() { 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 getEnum() { result = this.getEnumCaseClass().getEnum() }
+  }
+
+  /**
+   * 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 getEnum() { 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

@@ -38,6 +38,12 @@ predicate associatedTypeParameterInherited(
   name = a.getName()
 }
 
+private predicate isNestedType(ClassLikeDeclaration nested, ClassLikeDeclaration parent) {
+  exists(nested.getNameNode()) and
+  exists(parent.getNameNode()) and
+  parent.getAMemberIncludingExtensions() = nested
+}
+
 cached
 newtype TType =
   TClassLikeDeclarationType(ClassLikeDeclaration c) {
@@ -56,6 +62,7 @@ newtype TType =
   ) {
     associatedTypeParameter(c, a, inherited)
   } or
+  TParentTypeParameterType(ClassLikeDeclaration c) { isNestedType(c, _) } or
   TUnknownType() or
   TUnknownTypeTypeParameter(int i) { i in [0 .. 20] }
 
@@ -112,9 +119,11 @@ class ClassLikeDeclarationType extends TypeImpl, TClassLikeDeclarationType {
     result = super.getATypeParameter()
     or
     result = TAssociatedTypeParameterType(c, _, _)
+    or
+    result = TParentTypeParameterType(c)
   }
 
-  override string toString() { result = c.getName() }
+  override string toString() { result = c.getQualifiedName() }
 
   override Location getLocation() { result = c.getLocation() }
 }
@@ -156,6 +165,8 @@ class ClosureParameterPseudoType extends PseudoType, TClosureParameterPseudoType
 
 /** A type parameter. */
 abstract class TypeParameter extends TypeImpl {
+  abstract AstNode getScope();
+
   override TypeParameter getPositionalTypeParameter(int i) { none() }
 
   abstract AstNode getDeclaringItem();
@@ -178,6 +189,12 @@ class TypeParameterType extends TypeParameter, TTypeParameterType {
 
   Unified::TypeParameter getTypeParameter() { result = typeParam }
 
+  override AstNode getScope() {
+    result = typeParam.getParent()
+    or
+    result.(ClassLikeDeclaration).resolve().getEnclosingClassResolved*() = typeParam.getParent()
+  }
+
   override ClassLikeDeclaration getDeclaringItem() { typeParam = result.getATypeParameter() }
 
   override Identifier getNameNode() { result = typeParam.getNameNode() }
@@ -199,6 +216,8 @@ class AssociatedTypeParameterType extends TypeParameter, TAssociatedTypeParamete
 
   Unified::AssociatedTypeDeclaration getAssociatedTypeDeclaration() { result = assocTypeDecl }
 
+  override AstNode getScope() { result.(ClassLikeDeclaration).resolve() = c }
+
   override ClassLikeDeclaration getDeclaringItem() { result = c }
 
   override Identifier getNameNode() { result = assocTypeDecl.getNameNode() }
@@ -214,6 +233,27 @@ class AssociatedTypeParameterType extends TypeParameter, TAssociatedTypeParamete
   }
 }
 
+/**
+ * A special type parameter used to represent the parent type of a nested class.
+ */
+class ParentTypeParameterType extends TypeParameter, TParentTypeParameterType {
+  private Unified::ClassLikeDeclaration c;
+
+  ParentTypeParameterType() { this = TParentTypeParameterType(c) }
+
+  ClassLikeDeclaration getParent() { isNestedType(c, result) }
+
+  override AstNode getScope() { none() }
+
+  override ClassLikeDeclaration getDeclaringItem() { result = c }
+
+  override Identifier getNameNode() { none() }
+
+  override string toString() { result = "(parent:" + c.getName() + ")" }
+
+  override Location getLocation() { result = c.getLocation() }
+}
+
 /**
  * A type parameter of the special `UnknownType`.
  */
@@ -222,6 +262,8 @@ class UnknownTypeTypeParameter extends TypeParameter, TUnknownTypeTypeParameter
 
   UnknownTypeTypeParameter() { this = TUnknownTypeTypeParameter(i) }
 
+  override AstNode getScope() { none() }
+
   override AstNode getDeclaringItem() { none() }
 
   override Identifier getNameNode() { none() }

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()
+  }
 }