Swift: extract all output-producing source files, not only primary files

Author: AlexDenisov

swift/extractor/SwiftExtractor.cpp

@@ -6,8 +6,10 @@
 #include <sstream>
 #include <memory>
 #include <unistd.h>
+#include <unordered_set>
 
 #include <swift/AST/SourceFile.h>
+#include <swift/Basic/FileTypes.h>
 #include <llvm/ADT/SmallString.h>
 #include <llvm/Support/FileSystem.h>
 #include <llvm/Support/Path.h>
@@ -124,6 +126,17 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
 
 void codeql::extractSwiftFiles(const SwiftExtractorConfiguration& config,
                                swift::CompilerInstance& compiler) {
+  // The frontend can be called in many different ways.
+  // At each invocation we only extract system and builtin modules and any input source files that
+  // have an output associated with them.
+  std::unordered_set<std::string> sourceFiles;
+  auto inputFiles = compiler.getInvocation().getFrontendOptions().InputsAndOutputs.getAllInputs();
+  for (auto& input : inputFiles) {
+    if (input.getType() == swift::file_types::TY_Swift && !input.outputFilename().empty()) {
+      sourceFiles.insert(input.getFileName());
+    }
+  }
+
   for (auto& [_, module] : compiler.getASTContext().getLoadedModules()) {
     // We only extract system and builtin modules here as the other "user" modules can be built
     // during the build process and then re-used at a later stage. In this case, we extract the
@@ -135,12 +148,16 @@ void codeql::extractSwiftFiles(const SwiftExtractorConfiguration& config,
       // TODO: pass ModuleDecl directly when we have module extraction in place?
       extractDeclarations(config, decls, compiler, *module);
     } 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, primaryFile->getTopLevelDecls(), compiler, *module,
-                            primaryFile);
+      for (auto file : module->getFiles()) {
+        if (!llvm::isa<swift::SourceFile>(file)) {
+          continue;
+        }
+        auto sourceFile = llvm::cast<swift::SourceFile>(file);
+        if (sourceFiles.count(sourceFile->getFilename().str()) == 0) {
+          continue;
+        }
+        archiveFile(config, *sourceFile);
+        extractDeclarations(config, sourceFile->getTopLevelDecls(), compiler, *module, sourceFile);
       }
     }
   }

swift/integration-tests/frontend-invocations/A.swift

@@ -0,0 +1,6 @@
+| A.swift:0:0:0:0 | A.swift |
+| B.swift:0:0:0:0 | B.swift |
+| C.swift:0:0:0:0 | C.swift |
+| D.swift:0:0:0:0 | D.swift |
+| E.swift:0:0:0:0 | E.swift |
+| file://:0:0:0:0 |  |

swift/integration-tests/frontend-invocations/B.swift

@@ -0,0 +1,4 @@
+import swift
+
+from File f
+select f

swift/integration-tests/frontend-invocations/C.swift

@@ -0,0 +1,10 @@
+# TODO: Add linux
+SDK=$(shell xcrun -show-sdk-path)
+FRONTEND=$(shell xcrun -find swift-frontend)
+
+all:
+	$(FRONTEND) -frontend -c A.swift -sdk $(SDK)
+	$(FRONTEND) -frontend -c B.swift -o B.o -sdk $(SDK)
+	$(FRONTEND) -frontend -c -primary-file C.swift -sdk $(SDK)
+	$(FRONTEND) -frontend -c -primary-file D.swift -o D.o -sdk $(SDK)
+	$(FRONTEND) -frontend -c -primary-file E.swift Esup.swift -o E.o -sdk $(SDK)