Swift: extract location extraction into a separate entity

AlexDenisov · AlexDenisov · commit 680b7a16fb95 · 2022-11-08T14:41:10.000+01:00
diff --git a/swift/extractor/infra/SwiftDispatcher.h b/swift/extractor/infra/SwiftDispatcher.h
@@ -11,6 +11,7 @@
 #include "swift/extractor/infra/SwiftTagTraits.h"
 #include "swift/extractor/trap/generated/TrapClasses.h"
 #include "swift/extractor/infra/file/PathHash.h"
+#include "swift/extractor/infra/SwiftLocationExtractor.h"
 
 namespace codeql {
 
@@ -29,8 +30,7 @@ class SwiftDispatcher {
                                const swift::Expr*,
                                const swift::Pattern*,
                                const swift::TypeRepr*,
-                               const swift::TypeBase*,
-                               std::filesystem::path>;
+                               const swift::TypeBase*>;
 
   template <typename E>
   static constexpr bool IsStorable = std::is_constructible_v<Store::Handle, const E&>;
@@ -48,10 +48,11 @@ class SwiftDispatcher {
       : sourceManager{sourceManager},
         trap{trap},
         currentModule{currentModule},
-        currentPrimarySourceFile{currentPrimarySourceFile} {
+        currentPrimarySourceFile{currentPrimarySourceFile},
+        locationExtractor(trap) {
     if (currentPrimarySourceFile) {
       // we make sure the file is in the trap output even if the source is empty
-      fetchLabel(getFilePath(currentPrimarySourceFile->getFilename()));
+      locationExtractor.emitFile(currentPrimarySourceFile->getFilename());
     }
   }
 
@@ -325,20 +326,7 @@ class SwiftDispatcher {
   void attachLocation(swift::SourceLoc start,
                       swift::SourceLoc end,
                       TrapLabel<LocatableTag> locatableLabel) {
-    if (!start.isValid() || !end.isValid()) {
-      // invalid locations seem to come from entities synthesized by the compiler
-      return;
-    }
-    auto file = getFilePath(sourceManager.getDisplayNameForLoc(start));
-    DbLocation entry{{}};
-    entry.file = fetchLabel(file);
-    std::tie(entry.start_line, entry.start_column) = sourceManager.getLineAndColumnInBuffer(start);
-    std::tie(entry.end_line, entry.end_column) = sourceManager.getLineAndColumnInBuffer(end);
-    entry.id = trap.createLabel<DbLocationTag>('{', entry.file, "}:", entry.start_line, ':',
-                                               entry.start_column, ':', entry.end_line, ':',
-                                               entry.end_column);
-    emit(entry);
-    emit(LocatableLocationsTrap{locatableLabel, entry.id});
+    locationExtractor.attachLocation(sourceManager, start, end, locatableLabel);
   }
 
   template <typename Tag, typename... Ts>
@@ -391,19 +379,14 @@ class SwiftDispatcher {
   virtual void visit(const swift::TypeRepr* typeRepr, swift::Type type) = 0;
   virtual void visit(const swift::TypeBase* type) = 0;
 
-  void visit(const std::filesystem::path& file) {
-    auto entry = createEntry(file, file.string());
-    entry.name = file.string();
-    emit(entry);
-  }
-
   const swift::SourceManager& sourceManager;
   TrapDomain& trap;
   Store store;
   Store::Handle waitingForNewLabel{std::monostate{}};
   swift::ModuleDecl& currentModule;
   swift::SourceFile* currentPrimarySourceFile;
   std::unordered_set<swift::ModuleDecl*> encounteredModules;
+  SwiftLocationExtractor locationExtractor;
 };
 
 }  // namespace codeql
diff --git a/swift/extractor/infra/SwiftLocationExtractor.h b/swift/extractor/infra/SwiftLocationExtractor.h
@@ -0,0 +1,72 @@
+#pragma once
+
+#include <swift/AST/SourceFile.h>
+#include <swift/Basic/SourceManager.h>
+
+#include "swift/extractor/trap/TrapDomain.h"
+#include "swift/extractor/trap/generated/TrapEntries.h"
+#include "swift/extractor/trap/generated/TrapClasses.h"
+#include "swift/extractor/infra/file/PathHash.h"
+
+namespace codeql {
+
+class TrapDomain;
+
+class SwiftLocationExtractor {
+ public:
+  explicit SwiftLocationExtractor(TrapDomain& trap) : trap(trap) {}
+
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      swift::SourceLoc start,
+                      swift::SourceLoc end,
+                      TrapLabel<LocatableTag> locatableLabel) {
+    if (!start.isValid() || !end.isValid()) {
+      // invalid locations seem to come from entities synthesized by the compiler
+      return;
+    }
+    auto file = getFilePath(sourceManager.getDisplayNameForLoc(start));
+    DbLocation entry{{}};
+    entry.file = fetchFileLabel(file);
+    std::tie(entry.start_line, entry.start_column) = sourceManager.getLineAndColumnInBuffer(start);
+    std::tie(entry.end_line, entry.end_column) = sourceManager.getLineAndColumnInBuffer(end);
+    entry.id = trap.createLabel<DbLocationTag>('{', entry.file, "}:", entry.start_line, ':',
+                                               entry.start_column, ':', entry.end_line, ':',
+                                               entry.end_column);
+    trap.emit(entry);
+    trap.emit(LocatableLocationsTrap{locatableLabel, entry.id});
+  }
+
+  void emitFile(llvm::StringRef path) { fetchFileLabel(getFilePath(path)); }
+
+ private:
+  TrapLabel<FileTag> fetchFileLabel(const std::filesystem::path& file) {
+    if (store.count(file)) {
+      return store[file];
+    }
+
+    DbFile entry({});
+    entry.id = trap.createLabel<DbFileTag>(file.string());
+    entry.name = file.string();
+    trap.emit(entry);
+    store[file] = entry.id;
+    return entry.id;
+  }
+
+  static std::filesystem::path getFilePath(std::string_view path) {
+    // TODO: this needs more testing
+    // TODO: check canonicalization of names on a case insensitive filesystems
+    // TODO: make symlink resolution conditional on CODEQL_PRESERVE_SYMLINKS=true
+    std::error_code ec;
+    auto ret = std::filesystem::canonical(path, ec);
+    if (ec) {
+      std::cerr << "Cannot get real path: " << std::quoted(path) << ": " << ec.message() << "\n";
+      return {};
+    }
+    return ret;
+  }
+
+  TrapDomain& trap;
+  std::unordered_map<std::filesystem::path, TrapLabel<FileTag>> store;
+};
+
+}  // namespace codeql
diff --git a/swift/extractor/invocation/CodeQLDiagnosticsConsumer.cpp b/swift/extractor/invocation/CodeQLDiagnosticsConsumer.cpp
@@ -1,7 +1,5 @@
 #include "swift/extractor/invocation/CodeQLDiagnosticsConsumer.h"
-#include "swift/extractor/trap/TrapDomain.h"
 #include "swift/extractor/trap/generated/TrapEntries.h"
-#include "swift/extractor/trap/generated/TrapClasses.h"
 
 #include <swift/AST/DiagnosticEngine.h>
 #include <swift/Basic/SourceManager.h>
@@ -27,43 +25,6 @@ static int diagnosticsKind(const swift::DiagnosticInfo& diagInfo) {
   return 0;
 }
 
-static std::filesystem::path getFilePath(std::string_view path) {
-  // TODO: this needs more testing
-  // TODO: check canonicalization of names on a case insensitive filesystems
-  // TODO: make symlink resolution conditional on CODEQL_PRESERVE_SYMLINKS=true
-  std::error_code ec;
-  auto ret = std::filesystem::canonical(path, ec);
-  if (ec) {
-    std::cerr << "Cannot get real path: " << std::quoted(path) << ": " << ec.message() << "\n";
-    return {};
-  }
-  return ret;
-}
-
-static void attachLocation(TrapDomain& trap,
-                           swift::SourceManager& sourceManager,
-                           const swift::DiagnosticInfo& diagInfo,
-                           DiagnosticsTrap& locatable) {
-  auto loc = diagInfo.Loc;
-  if (!loc.isValid()) {
-    return;
-  }
-  auto filepath = getFilePath(sourceManager.getDisplayNameForLoc(loc));
-  DbFile file({});
-  file.id = trap.createLabel<DbFileTag>();
-  file.name = filepath;
-  trap.emit(file);
-
-  DbLocation location({});
-  location.id = trap.createLabel<DbLocationTag>();
-  location.file = file.id;
-  std::tie(location.start_line, location.start_column) =
-      sourceManager.getLineAndColumnInBuffer(loc);
-  std::tie(location.end_line, location.end_column) = sourceManager.getLineAndColumnInBuffer(loc);
-  trap.emit(location);
-  trap.emit(LocatableLocationsTrap{locatable.id, location.id});
-}
-
 void CodeQLDiagnosticsConsumer::handleDiagnostic(swift::SourceManager& sourceManager,
                                                  const swift::DiagnosticInfo& diagInfo) {
   auto message = getDiagMessage(sourceManager, diagInfo);
@@ -72,7 +33,7 @@ void CodeQLDiagnosticsConsumer::handleDiagnostic(swift::SourceManager& sourceMan
   diag.kind = diagnosticsKind(diagInfo);
   diag.text = message;
   trap.emit(diag);
-  attachLocation(trap, sourceManager, diagInfo, diag);
+  locationExtractor.attachLocation(sourceManager, diagInfo.Loc, diagInfo.Loc, diag.id);
 }
 
 std::string CodeQLDiagnosticsConsumer::getDiagMessage(swift::SourceManager& sourceManager,
diff --git a/swift/extractor/invocation/CodeQLDiagnosticsConsumer.h b/swift/extractor/invocation/CodeQLDiagnosticsConsumer.h
@@ -1,21 +1,24 @@
 #pragma once
 
 #include <swift/AST/DiagnosticConsumer.h>
+#include "swift/extractor/infra/SwiftLocationExtractor.h"
 
 namespace codeql {
 
 class TrapDomain;
 
 class CodeQLDiagnosticsConsumer : public swift::DiagnosticConsumer {
  public:
-  explicit CodeQLDiagnosticsConsumer(TrapDomain& targetFile) : trap(targetFile) {}
+  explicit CodeQLDiagnosticsConsumer(TrapDomain& targetFile)
+      : trap(targetFile), locationExtractor(targetFile) {}
   void handleDiagnostic(swift::SourceManager& sourceManager,
                         const swift::DiagnosticInfo& diagInfo) override;
 
  private:
   static std::string getDiagMessage(swift::SourceManager& sourceManager,
                                     const swift::DiagnosticInfo& diagInfo);
   TrapDomain& trap;
+  SwiftLocationExtractor locationExtractor;
 };
 
 }  // namespace codeql
