Unified: Add extensions convenience predicates

Author: hvitved

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

@@ -4,6 +4,7 @@
 overlay[local?]
 module;
 
+private import unified as U
 private import AstPlugin
 private import codeql.files.FileSystem
 private import codeql.unified.internal.NameBinding as NameBinding
@@ -37,7 +38,11 @@ module Unified {
       )
     }
 
-    /** Gets the nearest enclosing class declaration, if any. */
+    /**
+     * Gets the nearest syntactically enclosing class declaration, if any.
+     *
+     * For the semantically enclosing class, use `getParentClass()` instead.
+     */
     ClassLikeDeclaration getEnclosingClass() {
       exists(AstNode parent | parent = this.getParent() |
         result = parent
@@ -47,6 +52,33 @@ module Unified {
       )
     }
 
+    /**
+     * Gets the class that this node belongs to, if any.
+     *
+     * Unlike `getEnclosingClass`, this takes extensions into account:
+     *
+     * ```swift
+     * class A {
+     *   let x : Int // belongs to class A
+     * }
+     *
+     * extension A {
+     *   func foo() {} // also belongs to the class A
+     * }
+     * ```
+     */
+    ClassLikeDeclaration getParentClass() {
+      exists(ClassLikeDeclaration encl | encl = this.getEnclosingClass() |
+        not exists(encl.getExtensionTarget()) and
+        result = encl
+        or
+        exists(NameBinding::NamespaceNode ns |
+          ns.isStaticMemberNamespace(result) and
+          ns.getAnExtension().isStaticMemberNamespace(encl)
+        )
+      )
+    }
+
     private AstNode overrideEnclosingCallableParent() {
       exists(FunctionExpr func |
         // Capture declarations are evaluated as part of the outer context, and
@@ -205,17 +237,36 @@ module Unified {
     /**
      * Gets the name of this member qualified by the short names of each immediately-enclosing class.
      */
-    string getQualifiedName() {
-      exists(ClassLikeDeclaration cls |
-        this = cls.getAMember() and
-        result = cls.getQualifiedName() + "." + this.getNameOrFallback()
-      )
+    string getQualifiedName() { result = QualifiedName<getParentClass/1>::getQualifiedName(this) }
+  }
+
+  private signature ClassLikeDeclaration getDeclaringClassSig(Member m);
+
+  private module QualifiedName<getDeclaringClassSig/1 getDeclaringClass> {
+    string getQualifiedName(Member m) {
+      not m instanceof U::EnumCaseConstructor and
+      result = getQualifiedName(getDeclaringClass(m)) + "." + m.getNameOrFallback()
+      or
+      not exists(getDeclaringClass(m)) and
+      result = m.getNameOrFallback()
       or
-      not this = any(ClassLikeDeclaration cls).getAMember() and
-      result = this.getNameOrFallback()
+      result = getQualifiedName(m.(U::EnumCaseConstructor).getEnumCaseClass())
     }
   }
 
+  private ClassLikeDeclaration getParentClass(Member m) {
+    result.getAMemberIncludingExtensions() = m
+  }
+
+  private ClassLikeDeclaration getEnclosingClass(Member m) { result.getAMember() = m }
+
+  /**
+   * Similar to `Member.getQualifiedName`, but avoids a dependency on name binding, making
+   * it suitable for use in `toString` predicates.
+   */
+  private predicate getQualifiedNameForToString =
+    QualifiedName<getEnclosingClass/1>::getQualifiedName/1;
+
   class ClassLikeDeclaration extends G::ClassLikeDeclaration {
     /** Gets the name of this declaration. */
     string getName() { result = this.getNameNode().getValue() }
@@ -226,8 +277,24 @@ module Unified {
         NameBinding::getStaticBindingTargetFromRef(this.getABaseType().getType())
     }
 
+    /**
+     * Gets a member of this class, including members from extensions.
+     */
+    Member getAMemberIncludingExtensions() {
+      result = this.getAMember() and
+      not exists(this.getExtensionTarget())
+      or
+      exists(ClassLikeDeclaration ext, NameBinding::NamespaceNode ns |
+        ns.isStaticMemberNamespace(this) and
+        ns.getAnExtension().isStaticMemberNamespace(ext) and
+        result = ext.getAMember()
+      )
+    }
+
     override string toString() {
-      result = concat(getClassLikeDeclarationKeyword(this) + " ") + concat(this.getQualifiedName())
+      result =
+        concat(getClassLikeDeclarationKeyword(this) + " ") +
+          concat(getQualifiedNameForToString(this))
     }
   }
 
@@ -237,7 +304,8 @@ module Unified {
 
     override string toString() {
       result =
-        concat(getConstructorDeclarationKeyword(this) + " ") + concat(this.getQualifiedName())
+        concat(getConstructorDeclarationKeyword(this) + " ") +
+          concat(getQualifiedNameForToString(this))
     }
   }
 
@@ -251,7 +319,9 @@ module Unified {
     string getName() { result = this.getNameNode().getValue() }
 
     override string toString() {
-      result = concat(getFunctionDeclarationKeyword(this) + " ") + concat(this.getQualifiedName())
+      result =
+        concat(getFunctionDeclarationKeyword(this) + " ") +
+          concat(getQualifiedNameForToString(this))
     }
   }
 
@@ -263,7 +333,9 @@ module Unified {
     string getName() { result = this.getNameNode().getValue() }
 
     override string toString() {
-      result = concat(getVariableDeclarationKeyword(this) + " ") + concat(this.getQualifiedName())
+      result =
+        concat(getVariableDeclarationKeyword(this) + " ") +
+          concat(getQualifiedNameForToString(this))
     }
   }
 

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

@@ -1,3 +1,6 @@
+overlay[local?]
+module;
+
 private import unified
 private import codeql.util.Unit
 private import codeql.util.Option

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

@@ -1,6 +1,8 @@
 /**
  * Provides Swift-specific name binding rules.
  */
+overlay[local?]
+module;
 
 private import unified
 private import codeql.unified.internal.StaticNameBinding

unified/ql/lib/codeql/unified/internal/mad/LegacyMaD.qll

@@ -118,7 +118,7 @@ private predicate callableSelector(
   name = callable.getName() and
   argLabels = getArgLabelsFromCallable(callable) and
   exists(ClassLikeDeclaration cls |
-    callable = cls.getAMember() and
+    callable = cls.getAMemberIncludingExtensions() and
     isSubclassOfType(cls, type)
   )
 }

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

@@ -784,7 +784,7 @@ module Input3 implements InputSig3 {
     path.isEmpty()
     or
     exists(ClassLikeDeclaration c |
-      c = n.(SuperExpr).getEnclosingClass() and
+      c = n.(SuperExpr).getParentClass() and
       result = c.getBaseType(0).getType().(TypeMention).getTypeAt(path)
     )
     or

unified/ql/lib/utils/test/TestUtils.qll

@@ -4,15 +4,6 @@ private import codeql.unified.internal.NameBinding
 private import codeql.unified.internal.typeinference.TypeInference as TypeInference
 private import codeql.unified.internal.typeinference.TypeInferencePlugin
 
-private string deriveClassName(ClassLikeDeclaration cls) {
-  not exists(cls.getEnclosingClass()) and
-  result = cls.getName()
-  or
-  result = getStaticBindingTargetFromRef(cls.getExtensionTarget()).getValue()
-  or
-  result = deriveClassName(cls.getEnclosingClass()) + "." + cls.getName()
-}
-
 private string getTupleFieldName(NameBinding decl) {
   decl = any(TupleType tt).getField(result).getPattern()
 }
@@ -24,11 +15,11 @@ private string defaultName(NameBinding decl) {
     not exists(getTupleFieldName(decl)) and name = decl.getName()
   |
     exists(ClassLikeDeclaration cls |
-      decl.getDeclaration() = cls.getAMember() and
-      result = deriveClassName(cls) + "." + name
+      decl.getDeclaration() = cls.getAMemberIncludingExtensions() and
+      result = cls.getQualifiedName() + "." + name
     )
     or
-    not decl.getDeclaration() = any(ClassLikeDeclaration cls).getAMember() and
+    not decl.getDeclaration() = any(ClassLikeDeclaration cls).getAMemberIncludingExtensions() and
     result = name
   )
 }
@@ -47,25 +38,6 @@ predicate nameBinding(NameBinding v, string alias) {
   )
 }
 
-private string defaultCallableName(Callable c) {
-  exists(TypeInference::Input3::Callable callable | callable.asCallable() = c |
-    not c instanceof EnumConstructor and
-    exists(ClassLikeDeclaration cls |
-      c = cls.getAMember() and
-      result = deriveClassName(cls) + "." + callable.getName()
-    )
-    or
-    not c = any(ClassLikeDeclaration cls).getAMember() and
-    result = callable.getName()
-    or
-    exists(ClassLikeDeclaration enum, EnumConstructor ctor |
-      ctor = c and
-      enum = ctor.getEnum() and
-      result = deriveClassName(enum) + "." + ctor.getEnclosingClass().getName()
-    )
-  )
-}
-
 private predicate callableAt(Callable c, string filepath, int line) {
   c.getLocation().hasLocationInfo(filepath, line, _, _, _)
 }
@@ -76,6 +48,6 @@ predicate callableName(Callable c, string alias) {
     keyValueCommentAt(filepath, line, "name", alias)
     or
     not keyValueCommentAt(filepath, line, "name", _) and
-    alias = defaultCallableName(c)
+    alias = c.(Member).getQualifiedName()
   )
 }

unified/ql/test/library-tests/controlflow/basicblock-slices.expected

@@ -7,7 +7,7 @@
 | 8 | cfg.swift:8:1:8:6 | Double | 'Double -> Argument -V topLevelDecl -^ Double(...)' |
 | 10 | cfg.swift:10:1:13:1 | enum MyError | 'enum MyError -V MyError -> BaseType -V Error' |
 | 11 | cfg.swift:11:10:11:15 | MyError.error1 | 'MyError.error1 -V error1 -> MyError.error2 -V error2' |
-| 12 | cfg.swift:12:10:12:31 | MyError.error3 | 'MyError.error3 -V error3 -^ MyError.error3.(constructor)' |
+| 12 | cfg.swift:12:10:12:31 | MyError.error3 | 'MyError.error3 -V error3 -^ MyError.error3' |
 | 12 | cfg.swift:12:17:12:25 | withParam | 'withParam -? Block' |
 | 15 | cfg.swift:15:1:17:1 | func isZero | 'func isZero -V isZero' |
 | 15 | cfg.swift:15:13:15:13 | x | 'x -> Block' |