Swift: trap output rework

Firstly, this change reworks how inter-process races are resolved.
Moreover some responsability reorganization has led to merging
`TrapArena` and `TrapOutput` again into a `TrapDomain` class.

A `TargetFile` class is introduced, that is successfully created
only for the first process that starts processing a given trap output
file. From then on `TargetFile` simply wraps around `<<` stream
operations, dumping them to a temporary file. When `TargetFile::commit`
is called, the temporary file is moved on to the actual target trap
file.

Processes that lose the race can now just ignore the unneeded
extraction and go on, while previously all processes would carry out
all extractions overwriting each other at the end.

Some of the file system logic contained in `SwiftExtractor.cpp` has been
moved to this class, and two TODOs are solved:
* introducing a better inter process file collision avoidance strategy
* better error handling for trap output operations: if unable to write
  to the trap file (or carry out other basic file operations), we just
  abort.

The changes to `ExprVisitor` and `StmtVisitor` are due to wanting to
hide the raw `TrapDomain::createLabel` from them, and bring more
funcionality under the generic caching/dispatching mechanism.

Author: redsun82

swift/codegen/schema.yml

@@ -319,6 +319,7 @@ AppliedPropertyWrapperExpr:
   _extends: Expr
 
 Argument:
+  _extends: Locatable
   label: string
   _children:
     expr: Expr

swift/extractor/SwiftExtractor.cpp

@@ -16,8 +16,9 @@
 #include <llvm/Support/Path.h>
 
 #include "swift/extractor/trap/generated/TrapClasses.h"
-#include "swift/extractor/trap/TrapOutput.h"
+#include "swift/extractor/trap/TrapDomain.h"
 #include "swift/extractor/visitors/SwiftVisitor.h"
+#include "swift/extractor/infra/TargetFile.h"
 
 using namespace codeql;
 
@@ -52,7 +53,7 @@ static void archiveFile(const SwiftExtractorConfiguration& config, swift::Source
   }
 }
 
