Unified: Type inference for nested types

Author: hvitved

shared/typeinference/codeql/typeinference/internal/TypeInference.qll

@@ -2140,13 +2140,24 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
 
       /** An AST node. */
       class AstNode {
+        /** Gets the parent of this AST node, if any. */
+        AstNode getParent();
+
         /** Gets a textual representation of this AST node. */
         string toString();
 
         /** Gets the location of this AST node. */
         Location getLocation();
       }
 
+      /**
+       * Gets the scope of type parameter `tp`.
+       *
+       * Only used in a consistency check to ensure that type parameters do not
+       * escape their scope.
+       */
+      AstNode getTypeParameterScope(TypeParameter tp);
+
       /** An expression. */
       class Expr extends AstNode;
 
@@ -3744,6 +3755,11 @@ module Make1<LocationSig Location, InputSig1<Location> Input1> {
           Input3::inferTypeForDefaults(n, path) = t and
           not inferType(n, path) = t
         }
+
+        query predicate typeParameterEscape(AstNode n, TypePath path, TypeParameter tp) {
+          tp = inferType(n, path) and
+          not n.getParent*() = getTypeParameterScope(tp)
+        }
       }
 
       /**

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,12 +38,22 @@ predicate associatedTypeParameterInherited(
   name = a.getName()
 }
 
+private predicate isNestedType(Member m, Identifier nested, ClassLikeDeclaration parent) {
+  // exists(nested.getNameNode()) and
+  // exists(Member m |
+  exists(parent.getNameNode()) and
+  nested = m.getNameNode() and
+  m = parent.getAMemberIncludingExtensions()
+  // )
+}
+
 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 +66,14 @@ newtype TType =
   ) {
     associatedTypeParameter(c, a, inherited)
   } or
+  TParentTypeParameterType(Member m) {
+    isNestedType(m, _, _) and
+    (
+      exists(m.(ClassLikeDeclaration).getNameNode())
+      or
+      m instanceof TypeAliasDeclaration
+    )
+  } or
   TUnknownType() or
   TUnknownTypeTypeParameter(int i) { i in [0 .. 20] }
 
@@ -78,7 +96,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 +142,29 @@ 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() }
 }
 
+class AliasType extends TypeImpl, TAliasType {
+  TypeAliasDeclaration alias;
+
+  AliasType() { this = TAliasType(alias) }
+
+  TypeAliasDeclaration getTypeAliasDeclaration() { result = alias }
+
+  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.
  */
@@ -159,13 +202,16 @@ 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();
 }
 
 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 }
 
@@ -181,6 +227,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() }
@@ -202,6 +254,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() }
@@ -217,6 +271,35 @@ 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 Member m;
+
+  ParentTypeParameterType() { this = TParentTypeParameterType(m) }
+
+  ClassLikeDeclaration getParent() { isNestedType(m, _, result) }
+
+  Type getChild() {
+    result.(ClassLikeDeclarationType).getClassLikeDeclaration() = m
+    or
+    result.(AliasType).getTypeAliasDeclaration() = m
+  }
+
+  Member getMember() { result = m }
+
+  override AstNode getScope() { none() }
+
+  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`.
  */
@@ -225,6 +308,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()
+  }
 }