Reduce memory usage in AST parent map generation by lazily checking if nodes have been seen

Author: higher-performance

clang/lib/AST/ParentMapContext.cpp

@@ -60,6 +60,34 @@ class ParentMapContext::ParentMap {
 
   template <typename, typename...> friend struct ::MatchParents;
 
+  template <class T> struct IndirectDenseMapInfo {
+    using Ptr = T *;
+    using Base = llvm::DenseMapInfo<std::remove_cv_t<T>>;
+    static T &Verify(T &Val) {
+      assert(Val.getMemoizationData());
+      return Val;
+    }
+    static inline Ptr getEmptyKey() {
+      return static_cast<Ptr>(llvm::DenseMapInfo<void *>::getEmptyKey());
+    }
+    static inline Ptr getTombstoneKey() {
+      return static_cast<Ptr>(llvm::DenseMapInfo<void *>::getTombstoneKey());
+    }
+    static unsigned getHashValue(Ptr Val) {
+      return Val == getEmptyKey() || Val == getTombstoneKey()
+                 ? 0
+                 : Base::getHashValue(Verify(*Val));
+    }
+    static bool isEqual(Ptr LHS, Ptr RHS) {
+      if (LHS == getEmptyKey() || LHS == getTombstoneKey() ||
+          RHS == getEmptyKey() || RHS == getTombstoneKey() || LHS == RHS) {
+        return LHS == RHS;
+      }
+      return Base::isEqual(Verify(*LHS), Verify(*RHS));
+    }
+  };
+  using MapInfo = IndirectDenseMapInfo<const DynTypedNode>;
+
   /// Contains parents of a node.
   class ParentVector {
   public:
@@ -69,17 +97,70 @@ class ParentMapContext::ParentMap {
       for (; N > 0; --N)
         push_back(Value);
     }
-    bool contains(const DynTypedNode &Value) {
-      return Seen.contains(Value);
-    }
+    bool contains(const DynTypedNode &Value) { return CacheUntil(Value); }
     void push_back(const DynTypedNode &Value) {
-      if (!Value.getMemoizationData() || Seen.insert(Value).second)
+      if (!Value.getMemoizationData() || !FragileDedupCache.contains(&Value)) {
+        const size_t OldCapacity = Items.capacity();
         Items.push_back(Value);
+        // Danger: Clear the cache if its pointers might be invalidated.
+        if (OldCapacity != Items.capacity()) {
+          ItemsProcessed = 0;
+          FragileDedupCache.clear();
+        }
+      }
+
+      // Cache some nodes every time we attempt an insertion.
+      // This is important to guarantee that our (cheap) cache catches up with
+      // the (expensive) main vector upon repeated insertions, thus bounding
+      // element duplication.
+      // It also amortizes insertions to constant-time, preventing us from
+      // spending unnecessary time rebuilding the cache when it is invalidated.
+      CacheUntil(2);
+    }
+    llvm::ArrayRef<DynTypedNode> view() {
+      CacheUntil();
+      assert(ItemsProcessed == Items.size());
+      return Items;
     }
-    llvm::ArrayRef<DynTypedNode> view() const { return Items; }
+
   private:
+    // Places at most `MaxCount` items into the cache, stopping early if the
+    // node is seen.
+    bool CacheUntil(size_t MaxCount = std::numeric_limits<size_t>::max(),
+                    const DynTypedNode *OptionalNeedle = nullptr) {
+      assert(!OptionalNeedle || OptionalNeedle->getMemoizationData());
+
+      bool Found = OptionalNeedle && FragileDedupCache.contains(OptionalNeedle);
+      while (!Found && ItemsProcessed < Items.size() && MaxCount > 0) {
+        const auto *Item = &Items[ItemsProcessed];
+
+        if (Item->getMemoizationData()) {
+          FragileDedupCache.insert(Item);
+        }
+        ++ItemsProcessed;
+        --MaxCount;
+
+        assert(!Found);
+        Found = OptionalNeedle && Item->getMemoizationData() &&
+                MapInfo::isEqual(OptionalNeedle, Item);
+      }
+
+      return Found;
+    }
+    bool CacheUntil(const DynTypedNode &Needle) {
+      return CacheUntil(std::numeric_limits<size_t>::max(), &Needle);
+    }
+
+    // A partitioned vector of nodes, where the first `ItemsProcessed` elements
+    // are already processed into the cache, and the remaining have not.
+    // BE CAREFUL. Pointers into this container are stored in the
+    // `FragileDedupCache` set below.
     llvm::SmallVector<DynTypedNode, 2> Items;
-    llvm::SmallDenseSet<DynTypedNode, 2> Seen;
+    // This cache is fragile because it contains pointers that are invalidated
+    // when the vector capacity changes.
+    llvm::SmallDenseSet<const DynTypedNode *, 2, MapInfo> FragileDedupCache;
+    // Lazily tracks which items have been processed for the cache.
+    size_t ItemsProcessed = 0;
   };
 
   /// Maps from a node to its parents. This is used for nodes that have
@@ -117,7 +198,7 @@ class ParentMapContext::ParentMap {
     if (I == Map.end()) {
       return llvm::ArrayRef<DynTypedNode>();
     }
-    if (const auto *V = dyn_cast<ParentVector *>(I->second)) {
+    if (auto *V = dyn_cast<ParentVector *>(I->second)) {
       return V->view();
     }
     return getSingleDynTypedNodeFromParentMap(I->second);