Improve performance when batch invalidating/removing subgraphs

Sources/ComputeCxx/Collection/Set.h

@@ -0,0 +1,92 @@
+#pragma once
+
+#include <algorithm>
+#include <utility>
+
+#include "ComputeCxx/IAGBase.h"
+#include "Vector.h"
+
+IAG_ASSUME_NONNULL_BEGIN
+
+namespace IAG {
+
+/// A set of pointers stored in a vector sorted by address.
+template <typename T, typename size_type = std::size_t>
+    requires std::is_pointer_v<T> && std::unsigned_integral<size_type>
+class set {
+  private:
+    static constexpr uintptr_t removed = 0x1;
+
+    vector<uintptr_t, 0, size_type> _slots;
+    size_type _num_removed = 0;
+
+    uintptr_t *_Nonnull find_slot(const T &value) {
+        return std::lower_bound(_slots.begin(), _slots.end(), uintptr_t(value),
+                                [](uintptr_t slot, uintptr_t value) { return (slot & ~removed) < value; });
+    };
+
+  public:
+    class iterator {
+      private:
+        const uintptr_t *_Nonnull _slot;
+        const uintptr_t *_Nonnull _end;
+
+        void skip_removed() {
+            while (_slot != _end && (*_slot & removed)) {
+                ++_slot;
+            }
+        };
+
+      public:
+        iterator(const uintptr_t *_Nonnull slot, const uintptr_t *_Nonnull end) : _slot(slot), _end(end) {
+            skip_removed();
+        };
+
+        T _Nonnull operator*() const { return reinterpret_cast<T>(*_slot); };
+        iterator &operator++() {
+            ++_slot;
+            skip_removed();
+            return *this;
+        };
+        bool operator==(const iterator &other) const { return _slot == other._slot; };
+    };
+
+    iterator begin() const { return iterator(_slots.begin(), _slots.end()); };
+    iterator end() const { return iterator(_slots.end(), _slots.end()); };
+
+    bool empty() const { return _slots.size() == _num_removed; };
+
+    std::pair<iterator, bool> insert(const T &value) {
+        auto *slot = find_slot(value);
+        bool inserted = true;
+        if (slot == _slots.end() || (*slot & ~removed) != uintptr_t(value)) {
+            slot = _slots.insert(slot, uintptr_t(value));
+        } else if (*slot & removed) {
+            // A new element reused the address of a removed one
+            *slot = uintptr_t(value);
+            _num_removed -= 1;
+        } else {
+            inserted = false;
+        }
+        return {iterator(slot, _slots.end()), inserted};
+    };
+
+    size_type erase(const T &value) {
+        auto *slot = find_slot(value);
+        if (slot == _slots.end() || *slot != uintptr_t(value)) {
+            return 0;
+        }
+        *slot |= removed;
+        _num_removed += 1;
+        if (_num_removed > _slots.size() / 2) {
+            _slots.erase(std::remove_if(_slots.begin(), _slots.end(), [](uintptr_t slot) { return slot & removed; }),
+                         _slots.end());
+            _num_removed = 0;
+        }
+        return 1;
+    };
+};
+
+} // namespace IAG
+
+IAG_ASSUME_NONNULL_END

Sources/ComputeCxx/Comparison/Builder.h

@@ -2,11 +2,11 @@
 
 #include <platform/lock.h>
 
+#include "Collection/Vector.h"
 #include "ComputeCxx/IAGBase.h"
 #include "ComputeCxx/IAGComparison.h"
 #include "LayoutDescriptor.h"
 #include "Swift/MetadataVisitor.h"
-#include "Vector/Vector.h"
 
 IAG_ASSUME_NONNULL_BEGIN
 

Sources/ComputeCxx/Comparison/Compare.h

@@ -2,9 +2,9 @@
 
 #include "LayoutDescriptor.h"
 
+#include "Collection/Vector.h"
 #include "ComputeCxx/IAGBase.h"
 #include "ComputeCxx/IAGComparison.h"
-#include "Vector/Vector.h"
 
 IAG_ASSUME_NONNULL_BEGIN
 

Sources/ComputeCxx/Data/Table.h

@@ -8,8 +8,8 @@
 #include <platform/malloc.h>
 #include <platform/vm.h>
 
