Swift: make TrapLabelStore store untyped label internally

AlexDenisov · AlexDenisov · commit 039aaec6b799 · 2022-05-13T15:20:30.000+02:00
diff --git a/swift/extractor/trap/TrapLabel.h b/swift/extractor/trap/TrapLabel.h
@@ -13,9 +13,8 @@ class UntypedTrapLabel {
   uint64_t id_;
 
   friend class std::hash<UntypedTrapLabel>;
-
-  // we want to have access to the untyped, underlying id
-  friend class TrapLabelStore;
+  template <typename Tag>
+  friend class TrapLabel;
 
  protected:
   UntypedTrapLabel() : id_{0xffffffffffffffff} {}
@@ -44,6 +43,7 @@ class TrapLabel : public UntypedTrapLabel {
 
   // The caller is responsible for ensuring ID uniqueness.
   static TrapLabel unsafeCreateFromExplicitId(uint64_t id) { return {id}; }
+  static TrapLabel unsafeCreateFromUntyped(UntypedTrapLabel label) { return {label.id_}; }
 
   template <typename OtherTag>
   TrapLabel(const TrapLabel<OtherTag>& other) : UntypedTrapLabel(other) {
diff --git a/swift/extractor/trap/TrapLabelStore.h b/swift/extractor/trap/TrapLabelStore.h
@@ -41,19 +41,19 @@ class TrapLabelStore {
   template <typename T>
   std::optional<TrapLabel<ToTag<T>>> get(const T* e) {
     if (auto found = store_.find(getCanonicalPtr(e)); found != store_.end()) {
-      return TrapLabel<ToTag<T>>::unsafeCreateFromExplicitId(found->second);
+      return TrapLabel<ToTag<T>>::unsafeCreateFromUntyped(found->second);
     }
     return std::nullopt;
   }
 
   template <typename T>
   void insert(const T* e, TrapLabel<ToTag<T>> l) {
-    auto [_, inserted] = store_.emplace(getCanonicalPtr(e), l.id_);
+    auto [_, inserted] = store_.emplace(getCanonicalPtr(e), l);
     assert(inserted && "already inserted");
   }
 
  private:
-  std::unordered_map<const void*, uint64_t> store_;
+  std::unordered_map<const void*, UntypedTrapLabel> store_;
 };
 
 }  // namespace codeql
