Swift: attach NominalTypeDecl to the ExtensionDecl

Author: AlexDenisov

swift/codegen/schema.yml

@@ -36,6 +36,7 @@ IterableDeclContext:
   members: Decl*
 
 ExtensionDecl:
+  extended_type_decl: NominalTypeDecl
   _extends:
     - GenericContext
     - IterableDeclContext

swift/extractor/visitors/DeclVisitor.h

@@ -195,7 +195,8 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
 
   void visitExtensionDecl(swift::ExtensionDecl* decl) {
     auto label = dispatcher_.assignNewLabel(decl);
-    dispatcher_.emit(ExtensionDeclsTrap{label});
+    auto typeLabel = dispatcher_.fetchLabel(decl->getExtendedNominal());
+    dispatcher_.emit(ExtensionDeclsTrap{label, typeLabel});
     emitGenericContext(decl, label);
     emitIterableDeclContext(decl, label);
   }

swift/ql/lib/codeql/swift/generated/decl/ExtensionDecl.qll

@@ -2,7 +2,15 @@
 import codeql.swift.elements.decl.Decl
 import codeql.swift.elements.decl.GenericContext
 import codeql.swift.elements.decl.IterableDeclContext
+import codeql.swift.elements.decl.NominalTypeDecl
 
 class ExtensionDeclBase extends @extension_decl, Decl, GenericContext, IterableDeclContext {
   override string getAPrimaryQlClass() { result = "ExtensionDecl" }
+
+  NominalTypeDecl getExtendedTypeDecl() {
+    exists(NominalTypeDecl x |
+      extension_decls(this, x) and
+      result = x.resolve()
+    )
+  }
 }

swift/ql/lib/swift.dbscheme

@@ -98,7 +98,8 @@ iterable_decl_context_members(
 );
 
 extension_decls(
-  unique int id: @extension_decl
+  unique int id: @extension_decl,
+  int extended_type_decl: @nominal_type_decl ref
 );
 
 @nominal_type_decl =