Cached values

CompatibilityTesting/Package.swift

@@ -2,8 +2,6 @@
 
 import PackageDescription
 
-let swiftCheckoutPath = "\(Context.packageDirectory)/Checkouts/swift"
-
 var dependencies: [Package.Dependency] = [
     .package(name: "Compute", path: ".."),
     .package(url: "https://github.com/apple/swift-algorithms", from: "1.2.0"),

Sources/Compute/Attribute/Rule/Rule.swift

@@ -22,7 +22,6 @@ extension Rule {
         withUnsafePointer(to: initialValue) { initialValuePointer in
             Graph.setOutputValue(initialValuePointer)
         }
-
     }
 
     public static func _update(_ self: UnsafeMutableRawPointer, attribute: AnyAttribute) {
@@ -51,16 +50,50 @@ extension Rule {
 
 }
 
-public struct CachedValueOptions {}
+extension Graph {
+    @_extern(c, "AGGraphReadCachedAttribute")
+    static func readCachedAttribute(
+        hash: Int,
+        type: Metadata,
+        body: UnsafeRawPointer,
+        valueType: Metadata,
+        options: CachedValueOptions,
+        owner: AnyAttribute,
+        changed: UnsafeMutablePointer<Bool>?,
+        attributeTypeID: (AGUnownedGraphContextRef) -> UInt32
+    ) -> UnsafeRawPointer
+}
 
-extension Rule where Value: Hashable {
+extension Rule where Self: Hashable {
 
     public func cachedValue(options: CachedValueOptions, owner: AnyAttribute?) -> Value {
-        fatalError("not implemented")
+        return withUnsafePointer(to: self) { bodyPointer in
+            Self._cachedValue(
+                options: options,
+                owner: owner,
+                hashValue: hashValue,
+                bodyPtr: bodyPointer,
+                update: { Self._update }
+            ).pointee
+        }
     }
 
     public func cachedValueIfExists(options: CachedValueOptions, owner: AnyAttribute?) -> Value? {
-        fatalError("not implemented")
+        return withUnsafePointer(to: self) { bodyPointer in
+            let value = __AGGraphReadCachedAttributeIfExists(
+                hashValue,
+                Metadata(Self.self),
+                bodyPointer,
+                Metadata(Value.self),
+                options,
+                owner ?? .nil,
+                nil
+            )
+            guard let value else {
+                return nil
+            }
+            return value.assumingMemoryBound(to: Value.self).pointee
+        }
     }
 
     public static func _cachedValue(
@@ -70,7 +103,28 @@ extension Rule where Value: Hashable {
         bodyPtr: UnsafeRawPointer,
         update: () -> (UnsafeMutableRawPointer, AnyAttribute) -> Void
     ) -> UnsafePointer<Value> {
-        fatalError("not implemented")
+        let value = Graph.readCachedAttribute(
+            hash: hashValue,
+            type: Metadata(Self.self),
+            body: bodyPtr,
+            valueType: Metadata(Value.self),
+            options: options,
+            owner: owner ?? .nil,
+            changed: nil
+        ) { graph in
+            return graph.internAttributeType(type: Metadata(Self.self)) {
+                let attributeType = _AttributeType(
+                    selfType: Self.self,
+                    valueType: Value.self,
+                    flags: [],  // TODO: check flags are empty
+                    update: update()
+                )
+                let pointer = UnsafeMutablePointer<_AttributeType>.allocate(capacity: 1)
+                pointer.initialize(to: attributeType)
+                return UnsafePointer(pointer)
+            }
+        }
+        return value.assumingMemoryBound(to: Value.self)
     }
 
 }

Sources/Compute/Attribute/RuleContext/RuleContext.swift

@@ -6,11 +6,8 @@ public struct RuleContext<Value> {
         self.attribute = attribute
     }
 
-    @_extern(c, "AGGraphWithUpdate")
-    private static func withUpdate(_ attribute: AnyAttribute, body: @escaping () -> Void)
-
-    public func update(body: @escaping () -> Void) {
-        RuleContext.withUpdate(attribute.identifier, body: body)
+    public func update(body: () -> Void) {
+        AnyRuleContext(attribute: attribute.identifier).update(body: body)
     }
 
     public var value: Value {

Sources/ComputeCxx/Attribute/AttributeData/Node/Node.h

@@ -53,7 +53,7 @@ class Node {
     unsigned int _has_indirect_value : 1 = 0;
     unsigned int _input_edges_traverse_contexts : 1 = 0;
     unsigned int _needs_sort_input_edges : 1 = 0;
-    unsigned int _cacheable : 1 = 0;
+    unsigned int _cached : 1 = 0;
     unsigned int _main_ref : 1 = 0;
     unsigned int _self_modified : 1 = 0;
 
@@ -143,6 +143,9 @@ class Node {
         }
     }
 
+    bool is_cached() const { return _cached; }
+    void set_cached(bool value) { _cached = value; }
+
     bool is_main_ref() const { return _main_ref; }
     void set_main_ref(bool value) { _main_ref = value; }
 

Sources/ComputeCxx/Closure/ClosureFunction.h

@@ -16,6 +16,7 @@ template <typename ReturnType, typename... Args> class ClosureFunction {
     Context _context;
 
   public:
+    inline ClosureFunction(std::nullptr_t): _function(nullptr), _context(nullptr) {}
     inline ClosureFunction(Function function, Context context) noexcept : _function(function), _context(context) {
         void *mutable_context = const_cast<void *>(_context);
         ::swift::swift_retain(reinterpret_cast<::swift::HeapObject *>(mutable_context));
@@ -94,4 +95,8 @@ template <typename ReturnType, typename Arg>
     requires std::same_as<ReturnType, void>
 using ClosureFunctionPV = ClosureFunction<ReturnType, Arg>;
 
+template <typename ReturnType, typename Arg>
+    requires std::unsigned_integral<ReturnType>
+using ClosureFunctionCI = ClosureFunction<ReturnType, Arg>;
+
 } // namespace AG

Sources/ComputeCxx/Comparison/LayoutDescriptor.h

@@ -29,7 +29,7 @@ enum class HeapMode : uint16_t {
     GenericLocals = 1 << 2,
 };
 inline bool operator&(HeapMode a, HeapMode b) { return (uint16_t)a & (uint16_t)b; }
-inline HeapMode operator|(HeapMode a, HeapMode b) { return a | b; }
+inline HeapMode operator|(HeapMode a, HeapMode b) { return (HeapMode)((uint16_t)a | (uint16_t)b); }
 
 extern unsigned char base_address;
 

Sources/ComputeCxx/Graph/AGGraph.cpp

@@ -639,6 +639,84 @@ void AGGraphSetInvalidationCallback(AGGraphRef graph,
     graph_context->set_invalidation_callback(AG::ClosureFunctionAV<void, AGAttribute>(callback, callback_context));
 }
 
+#pragma mark - Cached value
+
+namespace {
+
+void *read_cached_attribute(size_t hash, const AG::swift::metadata &metadata, const void *body,
+                            const AG::swift::metadata &value_metadata, AGCachedValueOptions options,
+                            AG::AttributeID owner_id, AGChangedValueFlags *flags_out,
+                            AG::ClosureFunctionCI<uint32_t, AGUnownedGraphContextRef> get_attribute_type_id) {
+    auto update = AG::Graph::current_update();
+    auto update_stack = update.tag() == 0 ? update.get() : nullptr;
+
+    AG::Subgraph *subgraph = nullptr;
+    if (owner_id && !owner_id.is_nil()) {
+        owner_id.validate_data_offset();
+        subgraph = owner_id.subgraph();
+    } else {
+        if (update_stack != nullptr) {
+            subgraph = AG::AttributeID(update_stack->frames().back().attribute).subgraph();
+        } else {
+            subgraph = AG::Subgraph::current_subgraph();
+        }
+    }
+    if (subgraph == nullptr) {
+        AG::precondition_failure("no subgraph");
+    }
+
+    AG::data::ptr<AG::Node> cached_node = subgraph->cache_fetch(hash, metadata, body, get_attribute_type_id);
+    if (cached_node == nullptr) {
+        return nullptr;
+    }
+
+    if (update_stack == nullptr) {
+        void *value = subgraph->graph()->value_ref(AG::AttributeID(cached_node), 0, value_metadata, flags_out);
+        subgraph->cache_insert(cached_node); // TODO: when this becomes an input, is it removed from cache?
+        return value;
+    }
+
+    AGInputOptions input_options = options & AGCachedValueOptionsUnprefetched ? AGInputOptionsUnprefetched : AGInputOptionsNone;
+    return subgraph->graph()->input_value_ref(update_stack->frames().back().attribute, AG::AttributeID(cached_node), 0,
+                                              input_options, value_metadata, flags_out);
+}
+
+} // namespace
+
+void *AGGraphReadCachedAttribute(size_t hash, AGTypeID type, const void *body, AGTypeID value_type,
+                                 AGCachedValueOptions options, AGAttribute owner, bool *_Nullable changed_out,
+                                 uint32_t (*closure)(const void *context AG_SWIFT_CONTEXT,
+                                                     AGUnownedGraphContextRef graph_context) AG_SWIFT_CC(swift),
+                                 const void *closure_context) {
+    auto metadata = reinterpret_cast<const AG::swift::metadata *>(type);
+    auto value_metadata = reinterpret_cast<const AG::swift::metadata *>(value_type);
+    auto owner_id = AG::AttributeID(owner);
+
+    AGChangedValueFlags flags = 0;
+    void *value =
+        read_cached_attribute(hash, *metadata, body, *value_metadata, options, owner_id, &flags,
+                              AG::ClosureFunctionCI<uint32_t, AGUnownedGraphContextRef>(closure, closure_context));
+    if (changed_out) {
+        *changed_out = flags & AGChangedValueFlagsChanged ? true : false;
+    }
+    return value;
+}
+
+void *AGGraphReadCachedAttributeIfExists(size_t hash, AGTypeID type, const void *body, AGTypeID value_type,
+                                         AGCachedValueOptions options, AGAttribute owner, bool *_Nullable changed_out) {
+
+    auto metadata = reinterpret_cast<const AG::swift::metadata *>(type);
+    auto value_metadata = reinterpret_cast<const AG::swift::metadata *>(value_type);
+    auto owner_id = AG::AttributeID(owner);
+
+    AGChangedValueFlags flags = 0;
+    void *value = read_cached_attribute(hash, *metadata, body, *value_metadata, options, owner_id, &flags, nullptr);
+    if (changed_out) {
+        *changed_out = flags & AGChangedValueFlagsChanged ? true : false;
+    }
+    return value;
+}
+
 #pragma mark - Update
 
 void AGGraphSetUpdate(const void *update) {

Sources/ComputeCxx/Graph/Graph.cpp

@@ -40,7 +40,7 @@ Graph::Graph()
       _contexts_by_id(nullptr, nullptr, nullptr, nullptr, &_heap), _id(AGMakeUniqueID()) {
 
     static platform_once_t make_keys;
-          platform_once(&make_keys, []() {
+    platform_once(&make_keys, []() {
         pthread_key_create(&Graph::_current_update_key, 0);
         Subgraph::make_current_subgraph_key();
     });
@@ -188,6 +188,12 @@ void Graph::remove_subgraph(Subgraph &subgraph) {
         }
     }
 
+    if (subgraph.has_cached_nodes()) {
+        subgraph.set_has_cached_nodes(false);
+        auto iter = std::remove(_subgraphs_with_cached_nodes.begin(), _subgraphs_with_cached_nodes.end(), &subgraph);
+        _subgraphs_with_cached_nodes.erase(iter);
+    }
+
     _num_subgraphs -= 1;
 }
 
@@ -208,6 +214,19 @@ void Graph::invalidate_subgraphs() {
     }
 
     if (_main_handler == nullptr) {
+        auto iter = _subgraphs_with_cached_nodes.begin(), end = _subgraphs_with_cached_nodes.end();
+        while (iter != end) {
+            auto subgraph = *iter;
+            subgraph->set_graph_invalidating_subgraphs(true);
+            subgraph->cache_collect();
+            subgraph->set_graph_invalidating_subgraphs(false);
+
+            if (!subgraph->has_cached_nodes()) {
+                end = _subgraphs_with_cached_nodes.erase(iter);
+            } else {
+                ++iter;
+            }
+        }
         while (!_invalidating_subgraphs.empty()) {
             auto subgraph = _invalidating_subgraphs.back();
             _invalidating_subgraphs.pop_back();

Sources/ComputeCxx/Graph/Graph.h

@@ -105,6 +105,7 @@ class Graph {
 
     // Subgraphs
     vector<Subgraph *, 0, uint32_t> _subgraphs;
+    vector<Subgraph *, 0, uint32_t> _subgraphs_with_cached_nodes;
     vector<Subgraph *, 2, uint32_t> _invalidating_subgraphs;
     bool _deferring_subgraph_invalidation;
 
@@ -138,7 +139,6 @@ class Graph {
 
     void remove_input(data::ptr<Node> node, uint32_t index);
     void remove_input_edge(data::ptr<Node> node_ptr, Node &node, uint32_t index);
-    void remove_all_inputs(data::ptr<Node> node);
     void all_inputs_removed(data::ptr<Node> node);
 
     template <typename T> void add_output_edge(data::ptr<T> node, AttributeID output);
@@ -236,6 +236,8 @@ class Graph {
 
     void add_subgraph(Subgraph &subgraph);
     void remove_subgraph(Subgraph &subgraph);
+
+    void add_subgraphs_with_cached_nodes(Subgraph &subgraph) { _subgraphs_with_cached_nodes.push_back(&subgraph); };
 
     class without_invalidating {
       private:
@@ -292,6 +294,7 @@ class Graph {
     void remove_indirect_node(data::ptr<IndirectNode> node);
 
     uint32_t add_input(data::ptr<Node> node, AttributeID input, bool allow_nil, AGInputOptions options);
+    void remove_all_inputs(data::ptr<Node> node);
 
     void indirect_attribute_set(data::ptr<IndirectNode> indirect_node, AttributeID source);
     bool indirect_attribute_reset(data::ptr<IndirectNode> indirect_node, bool non_nil);

Sources/ComputeCxx/Subgraph/NodeCache.cpp

@@ -0,0 +1,52 @@
+#include "NodeCache.h"
+
+#include "Swift/SwiftShims.h"
+
+namespace AG {
+
+Subgraph::NodeCache::NodeCache() noexcept
+    : _heap(nullptr, 0, 0), _types(nullptr, nullptr, nullptr, nullptr, &_heap),
+      _items([](const ItemKey *item_key) -> uint64_t { return item_key->hash_and_age >> 8; },
+             [](const ItemKey *a, const ItemKey *b) -> bool {
+                 if (a == b) {
+                     return true;
+                 }
+
+                 if (a->type != b->type || (a->hash_and_age ^ b->hash_and_age) > 0xff) {
+                     return false;
+                 }
+
+                 const void *a_body = nullptr;
+                 if (a->node) {
+                     data::ptr<Node> a_node = a->node;
+                     auto a_attribute_type = AttributeID(a_node).subgraph()->graph()->attribute_type(a_node->type_id());
+                     a_body = a->node->get_self(a_attribute_type);
+                 } else {
+                     a_body = a->body;
+                 }
+
+                 const void *b_body = nullptr;
+                 if (b->node) {
+                     data::ptr<Node> b_node = b->node;
+                     auto b_attribute_type = AttributeID(b_node).subgraph()->graph()->attribute_type(b_node->type_id());
+                     b_body = b->node->get_self(b_attribute_type);
+                 } else {
+                     b_body = b->body;
+                 }
+
+                 return AGDispatchEquatable(a_body, b_body, a->type->type, a->type->equatable);
+             },
+             nullptr, nullptr, &_heap),
+      _items_by_node(nullptr, nullptr, nullptr, nullptr, &_heap) {}
+
+Subgraph::NodeCache::~NodeCache() noexcept {
+    _items_by_node.for_each(
+        [](data::ptr<Node> node, Item *item, void *context) {
+            if (item) {
+                delete item;
+            }
+        },
+        nullptr);
+}
+
+} // namespace AG