Merge pull request #9107 from redsun82/swift-arena

Swift: `TrapOutput`

Author: AlexDenisov

swift/extractor/SwiftExtractor.cpp

@@ -13,6 +13,8 @@
 #include <llvm/Support/Path.h>
 
 #include "swift/extractor/trap/TrapClasses.h"
+#include "swift/extractor/trap/TrapArena.h"
+#include "swift/extractor/trap/TrapOutput.h"
 
 using namespace codeql;
 
@@ -61,24 +63,28 @@ static void extractFile(const SwiftExtractorConfiguration& config, swift::Source
     return;
   }
 
-  std::ofstream trap(tempTrapPath.str().str());
-  if (!trap) {
+  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;
   }
-  trap << "// extractor-args: ";
+  trapStream << "// extractor-args: ";
   for (auto opt : config.frontendOptions) {
-    trap << std::quoted(opt) << " ";
+    trapStream << std::quoted(opt) << " ";
   }
-  trap << "\n\n";
-
-  File f;
-  f.id = TrapLabel<FileTag>{};
+  trapStream << "\n\n";
+
+  TrapOutput trap{trapStream};
+  TrapArena arena{};
+  auto label = arena.allocateLabel<FileTag>();
+  trap.assignStar(label);
+  File f{};
+  f.id = label;
   f.name = srcFilePath.str().str();
-  trap << f.id << "=*\n" << f;
+  trap.emit(f);
 
   // TODO: Pick a better name to avoid collisions
   std::string trapName = file.getFilename().str() + ".trap";

swift/extractor/trap/TrapArena.h

@@ -0,0 +1,23 @@
+#pragma once
+
+#include <iostream>
+#include <sstream>
+#include <vector>
+
+#include "swift/extractor/trap/TrapLabel.h"
+
+namespace codeql {
+
+// 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:
+  template <typename Tag>
+  TrapLabel<Tag> allocateLabel() {
+    return TrapLabel<Tag>::unsafeCreateFromExplicitId(id_++);
+  }
+};
+
+}  // namespace codeql

swift/extractor/trap/TrapLabel.h

@@ -39,6 +39,9 @@ class TrapLabel : public UntypedTrapLabel {
 
   TrapLabel() = default;
 
+  // The caller is responsible for ensuring ID uniqueness.
+  static TrapLabel unsafeCreateFromExplicitId(uint64_t id) { return {id}; }
+
   template <typename OtherTag>
   TrapLabel(const TrapLabel<OtherTag>& other) : UntypedTrapLabel(other) {
     // we temporarily need to bypass the label type system for unknown AST nodes and types

swift/extractor/trap/TrapOutput.h

@@ -0,0 +1,49 @@
+#pragma once
+
+#include <memory>
+#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
+// TODO better error handling
+class TrapOutput {
+  std::ostream& out_;
+
+ public:
+  explicit TrapOutput(std::ostream& out) : out_{out} {}
+
+  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:
+  template <typename... Args>
+  void print(const Args&... args) {
+    (out_ << ... << args) << '\n';
+  }
+};
+
+}  // namespace codeql