-static std::string getTrapFilename(swift::ModuleDecl& module, swift::SourceFile* primaryFile) {
+static std::string getFilename(swift::ModuleDecl& module, swift::SourceFile* primaryFile) {
   if (primaryFile) {
     return primaryFile->getFilename().str();
   }
@@ -80,56 +81,39 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
                                 swift::CompilerInstance& compiler,
                                 swift::ModuleDecl& module,
                                 swift::SourceFile* primaryFile = nullptr) {
+  auto filename = getFilename(module, primaryFile);
+
   // 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?
-  auto name = getTrapFilename(module, primaryFile);
-  llvm::StringRef filename(name);
-  std::string tempTrapName = filename.str() + '.' + std::to_string(getpid()) + ".trap";
-  llvm::SmallString<PATH_MAX> tempTrapPath(config.getTempTrapDir());
-  llvm::sys::path::append(tempTrapPath, tempTrapName);
-
-  llvm::StringRef tempTrapParent = llvm::sys::path::parent_path(tempTrapPath);
-  if (std::error_code ec = llvm::sys::fs::create_directories(tempTrapParent)) {
-    std::cerr << "Cannot create temp trap directory '" << tempTrapParent.str()
-              << "': " << ec.message() << "\n";
+  // the same input file(s). Using `TargetFile` the first process will win, and the following
+  // will just skip the work
+  TargetFile trapStream{filename + ".trap", config.trapDir, config.getTempTrapDir()};
+  if (!trapStream.good()) {
+    // another process arrived first, nothing to do for us
     return;
   }
 
-  std::ofstream trapStream(tempTrapPath.str().str());
-  if (!trapStream) {
-    std::error_code ec;
-    ec.assign(errno, std::generic_category());
-    std::cerr << "Cannot create temp trap file '" << tempTrapPath.str().str()
-              << "': " << ec.message() << "\n";
-    return;
-  }
   trapStream << "/* extractor-args:\n";
-  for (auto opt : config.frontendOptions) {
+  for (const auto& opt : config.frontendOptions) {
     trapStream << "  " << std::quoted(opt) << " \\\n";
   }
   trapStream << "\n*/\n";
 
   trapStream << "/* swift-frontend-args:\n";
-  for (auto opt : config.patchedFrontendOptions) {
+  for (const auto& opt : config.patchedFrontendOptions) {
     trapStream << "  " << std::quoted(opt) << " \\\n";
   }
   trapStream << "\n*/\n";
 
-  TrapOutput trap{trapStream};
-  TrapArena arena{};
+  TrapDomain trap{trapStream};
 
   // TODO: move default location emission elsewhere, possibly in a separate global trap file
-  auto unknownFileLabel = arena.allocateLabel<FileTag>();
   // the following cannot conflict with actual files as those have an absolute path starting with /
-  trap.assignKey(unknownFileLabel, "unknown");
+  auto unknownFileLabel = trap.createLabel<FileTag>("unknown");
+  auto unknownLocationLabel = trap.createLabel<LocationTag>("unknown");
   trap.emit(FilesTrap{unknownFileLabel});
-  auto unknownLocationLabel = arena.allocateLabel<LocationTag>();
-  trap.assignKey(unknownLocationLabel, "unknown");
   trap.emit(LocationsTrap{unknownLocationLabel, unknownFileLabel});
 
-  SwiftVisitor visitor(compiler.getSourceMgr(), arena, trap, module, primaryFile);
+  SwiftVisitor visitor(compiler.getSourceMgr(), trap, module, primaryFile);
   auto topLevelDecls = getTopLevelDecls(module, primaryFile);
   for (auto decl : topLevelDecls) {
     visitor.extract(decl);
@@ -142,28 +126,10 @@ static void extractDeclarations(const SwiftExtractorConfiguration& config,
     // fact that the file was extracted
     llvm::SmallString<PATH_MAX> name(filename);
     llvm::sys::fs::make_absolute(name);
-    auto fileLabel = arena.allocateLabel<FileTag>();
-    trap.assignKey(fileLabel, name.str().str());
+    auto fileLabel = trap.createLabel<FileTag>(name.str().str());
     trap.emit(FilesTrap{fileLabel, name.str().str()});
   }
-
-  // TODO: Pick a better name to avoid collisions
-  std::string trapName = filename.str() + ".trap";
-  llvm::SmallString<PATH_MAX> trapPath(config.trapDir);
-  llvm::sys::path::append(trapPath, trapName);
-
-  llvm::StringRef trapParent = llvm::sys::path::parent_path(trapPath);
-  if (std::error_code ec = llvm::sys::fs::create_directories(trapParent)) {
-    std::cerr << "Cannot create trap directory '" << trapParent.str() << "': " << ec.message()
-              << "\n";
-    return;
-  }
-
-  // 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";
-  }
+  trapStream.commit();
 }
 
 static std::unordered_set<std::string> collectInputFilenames(swift::CompilerInstance& compiler) {

swift/extractor/infra/BUILD.bazel

@@ -2,9 +2,11 @@ load("//swift:rules.bzl", "swift_cc_library")
 
 swift_cc_library(
     name = "infra",
+    srcs = glob(["*.cpp"]),
     hdrs = glob(["*.h"]),
     visibility = ["//swift:__subpackages__"],
     deps = [
         "//swift/extractor/trap",
+        "//swift/tools/prebuilt:swift-llvm-support",
     ],
 )

swift/extractor/infra/SwiftDispatcher.h

@@ -4,9 +4,8 @@
 #include <swift/Basic/SourceManager.h>
 #include <llvm/Support/FileSystem.h>
 
-#include "swift/extractor/trap/TrapArena.h"
 #include "swift/extractor/trap/TrapLabelStore.h"
-#include "swift/extractor/trap/TrapOutput.h"
+#include "swift/extractor/trap/TrapDomain.h"
 #include "swift/extractor/infra/SwiftTagTraits.h"
 #include "swift/extractor/trap/generated/TrapClasses.h"
 
@@ -22,12 +21,10 @@ class SwiftDispatcher {
   // all references and pointers passed as parameters to this constructor are supposed to outlive
   // the SwiftDispatcher
   SwiftDispatcher(const swift::SourceManager& sourceManager,
-                  TrapArena& arena,
-                  TrapOutput& trap,
+                  TrapDomain& trap,
                   swift::ModuleDecl& currentModule,
                   swift::SourceFile* currentPrimarySourceFile = nullptr)
       : sourceManager{sourceManager},
-        arena{arena},
         trap{trap},
         currentModule{currentModule},
         currentPrimarySourceFile{currentPrimarySourceFile} {}
@@ -77,6 +74,7 @@ class SwiftDispatcher {
     }
     waitingForNewLabel = e;
     visit(e);
+    // TODO when everything is moved to structured C++ classes, this should be moved to createEntry
     if (auto l = store.get(e)) {
       if constexpr (!std::is_base_of_v<swift::TypeBase, E>) {
         attachLocation(e, *l);
@@ -95,13 +93,17 @@ class SwiftDispatcher {
     return fetchLabelFromUnion<AstNodeTag>(node);
   }
 
+  TrapLabel<ConditionElementTag> fetchLabel(const swift::StmtConditionElement& element) {
+    return fetchLabel(&element);
+  }
+
   // 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, typename... Args>
   TrapLabelOf<E> assignNewLabel(E* e, Args&&... args) {
     assert(waitingForNewLabel == Store::Handle{e} && "assignNewLabel called on wrong entity");
-    auto label = createLabel<TrapTagOf<E>>(std::forward<Args>(args)...);
+    auto label = trap.createLabel<TrapTagOf<E>>(std::forward<Args>(args)...);
     store.insert(e, label);
     waitingForNewLabel = std::monostate{};
     return label;
@@ -118,18 +120,13 @@ class SwiftDispatcher {
     return TrapClassOf<E>{assignNewLabel(&e, std::forward<Args>(args)...)};
   }
 
-  template <typename Tag>
-  TrapLabel<Tag> createLabel() {
-    auto ret = arena.allocateLabel<Tag>();
-    trap.assignStar(ret);
-    return ret;
-  }
-
-  template <typename Tag, typename... Args>
-  TrapLabel<Tag> createLabel(Args&&... args) {
-    auto ret = arena.allocateLabel<Tag>();
-    trap.assignKey(ret, std::forward<Args>(args)...);
-    return ret;
+  // used to create a new entry for entities that should not be cached
+  // an example is swift::Argument, that are created on the fly and thus have no stable pointer
+  template <typename E, typename... Args, std::enable_if_t<!std::is_pointer_v<E>>* = nullptr>
+  auto createUncachedEntry(const E& e, Args&&... args) {
+    auto label = trap.createLabel<TrapTagOf<E>>(std::forward<Args>(args)...);
+    attachLocation(&e, label);
+    return TrapClassOf<E>{label};
   }
 
   template <typename Locatable>
@@ -213,6 +210,7 @@ class SwiftDispatcher {
   using Store = TrapLabelStore<swift::Decl,
                                swift::Stmt,
                                swift::StmtCondition,
+                               swift::StmtConditionElement,
                                swift::CaseLabelItem,
                                swift::Expr,
                                swift::Pattern,
@@ -227,13 +225,13 @@ class SwiftDispatcher {
       return;
     }
     std::string filepath = getFilepath(start);
-    auto fileLabel = createLabel<FileTag>(filepath);
+    auto fileLabel = trap.createLabel<FileTag>(filepath);
     // TODO: do not emit duplicate trap entries for Files
     trap.emit(FilesTrap{fileLabel, filepath});
     auto [startLine, startColumn] = sourceManager.getLineAndColumnInBuffer(start);
     auto [endLine, endColumn] = sourceManager.getLineAndColumnInBuffer(end);
-    auto locLabel = createLabel<LocationTag>('{', fileLabel, "}:", startLine, ':', startColumn, ':',
-                                             endLine, ':', endColumn);
+    auto locLabel = trap.createLabel<LocationTag>('{', fileLabel, "}:", startLine, ':', startColumn,
+                                                  ':', endLine, ':', endColumn);
     trap.emit(LocationsTrap{locLabel, fileLabel, startLine, startColumn, endLine, endColumn});
     trap.emit(LocatableLocationsTrap{locatableLabel, locLabel});
   }
@@ -275,16 +273,16 @@ class SwiftDispatcher {
   // which are to be introduced in follow-up PRs
   virtual void visit(swift::Decl* decl) = 0;
   virtual void visit(swift::Stmt* stmt) = 0;
-  virtual void visit(swift::StmtCondition* cond) = 0;
+  virtual void visit(const swift::StmtCondition* cond) = 0;
+  virtual void visit(const swift::StmtConditionElement* cond) = 0;
   virtual void visit(swift::CaseLabelItem* item) = 0;
   virtual void visit(swift::Expr* expr) = 0;
   virtual void visit(swift::Pattern* pattern) = 0;
   virtual void visit(swift::TypeRepr* type) = 0;
   virtual void visit(swift::TypeBase* type) = 0;
 
   const swift::SourceManager& sourceManager;
-  TrapArena& arena;
-  TrapOutput& trap;
+  TrapDomain& trap;
   Store store;
   Store::Handle waitingForNewLabel{std::monostate{}};
   swift::ModuleDecl& currentModule;

swift/extractor/infra/SwiftTagTraits.h

@@ -36,12 +36,14 @@ using SILBoxTypeReprTag = SilBoxTypeReprTag;
 
 MAP_TAG(Stmt);
 MAP_TAG(StmtCondition);
+MAP_TYPE_TO_TAG(StmtConditionElement, ConditionElementTag);
 MAP_TAG(CaseLabelItem);
 #define ABSTRACT_STMT(CLASS, PARENT) MAP_SUBTAG(CLASS##Stmt, PARENT)
 #define STMT(CLASS, PARENT) ABSTRACT_STMT(CLASS, PARENT)
 #include <swift/AST/StmtNodes.def>
 
 MAP_TAG(Expr);
+MAP_TAG(Argument);
 #define ABSTRACT_EXPR(CLASS, PARENT) MAP_SUBTAG(CLASS##Expr, PARENT)
 #define EXPR(CLASS, PARENT) ABSTRACT_EXPR(CLASS, PARENT)
 #include <swift/AST/ExprNodes.def>

swift/extractor/infra/TargetFile.cpp

@@ -0,0 +1,73 @@
+#include "swift/extractor/infra/TargetFile.h"
+
+#include <llvm/Support/FileSystem.h>
+#include <llvm/Support/Path.h>
+
+namespace codeql {
+namespace {
+[[noreturn]] void error(const char* action, const std::string& arg, std::error_code ec) {
+  std::cerr << "Unable to " << action << ": " << arg << " (" << ec.message() << ")\n";
+  std::abort();
+}
+
+[[noreturn]] void error(const char* action, const std::string& arg) {
+  error(action, arg, {errno, std::system_category()});
+}
+
+void ensureParentDir(const std::string& path) {
+  auto parent = llvm::sys::path::parent_path(path);
+  if (auto ec = llvm::sys::fs::create_directories(parent)) {
+    error("create directory", parent.str(), ec);
+  }
+}
+
+std::string initPath(std::string_view target, std::string_view dir) {
+  std::string ret{dir};
+  assert(!target.empty() && "target must be a non-empty path");
+  if (target[0] != '/') {
+    ret += '/';
+  }
+  ret.append(target);
+  ensureParentDir(ret);
+  return ret;
+}
+}  // namespace
+
+TargetFile::TargetFile(std::string_view target,
+                       std::string_view targetDir,
+                       std::string_view workingDir)
+    : workingPath{initPath(target, workingDir)}, targetPath{initPath(target, targetDir)} {
+  errno = 0;
+  // since C++17 "x" mode opens with O_EXCL (fails if file already exists)
+  if (auto f = std::fopen(targetPath.c_str(), "wx")) {
+    std::fclose(f);
+    out.open(workingPath);
+    checkOutput("open file for writing");
+  } else {
+    if (errno != EEXIST) {
+      error("open file for writing", targetPath);
+    }
+    // else we just lost the race, do nothing (good() will return false to signal this)
+  }
+}
+
+bool TargetFile::good() const {
+  return out && out.is_open();
+}
+
+void TargetFile::commit() {
+  assert(good());
+  out.close();
+  errno = 0;
+  if (std::rename(workingPath.c_str(), targetPath.c_str()) != 0) {
+    error("rename file", targetPath);
+  }
+}
+
+void TargetFile::checkOutput(const char* action) {
+  if (!out) {
+    error(action, workingPath);
+  }
+}
+
+}  // namespace codeql

swift/extractor/infra/TargetFile.h

@@ -0,0 +1,40 @@
+#pragma once
+
+#include <string>
+#include <fstream>
+#include <iostream>
+#include <cstdio>
+#include <cerrno>
+#include <system_error>
+#include <sstream>
+
+namespace codeql {
+
+// Only the first process trying to open an `TargetFile` for a given `target` is allowed to do
+// so, all others will have an instance with `good() == false` and failing on any other operation.
+// The content streamed to the `TargetFile` is written to `workingDir/target`, and is moved onto
+// `targetDir/target` only when `commit()` is called
+class TargetFile {
+  std::string workingPath;
+  std::string targetPath;
+  std::ofstream out;
+
+ public:
+  TargetFile(std::string_view target, std::string_view targetDir, std::string_view workingDir);
+
+  bool good() const;
+  void commit();
+
+  template <typename T>
+  TargetFile& operator<<(T&& value) {
+    errno = 0;
+    out << value;
+    checkOutput("write to file");
+    return *this;
+  }
+
+ private:
+  void checkOutput(const char* action);
+};
+
+}  // namespace codeql