Unified: Extend isInTypeContext

Author: hvitved

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

@@ -1,8 +1,9 @@
 private import unified
 private import NameBinding as NameBinding
+private import ExprPositionsPlugin
 
 /**
- * Holds if `expr` appears in the context of a type annotation.
+ * Holds if `expr` appears in a context where it must refer to a type.
  */
 predicate isInTypeContext(Expr expr) {
   expr = any(TypeCastExpr n).getType()
@@ -27,8 +28,22 @@ predicate isInTypeContext(Expr expr) {
   or
   expr = any(AssociatedTypeDeclaration n).getBound()
   or
+  expr = any(ClassLikeDeclaration c).getExtensionTarget()
+  or
+  expr = any(GenericTypeExpr gte).getATypeArgument()
+  or
   expr.getParent() instanceof TypeConstraint
   or
+  exists(Identifier id | id = NameBinding::getStaticBindingTarget(expr) |
+    id = any(ClassLikeDeclaration c).getNameNode()
+    or
+    id = any(TypeAliasDeclaration t).getNameNode()
+    or
+    id = any(TypeParameter t).getNameNode()
+  )
+  or
+  any(ExprPositionsPlugin p).isInTypeContext(expr)
+  or
   isInTypeContext(expr.getEnclosingExpr())
 }
 

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

@@ -0,0 +1,10 @@
+private import unified
+private import codeql.util.Unit
+
+private module Plugins {
+  private import ExprPositionsPluginSwift
+}
+
+class ExprPositionsPlugin extends Unit {
+  predicate isInTypeContext(Expr e) { none() }
+}

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

@@ -0,0 +1,6 @@
+private import unified
+private import ExprPositionsPlugin
+
+private class ExprPositionsPluginSwift extends ExprPositionsPlugin {
+  override predicate isInTypeContext(Expr e) { e = any(GenericTypeExpr g).getBase() }
+}