Merge pull request #8788 from AlexDenisov/alexdenisov/swift-first-extractor-test

Swift: file extraction

Author: redsun82

.bazelrc

@@ -1,3 +1,4 @@
-build --copt="-std=c++17"
+# -fno-rtti is required by LLVM/Swift
+build --repo_env=CC=clang --repo_env=CXX=clang++ --copt="-std=c++17" --copt="-fno-rtti"
 
 try-import %workspace%/local.bazelrc

swift/codegen/prefix.dbscheme

@@ -5,6 +5,3 @@ sourceLocationPrefix(
   string prefix: string ref
 );
 
-answer_to_life_the_universe_and_everything(
-   int answer: int ref
-);

swift/extractor/BUILD.bazel

@@ -8,7 +8,12 @@ alias(
 
 cc_binary(
     name = "extractor",
-    srcs = ["main.cpp"],
+    srcs = [
+        "SwiftExtractor.h",
+        "SwiftExtractor.cpp",
+        "SwiftExtractorConfiguration.h",
+        "main.cpp",
+    ],
     target_compatible_with = select({
         "@platforms//os:linux": [],
         "@platforms//os:macos": [],

swift/extractor/SwiftExtractor.cpp

@@ -0,0 +1,102 @@
+#include "SwiftExtractor.h"
+
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <sstream>
+#include <memory>
+#include <unistd.h>
+
+#include <swift/AST/SourceFile.h>
+#include <llvm/ADT/SmallString.h>
+#include <llvm/Support/FileSystem.h>
+#include <llvm/Support/Path.h>
+
+using namespace codeql;
+
+static void extractFile(const SwiftExtractorConfiguration& config, swift::SourceFile& file) {
+  if (std::error_code ec = llvm::sys::fs::create_directories(config.trapDir)) {
+    std::cerr << "Cannot create TRAP directory: " << ec.message() << "\n";
+    return;
+  }
+
+  if (std::error_code ec = llvm::sys::fs::create_directories(config.sourceArchiveDir)) {
+    std::cerr << "Cannot create source archive directory: " << ec.message() << "\n";
+    return;
+  }
+
+  llvm::SmallString<PATH_MAX> srcFilePath(file.getFilename());
+  llvm::sys::fs::make_absolute(srcFilePath);
+
+  llvm::SmallString<PATH_MAX> dstFilePath(config.sourceArchiveDir);
+  llvm::sys::path::append(dstFilePath, srcFilePath);
+
+  llvm::StringRef parent = llvm::sys::path::parent_path(dstFilePath);
+  if (std::error_code ec = llvm::sys::fs::create_directories(parent)) {
+    std::cerr << "Cannot create source archive destination directory '" << parent.str()
+              << "': " << ec.message() << "\n";
+    return;
+  }
+
+  if (std::error_code ec = llvm::sys::fs::copy_file(srcFilePath, dstFilePath)) {
+    std::cerr << "Cannot archive source file '" << srcFilePath.str().str() << "' -> '"
+              << dstFilePath.str().str() << "': " << ec.message() << "\n";
+    return;
+  }
+
+  // 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 = file.getFilename().str() + '.' + std::to_string(getpid()) + ".trap";
+  llvm::SmallString<PATH_MAX> tempTrapPath(config.trapDir);
+  llvm::sys::path::append(tempTrapPath, tempTrapName);
+
+  std::ofstream trap(tempTrapPath.str().str());
+  if (!trap) {
+    std::error_code ec;
+    ec.assign(errno, std::generic_category());
+    std::cerr << "Cannot create temp trap file '" << tempTrapPath.str().str()
+              << "': " << ec.message() << "\n";
+    return;
+  }
+  std::stringstream ss;
+  for (auto opt : config.frontendOptions) {
+    ss << std::quoted(opt) << " ";
+  }
+  ss << "\n";
+  trap << "// extractor-args: " << ss.str();
+
+  trap << "#0=*\n";
+  trap << "files(#0, " << std::quoted(srcFilePath.str().str()) << ")\n";
+
+  // TODO: Pick a better name to avoid collisions
+  std::string trapName = file.getFilename().str() + ".trap";
+  llvm::SmallString<PATH_MAX> trapPath(config.trapDir);
+  llvm::sys::path::append(trapPath, trapName);
+
+  // TODO: The last process wins. Should we do better than that?
+  if (std::error_code ec = llvm::sys::fs::rename(tempTrapPath, trapPath)) {
+    std::cerr << "Cannot rename temp trap file '" << tempTrapPath.str().str() << "' -> '"
+              << trapPath.str().str() << "': " << ec.message() << "\n";
+  }
+}
+
+void codeql::extractSwiftFiles(const SwiftExtractorConfiguration& config,
+                               swift::CompilerInstance& compiler) {
+  // Swift frontend can be called in several different modes, we are interested
+  // only in the cases when either a primary or a main source file is present
+  if (compiler.getPrimarySourceFiles().empty()) {
+    swift::ModuleDecl* module = compiler.getMainModule();
+    if (!module->getFiles().empty() &&
+        module->getFiles().front()->getKind() == swift::FileUnitKind::Source) {
+      // We can only call getMainSourceFile if the first file is of a Source kind
+      swift::SourceFile& file = module->getMainSourceFile();
+      extractFile(config, file);
+    }
+  } else {
+    for (auto s : compiler.getPrimarySourceFiles()) {
+      extractFile(config, *s);
+    }
+  }
+}

swift/extractor/SwiftExtractor.h

@@ -0,0 +1,14 @@
+#ifndef SWIFT_EXTRACTOR_H_
+#define SWIFT_EXTRACTOR_H_
+
+#include "SwiftExtractorConfiguration.h"
+#include <swift/AST/SourceFile.h>
+#include <swift/Frontend/Frontend.h>
+#include <memory>
+
+namespace codeql {
+void extractSwiftFiles(const SwiftExtractorConfiguration& config,
+                       swift::CompilerInstance& compiler);
+}  // namespace codeql
+
+#endif  // SWIFT_EXTRACTOR_H_

swift/extractor/SwiftExtractorConfiguration.h

@@ -0,0 +1,18 @@
+#ifndef SWIFT_EXTRACTOR_CONFIGURATION_H_
+#define SWIFT_EXTRACTOR_CONFIGURATION_H_
+
+#include <string>
+#include <vector>
+
+namespace codeql {
+struct SwiftExtractorConfiguration {
+  // The location for storing TRAP files to be imported by CodeQL engine.
+  std::string trapDir;
+  // The location for storing extracted source files.
+  std::string sourceArchiveDir;
+  // The arguments passed to the extractor. Used for debugging.
+  std::vector<std::string> frontendOptions;
+};
+}  // namespace codeql
+
+#endif  // SWIFT_EXTRACTOR_CONFIGURATION_H_

swift/extractor/main.cpp

@@ -3,15 +3,52 @@
 #include <stdlib.h>
 
 #include <swift/Basic/LLVMInitialize.h>
+#include <swift/FrontendTool/FrontendTool.h>
+
+#include "SwiftExtractor.h"
+
+using namespace std::string_literals;
+
+// This is part of the swiftFrontendTool interface, we hook into the
+// compilation pipeline and extract files after the Swift frontend performed
+// semantic analysys
+class Observer : public swift::FrontendObserver {
+ public:
+  explicit Observer(const codeql::SwiftExtractorConfiguration& config) : config{config} {}
+
+  void performedSemanticAnalysis(swift::CompilerInstance& compiler) override {
+    codeql::extractSwiftFiles(config, compiler);
+  }
+
+ private:
+  const codeql::SwiftExtractorConfiguration& config;
+};
+
+static std::string getenv_or(const char* envvar, const std::string& def) {
+  if (const char* var = getenv(envvar)) {
+    return var;
+  }
+  return def;
+}
 
 int main(int argc, char** argv) {
+  if (argc == 1) {
+    // TODO: print usage
+    return 1;
+  }
+  // Required by Swift/LLVM
   PROGRAM_START(argc, argv);
-  if (auto trapDir = getenv("CODEQL_EXTRACTOR_SWIFT_TRAP_DIR")) {
-    std::string file = trapDir;
-    file += "/my_first.trap";
-    if (std::ofstream out{file}) {
-      out << "answer_to_life_the_universe_and_everything(42)\n";
-    }
+  INITIALIZE_LLVM();
+
+  codeql::SwiftExtractorConfiguration configuration{};
+  configuration.trapDir = getenv_or("CODEQL_EXTRACTOR_SWIFT_TRAP_DIR", ".");
+  configuration.sourceArchiveDir = getenv_or("CODEQL_EXTRACTOR_SWIFT_SOURCE_ARCHIVE_DIR", ".");
+  std::vector<const char*> args;
+  for (int i = 1; i < argc; i++) {
+    args.push_back(argv[i]);
   }
-  return 0;
+  std::copy(std::begin(args), std::end(args), std::back_inserter(configuration.frontendOptions));
+  Observer observer(configuration);
+  int frontend_rc = swift::performFrontend(args, "swift-extractor", (void*)main, &observer);
+  return frontend_rc;
 }

swift/ql/lib/codeql/swift/elements/File.qll

@@ -1,4 +1,23 @@
-// generated by codegen/codegen.py, remove this comment if you wish to edit this file
 private import codeql.swift.generated.File
 
-class File extends FileBase { }
+class File extends FileBase {
+  /** toString */
+  override string toString() { result = getAbsolutePath() }
+
+  /** Gets the name of this file. */
+  override string getName() { files(this, result) }
+
+  /** Gets the absolute path of this file. */
+  string getAbsolutePath() { result = getName() }
+
+  /** Gets the full name of this file. */
+  string getFullName() { result = getAbsolutePath() }
+
+  /** Gets the URL of this file. */
+  string getURL() { result = "file://" + this.getAbsolutePath() + ":0:0:0:0" }
+
+  /** Gets the base name of this file. */
+  string getBaseName() {
+    result = this.getAbsolutePath().regexpCapture(".*/(([^/]*?)(?:\\.([^.]*))?)", 1)
+  }
+}

swift/ql/lib/swift.dbscheme

@@ -8,9 +8,6 @@ sourceLocationPrefix(
   string prefix: string ref
 );
 
-answer_to_life_the_universe_and_everything(
-   int answer: int ref
-);
 
 
 // from codegen/schema.yml

swift/ql/lib/swift.qll

@@ -0,0 +1,3 @@
+/** Top-level import for the Swift language pack */
+
+import codeql.swift.elements