Better Fabric ShadowNode memory safety

ReactCommon/react/renderer/components/text/BaseTextShadowNode.cpp

@@ -11,6 +11,7 @@
 #include <react/renderer/components/text/RawTextShadowNode.h>
 #include <react/renderer/components/text/TextProps.h>
 #include <react/renderer/components/text/TextShadowNode.h>
+#include <react/renderer/core/TraitCast.h>
 #include <react/renderer/mounting/ShadowView.h>
 
 namespace facebook::react {

ReactCommon/react/renderer/components/text/ParagraphShadowNode.cpp

@@ -13,6 +13,7 @@
 #include <react/renderer/attributedstring/AttributedStringBox.h>
 #include <react/renderer/components/view/ViewShadowNode.h>
 #include <react/renderer/components/view/conversions.h>
+#include <react/renderer/core/TraitCast.h>
 #include <react/renderer/graphics/rounding.h>
 #include <react/renderer/telemetry/TransactionTelemetry.h>
 
@@ -205,8 +206,8 @@ void ParagraphShadowNode::layout(LayoutContext layoutContext) {
     paragraphShadowNode =
         static_cast<ParagraphShadowNode *>(paragraphOwningShadowNode.get());
 
-    auto &layoutableShadowNode = const_cast<LayoutableShadowNode &>(
-        traitCast<LayoutableShadowNode const &>(*clonedShadowNode));
+    auto &layoutableShadowNode =
+        traitCast<LayoutableShadowNode &>(*clonedShadowNode);
 
     auto attachmentFrame = measurement.attachments[i].frame;
     auto attachmentSize = roundToPixel<&ceil>(

ReactCommon/react/renderer/components/text/RawTextShadowNode.h

@@ -34,29 +34,5 @@ class RawTextShadowNode : public ConcreteShadowNode<
   }
 };
 
-template <>
-inline RawTextShadowNode const &traitCast<RawTextShadowNode const &>(
-    ShadowNode const &shadowNode) {
-  bool castable =
-      shadowNode.getTraits().check(ShadowNodeTraits::Trait::RawText);
-  react_native_assert(castable);
-  (void)castable;
-  return static_cast<RawTextShadowNode const &>(shadowNode);
-}
-
-template <>
-inline RawTextShadowNode const *traitCast<RawTextShadowNode const *>(
-    ShadowNode const *shadowNode) {
-  if (!shadowNode) {
-    return nullptr;
-  }
-  bool castable =
-      shadowNode->getTraits().check(ShadowNodeTraits::Trait::RawText);
-  if (!castable) {
-    return nullptr;
-  }
-  return static_cast<RawTextShadowNode const *>(shadowNode);
-}
-
 } // namespace react
 } // namespace facebook

ReactCommon/react/renderer/components/text/TextShadowNode.h

@@ -58,27 +58,5 @@ class TextShadowNode : public ConcreteShadowNode<
 #endif
 };
 
-template <>
-inline TextShadowNode const &traitCast<TextShadowNode const &>(
-    ShadowNode const &shadowNode) {
-  bool castable = shadowNode.getTraits().check(ShadowNodeTraits::Trait::Text);
-  react_native_assert(castable);
-  (void)castable;
-  return static_cast<TextShadowNode const &>(shadowNode);
-}
-
-template <>
-inline TextShadowNode const *traitCast<TextShadowNode const *>(
-    ShadowNode const *shadowNode) {
-  if (!shadowNode) {
-    return nullptr;
-  }
-  bool castable = shadowNode->getTraits().check(ShadowNodeTraits::Trait::Text);
-  if (!castable) {
-    return nullptr;
-  }
-  return static_cast<TextShadowNode const *>(shadowNode);
-}
-
 } // namespace react
 } // namespace facebook

ReactCommon/react/renderer/components/view/YogaLayoutableShadowNode.cpp

@@ -13,6 +13,7 @@
 #include <react/renderer/components/view/conversions.h>
 #include <react/renderer/core/LayoutConstraints.h>
 #include <react/renderer/core/LayoutContext.h>
+#include <react/renderer/core/TraitCast.h>
 #include <react/renderer/debug/DebugStringConvertibleItem.h>
 #include <react/renderer/debug/SystraceSection.h>
 #include <yoga/Yoga.h>
