[clang][ExtractAPI] Implement Record removal from APISet

rdar://128092236

Author: daniel-grumberg

clang/include/clang/ExtractAPI/API.h

@@ -19,21 +19,13 @@
 #define LLVM_CLANG_EXTRACTAPI_API_H
 
 #include "clang/AST/Availability.h"
-#include "clang/AST/Decl.h"
 #include "clang/AST/DeclBase.h"
-#include "clang/AST/DeclObjC.h"
 #include "clang/AST/RawCommentList.h"
 #include "clang/Basic/SourceLocation.h"
-#include "clang/Basic/Specifiers.h"
 #include "clang/ExtractAPI/DeclarationFragments.h"
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/MapVector.h"
-#include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Casting.h"
-#include "llvm/Support/Compiler.h"
-#include "llvm/Support/ErrorHandling.h"
-#include "llvm/Support/raw_ostream.h"
 #include "llvm/TargetParser/Triple.h"
 #include <cstddef>
 #include <iterator>
@@ -328,6 +320,8 @@ class RecordContext {
   /// chain.
   void stealRecordChain(RecordContext &Other);
 
+  void removeFromRecordChain(APIRecord *Record);
+
   APIRecord::RecordKind getKind() const { return Kind; }
 
   struct record_iterator {
@@ -1426,10 +1420,14 @@ class APISet {
   typename std::enable_if_t<std::is_base_of_v<APIRecord, RecordTy>, RecordTy> *
   createRecord(StringRef USR, StringRef Name, CtorArgsContTy &&...CtorArgs);
 
-  ArrayRef<const APIRecord *> getTopLevelRecords() const {
-    return TopLevelRecords;
+  auto getTopLevelRecords() const {
+    return llvm::iterator_range<decltype(TopLevelRecords)::iterator>(TopLevelRecords);
   }
 
+  void removeRecord(StringRef USR);
+
+  void removeRecord(APIRecord *Record);
+
   APISet(const llvm::Triple &Target, Language Lang,
          const std::string &ProductName)
       : Target(Target), Lang(Lang), ProductName(ProductName) {}
@@ -1456,7 +1454,7 @@ class APISet {
   // lives in the BumpPtrAllocator.
   using APIRecordStoredPtr = std::unique_ptr<APIRecord, APIRecordDeleter>;
   llvm::DenseMap<StringRef, APIRecordStoredPtr> USRBasedLookupTable;
-  std::vector<const APIRecord *> TopLevelRecords;
+  llvm::SmallPtrSet<const APIRecord *, 32> TopLevelRecords;
 
 public:
   const std::string ProductName;
@@ -1482,7 +1480,7 @@ APISet::createRecord(StringRef USR, StringRef Name,
             dyn_cast_if_present<RecordContext>(Record->Parent.Record))
       ParentContext->addToRecordChain(Record);
     else
-      TopLevelRecords.push_back(Record);
+      TopLevelRecords.insert(Record);
   } else {
     Record = dyn_cast<RecordTy>(Result.first->second.get());
   }

clang/lib/ExtractAPI/API.cpp

@@ -13,8 +13,6 @@
 //===----------------------------------------------------------------------===//
 
 #include "clang/ExtractAPI/API.h"
-#include "clang/AST/RawCommentList.h"
-#include "clang/Index/USRGeneration.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/ErrorHandling.h"
 #include <memory>
@@ -60,6 +58,10 @@ bool RecordContext::IsWellFormed() const {
 
 void RecordContext::stealRecordChain(RecordContext &Other) {
   assert(IsWellFormed());
+  // Other's record chain is empty, nothing to do
+  if (Other.First == nullptr && Other.Last == nullptr)
+    return;
+
   // If we don't have an empty chain append Other's chain into ours.
   if (First)
     Last->NextInContext = Other.First;
@@ -68,6 +70,9 @@ void RecordContext::stealRecordChain(RecordContext &Other) {
 
   Last = Other.Last;
 
+  for (auto *StolenRecord = Other.First; StolenRecord != nullptr; StolenRecord = StolenRecord->getNextInContext())
+    StolenRecord->Parent = SymbolReference(cast<APIRecord>(this));
+
   // Delete Other's chain to ensure we don't accidentally traverse it.
   Other.First = nullptr;
   Other.Last = nullptr;
@@ -85,6 +90,22 @@ void RecordContext::addToRecordChain(APIRecord *Record) const {
   Last = Record;
 }
 
+void RecordContext::removeFromRecordChain(APIRecord *Record) {
+  APIRecord *Prev = nullptr;
+  for (APIRecord *Curr = First; Curr != Record; Curr = Curr->NextInContext)
+    Prev = Curr;
+
+  if (Prev)
+    Prev->NextInContext = Record->NextInContext;
+  else
+    First = Record->NextInContext;
+
+  if (Last == Record)
+    Last = Prev;
+
+  Record->NextInContext = nullptr;
+}
+
 APIRecord *APISet::findRecordForUSR(StringRef USR) const {
   if (USR.empty())
     return nullptr;
@@ -114,6 +135,38 @@ SymbolReference APISet::createSymbolReference(StringRef Name, StringRef USR,
   return SymbolReference(copyString(Name), copyString(USR), copyString(Source));
 }
 
+
+void APISet::removeRecord(StringRef USR) {
+  auto Result = USRBasedLookupTable.find(USR);
+  if (Result != USRBasedLookupTable.end()) {
+    auto *Record = Result->getSecond().get();
+    auto &ParentReference= Record->Parent;
+    auto *ParentRecord = const_cast<APIRecord *>(ParentReference.Record);
+    if (!ParentRecord)
+      ParentRecord = findRecordForUSR(ParentReference.USR);
+
+    if (auto *ParentCtx = llvm::cast_if_present<RecordContext>(ParentRecord)) {
+      ParentCtx->removeFromRecordChain(Record);
+      if (auto *RecordAsCtx=
+              llvm::dyn_cast<RecordContext>(Record))
+        ParentCtx->stealRecordChain(*RecordAsCtx);
+    } else {
+      TopLevelRecords.erase(Record);
+      if (auto *RecordAsCtx=
+              llvm::dyn_cast<RecordContext>(Record)) {
+        for (const auto *Child = RecordAsCtx->First; Child != nullptr;
+             Child = Child->getNextInContext())
+          TopLevelRecords.insert(Child);
+      }
+    }
+    USRBasedLookupTable.erase(Result);
+  }
+}
+
+void APISet::removeRecord(APIRecord *Record) {
+  removeRecord(Record->USR);
+}
+
 APIRecord::~APIRecord() {}
 TagRecord::~TagRecord() {}
 RecordRecord::~RecordRecord() {}