[ORC-RT][MachO] Add RecordSectionsTracker utility to track metadata sections.

Newly added sections can be processed by calling processNewSections. Calling
reset moves all sections back to the "new" state for reprocessing (expected to
be used by dlclose).

diff --git a/compiler-rt/lib/orc/macho_platform.cpp b/compiler-rt/lib/orc/macho_platform.cpp
index f17230c..e99fe87 100644 (file)
--- a/compiler-rt/lib/orc/macho_platform.cpp
+++ b/compiler-rt/lib/orc/macho_platform.cpp
@@ -128,6 +128,90 @@ private:
 
   using AtExitsVector = std::vector<AtExitEntry>;
 
+  /// Used to manage sections of fixed-sized metadata records (e.g. pointer
+  /// sections, selector refs, etc.)
+  template <typename RecordElement> class RecordSectionsTracker {
+  public:
+    /// Add a section to the "new" list.
+    void add(span<RecordElement> Sec) { New.push_back(std::move(Sec)); }
+
+    /// Returns true if there are new sections to process.
+    bool hasNewSections() const { return !New.empty(); }
+
+    /// Process all new sections.
+    template <typename ProcessSectionFunc>
+    std::enable_if_t<std::is_void_v<
+        std::invoke_result_t<ProcessSectionFunc, span<RecordElement>>>>
+    processNewSections(ProcessSectionFunc &&ProcessSection) {
+      for (auto &Sec : New)
+        ProcessSection(Sec);
+      moveNewToProcessed();
+    }
+
+    /// Proces all new sections with a fallible handler.
+    ///
+    /// Successfully handled sections will be moved to the Processed
+    /// list.
+    template <typename ProcessSectionFunc>
+    std::enable_if_t<
+        std::is_same_v<Error, std::invoke_result_t<ProcessSectionFunc,
+                                                   span<RecordElement>>>,
+        Error>
+    processNewSections(ProcessSectionFunc &&ProcessSection) {
+      for (size_t I = 0; I != New.size(); ++I) {
+        if (auto Err = ProcessSection(New[I])) {
+          for (size_t J = 0; J != I; ++J)
+            Processed.push_back(New[J]);
+          New.erase(New.begin(), New.begin() + I);
+          return Err;
+        }
+      }
+      moveNewToProcessed();
+      return Error::success();
+    }
+
+    /// Move all sections back to New for reprocessing.
+    void reset() {
+      moveNewToProcessed();
+      New = std::move(Processed);
+    }
+
+    /// Remove the section with the given range.
+    bool removeIfPresent(ExecutorAddrRange R) {
+      if (removeIfPresent(New, R))
+        return true;
+      return removeIfPresent(Processed, R);
+    }
+
+  private:
+    void moveNewToProcessed() {
+      if (Processed.empty())
+        Processed = std::move(New);
+      else {
+        Processed.reserve(Processed.size() + New.size());
+        std::copy(New.begin(), New.end(), std::back_inserter(Processed));
+        New.clear();
+      }
+    }
+
+    bool removeIfPresent(std::vector<span<RecordElement>> &V,
+                         ExecutorAddrRange R) {
+      auto RI = std::find_if(
+          V.rbegin(), V.rend(),
+          [RS = R.toSpan<RecordElement>()](const span<RecordElement> &E) {
+            return E.data() == RS.data();
+          });
+      if (RI != V.rend()) {
+        V.erase(std::next(RI).base());
+        return true;
+      }
+      return false;
+    }
+
+    std::vector<span<RecordElement>> Processed;
+    std::vector<span<RecordElement>> New;
+  };
+
   struct JITDylibState {
     std::string Name;
     void *Header = nullptr;
@@ -139,18 +223,12 @@ private:
     const objc_image_info *ObjCImageInfo = nullptr;
     std::unordered_map<void *, std::vector<char>> DataSectionContent;
     std::unordered_map<void *, size_t> ZeroInitRanges;
-    std::vector<span<void (*)()>> ModInitsSections;
-    std::vector<span<void (*)()>> ModInitsSectionsNew;
-    std::vector<span<uintptr_t>> ObjCClassListSections;
-    std::vector<span<uintptr_t>> ObjCClassListSectionsNew;
-    std::vector<span<uintptr_t>> ObjCSelRefsSections;
-    std::vector<span<uintptr_t>> ObjCSelRefsSectionsNew;
-    std::vector<span<char>> Swift5ProtocolsSections;
-    std::vector<span<char>> Swift5ProtocolsSectionsNew;
-    std::vector<span<char>> Swift5ProtocolConformancesSections;
-    std::vector<span<char>> Swift5ProtocolConformancesSectionsNew;
-    std::vector<span<char>> Swift5TypesSections;
-    std::vector<span<char>> Swift5TypesSectionsNew;
+    RecordSectionsTracker<void (*)()> ModInitsSections;
+    RecordSectionsTracker<void *> ObjCClassListSections;
+    RecordSectionsTracker<void *> ObjCSelRefsSections;
+    RecordSectionsTracker<char> Swift5ProtocolsSections;
+    RecordSectionsTracker<char> Swift5ProtocolConformancesSections;
+    RecordSectionsTracker<char> Swift5TypesSections;
 
     bool referenced() const {
       return LinkedAgainstRefCount != 0 || DlRefCount != 0;
@@ -356,18 +434,17 @@ Error MachOPlatformRuntimeState::registerObjectPlatformSections(
       if (auto Err = registerThreadDataSection(KV.second.toSpan<const char>()))
         return Err;
     } else if (KV.first == "__DATA,__objc_selrefs")
-      JDS->ObjCSelRefsSectionsNew.push_back(KV.second.toSpan<uintptr_t>());
+      JDS->ObjCSelRefsSections.add(KV.second.toSpan<void *>());
     else if (KV.first == "__DATA,__objc_classlist")
-      JDS->ObjCClassListSectionsNew.push_back(KV.second.toSpan<uintptr_t>());
+      JDS->ObjCClassListSections.add(KV.second.toSpan<void *>());
     else if (KV.first == "__TEXT,__swift5_protos")
-      JDS->Swift5ProtocolsSectionsNew.push_back(KV.second.toSpan<char>());
+      JDS->Swift5ProtocolsSections.add(KV.second.toSpan<char>());
     else if (KV.first == "__TEXT,__swift5_proto")
-      JDS->Swift5ProtocolConformancesSectionsNew.push_back(
-          KV.second.toSpan<char>());
+      JDS->Swift5ProtocolConformancesSections.add(KV.second.toSpan<char>());
     else if (KV.first == "__TEXT,__swift5_types")
-      JDS->Swift5TypesSectionsNew.push_back(KV.second.toSpan<char>());
+      JDS->Swift5TypesSections.add(KV.second.toSpan<char>());
     else if (KV.first == "__DATA,__mod_init_func")
-      JDS->ModInitsSectionsNew.push_back(KV.second.toSpan<void (*)()>());
+      JDS->ModInitsSections.add(KV.second.toSpan<void (*)()>());
     else {
       // Should this be a warning instead?
       return make_error<StringError>(
@@ -380,20 +457,6 @@ Error MachOPlatformRuntimeState::registerObjectPlatformSections(
   return Error::success();
 }
 
-// Remove the given range from the given vector if present.
-// Returns true if the range was removed, false otherwise.
-template <typename T>
-bool removeIfPresent(std::vector<span<T>> &V, ExecutorAddrRange R) {
-  auto RI = std::find_if(
-      V.rbegin(), V.rend(),
-      [RS = R.toSpan<T>()](const span<T> &E) { return E.data() == RS.data(); });
-  if (RI != V.rend()) {
-    V.erase(std::next(RI).base());
-    return true;
-  }
-  return false;
-}
-
 Error MachOPlatformRuntimeState::deregisterObjectPlatformSections(
     ExecutorAddr HeaderAddr,
     std::vector<std::pair<std::string_view, ExecutorAddrRange>> Secs) {
@@ -432,25 +495,19 @@ Error MachOPlatformRuntimeState::deregisterObjectPlatformSections(
       if (auto Err =
               deregisterThreadDataSection(KV.second.toSpan<const char>()))
         return Err;
-    } else if (KV.first == "__DATA,__objc_selrefs") {
-      if (!removeIfPresent(JDS->ObjCSelRefsSections, KV.second))
-        removeIfPresent(JDS->ObjCSelRefsSectionsNew, KV.second);
-    } else if (KV.first == "__DATA,__objc_classlist") {
-      if (!removeIfPresent(JDS->ObjCClassListSections, KV.second))
-        removeIfPresent(JDS->ObjCClassListSectionsNew, KV.second);
-    } else if (KV.first == "__TEXT,__swift5_protos") {
-      if (!removeIfPresent(JDS->Swift5ProtocolsSections, KV.second))
-        removeIfPresent(JDS->Swift5ProtocolsSectionsNew, KV.second);
-    } else if (KV.first == "__TEXT,__swift5_proto") {
-      if (!removeIfPresent(JDS->Swift5ProtocolConformancesSections, KV.second))
-        removeIfPresent(JDS->Swift5ProtocolConformancesSectionsNew, KV.second);
-    } else if (KV.first == "__TEXT,__swift5_types") {
-      if (!removeIfPresent(JDS->Swift5TypesSections, KV.second))
-        removeIfPresent(JDS->Swift5TypesSectionsNew, KV.second);
-    } else if (KV.first == "__DATA,__mod_init_func") {
-      if (!removeIfPresent(JDS->ModInitsSections, KV.second))
-        removeIfPresent(JDS->ModInitsSectionsNew, KV.second);
-    } else {
+    } else if (KV.first == "__DATA,__objc_selrefs")
+      JDS->ObjCSelRefsSections.removeIfPresent(KV.second);
+    else if (KV.first == "__DATA,__objc_classlist")
+      JDS->ObjCClassListSections.removeIfPresent(KV.second);
+    else if (KV.first == "__TEXT,__swift5_protos")
+      JDS->Swift5ProtocolsSections.removeIfPresent(KV.second);
+    else if (KV.first == "__TEXT,__swift5_proto")
+      JDS->Swift5ProtocolConformancesSections.removeIfPresent(KV.second);
+    else if (KV.first == "__TEXT,__swift5_types")
+      JDS->Swift5TypesSections.removeIfPresent(KV.second);
+    else if (KV.first == "__DATA,__mod_init_func")
+      JDS->ModInitsSections.removeIfPresent(KV.second);
+    else {
       // Should this be a warning instead?
       return make_error<StringError>(
           "Encountered unexpected section " +
@@ -591,42 +648,26 @@ MachOPlatformRuntimeState::lookupSymbolInJITDylib(void *DSOHandle,
   return Result;
 }
 
-template <typename T>
-static void moveAppendSections(std::vector<span<T>> &Dst,
-                               std::vector<span<T>> &Src) {
-  if (Dst.empty()) {
-    Dst = std::move(Src);
-    return;
-  }
-
-  Dst.reserve(Dst.size() + Src.size());
-  std::copy(Src.begin(), Src.end(), std::back_inserter(Dst));
-  Src.clear();
-}
-
 Error MachOPlatformRuntimeState::registerObjCSelectors(JITDylibState &JDS) {
-
-  if (JDS.ObjCSelRefsSectionsNew.empty())
+  if (!JDS.ObjCSelRefsSections.hasNewSections())
     return Error::success();
 
   if (ORC_RT_UNLIKELY(!sel_registerName))
     return make_error<StringError>("sel_registerName is not available");
 
-  for (const auto &ObjCSelRefs : JDS.ObjCSelRefsSectionsNew) {
-    for (uintptr_t &SelEntry : ObjCSelRefs) {
+  JDS.ObjCSelRefsSections.processNewSections([](span<void *> SelRefs) {
+    for (void *&SelEntry : SelRefs) {
       const char *SelName = reinterpret_cast<const char *>(SelEntry);
       auto Sel = sel_registerName(SelName);
       *reinterpret_cast<SEL *>(&SelEntry) = Sel;
     }
-  }
+  });
 
-  moveAppendSections(JDS.ObjCSelRefsSections, JDS.ObjCSelRefsSectionsNew);
   return Error::success();
 }
 
 Error MachOPlatformRuntimeState::registerObjCClasses(JITDylibState &JDS) {
-
-  if (JDS.ObjCClassListSectionsNew.empty())
+  if (!JDS.ObjCClassListSections.hasNewSections())
     return Error::success();
 
   if (ORC_RT_UNLIKELY(!objc_msgSend))
@@ -644,91 +685,87 @@ Error MachOPlatformRuntimeState::registerObjCClasses(JITDylibState &JDS) {
 
   auto ClassSelector = sel_registerName("class");
 
-  for (const auto &ObjCClassList : JDS.ObjCClassListSectionsNew) {
-    for (uintptr_t ClassPtr : ObjCClassList) {
-      auto *Cls = reinterpret_cast<Class>(ClassPtr);
-      auto *ClassCompiled = reinterpret_cast<ObjCClassCompiled *>(ClassPtr);
-      objc_msgSend(reinterpret_cast<id>(ClassCompiled->Parent), ClassSelector);
-      auto Registered = objc_readClassPair(Cls, JDS.ObjCImageInfo);
-
-      // FIXME: Improve diagnostic by reporting the failed class's name.
-      if (Registered != Cls)
-        return make_error<StringError>("Unable to register Objective-C class");
-    }
-  }
-
-  moveAppendSections(JDS.ObjCClassListSections, JDS.ObjCClassListSectionsNew);
-  return Error::success();
+  return JDS.ObjCClassListSections.processNewSections(
+      [&](span<void *> ClassPtrs) -> Error {
+        for (void *ClassPtr : ClassPtrs) {
+          auto *Cls = reinterpret_cast<Class>(ClassPtr);
+          auto *ClassCompiled = reinterpret_cast<ObjCClassCompiled *>(ClassPtr);
+          objc_msgSend(reinterpret_cast<id>(ClassCompiled->Parent),
+                       ClassSelector);
+          auto Registered = objc_readClassPair(Cls, JDS.ObjCImageInfo);
+          // FIXME: Improve diagnostic by reporting the failed class's name.
+          if (Registered != Cls)
+            return make_error<StringError>(
+                "Unable to register Objective-C class");
+        }
+        return Error::success();
+      });
 }
 
 Error MachOPlatformRuntimeState::registerSwift5Protocols(JITDylibState &JDS) {
 
-  if (JDS.Swift5ProtocolsSectionsNew.empty())
+  if (!JDS.Swift5ProtocolsSections.hasNewSections())
     return Error::success();
 
   if (ORC_RT_UNLIKELY(!swift_registerProtocols))
     return make_error<StringError>("swift_registerProtocols is not available");
 
-  for (const auto &Swift5Protocols : JDS.Swift5ProtocolsSectionsNew)
+  JDS.Swift5ProtocolsSections.processNewSections([](span<char> ProtoSec) {
     swift_registerProtocols(
-        reinterpret_cast<const ProtocolRecord *>(Swift5Protocols.data()),
-        reinterpret_cast<const ProtocolRecord *>(Swift5Protocols.data() +
-                                                 Swift5Protocols.size()));
+        reinterpret_cast<const ProtocolRecord *>(ProtoSec.data()),
+        reinterpret_cast<const ProtocolRecord *>(ProtoSec.data() +
+                                                 ProtoSec.size()));
+  });
 
-  moveAppendSections(JDS.Swift5ProtocolsSections,
-                     JDS.Swift5ProtocolsSectionsNew);
   return Error::success();
 }
 
 Error MachOPlatformRuntimeState::registerSwift5ProtocolConformances(
     JITDylibState &JDS) {
 
-  if (JDS.Swift5ProtocolConformancesSectionsNew.empty())
+  if (!JDS.Swift5ProtocolConformancesSections.hasNewSections())
     return Error::success();
 
   if (ORC_RT_UNLIKELY(!swift_registerProtocolConformances))
     return make_error<StringError>(
         "swift_registerProtocolConformances is not available");
 
-  for (const auto &ProtoConfSec : JDS.Swift5ProtocolConformancesSectionsNew)
-    swift_registerProtocolConformances(
-        reinterpret_cast<const ProtocolConformanceRecord *>(
-            ProtoConfSec.data()),
-        reinterpret_cast<const ProtocolConformanceRecord *>(
-            ProtoConfSec.data() + ProtoConfSec.size()));
+  JDS.Swift5ProtocolConformancesSections.processNewSections(
+      [](span<char> ProtoConfSec) {
+        swift_registerProtocolConformances(
+            reinterpret_cast<const ProtocolConformanceRecord *>(
+                ProtoConfSec.data()),
+            reinterpret_cast<const ProtocolConformanceRecord *>(
+                ProtoConfSec.data() + ProtoConfSec.size()));
+      });
 
-  moveAppendSections(JDS.Swift5ProtocolConformancesSections,
-                     JDS.Swift5ProtocolConformancesSectionsNew);
   return Error::success();
 }
 
 Error MachOPlatformRuntimeState::registerSwift5Types(JITDylibState &JDS) {
 
-  if (JDS.Swift5TypesSectionsNew.empty())
+  if (!JDS.Swift5TypesSections.hasNewSections())
     return Error::success();
 
   if (ORC_RT_UNLIKELY(!swift_registerTypeMetadataRecords))
     return make_error<StringError>(
         "swift_registerTypeMetadataRecords is not available");
 
-  for (const auto &TypeSec : JDS.Swift5TypesSectionsNew)
+  JDS.Swift5TypesSections.processNewSections([&](span<char> TypesSec) {
     swift_registerTypeMetadataRecords(
-        reinterpret_cast<const TypeMetadataRecord *>(TypeSec.data()),
-        reinterpret_cast<const TypeMetadataRecord *>(TypeSec.data() +
-                                                     TypeSec.size()));
+        reinterpret_cast<const TypeMetadataRecord *>(TypesSec.data()),
+        reinterpret_cast<const TypeMetadataRecord *>(TypesSec.data() +
+                                                     TypesSec.size()));
+  });
 
-  moveAppendSections(JDS.Swift5TypesSections, JDS.Swift5TypesSectionsNew);
   return Error::success();
 }
 
 Error MachOPlatformRuntimeState::runModInits(JITDylibState &JDS) {
-
-  for (const auto &ModInits : JDS.ModInitsSectionsNew) {
-    for (void (*Init)() : ModInits)
-      (*Init)();
-  }
-
-  moveAppendSections(JDS.ModInitsSections, JDS.ModInitsSectionsNew);
+  JDS.ModInitsSections.processNewSections([](span<void (*)()> Inits) {
+    for (auto *Init : Inits)
+      Init();
+  });
   return Error::success();
 }
 
@@ -891,8 +928,7 @@ Error MachOPlatformRuntimeState::dlcloseDeinitialize(JITDylibState &JDS) {
   runAtExits(JDS);
 
   // Reset mod-inits
-  moveAppendSections(JDS.ModInitsSections, JDS.ModInitsSectionsNew);
-  JDS.ModInitsSectionsNew = std::move(JDS.ModInitsSections);
+  JDS.ModInitsSections.reset();
 
   // Reset data section contents.
   for (auto &KV : JDS.DataSectionContent)