Swift: replace getCanonicalPointer with std::variant

This turned out easier than expected previously. `llvm::PointerUnion`
was also considered, which would have less memory footprint, but it
would require more effort as it is lacking the same implicit conversions
and operators that `std::variant` provides.

Also renamed `ToTag<E>` to `TrapTagOf<E>` and introduced a derived
convenience functor `TrapLabelOf<E>`.

Author: redsun82

swift/extractor/SwiftDispatcher.h

@@ -46,7 +46,7 @@ std::string inline getKindName<swift::TypeRepr, swift::TypeReprKind>(swift::Type
 
 }  // namespace detail
 
-// The main reponsibilities of the SwiftDispatcher are as follows:
+// The main responsibilities of the SwiftDispatcher are as follows:
 // * redirect specific AST node emission to a corresponding visitor (statements, expressions, etc.)
 // * storing TRAP labels for emitted AST nodes (in the TrapLabelStore) to avoid re-emission
 // Since SwiftDispatcher sees all the AST nodes, it also attaches a location to every 'locatable'
@@ -72,19 +72,28 @@ class SwiftDispatcher {
   }
 
  private:
+  // types to be supported by assignNewLabel/fetchLabel need to be listed here
+  using Store = TrapLabelStore<swift::Decl,
+                               swift::Stmt,
+                               swift::Expr,
+                               swift::Pattern,
+                               swift::TypeRepr,
+                               swift::TypeBase>;
+
   // This method gives a TRAP label for already emitted AST node.
   // If the AST node was not emitted yet, then the emission is dispatched to a corresponding
   // visitor (see `visit(T *)` methods below).
   template <typename E>