+#include "Collection/Vector.h"
 #include "ComputeCxx/IAGBase.h"
-#include "Vector/Vector.h"
 
 IAG_ASSUME_NONNULL_BEGIN
 

Sources/ComputeCxx/Debug/DebugServer.h

@@ -10,8 +10,8 @@
 
 #include <Utilities/ObjCPointer.h>
 
+#include "Collection/Vector.h"
 #include "ComputeCxx/IAGDebugServer.h"
-#include "Vector/Vector.h"
 
 IAG_ASSUME_NONNULL_BEGIN
 

Sources/ComputeCxx/Graph/Graph.cpp

@@ -243,16 +243,14 @@ void Graph::call_main_handler(void *context, void (*body)(void *)) {
 #pragma mark - Subgraphs
 
 void Graph::add_subgraph(Subgraph &subgraph) {
-    auto *pos = std::lower_bound(_subgraphs.begin(), _subgraphs.end(), &subgraph);
-    _subgraphs.insert(pos, &subgraph);
+    _subgraphs.insert(&subgraph);
 
     _num_subgraphs += 1;
     _num_subgraphs_total += 1;
 }
 
 void Graph::remove_subgraph(Subgraph &subgraph) {
-    auto *iter = std::remove(_subgraphs.begin(), _subgraphs.end(), &subgraph);
-    _subgraphs.erase(iter, _subgraphs.end());
+    _subgraphs.erase(&subgraph);
 
     if (auto *map = _tree_data_elements_by_subgraph.get()) {
         auto iter = map->find(&subgraph);

Sources/ComputeCxx/Graph/Graph.h

@@ -20,10 +20,11 @@
 #include "Attribute/AttributeID/AttributeID.h"
 #include "Attribute/AttributeType/AttributeType.h"
 #include "Closure/ClosureFunction.h"
+#include "Collection/Set.h"
+#include "Collection/Vector.h"
 #include "ComputeCxx/IAGGraphProfiling.h"
 #include "ComputeCxx/IAGGraphTracing.h"
 #include "Swift/Metadata.h"
-#include "Vector/Vector.h"
 
 IAG_ASSUME_NONNULL_BEGIN
 
@@ -191,7 +192,7 @@ class Graph {
     std::unique_ptr<KeyTable> _keys = nullptr;
 
     // Subgraphs
-    vector<Subgraph *, 0, uint32_t> _subgraphs;
+    set<Subgraph *, uint32_t> _subgraphs;
     vector<Subgraph *, 0, uint32_t> _subgraphs_with_cached_nodes;
     vector<Subgraph *, 2, uint32_t> _invalidating_subgraphs;
     bool _deferring_subgraph_invalidation = false;
@@ -337,7 +338,7 @@ class Graph {
 
     // MARK: Subgraphs
 
-    vector<Subgraph *, 0, uint32_t> &subgraphs() { return _subgraphs; };
+    set<Subgraph *, uint32_t> &subgraphs() { return _subgraphs; };
 
     void add_subgraph(Subgraph &subgraph);
     void remove_subgraph(Subgraph &subgraph);

Sources/ComputeCxx/Graph/Graph.mm

@@ -543,8 +543,9 @@ int trap_cycles() {
         NSMutableArray *subgraph_dicts = [NSMutableArray array];
 
         auto subgraph_indices = std::unordered_map<Subgraph *, uint64_t>();
-        for (uint32_t index = 0, iter = graph->subgraphs().size(); iter > 0; --iter, ++index) {
-            subgraph_indices.try_emplace(graph->subgraphs()[index], index);
+        uint64_t subgraph_index = 0;
+        for (auto subgraph : graph->subgraphs()) {
+            subgraph_indices.try_emplace(subgraph, subgraph_index++);
         }
 
         for (auto subgraph : graph->subgraphs()) {

Sources/ComputeCxx/Protobuf/Encoder.h

@@ -2,7 +2,7 @@
 
 #include "ComputeCxx/IAGBase.h"
 
-#include "Vector/Vector.h"
+#include "Collection/Vector.h"
 
 IAG_ASSUME_NONNULL_BEGIN
 

Sources/ComputeCxx/Subgraph/Subgraph.h

@@ -6,12 +6,12 @@
 
 #include "Attribute/AttributeID/AttributeID.h"
 #include "Closure/ClosureFunction.h"
+#include "Collection/IndirectPointerVector.h"
 #include "ComputeCxx/IAGBase.h"
 #include "ComputeCxx/IAGUniqueID.h"
 #include "Data/Zone.h"
 #include "Graph/Graph.h"
 #include "Graph/Tree/TreeElement.h"
-#include "Vector/IndirectPointerVector.h"
 
 IAG_ASSUME_NONNULL_BEGIN
 

Sources/ComputeCxx/Swift/ContextDescriptor.h

@@ -3,9 +3,9 @@
 #include <llvm/ADT/ArrayRef.h>
 #include <swift/Runtime/Metadata.h>
 
+#include "Collection/Vector.h"
 #include "ComputeCxx/IAGBase.h"
 #include "Metadata.h"
-#include "Vector/Vector.h"
 
 IAG_ASSUME_NONNULL_BEGIN
 

Tests/ComputeTests/Shared/Subgraph/SubgraphLifecycleTests.swift

@@ -89,4 +89,87 @@ struct SubgraphLifecycleTests {
             #expect(graph.counter(for: .createdSubgraphs) == 1)
         }
     }
+
+    @Suite
+    struct SubgraphBatchInvalidateTests {
+        @Test
+        func invalidateManySubgraphs() {
+            let graph = Graph()
+
+            var subgraphs: [Subgraph] = []
+            for _ in 0..<10_000 {
+                var parent = Subgraph(graph: graph)
+                subgraphs.append(parent)
+                for _ in 0..<3 {
+                    let child = Subgraph(graph: graph)
+                    parent.addChild(child)
+                    subgraphs.append(child)
+                    parent = child
+                }
+            }
+            #expect(graph.counter(for: .subgraphs) == 40_000)
+
+            graph.withoutSubgraphInvalidation {
+                for subgraph in subgraphs {
+                    subgraph.invalidate()
+                }
+            }
+            #expect(graph.counter(for: .subgraphs) == 0)
+
+            withExtendedLifetime(subgraphs) {}
+        }
+
+        #if canImport(Darwin)
+        @Test(arguments: [1, 3, 5])
+        func remainingSubgraphsAfterPartialInvalidate(keepEvery: Int) throws {
+            let graph = Graph()
+
+            let root = Subgraph(graph: graph)
+            var subgraphs: [Subgraph] = []
+            for _ in 0..<1_000 {
+                let subgraph = Subgraph(graph: graph)
+                root.addChild(subgraph)
+                subgraphs.append(subgraph)
+            }
+
+            var invalidated: [Subgraph] = []
+            var remaining: [Subgraph] = []
+            for (index, subgraph) in subgraphs.enumerated() {
+                if index % (keepEvery + 1) == 0 {
+                    remaining.append(subgraph)
+                } else {
+                    invalidated.append(subgraph)
+                }
+            }
+            graph.withoutSubgraphInvalidation {
+                for subgraph in invalidated {
+                    subgraph.invalidate()
+                }
+            }
+
+            // New subgraphs may reuse the addresses of invalidated ones
+            for _ in 0..<100 {
+                let subgraph = Subgraph(graph: graph)
+                root.addChild(subgraph)
+                remaining.append(subgraph)
+            }
+
+            #expect(graph.counter(for: .subgraphs) == remaining.count + 1)
+            #expect(root.childCount == remaining.count)
+
+            let description = try #require(graph.dictionaryDescription())
+            let graphDescription = try #require(description.graphs.first)
+            #expect(graphDescription.subgraphs.count == remaining.count + 1)
+
+            let rootIndex = try #require(graphDescription.subgraphs.firstIndex { $0.children != nil })
+            let rootChildren = try #require(graphDescription.subgraphs[rootIndex].children)
+            #expect(Set(rootChildren).count == remaining.count)
+            for childIndex in rootChildren {
+                #expect(graphDescription.subgraphs[childIndex].parents == [rootIndex])
+            }
+
+            withExtendedLifetime((root, subgraphs, remaining)) {}
+        }
+        #endif
+    }
 }