Merge pull request #9321 from github/alexdenisov/xref-decls

Swift: do not duplicate 'external' declarations

Author: AlexDenisov

swift/extractor/SwiftDispatcher.h

@@ -10,16 +10,29 @@
 
 namespace codeql {
 
+enum class SwiftExtractionMode { Module, PrimaryFile };
+
 // The main responsibilities of the SwiftDispatcher are as follows:
 // * redirect specific AST node emission to a corresponding visitor (statements, expressions, etc.)
 // * storing TRAP labels for emitted AST nodes (in the TrapLabelStore) to avoid re-emission
 // Since SwiftDispatcher sees all the AST nodes, it also attaches a location to every 'locatable'
 // node (AST nodes that are not types: declarations, statements, expressions, etc.).
 class SwiftDispatcher {
  public:
-  // sourceManager, arena, and trap are supposed to outlive the SwiftDispatcher
-  SwiftDispatcher(const swift::SourceManager& sourceManager, TrapArena& arena, TrapOutput& trap)
-      : sourceManager{sourceManager}, arena{arena}, trap{trap} {}
+  // all references passed as parameters to this constructor are supposed to outlive the
+  // SwiftDispatcher
+  SwiftDispatcher(const swift::SourceManager& sourceManager,
+                  TrapArena& arena,
+                  TrapOutput& trap,
+                  SwiftExtractionMode extractionMode,
+                  swift::ModuleDecl& currentModule,
+                  llvm::StringRef currentFileName)
+      : sourceManager{sourceManager},
+        arena{arena},
+        trap{trap},
+        extractionMode(extractionMode),
+        currentModule{currentModule},
+        currentFileName(currentFileName) {}
 
   template <typename Entry>
   void emit(const Entry& entry) {
@@ -74,10 +87,10 @@ class SwiftDispatcher {
   // Due to the lazy emission approach, we must assign a label to a corresponding AST node before
   // it actually gets emitted to handle recursive cases such as recursive calls, or recursive type
   // declarations
-  template <typename E>
-  TrapLabelOf<E> assignNewLabel(E* e) {
+  template <typename E, typename... Args>
+  TrapLabelOf<E> assignNewLabel(E* e, Args&&... args) {
     assert(waitingForNewLabel == Store::Handle{e} && "assignNewLabel called on wrong entity");
-    auto label = createLabel<TrapTagOf<E>>();
+    auto label = createLabel<TrapTagOf<E>>(std::forward<Args>(args)...);
     store.insert(e, label);
     waitingForNewLabel = std::monostate{};
     return label;
@@ -120,6 +133,35 @@ class SwiftDispatcher {
                    locatableLabel);
   }
 
+  // In order to not emit duplicated entries for declarations, we restrict emission to only
+  // Decls declared within the current "scope".
+  // Depending on the SwiftExtractionMode the scope is defined as follows:
+  //  - SwiftExtractionMode::Module: the current scope means the current module. This is used in
+  //    the case of system or builtin modules.
+  //  - SwiftExtractionMode::PrimaryFile: in this mode, we extract several files belonging to the
+  //    same module one by one. In this mode, we restrict emission only to the same file ignoring
+  //    all the other files.
+  // TODO this currently does not extract compiler-synthesized entities without a valid location,
+  // this will be fixed in an upcoming PR
+  bool shouldEmitDeclBody(swift::Decl* decl) {
+    switch (extractionMode) {
+      case SwiftExtractionMode::Module: {
+        return &currentModule == decl->getModuleContext();
+      } break;
+      case SwiftExtractionMode::PrimaryFile: {
+        swift::SourceLoc location = decl->getStartLoc();
+        if (!location.isValid()) {
+          return false;
+        }
+        auto declFileName = sourceManager.getDisplayNameForLoc(location).str();
+        return &currentModule == decl->getModuleContext() && declFileName == currentFileName;
+      } break;
+      default:
+        return false;
+    }
+    return false;
+  }
+
  private:
   // types to be supported by assignNewLabel/fetchLabel need to be listed here
   using Store = TrapLabelStore<swift::Decl,
@@ -199,6 +241,9 @@ class SwiftDispatcher {
   TrapOutput& trap;
   Store store;
   Store::Handle waitingForNewLabel{std::monostate{}};
+  SwiftExtractionMode extractionMode;
+  swift::ModuleDecl& currentModule;
+  llvm::StringRef currentFileName;
 };
 
 }  // namespace codeql

swift/extractor/SwiftExtractor.cpp

@@ -51,13 +51,15 @@ static void archiveFile(const SwiftExtractorConfiguration& config, swift::Source
 
 static void extractDeclarations(const SwiftExtractorConfiguration& config,
                                 swift::CompilerInstance& compiler,
-                                llvm::StringRef filename,
+                                SwiftExtractionMode extractionMode,
+                                swift::ModuleDecl& module,
+                                llvm::StringRef fileName,
                                 llvm::ArrayRef<swift::Decl*> topLevelDecls) {
   // The extractor can be called several times from different processes with
   // the same input file(s)
   // We are using PID to avoid concurrent access
   // TODO: find a more robust approach to avoid collisions?
-  std::string tempTrapName = filename.str() + '.' + std::to_string(getpid()) + ".trap";
+  std::string tempTrapName = fileName.str() + '.' + std::to_string(getpid()) + ".trap";
   llvm::SmallString<PATH_MAX> tempTrapPath(config.trapDir);
   llvm::sys::path::append(tempTrapPath, tempTrapName);
 
@@ -94,23 +96,22 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
   trap.assignKey(unknownLocationLabel, "unknown");
   trap.emit(LocationsTrap{unknownLocationLabel, unknownFileLabel});
 
-  SwiftVisitor visitor(compiler.getSourceMgr(), arena, trap);
-  for (swift::Decl* decl : topLevelDecls) {
+  SwiftVisitor visitor(compiler.getSourceMgr(), arena, trap, extractionMode, module, fileName);
+  for (auto decl : topLevelDecls) {
     visitor.extract(decl);
   }
   if (topLevelDecls.empty()) {
-    // In the case of emtpy files, the dispatcher is not called, but we still want to 'record' the
+    // In the case of empty files, the dispatcher is not called, but we still want to 'record' the
     // fact that the file was extracted
-    // TODO: to be moved elsewhere
-    llvm::SmallString<PATH_MAX> srcFilePath(filename);
-    llvm::sys::fs::make_absolute(srcFilePath);
+    llvm::SmallString<PATH_MAX> name(fileName);
+    llvm::sys::fs::make_absolute(name);
     auto fileLabel = arena.allocateLabel<FileTag>();
-    trap.assignKey(fileLabel, srcFilePath.str().str());
-    trap.emit(FilesTrap{fileLabel, srcFilePath.str().str()});
+    trap.assignKey(fileLabel, name.str().str());
+    trap.emit(FilesTrap{fileLabel, name.str().str()});
   }
 
   // TODO: Pick a better name to avoid collisions
-  std::string trapName = filename.str() + ".trap";
+  std::string trapName = fileName.str() + ".trap";
   llvm::SmallString<PATH_MAX> trapPath(config.trapDir);
   llvm::sys::path::append(trapPath, trapName);
 
@@ -132,14 +133,15 @@ void codeql::extractSwiftFiles(const SwiftExtractorConfiguration& config,
       llvm::SmallVector<swift::Decl*> decls;
       module->getTopLevelDecls(decls);
       // TODO: pass ModuleDecl directly when we have module extraction in place?
-      extractDeclarations(config, compiler, module->getModuleFilename(), decls);
+      extractDeclarations(config, compiler, SwiftExtractionMode::Module, *module,
+                          module->getModuleFilename(), decls);
     } else {
       // The extraction will only work if one (or more) `-primary-file` CLI option is provided,
       // which is what always happens in case of `swift build` and `xcodebuild`
       for (auto primaryFile : module->getPrimarySourceFiles()) {
         archiveFile(config, *primaryFile);
-        extractDeclarations(config, compiler, primaryFile->getFilename(),
-                            primaryFile->getTopLevelDecls());
+        extractDeclarations(config, compiler, SwiftExtractionMode::PrimaryFile, *module,
+                            primaryFile->getFilename(), primaryFile->getTopLevelDecls());
       }
     }
   }

swift/extractor/trap/TrapOutput.h

@@ -39,6 +39,12 @@ class TrapOutput {
     print(e);
   }
 
+  template <typename... Args>
+  void debug(const Args&... args) {
+    out_ << "// DEBUG: ";
+    (out_ << ... << args) << '\n';
+  }
+
  private:
   template <typename... Args>
   void print(const Args&... args) {

swift/extractor/visitors/DeclVisitor.h

@@ -3,6 +3,7 @@
 #include <swift/AST/Decl.h>
 #include <swift/AST/GenericParamList.h>
 #include <swift/AST/ParameterList.h>
+#include <swift/AST/ASTMangler.h>
 
 #include "swift/extractor/visitors/VisitorBase.h"
 
@@ -20,20 +21,29 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
   using AstVisitorBase<DeclVisitor>::AstVisitorBase;
 
   void visitFuncDecl(swift::FuncDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(ConcreteFuncDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitAbstractFunctionDecl(decl, label);
   }
 
   void visitConstructorDecl(swift::ConstructorDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(ConstructorDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitConstructorDecl(decl, label);
   }
 
   void visitDestructorDecl(swift::DestructorDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(DestructorDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitDestructorDecl(decl, label);
   }
 
@@ -64,6 +74,7 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
   }
 
   void visitParamDecl(swift::ParamDecl* decl) {
+    // TODO: deduplicate
     auto label = dispatcher_.assignNewLabel(decl);
     dispatcher_.emit(ParamDeclsTrap{label});
     if (decl->isInOut()) {
@@ -92,33 +103,46 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
   }
 
   void visitVarDecl(swift::VarDecl* decl) {
+    // TODO: deduplicate all non-local variables
     auto label = dispatcher_.assignNewLabel(decl);
     auto introducer = static_cast<uint8_t>(decl->getIntroducer());
     dispatcher_.emit(ConcreteVarDeclsTrap{label, introducer});
     emitVarDecl(decl, label);
   }
 
   void visitStructDecl(swift::StructDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(StructDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitNominalTypeDecl(decl, label);
   }
 
   void visitClassDecl(swift::ClassDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(ClassDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitNominalTypeDecl(decl, label);
   }
 
   void visitEnumDecl(swift::EnumDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(EnumDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitNominalTypeDecl(decl, label);
   }
 
   void visitProtocolDecl(swift::ProtocolDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(ProtocolDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitNominalTypeDecl(decl, label);
   }
 
@@ -132,8 +156,11 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
   }
 
   void visitEnumElementDecl(swift::EnumElementDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(EnumElementDeclsTrap{label, decl->getNameStr().str()});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     if (decl->hasParameterList()) {
       auto i = 0u;
       for (auto p : *decl->getParameterList()) {
@@ -143,26 +170,36 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
   }
 
   void visitGenericTypeParamDecl(swift::GenericTypeParamDecl* decl) {
+    // TODO: deduplicate
     auto label = dispatcher_.assignNewLabel(decl);
     dispatcher_.emit(GenericTypeParamDeclsTrap{label});
     emitTypeDecl(decl, label);
   }
 
   void visitAssociatedTypeDecl(swift::AssociatedTypeDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(AssociatedTypeDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitTypeDecl(decl, label);
   }
 
   void visitTypeAliasDecl(swift::TypeAliasDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(TypeAliasDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     emitTypeDecl(decl, label);
   }
 
   void visitAccessorDecl(swift::AccessorDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
     dispatcher_.emit(AccessorDeclsTrap{label});
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     switch (decl->getAccessorKind()) {
       case swift::AccessorKind::Get:
         dispatcher_.emit(AccessorDeclIsGetterTrap{label});
@@ -181,7 +218,10 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
   }
 
   void visitSubscriptDecl(swift::SubscriptDecl* decl) {
-    auto label = dispatcher_.assignNewLabel(decl);
+    auto label = dispatcher_.assignNewLabel(decl, mangledName(decl));
+    if (!dispatcher_.shouldEmitDeclBody(decl)) {
+      return;
+    }
     auto elementTypeLabel = dispatcher_.fetchLabel(decl->getElementInterfaceType());
     dispatcher_.emit(SubscriptDeclsTrap{label, elementTypeLabel});
     if (auto indices = decl->getIndices()) {
@@ -202,6 +242,11 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
   }
 
  private:
+  std::string mangledName(swift::ValueDecl* decl) {
+    // prefix adds a couple of special symbols, we don't necessary need them
+    return mangler.mangleAnyDecl(decl, /* prefix = */ false);
+  }
+
   void emitConstructorDecl(swift::ConstructorDecl* decl, TrapLabel<ConstructorDeclTag> label) {
     emitAbstractFunctionDecl(decl, label);
   }
@@ -309,6 +354,9 @@ class DeclVisitor : public AstVisitorBase<DeclVisitor> {
     }
     emitValueDecl(decl, label);
   }
+
+ private:
+  swift::Mangle::ASTMangler mangler;
 };
 
 }  // namespace codeql

swift/integration-tests/.gitignore

@@ -1,7 +1,8 @@
 .DS_Store
-/.build
-/Packages
-/*.xcodeproj
+.build
+Packages
+*.xcodeproj
 xcuserdata/
 DerivedData/
 .swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata
+*.actual

swift/integration-tests/cross-references/Classes.expected

@@ -0,0 +1 @@
+| Sources/cross-references/lib.swift:3:1:6:1 | C |

swift/integration-tests/cross-references/Classes.ql

@@ -0,0 +1,5 @@
+import swift
+
+from ClassDecl d
+where d.getLocation().getFile().getBaseName() != "Package.swift"
+select d

swift/integration-tests/cross-references/Constructors.expected

@@ -0,0 +1,3 @@
+| Sources/cross-references/lib.swift:4:5:4:13 | init |
+| Sources/cross-references/lib.swift:14:8:14:8 | init |
+| Sources/cross-references/main.swift:11:8:11:8 | init |

swift/integration-tests/cross-references/Constructors.ql

@@ -0,0 +1,5 @@
+import swift
+
+from ConstructorDecl d
+where d.getLocation().getFile().getBaseName() != "Package.swift"
+select d

swift/integration-tests/cross-references/Destructors.expected

@@ -0,0 +1 @@
+| Sources/cross-references/lib.swift:5:5:5:13 | deinit |