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 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,11 @@ 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,
+ * and where omission of type arguments implicitly refer to the type
+ * parameters of the type being accessed (provided that they are in scope).
  */
-predicate isInTypeContext(Expr expr) {
+predicate isInTypeContextWithImplicitTypeArgs(Expr expr) {
   expr = any(TypeCastExpr n).getType()
   or
   expr = any(TypeTestExpr n).getType()
@@ -36,11 +38,26 @@ predicate isInTypeContext(Expr expr) {
   or
   expr.getParent() instanceof TypeConstraint
   or
+  any(ExprPositionsPlugin p).isInTypeContextWithImplicitTypeArgs(expr)
+  or
+  isInTypeContextWithImplicitTypeArgs(expr.getEnclosingExpr())
+}
+
+private predicate isInTypeContextWithoutImplicitTypeArgs(Expr expr) {
   exists(NameBinding::getTypeFromRef(expr))
   or
   any(ExprPositionsPlugin p).isInTypeContext(expr)
   or
-  isInTypeContext(expr.getEnclosingExpr())
+  isInTypeContextWithoutImplicitTypeArgs(expr.getEnclosingExpr())
+}
+
+/**
+ * Holds if `expr` appears in a context where it must refer to a type.
+ */
+predicate isInTypeContext(Expr expr) {
+  isInTypeContextWithImplicitTypeArgs(expr)
+  or
+  isInTypeContextWithoutImplicitTypeArgs(expr)
 }
 
 /** Holds if `e` appears in a name-binding position inside `declaration` */

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

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

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

@@ -2,7 +2,9 @@ private import unified
 private import ExprPositionsPlugin
 
 private class ExprPositionsPluginSwift extends ExprPositionsPlugin {
-  override predicate isInTypeContext(Expr e) { e = any(GenericTypeExpr g).getBase() }
+  override predicate isInTypeContextWithImplicitTypeArgs(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() }
 }
@@ -214,6 +223,25 @@ 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 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`.
  */

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

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

@@ -64,6 +64,13 @@ private module Input1 implements InputSig1<Location> {
         kind = 3 and
         tp0 = TUnknownTypeTypeParameter(id1) and
         id2 = 0
+        or
+        kind = 4 and
+        exists(ClassLikeDeclaration c |
+          tp0 = TParentTypeParameterType(c) and
+          id1 = idOfTypeParameterAstNode(c) and
+          id2 = 0
+        )
       |
         tp0 order by kind, id1, id2
       )
@@ -389,9 +396,6 @@ module Input3 implements InputSig3 {
     Unified::Callable getCallable() { result = c }
 
     override TypeMention getDeclaringType() {
-      result = c.(EnumConstructor).getEnum().getNameNode()
-      or
-      not c instanceof EnumConstructor and
       exists(ClassLikeDeclaration cls |
         c = cls.getAMember() and
         result = getClassLikeDeclarationTypeMention(cls)
@@ -405,9 +409,9 @@ module Input3 implements InputSig3 {
       or
       result = c.(FunctionDeclaration).getNameNode()
       or
-      result = c.(EnumConstructor).getEnclosingClass().getNameNode()
+      result = c.(EnumCaseConstructor).getEnumCaseClass().getNameNode()
       or
-      not c instanceof EnumConstructor and
+      not c instanceof EnumCaseConstructor and
       result = c.(ConstructorDeclaration).getNameNode()
     }
 
@@ -492,15 +496,21 @@ module Input3 implements InputSig3 {
     }
 
     Callable getTargetViaStaticNameBinding() {
-      exists(NameBinding b | b = getStaticBindingTargetFromRef(this.getCallee()) |
-        // object creation (including enum constructors): `String(42)`, `Optional.Some(42)`
-        exists(ClassLikeDeclaration cls |
-          cls.getNameNode() = b and
-          result.isInstanceMemberOf(cls, _) and
-          result.asCallable() instanceof ConstructorDeclaration
-        )
-        or
-        // call without explicit receiver: `foo(42)`
+      // object creation (including enum constructors): `String(42)`, `Optional.Some(42)`
+      exists(ClassLikeDeclaration cls |
+        cls =
+          this.getCallee()
+              .(TypeMention)
+              .getType()
+              .(ClassLikeDeclarationType)
+              .getClassLikeDeclaration() and
+        result.isInstanceMemberOf(cls, _) and
+        result.asCallable() instanceof ConstructorDeclaration
+      )
+      or
+      // call without explicit receiver: `foo(42)`
+      exists(NameBinding b |
+        b = getStaticBindingTargetFromRef(this.getCallee()) and
         result.getNameNode() = b
       )
     }
@@ -659,10 +669,20 @@ module Input3 implements InputSig3 {
       result = getQualifierTypeMention(invocation)
     }
 
-    private predicate constructionReceiver(Invocation i, Callable constructor, Expr receiver) {
-      constructor = i.getTargetImpl(_) and
-      constructor.asCallable() instanceof ConstructorDeclaration and
-      receiver = i.getArgument(0)
+    private predicate constructionReceiver(
+      Invocation i, Callable constructor, Expr receiver, TypePath enumPrefixPath
+    ) {
+      exists(ConstructorDeclaration cd |
+        constructor = i.getTargetImpl(_) and
+        cd = constructor.asCallable() and
+        receiver = i.getArgument(0)
+      |
+        enumPrefixPath =
+          TypePath::singleton(TParentTypeParameterType(cd.(EnumCaseConstructor).getEnumCaseClass()))
+        or
+        not cd instanceof EnumCaseConstructor and
+        enumPrefixPath.isEmpty()
+      )
     }
 
     /**
@@ -696,11 +716,15 @@ module Input3 implements InputSig3 {
       M3::ContextualTyping::ContextualReturnType<ContextualConstructorReturnTypeInput>;
 
     private predicate needsContextualTyping(Invocation i, TypePath path) {
-      exists(Callable constructor, Expr receiver, ClassLikeDeclarationType receiverType |
-        constructionReceiver(i, constructor, receiver) and
-        receiverType = inferType(receiver)
+      exists(
+        Callable constructor, Expr receiver, TypePath enumPrefixPath,
+        ClassLikeDeclarationType receiverType, TypePath path0
+      |
+        constructionReceiver(i, constructor, receiver, enumPrefixPath) and
+        receiverType = inferType(receiver) and
+        path = path0.stripPrefix(enumPrefixPath)
       |
-        ContextualConstructorReturnType::needsContextualTyping(i, path) and
+        ContextualConstructorReturnType::needsContextualTyping(i, path0) and
         receiverType = constructor.getDeclaringType().getType()
         or
         // inherited constructor
@@ -710,18 +734,19 @@ module Input3 implements InputSig3 {
         |
           declType = constructor.getDeclaringType().getType() and
           BaseTypes::rootTypesSatisfaction(receiverType, declType, abs, _, constraint) and
-          path = TypePath::singleton(tp)
+          tp =
+            getClassLikeDeclarationTypeMention(receiverType.getClassLikeDeclaration())
+                .getTypeAt(path0)
         |
           // `A<T> : B<T>`
-          exists(TypePath path0 |
-            BaseTypes::conditionSatisfiesConstraintTypeAt(abs, _, constraint, path0, tp) and
-            ContextualConstructorReturnType::needsContextualTyping(i, path0)
+          exists(TypePath path1 |
+            BaseTypes::conditionSatisfiesConstraintTypeAt(abs, _, constraint, path1, tp) and
+            ContextualConstructorReturnType::needsContextualTyping(i, path1)
           )
           or
           // `A<T> : B`
-          tp = receiverType.getATypeParameter() and
           not BaseTypes::conditionSatisfiesConstraintTypeAt(abs, _, constraint, _, tp) and
-          not exists(Type t | t = i.getTypeQualifier(path) and not t instanceof PseudoType)
+          not exists(Type t | t = i.getTypeQualifier(path0) and not t instanceof PseudoType)
         )
       )
     }
@@ -734,8 +759,7 @@ module Input3 implements InputSig3 {
      * as well as calls receivers, `self.init(0)` or `super.init(0)`.
      */
     predicate step(AstNode n1, TypePath prefix1, AstNode n2, TypePath prefix2) {
-      constructionReceiver(n2, _, n1) and
-      prefix1.isEmpty() and
+      constructionReceiver(n2, _, n1, prefix1) and
       prefix2.isEmpty()
     }
 

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

@@ -10,6 +10,7 @@ import TypeInference::Consistency
 
 query predicate illFormedTypeMention(TypeMention tm) {
   Consistency::illFormedTypeMention(tm) and
+  not tm.(AliasExprTypeMention).isTypeParameterAlias(_) and
   tm.fromSource()
 }