Swift: move most of TrapArena to TrapFile

swift/extractor/SwiftExtractor.cpp

@@ -14,6 +14,7 @@
 
 #include "swift/extractor/trap/TrapClasses.h"
 #include "swift/extractor/trap/TrapArena.h"
+#include "swift/extractor/trap/TrapOutput.h"
 
 using namespace codeql;
 
@@ -47,51 +48,16 @@ static void extractFile(const SwiftExtractorConfiguration& config, swift::Source
     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);
+  if (TrapOutput trap{config, file.getFilename().str()}) {
+    TrapArena arena{};
 
-  llvm::StringRef trapParent = llvm::sys::path::parent_path(tempTrapPath);
-  if (std::error_code ec = llvm::sys::fs::create_directories(trapParent)) {
-    std::cerr << "Cannot create trap directory '" << trapParent.str() << "': " << ec.message()
-              << "\n";
-    return;
-  }
+    auto label = arena.allocateLabel<File::Tag>();
+    trap.assignStar(label);
 
-  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;
-  }
-  trap << "// extractor-args: ";
-  for (auto opt : config.frontendOptions) {
-    trap << std::quoted(opt) << " ";
-  }
-  trap << "\n\n";
-
-  TrapArena arena{trap};
-  // arena will be passed to another class in a later PR, the next block of code is only an example
-  File f;
-  f.id = arena.getLabel<FileTag>();
-  f.name = srcFilePath.str().str();
-  arena.emit(f);
-
-  // 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";
+    File f{};
+    f.id = label;
+    f.name = srcFilePath.str().str();
+    trap.emit(f);
   }
 }
 

swift/extractor/trap/TrapArena.h

@@ -1,7 +1,3 @@
-//
-// Created by redsun82 on 18.01.22.
-//
-
 #pragma once
 
 #include <iostream>
@@ -12,59 +8,16 @@
 
 namespace codeql {
 
-// TrapArena has the responsibilities to
-// * allocate distinct trap #-labels outputting their assignment to '*' or a @-keys to a trap output
-// * forwarding trap entries to said output
-// TODO: split or move these responsibilities into separate classes
+// TrapArena has the responsibilities to allocate distinct trap #-labels
+// TODO this is now a small functionality that will be moved to code upcoming from other PRs
 class TrapArena {
+  uint64_t id_{0};
+
  public:
-  explicit TrapArena(std::ostream& out) : out_{out} {}
-
-  // get a new label for *
-  template <typename Tag>
-  TrapLabel<Tag> getLabel() {
-    auto ret = allocateLabel<Tag>();
-    print(ret, "=*");
-    return ret;
-  }
-
-  // get a new label for an @-key
-  template <typename Tag>
-  TrapLabel<Tag> getLabel(const std::string& key) {
-    auto ret = allocateLabel<Tag>();
-    // prefix the key with the id to guarantee the same key is not used wrongly with different tags
-    auto prefixed = std::string(Tag::prefix) + '_' + key;
-    print(ret, "=@", quoted(prefixed));
-    return ret;
-  }
-
-  // same as getLabel(const std::string&) above, concatenating keyParts
-  template <typename Tag, typename... Args>
-  TrapLabel<Tag> getLabel(const Args&... keyParts) {
-    std::ostringstream oss;
-    (oss << ... << keyParts);
-    return getLabel<Tag>(oss.str());
-  }
-
-  // emit a trap entry
-  template <typename Entry>
-  void emit(const Entry& e) {
-    print(e);
-  }
-
- private:
-  template <typename... Args>
-  void print(const Args&... args) {
-    (out_ << ... << args) << '\n';
-  }
-
   template <typename Tag>
   TrapLabel<Tag> allocateLabel() {
     return {id_++};
   }
-
-  uint64_t id_{0};
-  std::ostream& out_;
 };
 
 }  // namespace codeql

swift/extractor/trap/TrapLabel.h

