Fix stale graph edges during teardown

Author: Kyle-Ye

Sources/ComputeCxx/Attribute/AttributeID/WeakAttributeID.cpp

@@ -10,7 +10,7 @@ bool WeakAttributeID::expired() const {
     uint64_t raw_page_seed = data::table::shared().raw_page_seed(_identifier.page_ptr());
     if (raw_page_seed & 0xff00000000) {
         auto zone_info = data::zone::info::from_raw_value(uint32_t(raw_page_seed));
-        if (zone_info.zone_id() == _seed) {
+        if (zone_info.zone_id() == _seed && !zone_info.is_deleted()) {
             return false;
         }
     }

Sources/ComputeCxx/Data/Zone.h

@@ -25,10 +25,19 @@ class zone {
         info(uint32_t zone_id) : _value(zone_id & zone_id_mask) {};
 
         uint32_t zone_id() const { return _value & zone_id_mask; };
-        info with_deleted() const { return info(_value | deleted); };
+        bool is_deleted() const { return (_value & deleted) != 0; };
+        info with_deleted() const {
+            auto result = *this;
+            result._value |= deleted;
+            return result;
+        };
 
         uint32_t to_raw_value() const { return _value; };
-        static info from_raw_value(uint32_t value) { return info(value); };
+        static info from_raw_value(uint32_t value) {
+            info result(0);
+            result._value = value;
+            return result;
+        };
     };
 
   private:

Sources/ComputeCxx/Graph/Graph.cpp

@@ -36,6 +36,28 @@ platform_lock Graph::_all_graphs_lock = PLATFORM_LOCK_INIT;
 
 pthread_key_t Graph::_current_update_key = 0;
 
+namespace {
+
+bool attribute_page_is_allocated(AttributeID attribute) {
+    if (!attribute || attribute.is_nil()) {
+        return false;
+    }
+    return data::table::shared().raw_page_seed(attribute.page_ptr()) & 0xff00000000;
+}
+
+bool attribute_zone_is_live(AttributeID attribute) {
+    if (!attribute || attribute.is_nil()) {
+        return false;
+    }
+    auto raw_page_seed = data::table::shared().raw_page_seed(attribute.page_ptr());
+    if (!(raw_page_seed & 0xff00000000)) {
+        return false;
+    }
+    return !data::zone::info::from_raw_value(uint32_t(raw_page_seed)).is_deleted();
+}
+
+} // namespace
+
 Graph::Graph()
     : _heap(nullptr, 0, 0), _interned_types(nullptr, nullptr, nullptr, nullptr, &_heap),
       _contexts_by_id(nullptr, nullptr, nullptr, nullptr, &_heap), _id(IAGMakeUniqueID()) {
@@ -552,23 +574,40 @@ void Graph::remove_indirect_node(data::ptr<IndirectNode> indirect_node) {
 }
 
 void Graph::remove_removed_input(AttributeID attribute, AttributeID input) {
+    // An indirect input in the current invalidation batch can have a deleted
+    // zone while still resolving to a live source, so only require its page to
+    // remain allocated until after resolution.
+    if (!attribute_page_is_allocated(input)) {
+        return;
+    }
+
     auto resolved_input =
         input.resolve(TraversalOptions::SkipMutableReference | TraversalOptions::EvaluateWeakReferences).attribute();
+    if (!attribute_zone_is_live(resolved_input)) {
+        return;
+    }
+
+    auto input_subgraph = resolved_input.subgraph();
+    if (!input_subgraph || input_subgraph->graph() != this || !input_subgraph->is_valid()) {
+        return;
+    }
+
     if (auto input_node = resolved_input.get_node()) {
-        if (!resolved_input.subgraph()->is_invalidated()) {
-            remove_output_edge(input_node, attribute);
-        }
+        remove_output_edge(input_node, attribute);
     } else if (auto input_indirect_node = resolved_input.get_indirect_node()) {
-        if (!resolved_input.subgraph()->is_invalidated()) {
-            if (input_indirect_node->is_mutable()) {
-                remove_output_edge(input_indirect_node.unsafe_cast<MutableIndirectNode>(), attribute);
-            }
+        if (input_indirect_node->is_mutable()) {
+            remove_output_edge(input_indirect_node.unsafe_cast<MutableIndirectNode>(), attribute);
         }
     }
 }
 
 bool Graph::remove_removed_output(AttributeID attribute, AttributeID output, bool option) {
-    if (output.subgraph()->is_invalidated()) {
+    if (!attribute_zone_is_live(output)) {
+        return false;
+    }
+
+    auto output_subgraph = output.subgraph();
+    if (!output_subgraph || output_subgraph->graph() != this || !output_subgraph->is_valid()) {
         return false;
     }
 

Tests/ComputeTests/Shared/Attribute/AttributeBodyTests.swift

@@ -53,6 +53,51 @@ struct AttributeBodyTests {
             }
         }
 
+        @Test
+        func weakAttributeExpiresBeforeBodyIsDestroyed() {
+            struct TestBody: _AttributeBody {
+                var weakAttribute: WeakAttribute<Int>
+                var weakAttributeExpired: UnsafeMutablePointer<Bool>
+
+                static func _destroySelf(_ self: UnsafeMutableRawPointer) {
+                    let testBody = self.assumingMemoryBound(to: Self.self)
+                    testBody.pointee.weakAttributeExpired.pointee = testBody.pointee.weakAttribute.attribute == nil
+                }
+
+                static var _hasDestroySelf: Bool {
+                    return true
+                }
+            }
+
+            let graph = Graph()
+            let subgraph = Subgraph(graph: graph)
+            let oldSubgraph = Subgraph.current
+            defer {
+                Subgraph.current = oldSubgraph
+            }
+            Subgraph.current = subgraph
+
+            let attribute = Attribute(value: 0)
+            var weakAttributeExpired = false
+            withUnsafeMutablePointer(to: &weakAttributeExpired) { weakAttributeExpiredPointer in
+                let body = TestBody(
+                    weakAttribute: WeakAttribute(attribute),
+                    weakAttributeExpired: weakAttributeExpiredPointer
+                )
+                let _ = withUnsafePointer(to: body) { bodyPointer in
+                    Attribute<Int>(body: bodyPointer, value: nil, flags: []) {
+                        return { _, _ in }
+                    }
+                }
+
+                #expect(weakAttributeExpiredPointer.pointee == false)
+
+                subgraph.invalidate()
+
+                #expect(weakAttributeExpiredPointer.pointee == true)
+            }
+        }
+
         @Test
         func doesNotCallDestroySelfWhenHasDestroySelfIsFalse() {
             struct TestBody: _AttributeBody {