From cd4a4a84272da7f91f5cb9553a4f85ec82cb182d Mon Sep 17 00:00:00 2001
From: Sidharth Guglani <sidharthguglani@fb.com>
Date: Wed, 15 Sep 2021 09:53:35 -0700
Subject: [PATCH] Pass StateHandler to LayoutState.calculate

Reviewed By: adityasharat

Differential Revision: D30934733

fbshipit-source-id: e354ab1cee4a7a5817a4763f919632f2702aba23
---
 .../src/main/java/com/facebook/litho/ComponentTree.java      | 5 +++--
 litho-core/src/main/java/com/facebook/litho/LayoutState.java | 3 ++-
 .../litho/LayoutStateCalculateVisibilityOutputsTest.java     | 1 +
 .../src/test/java/com/facebook/litho/TreeDiffingTest.java    | 2 ++
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/litho-core/src/main/java/com/facebook/litho/ComponentTree.java b/litho-core/src/main/java/com/facebook/litho/ComponentTree.java
index 67af6382092..f5fda266a49 100644
--- a/litho-core/src/main/java/com/facebook/litho/ComponentTree.java
+++ b/litho-core/src/main/java/com/facebook/litho/ComponentTree.java
@@ -2949,9 +2949,9 @@ private LayoutState calculateLayoutStateInternal() {
       final ComponentContext contextWithStateHandler;
       final LayoutState previousLayoutState;
 
+      final StateHandler stateHandler;
       synchronized (ComponentTree.this) {
-        final StateHandler stateHandler =
-            StateHandler.createNewInstance(ComponentTree.this.mStateHandler);
+        stateHandler = StateHandler.createNewInstance(ComponentTree.this.mStateHandler);
 
         previousLayoutState = mCommittedLayoutState;
         contextWithStateHandler = new ComponentContext(context, stateHandler, treeProps, null);
@@ -2962,6 +2962,7 @@ private LayoutState calculateLayoutStateInternal() {
           contextWithStateHandler,
           root,
           layoutStateFuture,
+          stateHandler,
           ComponentTree.this.mId,
           widthSpec,
           heightSpec,
diff --git a/litho-core/src/main/java/com/facebook/litho/LayoutState.java b/litho-core/src/main/java/com/facebook/litho/LayoutState.java
index e28519d4bba..ab1fdf8d1d1 100644
--- a/litho-core/src/main/java/com/facebook/litho/LayoutState.java
+++ b/litho-core/src/main/java/com/facebook/litho/LayoutState.java
@@ -1517,6 +1517,7 @@ static LayoutState calculate(
         c,
         component,
         null,
+        new StateHandler(),
         componentTreeId,
         widthSpec,
         heightSpec,
@@ -1531,6 +1532,7 @@ static LayoutState calculate(
       ComponentContext c,
       Component component,
       @Nullable LayoutStateFuture layoutStateFuture,
+      StateHandler stateHandler,
       int componentTreeId,
       int widthSpec,
       int heightSpec,
@@ -1539,7 +1541,6 @@ static LayoutState calculate(
       @Nullable LayoutState currentLayoutState,
       @CalculateLayoutSource int source,
       @Nullable String extraAttribution) {
-    final StateHandler stateHandler = c.getStateHandler();
 
     final ComponentsLogger logger = c.getLogger();
 
diff --git a/litho-it/src/test/java/com/facebook/litho/LayoutStateCalculateVisibilityOutputsTest.java b/litho-it/src/test/java/com/facebook/litho/LayoutStateCalculateVisibilityOutputsTest.java
index 0bebbc9f4b0..c0a9a24fbe8 100644
--- a/litho-it/src/test/java/com/facebook/litho/LayoutStateCalculateVisibilityOutputsTest.java
+++ b/litho-it/src/test/java/com/facebook/litho/LayoutStateCalculateVisibilityOutputsTest.java
@@ -254,6 +254,7 @@ private static LayoutState calculateLayoutState(
         context,
         component,
         null,
+        new StateHandler(),
         componentTreeId,
         widthSpec,
         heightSpec,
diff --git a/litho-it/src/test/java/com/facebook/litho/TreeDiffingTest.java b/litho-it/src/test/java/com/facebook/litho/TreeDiffingTest.java
index b19e60c9cc7..e4525095da0 100644
--- a/litho-it/src/test/java/com/facebook/litho/TreeDiffingTest.java
+++ b/litho-it/src/test/java/com/facebook/litho/TreeDiffingTest.java
@@ -771,6 +771,7 @@ private static LayoutState calculateLayoutState(
         context,
         component,
         null,
+        new StateHandler(),
         -1,
         widthSpec,
         heightSpec,
@@ -791,6 +792,7 @@ private static LayoutState calculateLayoutStateWithDiffing(
         context,
         component,
         null,
+        new StateHandler(),
         -1,
         widthSpec,
         heightSpec,