Unified: Type inference for nested types

Author: hvitved

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

@@ -2272,7 +2272,11 @@ private module Input3 implements InputSig3 {
     BoolType() { this.getTypeItem() instanceof Builtins::Bool }
   }
 
-  class AstNode = Rust::AstNode;
+  class AstNode extends Rust::AstNode {
+    AstNode getParent() { result = this.getParentNode() }
+  }
+
+  AstNode getTypeParameterScope(TypeParameter tp) { result = tp.getDeclaringItem() }
 
   final class Expr = ExprImpl;
 
@@ -2286,13 +2290,13 @@ private module Input3 implements InputSig3 {
     TypeMention getType() { result = this.getTypeRepr() }
   }
 
-  class Switch extends Rust::MatchExpr {
+  class Switch extends Expr, Rust::MatchExpr {
     Expr getExpr() { result = this.getScrutinee() }
 
     Case getCase(int index) { result = this.getArm(index) }
   }
 
-  class Case extends Rust::MatchArm {
+  class Case extends AstNode, Rust::MatchArm {
     AstNode getAPattern() { result = this.getPat() }
 
     AstNode getBody() { result = this.getExpr() }
@@ -2453,8 +2457,8 @@ private module Input3 implements InputSig3 {
       // mutual recursion; resolving fields requires resolving types and vice versa
       result =
         [
-          resolveStructFieldExpr(this, _).(AstNode),
-          resolveTupleFieldExpr(this, _)
+          resolveStructFieldExpr(this, _).(Field),
+          resolveTupleFieldExpr(this, _).(Field)
         ]
     }
   }
@@ -2484,7 +2488,7 @@ private module Input3 implements InputSig3 {
     )
   }
 
-  class Return extends ReturnExpr {
+  class Return extends Expr, ReturnExpr {
     Expr getExpr() { result = super.getExpr() }
   }
 
@@ -2587,15 +2591,15 @@ private module Input3 implements InputSig3 {
     override TypeItem getTypeItem() { result = this }
 
     override Parameter getParameter(int i) {
-      result = [super.getTupleField(i).(AstNode), super.getNthStructField(i)]
+      result = [super.getTupleField(i).(Parameter), super.getNthStructField(i).(Parameter)]
     }
   }
 
   private class VariantConstructor extends ConstructorImpl instanceof Variant {
     override TypeItem getTypeItem() { result = super.getEnum() }
 
     override Parameter getParameter(int i) {
-      result = [super.getTupleField(i).(AstNode), super.getNthStructField(i)]
+      result = [super.getTupleField(i).(Parameter), super.getNthStructField(i).(Parameter)]
     }
   }
 
@@ -3161,7 +3165,7 @@ module Consistency = M3::Consistency;
 
 module CachedStage = M3::CachedStage;
 
-private predicate typeTestAstNodeRepr(AstNode n, string repr) {
+private predicate typeTestAstNodeRepr(Input3::AstNode n, string repr) {
   repr = [n.toString(), n.(IdentPat).getName().getText()]
 }
 

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;
 
@@ -3754,6 +3765,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