Swift: move location extraction logic into a separate class

AlexDenisov · AlexDenisov · commit b1aef82117be · 2023-03-01T11:32:50.000+01:00
diff --git a/swift/extractor/SwiftExtractor.cpp b/swift/extractor/SwiftExtractor.cpp
@@ -12,6 +12,7 @@
 #include "swift/extractor/infra/TargetDomains.h"
 #include "swift/extractor/SwiftBuiltinSymbols.h"
 #include "swift/extractor/infra/file/Path.h"
+#include "swift/extractor/infra/SwiftLocationExtractor.h"
 
 using namespace codeql;
 using namespace std::string_literals;
@@ -139,7 +140,9 @@ static std::unordered_set<swift::ModuleDecl*> extractDeclarations(
     }
   }
 
-  SwiftVisitor visitor(compiler.getSourceMgr(), *trap, module, primaryFile);
+  SwiftLocationExtractor locationExtractor(*trap);
+  locationExtractor.emitFile(primaryFile);
+  SwiftVisitor visitor(compiler.getSourceMgr(), *trap, locationExtractor, module, primaryFile);
   auto topLevelDecls = getTopLevelDecls(module, primaryFile);
   for (auto decl : topLevelDecls) {
     visitor.extract(decl);
diff --git a/swift/extractor/infra/SwiftDispatcher.h b/swift/extractor/infra/SwiftDispatcher.h
@@ -45,18 +45,14 @@ class SwiftDispatcher {
   // the SwiftDispatcher
   SwiftDispatcher(const swift::SourceManager& sourceManager,
                   TrapDomain& trap,
+                  SwiftLocationExtractor& locationExtractor,
                   swift::ModuleDecl& currentModule,
                   swift::SourceFile* currentPrimarySourceFile = nullptr)
       : sourceManager{sourceManager},
         trap{trap},
+        locationExtractor{locationExtractor},
         currentModule{currentModule},
-        currentPrimarySourceFile{currentPrimarySourceFile},
-        locationExtractor(trap) {
-    if (currentPrimarySourceFile) {
-      // we make sure the file is in the trap output even if the source is empty
-      locationExtractor.emitFile(currentPrimarySourceFile->getFilename());
-    }
-  }
+        currentPrimarySourceFile{currentPrimarySourceFile} {}
 
   const std::unordered_set<swift::ModuleDecl*> getEncounteredModules() && {
     return std::move(encounteredModules);
@@ -157,7 +153,7 @@ class SwiftDispatcher {
     // TODO when everything is moved to structured C++ classes, this should be moved to createEntry
     if (auto l = store.get(e)) {
       if constexpr (IsLocatable<E>) {
-        attachLocation(e, *l);
+        locationExtractor.attachLocation(sourceManager, e, *l);
       }
       return *l;
     }
@@ -206,50 +202,10 @@ class SwiftDispatcher {
   template <typename E, typename... Args>
   auto createUncachedEntry(const E& e, Args&&... args) {
     auto label = trap.createLabel<TrapTagOf<E>>(std::forward<Args>(args)...);
-    attachLocation(&e, label);
+    locationExtractor.attachLocation(sourceManager, &e, label);
     return TrapClassOf<E>{label};
   }
 
-  template <typename Locatable>
-  void attachLocation(Locatable locatable, TrapLabel<LocatableTag> locatableLabel) {
-    attachLocation(&locatable, locatableLabel);
-  }
-
-  // Emits a Location TRAP entry and attaches it to a `Locatable` trap label
-  template <typename Locatable>
-  void attachLocation(Locatable* locatable, TrapLabel<LocatableTag> locatableLabel) {
-    attachLocation(locatable->getStartLoc(), locatable->getEndLoc(), locatableLabel);
-  }
-
-  void attachLocation(const swift::CapturedValue* capture, TrapLabel<LocatableTag> locatableLabel) {
-    attachLocation(capture->getLoc(), locatableLabel);
-  }
-
-  void attachLocation(const swift::IfConfigClause* clause, TrapLabel<LocatableTag> locatableLabel) {
-    attachLocation(clause->Loc, clause->Loc, locatableLabel);
-  }
-
-  void attachLocation(swift::AvailabilitySpec* spec, TrapLabel<LocatableTag> locatableLabel) {
-    attachLocation(spec->getSourceRange().Start, spec->getSourceRange().End, locatableLabel);
-  }
-
-  // Emits a Location TRAP entry and attaches it to a `Locatable` trap label for a given `SourceLoc`
-  void attachLocation(swift::SourceLoc loc, TrapLabel<LocatableTag> locatableLabel) {
-    attachLocation(loc, loc, locatableLabel);
-  }
-
-  // Emits a Location TRAP entry for a list of swift entities and attaches it to a `Locatable` trap
-  // label
-  template <typename Locatable>
-  void attachLocation(llvm::MutableArrayRef<Locatable>* locatables,
-                      TrapLabel<LocatableTag> locatableLabel) {
-    if (locatables->empty()) {
-      return;
-    }
-    attachLocation(locatables->front().getStartLoc(), locatables->back().getEndLoc(),
-                   locatableLabel);
-  }
-
   // return `std::optional(fetchLabel(arg))` if arg converts to true, otherwise std::nullopt
   // universal reference `Arg&&` is used to catch both temporary and non-const references, not
   // for perfect forwarding
@@ -317,7 +273,7 @@ class SwiftDispatcher {
   void emitComment(swift::Token& comment) {
     CommentsTrap entry{trap.createLabel<CommentTag>(), comment.getRawText().str()};
     trap.emit(entry);
-    attachLocation(comment.getRange().getStart(), comment.getRange().getEnd(), entry.id);
+    locationExtractor.attachLocation(sourceManager, comment, entry.id);
   }
 
  private:
@@ -333,12 +289,6 @@ class SwiftDispatcher {
   template <typename T>
   struct HasId<T, decltype(std::declval<T>().id, void())> : std::true_type {};
 
-  void attachLocation(swift::SourceLoc start,
-                      swift::SourceLoc end,
-                      TrapLabel<LocatableTag> locatableLabel) {
-    locationExtractor.attachLocation(sourceManager, start, end, locatableLabel);
-  }
-
   template <typename Tag, typename... Ts>
   TrapLabel<Tag> fetchLabelFromUnion(const llvm::PointerUnion<Ts...> u) {
     TrapLabel<Tag> ret{};
@@ -382,11 +332,11 @@ class SwiftDispatcher {
   const swift::SourceManager& sourceManager;
   TrapDomain& trap;
   Store store;
+  SwiftLocationExtractor& locationExtractor;
   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.cpp b/swift/extractor/infra/SwiftLocationExtractor.cpp
@@ -1,5 +1,6 @@
 #include <swift/AST/SourceFile.h>
 #include <swift/Basic/SourceManager.h>
+#include <swift/Parse/Token.h>
 
 #include "swift/extractor/trap/TrapDomain.h"
 #include "swift/extractor/trap/generated/TrapEntries.h"
@@ -29,8 +30,42 @@ void SwiftLocationExtractor::attachLocation(const swift::SourceManager& sourceMa
   trap.emit(LocatableLocationsTrap{locatableLabel, entry.id});
 }
 
-void SwiftLocationExtractor::emitFile(llvm::StringRef path) {
-  fetchFileLabel(resolvePath(path));
+void SwiftLocationExtractor::emitFile(swift::SourceFile* file) {
+  if (file) {
+    fetchFileLabel(resolvePath(file->getFilename()));
+  }
+}
+
+void SwiftLocationExtractor::attachLocation(const swift::SourceManager& sourceManager,
+                                            const swift::CapturedValue* capture,
+                                            TrapLabel<LocatableTag> locatableLabel) {
+  attachLocation(sourceManager, capture->getLoc(), locatableLabel);
+}
+
+void SwiftLocationExtractor::attachLocation(const swift::SourceManager& sourceManager,
+                                            const swift::IfConfigClause* clause,
+                                            TrapLabel<LocatableTag> locatableLabel) {
+  attachLocation(sourceManager, clause->Loc, clause->Loc, locatableLabel);
+}
+
+void SwiftLocationExtractor::attachLocation(const swift::SourceManager& sourceManager,
+                                            swift::AvailabilitySpec* spec,
+                                            TrapLabel<LocatableTag> locatableLabel) {
+  attachLocation(sourceManager, spec->getSourceRange().Start, spec->getSourceRange().End,
+                 locatableLabel);
+}
+
+void SwiftLocationExtractor::attachLocation(const swift::SourceManager& sourceManager,
+                                            swift::Token& token,
+                                            TrapLabel<LocatableTag> locatableLabel) {
+  attachLocation(sourceManager, token.getRange().getStart(), token.getRange().getEnd(),
+                 locatableLabel);
+}
+
+void SwiftLocationExtractor::attachLocation(const swift::SourceManager& sourceManager,
+                                            swift::SourceLoc loc,
+                                            TrapLabel<LocatableTag> locatableLabel) {
+  attachLocation(sourceManager, loc, loc, locatableLabel);
 }
 
 TrapLabel<FileTag> SwiftLocationExtractor::fetchFileLabel(const std::filesystem::path& file) {
diff --git a/swift/extractor/infra/SwiftLocationExtractor.h b/swift/extractor/infra/SwiftLocationExtractor.h
@@ -1,5 +1,8 @@
 #pragma once
 
+#include <swift/AST/ASTAllocated.h>
+#include <swift/AST/AvailabilitySpec.h>
+#include <swift/AST/SourceFile.h>
 #include <swift/Basic/SourceManager.h>
 #include <unordered_map>
 #include <filesystem>
@@ -15,12 +18,60 @@ class SwiftLocationExtractor {
  public:
   explicit SwiftLocationExtractor(TrapDomain& trap) : trap(trap) {}
 
+  void emitFile(swift::SourceFile* file);
+
+  template <typename Locatable>
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      Locatable locatable,
+                      TrapLabel<LocatableTag> locatableLabel) {
+    attachLocation(sourceManager, &locatable, locatableLabel);
+  }
+
+  // Emits a Location TRAP entry and attaches it to a `Locatable` trap label
+  template <typename Locatable>
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      Locatable* locatable,
+                      TrapLabel<LocatableTag> locatableLabel) {
+    attachLocation(sourceManager, locatable->getStartLoc(), locatable->getEndLoc(), locatableLabel);
+  }
+
+  // Emits a Location TRAP entry for a list of swift entities and attaches it to a `Locatable` trap
+  // label
+  template <typename Locatable>
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      llvm::MutableArrayRef<Locatable>* locatables,
+                      TrapLabel<LocatableTag> locatableLabel) {
+    if (locatables->empty()) {
+      return;
+    }
+    attachLocation(sourceManager, locatables->front().getStartLoc(), locatables->back().getEndLoc(),
+                   locatableLabel);
+  }
+
   void attachLocation(const swift::SourceManager& sourceManager,
                       swift::SourceLoc start,
                       swift::SourceLoc end,
                       TrapLabel<LocatableTag> locatableLabel);
 
-  void emitFile(llvm::StringRef path);
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      swift::SourceLoc loc,
+                      TrapLabel<LocatableTag> locatableLabel);
+
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      const swift::CapturedValue* capture,
+                      TrapLabel<LocatableTag> locatableLabel);
+
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      const swift::IfConfigClause* clause,
+                      TrapLabel<LocatableTag> locatableLabel);
+
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      swift::AvailabilitySpec* spec,
+                      TrapLabel<LocatableTag> locatableLabel);
+
+  void attachLocation(const swift::SourceManager& sourceManager,
+                      swift::Token& token,
+                      TrapLabel<LocatableTag> locatableLabel);
 
  private:
   TrapLabel<FileTag> fetchFileLabel(const std::filesystem::path& file);
