react-native-unistyles 3.3.0 → 3.5.0

package/cxx/core/UnistylesRegistry.cpp

@@ -1,6 +1,7 @@
 #include "UnistylesRegistry.h"
 #include "UnistylesState.h"
 #include "Parser.h"
+#include "NativeProps.h"
 
 using namespace margelo::nitro::unistyles;
 using namespace facebook;
@@ -8,7 +9,13 @@ using namespace facebook::react;
 
 std::atomic<int> core::UnistylesRegistry::_nextStyleSheetTag{0};
 
-void core::UnistylesRegistry::registerTheme(jsi::Runtime& rt, std::string name, jsi::Value& theme) {
+void core::UnistylesRegistry::registerTheme(jsi::Runtime& rt, uint64_t generation, std::string name, jsi::Value& theme) {
+    std::lock_guard<std::mutex> lock(this->_ownershipMutex);
+
+    if (!this->isCurrent(generation)) {
+        throw std::runtime_error(helpers::RUNTIME_REPLACED_ERROR);
+    }
+
     auto& state = this->getState();
 
     state._jsThemes.emplace(name, std::move(theme));
@@ -45,7 +52,7 @@ void core::UnistylesRegistry::createState() {
     this->_state = std::make_unique<UnistylesState>();
 }
 
-void core::UnistylesRegistry::updateTheme(jsi::Runtime& rt, std::string& themeName, jsi::Function&& callback) {
+void core::UnistylesRegistry::updateTheme(jsi::Runtime& rt, uint64_t generation, std::string& themeName, jsi::Function&& callback) {
     auto& state = this->getState();
     auto it = state._jsThemes.find(themeName);
 
@@ -55,42 +62,103 @@ void core::UnistylesRegistry::updateTheme(jsi::Runtime& rt, std::string& themeNa
 
     helpers::assertThat(rt, result.isObject(), "Unistyles: Returned theme is not an object. Please check your updateTheme function.");
 
-    it->second = result.asObject(rt);
+    auto theme = result.asObject(rt);
+    std::lock_guard<std::mutex> lock(this->_ownershipMutex);
+
+    if (!this->isCurrent(generation)) {
+        throw std::runtime_error(helpers::RUNTIME_REPLACED_ERROR);
+    }
+
+    it->second = std::move(theme);
 }
 
 void core::UnistylesRegistry::linkShadowNodeWithUnistyle(
     jsi::Runtime& rt,
-    const ShadowNodeFamily* shadowNodeFamily,
+    uint64_t generation,
+    const std::shared_ptr<const ShadowNodeFamily>& shadowNodeFamily,
     std::vector<std::shared_ptr<UnistyleData>>& unistylesData,
     std::optional<folly::dynamic> initialScopedUpdate
 ) {
-    this->trafficController.withLock([this, &rt, &unistylesData, shadowNodeFamily, &initialScopedUpdate](){
+    // released after unlocking, so a family destructor never runs within the lock
+    ReleasedFamilies releasedFamilies;
+    std::lock_guard<std::mutex> lock(this->_ownershipMutex);
+
+    if (!this->isCurrent(generation)) {
+        return;
+    }
+
+    this->trafficController.withLock([this, &unistylesData, &shadowNodeFamily, &initialScopedUpdate, &releasedFamilies](){
+        auto family = shadowNodeFamily.get();
+        auto it = this->_shadowRegistry.find(family);
+
         // Clear suspension state if this family was previously suspended
-        if (_suspendedFamilies.erase(shadowNodeFamily) > 0) {
-            auto* mutableFamily = const_cast<ShadowNodeFamily*>(shadowNodeFamily);
-            mutableFamily->nativeProps_DEPRECATED.reset();
+        // native props stay, they hold values from theme changes that React doesn't know about
+        if (it != this->_shadowRegistry.end() && it->second.isSuspended) {
             // Clear old registry entries to prevent stale UnistyleData accumulation
-            this->_shadowRegistry.erase(shadowNodeFamily);
-            // Remove any stale traffic controller entry (e.g. from a theme change during suspension)
-            this->trafficController.removeShadowNode(shadowNodeFamily);
+            // and any stale traffic controller entry (e.g. from a theme change during suspension)
+            this->forgetFamilyUnsafe(family, releasedFamilies);
         }
 
-        std::for_each(unistylesData.begin(), unistylesData.end(), [this, shadowNodeFamily](std::shared_ptr<UnistyleData> unistyleData){
-            this->_shadowRegistry[shadowNodeFamily].emplace_back(unistyleData);
-        });
+        auto& linkedFamily = this->_shadowRegistry[family];
+
+        linkedFamily.family = shadowNodeFamily;
+        linkedFamily.unistyles.insert(linkedFamily.unistyles.end(), unistylesData.begin(), unistylesData.end());
 
         // Required for scoped themes to apply on initial mount
         if (initialScopedUpdate.has_value()) {
-            shadow::ShadowLeafUpdates updates;
+            this->trafficController.setUpdate(shadowNodeFamily, std::move(*initialScopedUpdate));
+        }
 
-            updates.emplace(shadowNodeFamily, std::move(*initialScopedUpdate));
-            this->trafficController.setUpdates(updates);
+        // families unmounted without unlink (eg. while frozen) are pinned by us only
+        // sweep them once the registry doubles, so it's amortized O(1) per link
+        if (this->_shadowRegistry.size() >= this->_sweepThreshold) {
+            this->sweepUnownedFamiliesUnsafe(releasedFamilies);
+            this->_sweepThreshold = std::max(MIN_SWEEP_THRESHOLD, this->_shadowRegistry.size() * 2);
         }
     });
 }
 
+bool core::UnistylesRegistry::isOwnedOnlyByUnistyles(const LinkedFamily& linkedFamily) noexcept {
+    // shadow nodes own their family, so if nothing but our pin owns it, the family is not part of any
+    // shadow tree anymore. React Native can't get it back, as it only keeps weak references beside shadow nodes
+    return linkedFamily.family.use_count() == 1;
+}
+
+void core::UnistylesRegistry::forgetFamilyUnsafe(const ShadowNodeFamily* shadowNodeFamily, ReleasedFamilies& releasedFamilies) {
+    auto it = this->_shadowRegistry.find(shadowNodeFamily);
+
+    if (it != this->_shadowRegistry.end()) {
+        releasedFamilies.emplace_back(std::move(it->second.family));
+        this->_shadowRegistry.erase(it);
+    }
+
+    if (auto pendingFamily = this->trafficController.removeShadowNode(shadowNodeFamily)) {
+        releasedFamilies.emplace_back(std::move(pendingFamily));
+    }
+}
+
+void core::UnistylesRegistry::sweepUnownedFamiliesUnsafe(ReleasedFamilies& releasedFamilies) {
+    std::vector<const ShadowNodeFamily*> unownedFamilies;
+
+    for (const auto& [family, linkedFamily] : this->_shadowRegistry) {
+        if (isOwnedOnlyByUnistyles(linkedFamily)) {
+            unownedFamilies.emplace_back(family);
+        }
+    }
+
+    for (const auto* family : unownedFamilies) {
+        this->forgetFamilyUnsafe(family, releasedFamilies);
+    }
+}
+
 void core::UnistylesRegistry::removeDuplicatedUnistyles(const ShadowNodeFamily *shadowNodeFamily, std::vector<core::Unistyle::Shared>& unistyles) {
-    auto targetFamilyUnistyles = this->_shadowRegistry[shadowNodeFamily];
+    auto it = this->_shadowRegistry.find(shadowNodeFamily);
+
+    if (it == this->_shadowRegistry.end()) {
+        return;
+    }
+
+    auto& targetFamilyUnistyles = it->second.unistyles;
 
     unistyles.erase(
         std::remove_if(
@@ -111,26 +179,36 @@ void core::UnistylesRegistry::removeDuplicatedUnistyles(const ShadowNodeFamily *
 }
 
 void core::UnistylesRegistry::unlinkShadowNodeWithUnistyles(const ShadowNodeFamily* shadowNodeFamily) {
-    this->trafficController.withLock([this, shadowNodeFamily](){
-        this->_shadowRegistry.erase(shadowNodeFamily);
-        this->_suspendedFamilies.erase(shadowNodeFamily);
-        this->trafficController.removeShadowNode(shadowNodeFamily);
+    ReleasedFamilies releasedFamilies;
+
+    this->trafficController.withLock([this, shadowNodeFamily, &releasedFamilies](){
+        this->forgetFamilyUnsafe(shadowNodeFamily, releasedFamilies);
     });
 }
 
 void core::UnistylesRegistry::suspendShadowNode(const ShadowNodeFamily* shadowNodeFamily) {
     this->trafficController.withLock([this, shadowNodeFamily](){
-        if (this->_shadowRegistry.contains(shadowNodeFamily)) {
-            this->_suspendedFamilies.insert(shadowNodeFamily);
+        auto it = this->_shadowRegistry.find(shadowNodeFamily);
+
+        if (it != this->_shadowRegistry.end()) {
+            it->second.isSuspended = true;
         }
     });
 }
 
 bool core::UnistylesRegistry::isSuspended(const ShadowNodeFamily* family) const noexcept {
-    return _suspendedFamilies.count(family) > 0;
+    auto it = this->_shadowRegistry.find(family);
+
+    return it != this->_shadowRegistry.end() && it->second.isSuspended;
 }
 
-std::shared_ptr<core::StyleSheet> core::UnistylesRegistry::addStyleSheet(jsi::Runtime& rt, core::StyleSheetType type, jsi::Object&& rawValue) {
+std::shared_ptr<core::StyleSheet> core::UnistylesRegistry::addStyleSheet(jsi::Runtime& rt, uint64_t generation, core::StyleSheetType type, jsi::Object&& rawValue) {
+    std::lock_guard<std::mutex> lock(this->_ownershipMutex);
+
+    if (!this->isCurrent(generation)) {
+        throw std::runtime_error(helpers::RUNTIME_REPLACED_ERROR);
+    }
+
     int tag = _nextStyleSheetTag.fetch_add(1);
 
     auto sheet = std::make_shared<core::StyleSheet>(tag, type, std::move(rawValue));
@@ -143,36 +221,56 @@ core::DependencyMap core::UnistylesRegistry::buildDependencyMap(std::vector<Unis
     core::DependencyMap dependencyMap;
 
     std::unordered_set<UnistyleDependency> uniqueDependencies(deps.begin(), deps.end());
+    ReleasedFamilies releasedFamilies;
+
+    // families stay pinned by the registry until the shadow tree update is committed
+    this->trafficController.withLock([this, &dependencyMap, &uniqueDependencies, &releasedFamilies](){
+        // we iterate whole registry anyway, so it's a free moment to drop families unmounted without unlink
+        this->sweepUnownedFamiliesUnsafe(releasedFamilies);
+
+        for (const auto& [family, linkedFamily] : this->_shadowRegistry) {
+            bool hasAnyOfDependencies = false;
+
+            // Check if any dependency matches
+            for (const auto& unistyleData : linkedFamily.unistyles) {
+                for (const auto& dep : unistyleData->unistyle->dependencies) {
+                    if (uniqueDependencies.count(dep)) {
+                        hasAnyOfDependencies = true;
+                        break;
+                    }
+                }
 
-    for (const auto& [family, unistyles] : this->_shadowRegistry) {
-        bool hasAnyOfDependencies = false;
-
-        // Check if any dependency matches
-        for (const auto& unistyleData : unistyles) {
-            for (const auto& dep : unistyleData->unistyle->dependencies) {
-                if (uniqueDependencies.count(dep)) {
-                    hasAnyOfDependencies = true;
+                if (hasAnyOfDependencies) {
                     break;
-                }
+                };
+            }
+
+            if (!hasAnyOfDependencies) {
+                continue;
             }
 
-            if (hasAnyOfDependencies) {
-                break;
-            };
+            dependencyMap[family].insert(
+                dependencyMap[family].end(),
+                linkedFamily.unistyles.begin(),
+                linkedFamily.unistyles.end()
+            );
         }
+    });
+
+    return dependencyMap;
+}
+
+void core::UnistylesRegistry::queueShadowLeafUpdatesUnsafe(shadow::ShadowLeafUpdates& updates) {
+    for (auto& [family, props] : updates) {
+        auto it = this->_shadowRegistry.find(family);
 
-        if (!hasAnyOfDependencies) {
+        // family was unlinked in the meantime, there is nothing to update
+        if (it == this->_shadowRegistry.end()) {
             continue;
         }
 
-        dependencyMap[family].insert(
-            dependencyMap[family].end(),
-            unistyles.begin(),
-            unistyles.end()
-        );
+        this->trafficController.setUpdate(it->second.family, std::move(props));
     }
-
-    return dependencyMap;
 }
 
 // called from proxied function only, we don't know host
@@ -185,15 +283,15 @@ void core::UnistylesRegistry::shadowLeafUpdateFromUnistyle(jsi::Runtime& rt, Uni
             ? std::make_optional(maybePressableId.asString(rt).utf8(rt))
             : std::nullopt;
 
-        for (const auto& [family, unistyles] : this->_shadowRegistry) {
-            for (const auto& unistyleData : unistyles) {
+        for (const auto& [family, linkedFamily] : this->_shadowRegistry) {
+            for (const auto& unistyleData : linkedFamily.unistyles) {
                 if (unistyleData->unistyle == unistyle) {
                     updates[family] = parser.parseStylesToShadowTreeStyles(rt, { unistyleData });
                 }
             }
         }
 
-        this->trafficController.setUpdates(updates);
+        this->queueShadowLeafUpdatesUnsafe(updates);
     });
 }
 
@@ -262,11 +360,85 @@ void core::UnistylesRegistry::setScopedTheme(std::optional<std::string> themeNam
     this->_scopedTheme = std::move(themeName);
 }
 
-void core::UnistylesRegistry::destroy() {
-    this->_state.reset();
-    this->_styleSheetRegistry.clear();
-    this->_shadowRegistry.clear();
-    this->_suspendedFamilies.clear();
+uint64_t core::UnistylesRegistry::takeOwnership() {
+    std::lock_guard<std::mutex> lock(this->_ownershipMutex);
+    auto previousGeneration = this->_generation.load();
+
+    if (previousGeneration != 0) {
+        this->_replacedStates[previousGeneration] = this->takeOwnedState();
+    }
+
+    this->_generation.store(++this->_lastGeneration);
+
+    return this->_lastGeneration;
+}
+
+void core::UnistylesRegistry::releaseOwnership(uint64_t generation) {
+    // called when module is invalidated, before React Native destroys the runtime
+    // declared before the lock, so jsi values and pins are released after unlocking
+    OwnedState releasedState;
+
+    {
+        std::lock_guard<std::mutex> lock(this->_ownershipMutex);
+
+        if (this->isCurrent(generation)) {
+            releasedState = this->takeOwnedState();
+            this->_generation.store(0);
+        } else if (auto it = this->_replacedStates.find(generation); it != this->_replacedStates.end()) {
+            releasedState = std::move(it->second);
+            this->_replacedStates.erase(it);
+        }
+    }
+
+    for (auto& [_, styleSheet] : releasedState.styleSheets) {
+        styleSheet->unistyles.clear();
+    }
+}
+
+core::UnistylesRegistry::OwnedState core::UnistylesRegistry::takeOwnedState() {
+    OwnedState ownedState;
+
+    this->trafficController.withLock([this, &ownedState](){
+        ownedState.shadowRegistry = std::exchange(this->_shadowRegistry, {});
+        ownedState.updates = this->trafficController.takeUpdates();
+        this->_sweepThreshold = MIN_SWEEP_THRESHOLD;
+    });
+
+    ownedState.styleSheets = std::exchange(this->_styleSheetRegistry, {});
+    ownedState.state = std::move(this->_state);
     this->_scopedTheme = std::nullopt;
+    this->_committedTags.clear();
     _nextStyleSheetTag.store(0);
+
+    return ownedState;
+}
+
+std::vector<core::LinkedFamilySnapshot> core::UnistylesRegistry::getLinkedFamiliesSnapshot() {
+    std::vector<LinkedFamilySnapshot> snapshot;
+
+    this->trafficController.withLock([this, &snapshot](){
+        snapshot.reserve(this->_shadowRegistry.size());
+
+        for (const auto& [_, linkedFamily] : this->_shadowRegistry) {
+            const bool isOwnedOnlyByUnistyles = this->isOwnedOnlyByUnistyles(linkedFamily);
+
+            snapshot.push_back({linkedFamily.family, linkedFamily.unistyles, linkedFamily.isSuspended, isOwnedOnlyByUnistyles});
+        }
+    });
+
+    return snapshot;
+}
+
+size_t core::UnistylesRegistry::getPendingUpdatesCount() {
+    return this->trafficController.withLock([this](){
+        return this->trafficController.getUpdatesCount();
+    });
+}
+
+void core::UnistylesRegistry::setCommittedTags(std::vector<Tag>&& tags) {
+    this->_committedTags = std::move(tags);
+}
+
+std::vector<Tag> core::UnistylesRegistry::takeCommittedTags() {
+    return std::exchange(this->_committedTags, {});
 }

package/cxx/core/UnistylesRegistry.h

@@ -2,6 +2,7 @@
 
 #include "set"
 #include <atomic>
+#include <mutex>
 #include <jsi/jsi.h>
 #include <folly/dynamic.h>
 #include <react/renderer/uimanager/UIManager.h>
@@ -23,6 +24,19 @@ using namespace facebook::react;
 
 using DependencyMap = std::unordered_map<const ShadowNodeFamily*, std::vector<std::shared_ptr<UnistyleData>>>;
 
+struct LinkedFamily {
+    std::shared_ptr<const ShadowNodeFamily> family;
+    std::vector<std::shared_ptr<UnistyleData>> unistyles;
+    bool isSuspended = false;
+};
+
+struct LinkedFamilySnapshot {
+    std::shared_ptr<const ShadowNodeFamily> family;
+    std::vector<std::shared_ptr<UnistyleData>> unistyles;
+    bool isSuspended = false;
+    bool isOwnedOnlyByUnistyles = false;
+};
+
 struct UnistylesRegistry: public StyleSheetRegistry {
     static UnistylesRegistry& get();
 
@@ -31,38 +45,81 @@ struct UnistylesRegistry: public StyleSheetRegistry {
 
     bool shouldUsePointsForBreakpoints = false;
 
-    void registerTheme(jsi::Runtime& rt, std::string name, jsi::Value& theme);
+    void registerTheme(jsi::Runtime& rt, uint64_t generation, std::string name, jsi::Value& theme);
     void registerBreakpoints(std::vector<std::pair<std::string, double>>& sortedBreakpoints);
     void setPrefersAdaptiveThemes(bool prefersAdaptiveThemes);
     void setInitialThemeName(std::string themeName);
-    void updateTheme(jsi::Runtime& rt, std::string& themeName, jsi::Function&& callback);
+    void updateTheme(jsi::Runtime& rt, uint64_t generation, std::string& themeName, jsi::Function&& callback);
 
     UnistylesState& getState();
     void createState();
     std::vector<std::shared_ptr<core::StyleSheet>> getStyleSheetsToRefresh(std::vector<UnistyleDependency>& unistylesDependencies);
-    void linkShadowNodeWithUnistyle(jsi::Runtime& rt, const ShadowNodeFamily*, std::vector<std::shared_ptr<UnistyleData>>& unistylesData, std::optional<folly::dynamic> initialScopedUpdate = std::nullopt);
+    void linkShadowNodeWithUnistyle(jsi::Runtime& rt, uint64_t generation, const std::shared_ptr<const ShadowNodeFamily>& shadowNodeFamily, std::vector<std::shared_ptr<UnistyleData>>& unistylesData, std::optional<folly::dynamic> initialScopedUpdate = std::nullopt);
     void unlinkShadowNodeWithUnistyles(const ShadowNodeFamily*);
     void suspendShadowNode(const ShadowNodeFamily*);
     bool isSuspended(const ShadowNodeFamily*) const noexcept;
-    std::shared_ptr<core::StyleSheet> addStyleSheet(jsi::Runtime& rt, core::StyleSheetType type, jsi::Object&& rawValue);
+    std::shared_ptr<core::StyleSheet> addStyleSheet(jsi::Runtime& rt, uint64_t generation, core::StyleSheetType type, jsi::Object&& rawValue);
     DependencyMap buildDependencyMap(std::vector<UnistyleDependency>& deps);
     void shadowLeafUpdateFromUnistyle(jsi::Runtime& rt, Unistyle::Shared unistyle, jsi::Value& maybePressableId);
+    // call it only within trafficController.withLock!
+    void queueShadowLeafUpdatesUnsafe(shadow::ShadowLeafUpdates& updates);
     shadow::ShadowTrafficController trafficController{};
     const std::optional<std::string> getScopedTheme();
     void removeDuplicatedUnistyles(const ShadowNodeFamily* shadowNodeFamily, std::vector<core::Unistyle::Shared>& unistyles);
     void setScopedTheme(std::optional<std::string> themeName);
     core::Unistyle::Shared getUnistyleById(std::string unistyleID);
-    void destroy();
+    std::vector<LinkedFamilySnapshot> getLinkedFamiliesSnapshot();
+    size_t getPendingUpdatesCount();
+    void setCommittedTags(std::vector<Tag>&& tags);
+    std::vector<Tag> takeCommittedTags();
+
+    // newest installed runtime owns the registry, generation 0 means no owner
+    // previous runtime's state is kept aside until it releases ownership, as it may still run JS
+    uint64_t takeOwnership();
+    void releaseOwnership(uint64_t generation);
+
+    inline bool isCurrent(uint64_t generation) const noexcept {
+        return generation != 0 && generation == this->_generation.load();
+    }
 
 private:
+    using ReleasedFamilies = std::vector<std::shared_ptr<const ShadowNodeFamily>>;
+
+    struct OwnedState {
+        std::unordered_map<const ShadowNodeFamily*, LinkedFamily> shadowRegistry;
+        shadow::PinnedShadowLeafUpdates updates;
+        std::unordered_map<int, std::shared_ptr<core::StyleSheet>> styleSheets;
+        std::unique_ptr<UnistylesState> state;
+    };
+
     UnistylesRegistry() = default;
 
+    // call it only within _ownershipMutex!
+    // destructs nothing, so it can run on the next runtime's thread
+    OwnedState takeOwnedState();
+
+    // call it only within trafficController.withLock!
+    // pins are moved to releasedFamilies, so they can be released after unlocking
+    void forgetFamilyUnsafe(const ShadowNodeFamily* shadowNodeFamily, ReleasedFamilies& releasedFamilies);
+    void sweepUnownedFamiliesUnsafe(ReleasedFamilies& releasedFamilies);
+
+    // React Native doesn't own the family anymore, it was unmounted without unlink
+    static bool isOwnedOnlyByUnistyles(const LinkedFamily& linkedFamily) noexcept;
+
+    static constexpr size_t MIN_SWEEP_THRESHOLD = 64;
+
+    // lock order is _ownershipMutex -> trafficController, never call JS within
+    std::mutex _ownershipMutex;
+    std::atomic<uint64_t> _generation{0};
+    uint64_t _lastGeneration = 0;
+    std::unordered_map<uint64_t, OwnedState> _replacedStates{};
     static std::atomic<int> _nextStyleSheetTag;
     std::optional<std::string> _scopedTheme{};
     std::unique_ptr<UnistylesState> _state{};
     std::unordered_map<int, std::shared_ptr<core::StyleSheet>> _styleSheetRegistry{};
-    std::unordered_map<const ShadowNodeFamily*, std::vector<std::shared_ptr<UnistyleData>>> _shadowRegistry{};
-    std::unordered_set<const ShadowNodeFamily*> _suspendedFamilies{};
+    std::unordered_map<const ShadowNodeFamily*, LinkedFamily> _shadowRegistry{};
+    size_t _sweepThreshold = MIN_SWEEP_THRESHOLD;
+    std::vector<Tag> _committedTags{};
 };
 
 inline UnistylesRegistry& UnistylesRegistry::get() {