@@ -205,7 +206,7 @@ void YogaLayoutableShadowNode::adoptYogaChild(size_t index) {
     // TODO: At this point, React has wrong reference to the node. (T138668036)
     auto clonedChildNode = childNode.clone({});
     auto &layoutableClonedChildNode =
-        traitCast<YogaLayoutableShadowNode const &>(*clonedChildNode);
+        traitCast<YogaLayoutableShadowNode &>(*clonedChildNode);
 
     // The owner must be nullptr for a newly cloned node.
     react_native_assert(
@@ -298,7 +299,7 @@ void YogaLayoutableShadowNode::updateYogaChildren() {
     adoptYogaChild(i);
 
     if (isClean) {
-      auto &oldYogaChildNode = *oldYogaChildren[i];
+      auto &oldYogaChildNode = *oldYogaChildren.at(i);
       auto &newYogaChildNode =
           traitCast<YogaLayoutableShadowNode const &>(*getChildren().at(i))
               .yogaNode_;
@@ -550,8 +551,7 @@ void YogaLayoutableShadowNode::layout(LayoutContext layoutContext) {
 
   auto contentFrame = Rect{};
   for (auto childYogaNode : yogaNode_.getChildren()) {
-    auto &childNode =
-        *static_cast<YogaLayoutableShadowNode *>(childYogaNode->getContext());
+    auto &childNode = shadowNodeFromContext(childYogaNode);
 
     // Verifying that the Yoga node belongs to the ShadowNode.
     react_native_assert(&childNode.yogaNode_ == childYogaNode);
@@ -635,17 +635,15 @@ YGNode *YogaLayoutableShadowNode::yogaNodeCloneCallbackConnector(
 
   // At this point it is guaranteed that all shadow nodes associated with yoga
   // nodes are `YogaLayoutableShadowNode` subclasses.
-  auto parentNode =
-      static_cast<YogaLayoutableShadowNode *>(parentYogaNode->getContext());
-  auto oldNode =
-      static_cast<YogaLayoutableShadowNode *>(oldYogaNode->getContext());
+  auto &parentNode = shadowNodeFromContext(parentYogaNode);
+  auto &oldNode = shadowNodeFromContext(oldYogaNode);
 
-  auto clonedNode = oldNode->clone(
+  auto clonedNode = oldNode.clone(
       {ShadowNodeFragment::propsPlaceholder(),
        ShadowNodeFragment::childrenPlaceholder(),
-       oldNode->getState()});
-  parentNode->replaceChild(*oldNode, clonedNode, childIndex);
-  return &static_cast<YogaLayoutableShadowNode &>(*clonedNode).yogaNode_;
+       oldNode.getState()});
+  parentNode.replaceChild(oldNode, clonedNode, childIndex);
+  return &traitCast<YogaLayoutableShadowNode &>(*clonedNode).yogaNode_;
 }
 
 YGSize YogaLayoutableShadowNode::yogaNodeMeasureCallbackConnector(
@@ -657,8 +655,7 @@ YGSize YogaLayoutableShadowNode::yogaNodeMeasureCallbackConnector(
   SystraceSection s(
       "YogaLayoutableShadowNode::yogaNodeMeasureCallbackConnector");
 
-  auto shadowNodeRawPtr =
-      static_cast<YogaLayoutableShadowNode *>(yogaNode->getContext());
+  auto &shadowNode = shadowNodeFromContext(yogaNode);
 
   auto minimumSize = Size{0, 0};
   auto maximumSize = Size{
@@ -689,13 +686,19 @@ YGSize YogaLayoutableShadowNode::yogaNodeMeasureCallbackConnector(
       break;
   }
 
-  auto size = shadowNodeRawPtr->measureContent(
+  auto size = shadowNode.measureContent(
       threadLocalLayoutContext, {minimumSize, maximumSize});
 
   return YGSize{
       yogaFloatFromFloat(size.width), yogaFloatFromFloat(size.height)};
 }
 
+YogaLayoutableShadowNode &YogaLayoutableShadowNode::shadowNodeFromContext(
+    YGNode *yogaNode) {
+  return traitCast<YogaLayoutableShadowNode &>(
+      *static_cast<ShadowNode *>(yogaNode->getContext()));
+}
+
 YGConfig &YogaLayoutableShadowNode::initializeYogaConfig(YGConfig &config) {
   config.setCloneNodeCallback(
       YogaLayoutableShadowNode::yogaNodeCloneCallbackConnector);

ReactCommon/react/renderer/components/view/YogaLayoutableShadowNode.h

@@ -139,6 +139,7 @@ class YogaLayoutableShadowNode : public LayoutableShadowNode {
       YGMeasureMode widthMode,
       float height,
       YGMeasureMode heightMode);
+  static YogaLayoutableShadowNode &shadowNodeFromContext(YGNode *yogaNode);
 
 #pragma mark - RTL Legacy Autoflip
 
@@ -190,29 +191,5 @@ class YogaLayoutableShadowNode : public LayoutableShadowNode {
   void ensureYogaChildrenLookFine() const;
 };
 
-template <>
-inline YogaLayoutableShadowNode const &
-traitCast<YogaLayoutableShadowNode const &>(ShadowNode const &shadowNode) {
-  bool castable =
-      shadowNode.getTraits().check(ShadowNodeTraits::Trait::YogaLayoutableKind);
-  react_native_assert(castable);
-  (void)castable;
-  return static_cast<YogaLayoutableShadowNode const &>(shadowNode);
-}
-
-template <>
-inline YogaLayoutableShadowNode const *
-traitCast<YogaLayoutableShadowNode const *>(ShadowNode const *shadowNode) {
-  if (!shadowNode) {
-    return nullptr;
-  }
-  bool castable = shadowNode->getTraits().check(
-      ShadowNodeTraits::Trait::YogaLayoutableKind);
-  if (!castable) {
-    return nullptr;
-  }
-  return static_cast<YogaLayoutableShadowNode const *>(shadowNode);
-}
-
 } // namespace react
 } // namespace facebook

ReactCommon/react/renderer/core/LayoutableShadowNode.cpp

@@ -11,6 +11,7 @@
 #include <react/renderer/core/LayoutContext.h>
 #include <react/renderer/core/LayoutMetrics.h>
 #include <react/renderer/core/ShadowNode.h>
+#include <react/renderer/core/TraitCast.h>
 #include <react/renderer/debug/DebugStringConvertibleItem.h>
 #include <react/renderer/graphics/conversions.h>
 

ReactCommon/react/renderer/core/LayoutableShadowNode.h

@@ -164,29 +164,5 @@ class LayoutableShadowNode : public ShadowNode {
   LayoutMetrics layoutMetrics_;
 };
 
-template <>
-inline LayoutableShadowNode const &traitCast<LayoutableShadowNode const &>(
-    ShadowNode const &shadowNode) {
-  bool castable =
-      shadowNode.getTraits().check(ShadowNodeTraits::Trait::LayoutableKind);
-  react_native_assert(castable);
-  (void)castable;
-  return static_cast<LayoutableShadowNode const &>(shadowNode);
-}
-
-template <>
-inline LayoutableShadowNode const *traitCast<LayoutableShadowNode const *>(
-    ShadowNode const *shadowNode) {
-  if (!shadowNode) {
-    return nullptr;
-  }
-  bool castable =
-      shadowNode->getTraits().check(ShadowNodeTraits::Trait::LayoutableKind);
-  if (!castable) {
-    return nullptr;
-  }
-  return static_cast<LayoutableShadowNode const *>(shadowNode);
-}
-
 } // namespace react
 } // namespace facebook

ReactCommon/react/renderer/core/ShadowNode.h

@@ -224,21 +224,5 @@ class ShadowNode : public Sealable, public DebugStringConvertible {
   ShadowNodeTraits traits_;
 };
 
-/*
- * Template declarations for future specializations in concrete classes.
- * `traitCast` checks for a trait that corresponds to the provided type and
- * performs `static_cast`. Practically, the behavior is identical to
- * `dynamic_cast` with very little runtime overhead.
- */
-template <typename ShadowNodeReferenceT>
-ShadowNodeReferenceT traitCast(ShadowNode const &shadowNode);
-
-template <typename ShadowNodePointerT>
-ShadowNodePointerT traitCast(ShadowNode const *shadowNode);
-
-template <typename ShadowNodePointerT>
-std::shared_ptr<ShadowNodePointerT const> traitCast(
-    std::shared_ptr<ShadowNode const> shadowNode);
-
 } // namespace react
 } // namespace facebook

ReactCommon/react/renderer/core/ShadowNodeTraits.h

@@ -92,17 +92,8 @@ class ShadowNodeTraits {
     // Temporary (?) to indicate MapBuffer support on Android
     AndroidMapBufferPropsSupported = 1 << 15,
 
-    // Reserved
-    ReservedTrait0 = 1 << 16,
-    ReservedTrait1 = 1 << 17,
-    ReservedTrait2 = 1 << 18,
-    ReservedTrait3 = 1 << 19,
-    ReservedTrait4 = 1 << 20,
-    ReservedTrait5 = 1 << 21,
-    ReservedTrait6 = 1 << 22,
-
-    // Unserved - alias these for local usage
-    UnreservedTrait1 = 1 << 23
+    // Inherits 'ArtBaseShadowNode' (previously built into RN)
+    Art = 1 << 16,
   };
 
   /*
@@ -120,6 +111,10 @@ class ShadowNodeTraits {
     return ShadowNodeTraits::Trait(traits_ & traits) == traits;
   }
 
+  inline Trait get() const {
+    return traits_;
+  }
+
  private:
   Trait traits_{Trait::None};
 };

ReactCommon/react/renderer/core/TraitCast.h

@@ -0,0 +1,90 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <glog/logging.h>
+#include <react/renderer/core/ShadowNode.h>
+#include <memory>
+
+namespace facebook::react::details {
+template <typename ShadowNodePointerT, typename ParamT>
+ShadowNodePointerT traitCastPointer(ParamT shadowNode) {
+  auto expectedTraits =
+      std::remove_cv_t<std::remove_pointer_t<ShadowNodePointerT>>::BaseTraits();
+  if (!shadowNode || !shadowNode->getTraits().check(expectedTraits.get())) {
+    return nullptr;
+  }
+
+  return static_cast<ShadowNodePointerT>(shadowNode);
+}
+
+template <typename ShadowNodeRefT, typename ParamT>
+ShadowNodeRefT traitCastRef(ParamT &&shadowNode) {
+  auto expectedTraits =
+      std::remove_cv_t<std::remove_reference_t<ShadowNodeRefT>>::BaseTraits();
+  if (!shadowNode.getTraits().check(expectedTraits.get())) {
+    LOG(FATAL) << "Invalid ShadowNode cast\n"
+               << "Expected traits: " << std::hex
+               << static_cast<int32_t>(expectedTraits.get()) << "\n"
+               << "Actual traits: " << std::hex
+               << static_cast<int32_t>(shadowNode.getTraits().get()) << "\n";
+  }
+
+  return static_cast<ShadowNodeRefT>(shadowNode);
+}
+
+template <typename ShadowNodeT, typename ParamT>
+std::shared_ptr<ShadowNodeT> traitCastShared(
+    const std::shared_ptr<ParamT> &shadowNode) {
+  auto expectedTraits = std::remove_cv_t<ShadowNodeT>::BaseTraits();
+  if (!shadowNode || !shadowNode->getTraits().check(expectedTraits.get())) {
+    return nullptr;
+  }
+
+  return std::static_pointer_cast<ShadowNodeT>(shadowNode);
+}
+} // namespace facebook::react::details
+
+namespace facebook::react {
+
+// Cast from one ShadowNode reference to another, terminating if the cast is
+// invalid.
+template <typename ShadowNodeReferenceT>
+ShadowNodeReferenceT traitCast(ShadowNode const &shadowNode) {
+  return details::traitCastRef<ShadowNodeReferenceT>(shadowNode);
+}
+template <typename ShadowNodeReferenceT>
+ShadowNodeReferenceT traitCast(ShadowNode &shadowNode) {
+  return details::traitCastRef<ShadowNodeReferenceT>(shadowNode);
+}
+
+// Cast from one ShadowNode pointer to another, returning nullptr if the cast is
+// invalid.
+template <typename ShadowNodePointerT>
+ShadowNodePointerT traitCast(ShadowNode const *shadowNode) {
+  return details::traitCastPointer<ShadowNodePointerT>(shadowNode);
+}
+template <typename ShadowNodePointerT>
+ShadowNodePointerT traitCast(ShadowNode *shadowNode) {
+  return details::traitCastPointer<ShadowNodePointerT>(shadowNode);
+}
+
+// Cast from one ShadowNode shared_pointer to another, returning nullptr if the
+// cast is invalid.
+template <typename ShadowNodeT, typename ParamT>
+std::shared_ptr<ShadowNodeT> traitCast(
+    const std::shared_ptr<ParamT> &shadowNode) {
+  return details::traitCastShared<ShadowNodeT>(shadowNode);
+}
+template <typename ShadowNodeT, typename ParamT>
+std::shared_ptr<ShadowNodeT const> traitCast(
+    const std::shared_ptr<ParamT const> &shadowNode) {
+  return details::traitCastShared<ShadowNodeT const>(shadowNode);
+}
+
+} // namespace facebook::react