Unified: Basic support for extensions

unified/extractor/ast_types.yml

@@ -478,6 +478,7 @@ named:
   class_like_declaration:
     modifier*: modifier
     name_node?: identifier
+    extension_target?: expr
     type_parameter*: type_parameter
     type_constraint*: type_constraint
     base_type*: base_type

unified/extractor/src/languages/swift/swift.rs

@@ -1266,22 +1266,19 @@ fn translation_rules() -> Vec<Rule<SwiftContext>> {
                 base_type: {bases.into_iter().map(|ty| tree!((base_type type: {ty})))}
                 member: {members})
         ),
-        // An `extension Foo { … }` is likewise a `class_like_declaration`, named
-        // by the extended type. The extended type is captured opaquely (as its
-        // source text) so that qualified names (`extension String.Interpolation`,
-        // a `memberType`) name the declaration just like simple ones.
+        // An `extension Foo.Bar { … }` is likewise a `class_like_declaration`.
         rule!(
             (extensionDecl
                 extensionKeyword: @kind
                 modifiers: _* @mods
-                extendedType: @@name
+                extendedType: @extendedType
                 inheritanceClause: (inheritanceClause inheritedTypes: (inheritedType type: @bases)*)?
                 memberBlock: (memberBlock members: _* @members))
             =>
             (class_like_declaration
                 modifier: (modifier #{kind})
                 modifier: {mods}
-                name_node: (identifier #{name})
+                extension_target: {extendedType}
                 base_type: {bases.into_iter().map(|ty| tree!((base_type type: {ty})))}
                 member: {members})
         ),

unified/extractor/tests/corpus/swift/types/extension.output

@@ -70,7 +70,7 @@ top_level
       stmt:
         class_like_declaration
           modifier: modifier "extension"
-          name_node: identifier "Int"
+          extension_target: identifier "Int"
           member:
             function_declaration
               name_node: identifier "squared"

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

@@ -418,6 +418,11 @@ module Unified {
     /** Gets the node corresponding to the field `base_type`. */
     final F::BaseType getABaseType() { result = this.getBaseType(_) }
 
+    /** Gets the node corresponding to the field `extension_target`. */
+    final F::Expr getExtensionTarget() {
+      unified_class_like_declaration_extension_target(this, result)
+    }
+
     /** Gets the node corresponding to the field `member`. */
     final F::Member getMember(int i) { unified_class_like_declaration_member(this, i, result) }
 
@@ -454,6 +459,7 @@ module Unified {
     /** Gets a field or child node of this node. */
     final override F::AstNode getAFieldOrChild() {
       unified_class_like_declaration_base_type(this, _, result) or
+      unified_class_like_declaration_extension_target(this, result) or
       unified_class_like_declaration_member(this, _, result) or
       unified_class_like_declaration_modifier(this, _, result) or
       unified_class_like_declaration_name_node(this, result) or
@@ -1614,6 +1620,10 @@ module Unified {
       or
       result = node.(ClassLikeDeclaration).getBaseType(i) and name = "getBaseType"
       or
+      result = node.(ClassLikeDeclaration).getExtensionTarget() and
+      i = -1 and
+      name = "getExtensionTarget"
+      or
       result = node.(ClassLikeDeclaration).getMember(i) and name = "getMember"
       or
       result = node.(ClassLikeDeclaration).getModifier(i) and name = "getModifier"

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

@@ -262,6 +262,13 @@ predicate valueStep(NameBindingNode node1, NameBindingNode node2) {
     node2 = getNodeFromRef(p.getSubPattern())
   )
   or
+  // Extensions have access to the members of the entity they extend.
+  // TODO: The type parameters of the target type should also be in the local scope (for Swift).
+  exists(ClassLikeDeclaration extension |
+    node1 = getNodeFromRef(extension.getExtensionTarget()) and
+    node2.isLocalNamespace(extension)
+  )
+  or
   FolderHeuristic::valueStep(node1, node2)
 }
 
@@ -282,6 +289,13 @@ predicate inheritanceStep(NameBindingNode supertype, NameBindingNode subtype) {
   )
 }
 
+predicate extensionStep(NameBindingNode extension, NameBindingNode targetClass) {
+  exists(ClassLikeDeclaration cls |
+    targetClass = getNodeFromRef(cls.getExtensionTarget()) and
+    extension.isStaticMemberNamespace(cls)
+  )
+}
+
 signature module TrackInputSig {
   /** Holds if the forward-flow of `node` should be tracked. */
   predicate shouldTrack(NameBindingNode node);
@@ -362,6 +376,17 @@ class NamespaceNode extends NameBindingNode {
   /** If this is the instance namespace for a class, gets the corresponding static namespace. */
   NamespaceNode toStaticNamespace() { result.toInstanceNamespace() = this }
 
+  private NamespaceNode getAnExtension1() { extensionStep(result, this.ref()) }
+
+  /** Gets a namespace that is an extension (i.e. containing extension methods) of this node. */
+  NamespaceNode getAnExtension() {
+    result = this.getAnExtension1()
+    or
+    // `extensionStep` connects the static namespaces of classes.
+    // Add the corresponding extension relation between the instance namespaces.
+    result = this.toStaticNamespace().getAnExtension1().toInstanceNamespace()
+  }
+
   private NamespaceNode getAnInheritanceParent1() { inheritanceStep(result.ref(), this) }
 
   /** Gets a namespace from which this namespace inherits directly. */
@@ -384,6 +409,8 @@ class NamespaceNode extends NameBindingNode {
     not this.hasOwnMember(name) and
     result = this.getAnInheritanceParent().getMember(name) and
     isInheritableMemberNode(result)
+    or
+    result = this.getAnExtension().getMember(name)
   }
 }
 
@@ -480,6 +507,9 @@ module DebugGraph<relevantNodeSig/1 relevantNode> {
       or
       inheritanceStep(node1, node2) and
       value = "inheritedBy"
+      or
+      extensionStep(node1, node2) and
+      value = "extensionOf"
     )
   }
 }
@@ -576,6 +606,17 @@ private module FolderHeuristic {
   }
 }
 
+private ClassLikeDeclaration resolveExtensionTarget(ClassLikeDeclaration cls) {
+  trackNameBinding(result.getNameNode()) = getNodeFromRef(cls.getExtensionTarget())
+}
+
+private ClassLikeDeclaration tryResolveExtensionTarget(ClassLikeDeclaration cls) {
+  result = resolveExtensionTarget(cls)
+  or
+  not exists(resolveExtensionTarget(cls)) and
+  result = cls
+}
+
 /**
  * Holds if `access` may resolve to `target` through the enclosing `accessingClass`.
  *
@@ -604,7 +645,8 @@ private predicate unqualifiedMemberAccessCand(
     // Resolved in an uncertain scope
     exists(NamespaceNode namespace, string name |
       name = access.getName() and
-      accessingClass = LocalNameBindingOutput::getAnUncertainScope(access, name)
+      accessingClass =
+        tryResolveExtensionTarget(LocalNameBindingOutput::getAnUncertainScope(access, name))
     |
       instanceAccess = true and
       namespace.isInstanceMemberNamespace(accessingClass) and

unified/ql/lib/unified.dbscheme

@@ -318,6 +318,11 @@ unified_class_like_declaration_base_type(
   unique int base_type: @unified_base_type ref
 );
 
+unified_class_like_declaration_extension_target(
+  unique int unified_class_like_declaration: @unified_class_like_declaration ref,
+  unique int extension_target: @unified_expr ref
+);
+
 #keyset[unified_class_like_declaration, index]
 unified_class_like_declaration_member(
   int unified_class_like_declaration: @unified_class_like_declaration ref,

unified/ql/test/library-tests/static-name-binding/extensions.swift

@@ -0,0 +1,79 @@
+class A {
+    func ownMethod() {
+        ownMethod() // $ access=A.ownMethod
+        extensionMethod1() // $ access=A.extensionMethod1
+        extensionMethod2() // $ access=A.extensionMethod2
+    }
+}
+
+extension A { // $ access=A
+    func extensionMethod1() { // name=A.extensionMethod1
+        ownMethod() // $ access=A.ownMethod
+        extensionMethod1() // $ access=A.extensionMethod1
+        extensionMethod2() // $ access=A.extensionMethod2
+    }
+}
+
+extension A { // $ access=A
+    func extensionMethod2() { // name=A.extensionMethod2
+        ownMethod() // $ access=A.ownMethod
+        extensionMethod1() // $ access=A.extensionMethod1
+        extensionMethod2() // $ access=A.extensionMethod2
+    }
+}
+
+class B {
+}
+
+extension B { // $ access=B
+    class C { // name=B.C
+        class D {} // name=B.C.D
+    }
+}
+extension B { // $ access=B
+    class Nested : C { // $ access=B.C
+        let x : D // $ access=B.C.D
+    }
+}
+
+// Protocol conformance through extension
+protocol Base {
+    func baseMethod();
+    func baseMethodNoImpl();
+    func baseMethodDefaultImpl();
+}
+extension Base { // $ access=Base
+    func baseMethodExt() {} // name=BaseImpl.baseMethodExt
+    func baseMethodDefaultImpl() {} // name=BaseImpl.baseMethodDefaultImpl
+}
+class X {
+    func xMethod() {
+        baseMethod() // $ access=X.baseMethod
+        baseMethodNoImpl() // $ access=Base.baseMethodNoImpl // with no visible implementation, just resolve to the signature
+        baseMethodExt() // $ access=BaseImpl.baseMethodExt
+
+        // Static name binding may find multiple targets. Type inference should disambiguate.
+        baseMethodDefaultImpl() // $ access=Base.baseMethodDefaultImpl access=BaseImpl.baseMethodDefaultImpl
+    }
+}
+
+extension X : Base { // $ access=X access=Base
+    func baseMethod() {} // name=X.baseMethod
+}
+
+class Y {
+    func yMethod() {
+        baseMethod() // $ access=Y.baseMethod
+        baseMethodDefaultImpl() // $ access=Y.baseMethodDefaultImpl
+    }
+}
+extension Y : Base { // $ access=Y access=Base
+    func baseMethod() {} // name=Y.baseMethod
+    func baseMethodDefaultImpl() {} // name=Y.baseMethodDefaultImpl
+}
+
+// Type parameters of the extended type should be in scope in the extension.
+class GenericExtensionTarget<ExtensionTypeParameter> {}
+extension GenericExtensionTarget { // $ access=GenericExtensionTarget
+    func useTypeParameter(_: ExtensionTypeParameter) {} // $ MISSING: access=ExtensionTypeParameter
+}

unified/ql/test/library-tests/static-name-binding/package1/Sources/Target1/File1.swift

@@ -1,2 +1,8 @@
 let x: A; // $ access=Target1.A
 let y: Target2.A; // not a valid reference
+
+public class ScopedExtensionTarget {
+    func useExtensionFromUnimportedModule() {
+        target2ExtensionMethod() // $ SPURIOUS: access=Target1.ScopedExtensionTarget.target2ExtensionMethod
+    }
+}

unified/ql/test/library-tests/static-name-binding/package1/Sources/Target2/File3.swift

@@ -1,5 +1,11 @@
+import Target1
+
 public class A {} // name=Target2.A
 
 public class B { // name=Target2.B
     public class C {} // name=Target2.B.C
 }
+
+extension ScopedExtensionTarget { // $ access=ScopedExtensionTarget
+    func target2ExtensionMethod() {} // name=Target1.ScopedExtensionTarget.target2ExtensionMethod
+}

unified/ql/test/library-tests/static-name-binding/test.swift

@@ -68,5 +68,5 @@ protocol P { }
 extension H // $ access=H1
   : P { } // $ access=P
 
-extension A.B.C // $ MISSING: access=A access=A.B access=A.B.C (`A.B.C` is currently parsed as a single identifier)
-  : P { } // $ access=P
+extension A.B.C // $ access=A access=A.B access=A.B.C
+  : P { } // $ access=P