@@ -34,7 +34,9 @@ class TrapLabel : public UntypedTrapLabel {
 
   using UntypedTrapLabel::UntypedTrapLabel;
 
-  // we want only TrapArena to create non-default labels
+  // we want one authority tasked with creating labels to avoid conflicts, having access to the
+  // private constructor
+  // this is the TrapArena class for the moment
   friend class TrapArena;
 
  public:
@@ -59,8 +61,10 @@ inline auto trapQuoted(const std::string& s) {
   return std::quoted(s, '"', '"');
 }
 
-template <typename Tag>
+template <typename TagParam>
 struct Binding {
+  using Tag = TagParam;
+
   TrapLabel<Tag> id;
 };
 

swift/extractor/trap/TrapOutput.h

@@ -0,0 +1,119 @@
+#pragma once
+
+#include <memory>
+#include <llvm/ADT/SmallString.h>
+#include <llvm/Support/FileSystem.h>
+#include <llvm/Support/Path.h>
+
+#include "swift/extractor/SwiftExtractorConfiguration.h"
+#include "swift/extractor/trap/TrapLabel.h"
+
+namespace codeql {
+
+// Sink for trap emissions and label assignments. This abstracts away `ofstream` operations
+// like `ofstream`, an explicit bool operator is provided, that return false if something
+// went wrong
+class TrapOutput {
+  std::ofstream out_{};
+  std::string target_;
+
+ public:
+  // open a output stream. Internally a temporary file is created which is then moved into place
+  // on `close()` or destruction
+  TrapOutput(const SwiftExtractorConfiguration& config, const std::string& target) {
+    // TODO: Pick a better name to avoid collisions
+    llvm::SmallString<PATH_MAX> trapPath(config.trapDir);
+    llvm::sys::path::append(trapPath, target);
+    target_ = trapPath.str().str();
+
+    auto trapDir = llvm::sys::path::parent_path(trapPath);
+    if (std::error_code ec = llvm::sys::fs::create_directories(trapDir)) {
+      std::cerr << "Cannot create trap directory '" << trapDir.str() << "': " << ec.message()
+                << "\n";
+      return;
+    }
+
+    auto tempFile = getTempFile();
+    out_ = std::ofstream{tempFile};
+    if (!out_) {
+      std::error_code ec;
+      ec.assign(errno, std::generic_category());
+      std::cerr << "Cannot create temp trap file '" << tempFile << "': " << ec.message() << "\n";
+      return;
+    }
+    out_ << "// extractor-args: ";
+    for (auto opt : config.frontendOptions) {
+      out_ << std::quoted(opt) << " ";
+    }
+    out_ << "\n\n";
+  }
+
+  ~TrapOutput() { close(); }
+
+  TrapOutput(const TrapOutput& other) = delete;
+  TrapOutput& operator=(const TrapOutput& other) = delete;
+
+  TrapOutput(TrapOutput&& other) : out_{std::move(other.out_)}, target_{std::move(other.target_)} {
+    assert(!other.out_.is_open());
+  }
+
+  TrapOutput& operator=(TrapOutput&& other) {
+    close();
+    out_ = std::move(other.out_);
+    target_ = std::move(other.target_);
+    return *this;
+  }
+
+  explicit operator bool() const { return bool{out_}; }
+
+  void close() {
+    // TODO: The last process wins. Should we do better than that?
+    if (out_.is_open()) {
+      out_.close();
+      auto tempFile = getTempFile();
+      auto targetFile = getTargetFile();
+      if (std::error_code ec = llvm::sys::fs::rename(tempFile, targetFile)) {
+        std::cerr << "Cannot rename temp trap file '" << tempFile << "' -> '" << targetFile
+                  << "': " << ec.message() << "\n";
+      }
+    }
+  }
+
+  template <typename Tag>
+  void assignStar(TrapLabel<Tag> label) {
+    print(label, "=*");
+  }
+
+  template <typename Tag>
+  void assignKey(TrapLabel<Tag> label, const std::string& key) {
+    // prefix the key with the id to guarantee the same key is not used wrongly with different tags
+    auto prefixed = std::string(Tag::prefix) + '_' + key;
+    print(label, "=@", trapQuoted(prefixed));
+  }
+
+  template <typename Tag, typename... Args>
+  void assignKey(TrapLabel<Tag> label, const Args&... keyParts) {
+    std::ostringstream oss;
+    (oss << ... << keyParts);
+    assignKey(label, oss.str());
+  }
+
+  template <typename Entry>
+  void emit(const Entry& e) {
+    print(e);
+  }
+
+ private:
+  std::string getTargetFile() { return target_ + ".trap"; }
+  std::string getTempFile() {
+    // TODO: find a more robust approach to avoid collisions?
+    return target_ + '.' + std::to_string(getpid()) + ".trap";
+  }
+
+  template <typename... Args>
+  void print(const Args&... args) {
+    (out_ << ... << args) << '\n';
+  }
+};
+
+}  // namespace codeql