Integrate StateProvider with State

Summary:
Incorporates several changes into State when read tracking is enabled
- State will now read its value via the StateProvider
- State equality and hashCode also now ignore the fallback value

A number of changes might need to be made to prepare the codebase for these changes (especially with the identity of state), but these will happen in a separate diff stack

[RFC](https://docs.google.com/document/d/1aFOON5LBtuTA9xibwv5WGqPtZi5SU3dquIMso7bw7hY/edit?pli=1&tab=t.0)

Reviewed By: astreet

Differential Revision: D72053550

fbshipit-source-id: 2b6eb759180a07460b76f31b1825e6c7fc67f7ce

Author: kingsleyadio

litho-core/src/main/java/com/facebook/litho/ComponentContext.java

@@ -82,7 +82,7 @@ public class ComponentContext {
   private boolean mIsParentTreePropsCloned;
 
   @ThreadConfined(ThreadConfined.ANY)
-  private @Nullable LithoTree mLithoTree;
+  private @Nullable final LithoTree mLithoTree;
 
   private @Nullable LithoVisibilityEventsController mLithoVisibilityEventsController;
 

litho-core/src/main/java/com/facebook/litho/ComponentTree.java

@@ -85,6 +85,7 @@
 import com.facebook.litho.debug.LithoDebugEventAttributes;
 import com.facebook.litho.lifecycle.LifecycleOwnerWrapper;
 import com.facebook.litho.perfboost.LithoPerfBooster;
+import com.facebook.litho.state.TreeStateProvider;
 import com.facebook.litho.stats.LithoStats;
 import com.facebook.rendercore.LogLevel;
 import com.facebook.rendercore.MountContentPools;
@@ -119,6 +120,7 @@
 public class ComponentTree
     implements LithoVisibilityEventsListener,
         StateUpdater,
+        TreeStateProvider,
         MountedViewReference,
         ErrorComponentReceiver,
         LithoTreeLifecycleProvider {

litho-core/src/main/java/com/facebook/litho/KState.kt

@@ -18,6 +18,10 @@ package com.facebook.litho
 
 import com.facebook.litho.annotations.Hook
 import com.facebook.litho.state.ComponentState
+import com.facebook.litho.state.StateId
+import com.facebook.litho.state.StateProvider
+import com.facebook.litho.state.StateReadRecorder
+import java.util.Objects
 
 /**
  * Declares a state variable within a Component. The initializer will provide the initial value if
@@ -33,8 +37,7 @@ fun <T> ComponentScope.useState(initializer: () -> T): State<T> {
   val treeState: TreeState =
       resolveContext?.treeState
           ?: throw IllegalStateException("Cannot create state outside of layout calculation")
-  val stateUpdater: StateUpdater =
-      context.stateUpdater ?: throw IllegalStateException("LithoTree is null")
+  val lithoTree = context.lithoTree ?: error("LithoTree is null")
 
   val isNestedTreeContext = context.isNestedTreeContext
   val kState =
@@ -54,11 +57,13 @@ fun <T> ComponentScope.useState(initializer: () -> T): State<T> {
     context.scopedComponentInfo.state = state
 
     return State(
-        stateUpdater,
+        lithoTree.stateProvider,
+        lithoTree.stateUpdater,
         hookIndex,
         globalKey,
         isNestedTreeContext,
         context.componentScope,
+        lithoTree.isReadTrackingEnabled,
         state.value.states[hookIndex] as T)
   } else {
     context.scopedComponentInfo.state = kState
@@ -70,25 +75,38 @@ fun <T> ComponentScope.useState(initializer: () -> T): State<T> {
   }
 
   return State(
-      stateUpdater,
+      lithoTree.stateProvider,
+      lithoTree.stateUpdater,
       hookIndex,
       globalKey,
       isNestedTreeContext,
       context.componentScope,
+      lithoTree.isReadTrackingEnabled,
       kState.value.states[hookIndex] as T)
 }
 
 /** Interface with which a component gets the value from a state or updates it. */
 class State<T>
 internal constructor(
+    private val stateProvider: StateProvider,
     private val stateUpdater: StateUpdater,
     private val hookStateIndex: Int,
     private val globalKey: String,
-    private val isNestedTreeContext: Boolean,
+    internal val isNestedTreeContext: Boolean,
     private val componentScope: Component?,
-    val value: T
+    private val isReadTrackingEnabled: Boolean,
+    private val fallback: T
 ) {
 
+  internal val stateId: StateId = StateId(stateProvider.treeId, globalKey, hookStateIndex)
+
+  val value: T
+    get() {
+      if (!isReadTrackingEnabled) return fallback
+      StateReadRecorder.read(stateId)
+      return stateProvider.getValue(this)
+    }
+
   /**
    * Updates this state value and enqueues a new layout calculation reflecting it to execute in the
    * background.
@@ -215,10 +233,12 @@ internal constructor(
 
     return globalKey == other.globalKey &&
         hookStateIndex == other.hookStateIndex &&
-        value == other.value
+        if (isReadTrackingEnabled) stateId.treeId == other.stateId.treeId
+        else fallback == other.fallback
   }
 
   override fun hashCode(): Int {
-    return globalKey.hashCode() * 17 + (value?.hashCode() ?: 0) * 11 + hookStateIndex
+    return Objects.hash(
+        globalKey, hookStateIndex, if (isReadTrackingEnabled) stateId.treeId else fallback)
   }
 }

litho-core/src/main/java/com/facebook/litho/LithoTree.kt

@@ -17,18 +17,28 @@
 package com.facebook.litho
 
 import com.facebook.infer.annotation.ThreadConfined
+import com.facebook.litho.config.ComponentsConfiguration
+import com.facebook.litho.state.StateProvider
+import com.facebook.litho.state.StateProviderImpl
+import com.facebook.litho.state.TreeStateProvider
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.atomic.AtomicReference
 
 /** Represents a pointer to the Tree that a ComponentContext is attached to */
 class LithoTree
-constructor(
+internal constructor(
+    treeStateProvider: TreeStateProvider,
     @field:ThreadConfined(ThreadConfined.ANY) val stateUpdater: StateUpdater,
     @field:ThreadConfined(ThreadConfined.UI) val mountedViewReference: MountedViewReference,
     val errorComponentReceiver: ErrorComponentReceiver,
     val lithoTreeLifecycleProvider: LithoTreeLifecycleProvider,
     val id: Int
 ) {
+  @field:ThreadConfined(ThreadConfined.ANY)
+  val stateProvider: StateProvider = StateProviderImpl(id, treeStateProvider)
+
+  val isReadTrackingEnabled: Boolean =
+      ComponentsConfiguration.defaultInstance.enableStateReadTracking
 
   // Used to lazily store a CoroutineScope, if coroutine helper methods are used.
   @JvmField val internalScopeRef: AtomicReference<Any> = AtomicReference<Any>()
@@ -39,6 +49,7 @@ constructor(
 
     fun create(componentTree: ComponentTree, stateUpdater: StateUpdater): LithoTree =
         LithoTree(
+            componentTree,
             stateUpdater,
             componentTree,
             componentTree,

litho-core/src/main/java/com/facebook/litho/NestedLithoPrimitive.kt

@@ -180,16 +180,14 @@ fun createdNestedTreeComponentContext(
     errorComponent: (Component?) -> Unit,
     lifecycleProvider: NestedLithoTreeLifecycleProvider,
 ): ComponentContext {
+  val stateUpdater = NestedStateUpdater(state = updatedState, updater = updater)
   return ComponentContext(
       androidContext,
       treeProps,
       lithoConfiguration,
       LithoTree(
-          stateUpdater =
-              NestedStateUpdater(
-                  state = updatedState,
-                  updater = updater,
-              ),
+          treeStateProvider = stateUpdater,
+          stateUpdater = stateUpdater,
           mountedViewReference = rootHostReference,
           errorComponentReceiver = errorComponent,
           lithoTreeLifecycleProvider = lifecycleProvider,

litho-core/src/main/java/com/facebook/litho/NestedLithoTree.kt

@@ -20,6 +20,7 @@ import android.view.View
 import androidx.annotation.UiThread
 import com.facebook.kotlin.compilerplugins.dataclassgenerate.annotation.DataClassGenerate
 import com.facebook.kotlin.compilerplugins.dataclassgenerate.annotation.Mode
+import com.facebook.litho.state.TreeStateProvider
 import com.facebook.rendercore.SizeConstraints
 import com.facebook.rendercore.thread.utils.ThreadUtils.assertMainThread
 
@@ -154,14 +155,17 @@ object NestedLithoTree {
 class NestedStateUpdater(
     private val state: TreeState,
     private val updater: StateUpdateRequester,
-) : StateUpdater {
+) : StateUpdater, TreeStateProvider {
 
   override var isFirstMount: Boolean
     get() = state.mountInfo.isFirstMount
     set(value) {
       state.mountInfo.isFirstMount = value
     }
 
+  override val treeState: TreeState
+    get() = state
+
   override fun updateHookStateAsync(
       globalKey: String,
       updateBlock: HookUpdater,

litho-core/src/main/java/com/facebook/litho/PrimitiveComponentScope.kt

@@ -281,14 +281,15 @@ fun PrimitiveComponentScope.useNestedTree(
   }
 
   val errorComponentRef = useState { AtomicReference<Component?>(null) }
-
+  val stateUpdater = NestedStateUpdater(state = newState, updater = onStateUpdate)
   val componentContext =
       ComponentContext(
           androidContext,
           treeProps,
           lithoConfig,
           LithoTree(
-              stateUpdater = NestedStateUpdater(state = newState, updater = onStateUpdate),
+              treeStateProvider = stateUpdater,
+              stateUpdater = stateUpdater,
               mountedViewReference = nestedTreeState.mountedViewReference,
               errorComponentReceiver = { errorComponentRef.update(AtomicReference(it)) },
               lithoTreeLifecycleProvider = nestedTreeState.treeLifecycleProvider,

litho-core/src/main/java/com/facebook/litho/TreeState.kt

@@ -268,6 +268,16 @@ class TreeState {
     return false
   }
 
+  @Suppress("UNCHECKED_CAST")
+  internal fun <T> getHookStateValue(
+      globalKey: String,
+      hookStateIndex: Int,
+      isNestedTree: Boolean
+  ): T {
+    val stateContainer = getState(globalKey, isNestedTree)?.value as KStateContainer
+    return stateContainer.states[hookStateIndex] as T
+  }
+
   val pendingStateUpdateTransitions: List<Transition>
     get() {
       val updateStateTransitions: MutableList<Transition> = ArrayList()

litho-core/src/main/java/com/facebook/litho/state/StateProvider.kt

@@ -91,6 +91,11 @@ internal class StateProviderImpl(
   }
 
   override fun <T> getValue(state: State<T>): T {
-    error("Not yet implemented")
+    val stateId = state.stateId
+    require(stateId.treeId == treeId) {
+      "State tree (id=${stateId.treeId}) does not match StateProvider tree (id=$treeId)"
+    }
+    val source = currentSource.get() ?: treeStateProvider.treeState
+    return source.getHookStateValue(stateId.globalKey, stateId.index, state.isNestedTreeContext)
   }
 }