-  TrapLabel<ToTag<E>> fetchLabel(E* e) {
+  TrapLabelOf<E> fetchLabel(E* e) {
     // this is required so we avoid any recursive loop: a `fetchLabel` during the visit of `e` might
     // end up calling `fetchLabel` on `e` itself, so we want the visit of `e` to call `fetchLabel`
-    // only after having called `assignNewLabel` on `e`
-    assert(!waitingForNewLabel && "fetchLabel called before assignNewLabel");
+    // only after having called `assignNewLabel` on `e`. `index() == 0` indicates the monostate of
+    // `Store::Handle`, i.e. its default construction
+    assert(waitingForNewLabel.index() == 0 && "fetchLabel called before assignNewLabel");
     if (auto l = store.get(e)) {
       return *l;
     }
-    waitingForNewLabel = getCanonicalPointer(e);
+    waitingForNewLabel = e;
     visit(e);
     if (auto l = store.get(e)) {
       if constexpr (!std::is_base_of_v<swift::TypeBase, E>) {
@@ -100,12 +109,12 @@ class SwiftDispatcher {
   // it actually gets emitted to handle recursive cases such as recursive calls, or recursive type
   // declarations
   template <typename E>
-  TrapLabel<ToTag<E>> assignNewLabel(E* e) {
-    assert(waitingForNewLabel == getCanonicalPointer(e) && "assignNewLabel called on wrong entity");
-    auto label = getLabel<ToTag<E>>();
+  TrapLabelOf<E> assignNewLabel(E* e) {
+    assert(waitingForNewLabel == Store::Handle{e} && "assignNewLabel called on wrong entity");
+    auto label = getLabel<TrapTagOf<E>>();
     trap.assignStar(label);
     store.insert(e, label);
-    waitingForNewLabel = nullptr;
+    waitingForNewLabel = {};
     return label;
   }
 
@@ -167,8 +176,8 @@ class SwiftDispatcher {
   const swift::SourceManager& sourceManager;
   TrapArena& arena;
   TrapOutput& trap;
-  TrapLabelStore store;
-  const void* waitingForNewLabel{nullptr};
+  Store store;
+  Store::Handle waitingForNewLabel{};
 };
 
 }  // namespace codeql

swift/extractor/SwiftTagTraits.h

@@ -1,6 +1,7 @@
 #pragma once
 
-// This file implements the mapping needed by the API defined in the TrapTagTraits.h
+// This file implements the mapping needed by the API defined in the TrapTagTraits.h, so that
+// TrapTagOf/TrapLabelOf provide the tags/labels for specific swift entity types.
 #include <swift/AST/ASTVisitor.h>
 #include "swift/extractor/trap/TrapTagTraits.h"
 #include "swift/extractor/trap/generated/TrapTags.h"

swift/extractor/trap/TrapLabel.h

@@ -4,7 +4,6 @@
 #include <iostream>
 #include <string>
 
-#include "swift/extractor/trap/TrapTagTraits.h"
 #include "swift/extractor/trap/generated/TrapTags.h"
 
 namespace codeql {

swift/extractor/trap/TrapLabelStore.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <variant>
 #include <cassert>
 #include <optional>
 #include <unordered_map>
@@ -12,50 +13,30 @@
 
 namespace codeql {
 
-// The following is needed to avoid the problem of subclass pointers not necessarily coinciding
-// with superclass ones in case of multiple inheritance
-// The interesting part here is implicit conversion from a derived class pointer to the parameter
-inline const swift::Decl* getCanonicalPointer(const swift::Decl* e) {
-  return e;
-}
-inline const swift::Stmt* getCanonicalPointer(const swift::Stmt* e) {
-  return e;
-}
-inline const swift::Expr* getCanonicalPointer(const swift::Expr* e) {
-  return e;
-}
-inline const swift::Pattern* getCanonicalPointer(const swift::Pattern* e) {
-  return e;
-}
-inline const swift::TypeRepr* getCanonicalPointer(const swift::TypeRepr* e) {
-  return e;
-}
-inline const swift::TypeBase* getCanonicalPointer(const swift::TypeBase* e) {
-  return e;
-}
-
 // The extraction is done in a lazy/on-demand fashion:
 // Each emitted TRAP entry for an AST node gets a TRAP label assigned to it.
 // To avoid re-emission, we store the "AST node <> label" entry in the TrapLabelStore.
+// The template parameters `Ts...` indicate to what entity types we restrict the storage
+template <typename... Ts>
 class TrapLabelStore {
  public:
+  using Handle = std::variant<std::monostate, const Ts*...>;
   template <typename T>
-  std::optional<TrapLabel<ToTag<T>>> get(const T* e) {
-    if (auto found = store_.find(getCanonicalPointer(e)); found != store_.end()) {
-      return TrapLabel<ToTag<T>>::unsafeCreateFromUntyped(found->second);
+  std::optional<TrapLabelOf<T>> get(const T* e) {
+    if (auto found = store_.find(e); found != store_.end()) {
+      return TrapLabelOf<T>::unsafeCreateFromUntyped(found->second);
     }
     return std::nullopt;
   }
 
   template <typename T>
-  void insert(const T* e, TrapLabel<ToTag<T>> l) {
-    auto [_, inserted] = store_.emplace(getCanonicalPointer(e), l);
+  void insert(const T* e, TrapLabelOf<T> l) {
+    auto [_, inserted] = store_.emplace(e, l);
     assert(inserted && "already inserted");
   }
 
  private:
-  // TODO: consider std::variant or llvm::PointerUnion instead of `void *`
-  std::unordered_map<const void*, UntypedTrapLabel> store_;
+  std::unordered_map<Handle, UntypedTrapLabel> store_;
 };
 
 }  // namespace codeql

swift/extractor/trap/TrapTagTraits.h

@@ -16,6 +16,9 @@ struct ToTagOverride : ToTagFunctor<T> {};
 }  // namespace detail
 
 template <typename T>
-using ToTag = typename detail::ToTagOverride<std::remove_const_t<T>>::type;
+using TrapTagOf = typename detail::ToTagOverride<std::remove_const_t<T>>::type;
+
+template <typename T>
+using TrapLabelOf = TrapLabel<TrapTagOf<T>>;
 
 }  // namespace codeql