Track state reads in Draw phase

Summary:
With the necessary plugs and buttons now in place. Tracking state reads in Draw phase can finally happen.

All state reads in the Draw phase will be committed into the UiStateReadRecords (similar to binders). These shall be used later to optimize the drawing phase by efficiently determining which draw hosts need to be invalidated

Reviewed By: adityasharat

Differential Revision: D77378261

fbshipit-source-id: 8aa53e57f751869ca1e5ba156d66190dab62257f

Author: kingsleyadio

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

@@ -136,7 +136,7 @@ constructor(context: ComponentContext, attrs: AttributeSet? = null) :
 
   private var onDirtyMountListener: OnDirtyMountListener? = null
 
-  private val uiStateReadRecords = UiStateReadRecords(this)
+  @get:JvmName("getUiStateReadRecords") internal val uiStateReadRecords = UiStateReadRecords(this)
 
   private val mountState = MountState(this, systrace, uiStateReadRecords)
 

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

@@ -19,6 +19,7 @@
 package com.facebook.litho
 
 import com.facebook.litho.state.StateReadRecorder
+import com.facebook.litho.state.UiStateReadRecords
 
 /**
  * Wrapper around the Component-specific render implementation.
@@ -49,3 +50,9 @@ internal val ComponentContext.isReadTrackingEnabled: Boolean
   // Workaround for package-private calls directly in an inline function that seem
   // to generate bytecode that's incompatible with Java call-sites
   get() = lithoTree?.isReadTrackingEnabled == true
+
+internal val ComponentContext.uiStateReadRecords: UiStateReadRecords
+  get() =
+      checkNotNull(lithoTree?.uiStateReadRecordsProvider?.getUiStateReadRecords()) {
+        "Could not retrieve the UI state read records. This is likely because the LithoTree is not initialized."
+      }

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

@@ -79,6 +79,8 @@
 import com.facebook.litho.perfboost.LithoPerfBooster;
 import com.facebook.litho.state.StateId;
 import com.facebook.litho.state.TreeStateProvider;
+import com.facebook.litho.state.UiStateReadRecords;
+import com.facebook.litho.state.UiStateReadRecordsProvider;
 import com.facebook.litho.stats.LithoStats;
 import com.facebook.rendercore.LogLevel;
 import com.facebook.rendercore.MountContentPools;
@@ -114,6 +116,7 @@ public class ComponentTree
     implements LithoVisibilityEventsListener,
         StateUpdater,
         TreeStateProvider,
+        UiStateReadRecordsProvider,
         MountedViewReference,
         ErrorComponentReceiver,
         LithoTreeLifecycleProvider {
@@ -2580,6 +2583,15 @@ private void debugLog(String eventName, String info) {
     }
   }
 
+  @Override
+  public UiStateReadRecords getUiStateReadRecords() {
+    if (mLithoView == null) {
+      throw new IllegalStateException(
+          "Trying to get UI state read records without a set host view");
+    }
+    return mLithoView.getUiStateReadRecords();
+  }
+
   private static void bindHandlesToComponentTree(
       ComponentTree componentTree, LayoutState layoutState) {
     for (Handle handle : layoutState.getComponentHandles()) {

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

@@ -60,7 +60,13 @@ class LithoRenderer(
           stateUpdater = this,
           rootHost = this,
           errorComponentReceiver = this,
-          treeLifecycle = NoOpLifecycleProvider())
+          treeLifecycle = NoOpLifecycleProvider(),
+          uiStateReadRecordsProvider = {
+            val host =
+                rootHost as? BaseMountingView
+                    ?: error("Trying to get UI state read records without a set host")
+            host.uiStateReadRecords
+          })
   private val errorComponentRef = AtomicReference<Component?>(null)
 
   /** The root component that this renderer is responsible for rendering. */

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

@@ -20,6 +20,7 @@ import com.facebook.infer.annotation.ThreadConfined
 import com.facebook.litho.config.ComponentsConfiguration
 import com.facebook.litho.state.StateProviderImpl
 import com.facebook.litho.state.TreeStateProvider
+import com.facebook.litho.state.UiStateReadRecordsProvider
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.atomic.AtomicReference
 
@@ -37,7 +38,8 @@ internal constructor(
     @field:ThreadConfined(ThreadConfined.ANY) val stateUpdater: StateUpdater,
     @field:ThreadConfined(ThreadConfined.UI) val rootHost: RootHostProvider,
     val errorComponentReceiver: ErrorComponentReceiver,
-    val treeLifecycle: LithoRendererLifecycleProvider
+    val treeLifecycle: LithoRendererLifecycleProvider,
+    internal val uiStateReadRecordsProvider: UiStateReadRecordsProvider,
 ) {
 
   val isReadTrackingEnabled: Boolean =
@@ -61,7 +63,8 @@ internal constructor(
             stateUpdater = stateUpdater,
             rootHost = componentTree,
             errorComponentReceiver = componentTree,
-            treeLifecycle = componentTree)
+            treeLifecycle = componentTree,
+            uiStateReadRecordsProvider = componentTree)
 
     @JvmStatic
     fun generateComponentTreeId(): Int {

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

@@ -278,7 +278,14 @@ fun PrimitiveComponentScope.useNestedTree(
             stateUpdater = stateUpdater,
             rootHost = nestedTreeState.mountedViewReference,
             errorComponentReceiver = { errorComponentRef.update(AtomicReference(it)) },
-            treeLifecycle = nestedTreeState.treeLifecycleProvider)
+            treeLifecycle = nestedTreeState.treeLifecycleProvider,
+            uiStateReadRecordsProvider = {
+              val host = nestedTreeState.mountedViewReference.mountedView
+              check(host is BaseMountingView) {
+                "Trying to get UI state read records without a set host"
+              }
+              host.uiStateReadRecords
+            })
       }
   val componentContext =
       ComponentContext(

litho-core/src/main/java/com/facebook/litho/graphics/drawscope/DrawingStyles.kt

@@ -19,24 +19,36 @@ package com.facebook.litho.graphics.drawscope
 import androidx.core.graphics.withSave
 import com.facebook.litho.ComponentHost
 import com.facebook.litho.Style
+import com.facebook.litho.binders.ModelWithContext
 import com.facebook.litho.binders.onBindHostView
+import com.facebook.litho.uiStateReadRecords
 import com.facebook.rendercore.graphics.drawscope.CanvasDrawScope
 import com.facebook.rendercore.graphics.drawscope.DrawScope
 
 fun Style.drawBehind(block: DrawScope.() -> Unit): Style {
   return this +
       onBindHostView(Unit) { content ->
+        val model = binderModel as ModelWithContext
+        val uiStateReadsRecords = model.scopedContext.uiStateReadRecords
+        // Retrieve the binderId during binder execution so that we can use it later in Draw
+        val binderId = binderId
+
         content as ComponentHost
         content.drawBehind = { canvas ->
-          canvas.withSave {
-            CanvasDrawScope(
-                    // canvas = canvas,
-                    // context = draw context,
-                    // size = Size(w, h),
-                    )
-                .block()
+          uiStateReadsRecords.recordOnDraw(binderId.renderUnitId) {
+            canvas.withSave {
+              CanvasDrawScope(
+                      // canvas = canvas,
+                      // context = draw context,
+                      // size = Size(w, h),
+                      )
+                  .block()
+            }
           }
         }
-        onUnbind { content.drawBehind = null }
+        onUnbind {
+          content.drawBehind = null
+          uiStateReadsRecords.removeDrawScope(binderId.renderUnitId)
+        }
       }
 }

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

@@ -17,8 +17,10 @@
 package com.facebook.litho.state
 
 import androidx.annotation.MainThread
+import androidx.collection.MutableLongSet
 import androidx.collection.MutableScatterSet
 import androidx.collection.ScatterSet
+import androidx.collection.mutableLongSetOf
 import androidx.collection.mutableScatterMapOf
 import androidx.collection.mutableScatterSetOf
 import com.facebook.rendercore.BinderId
@@ -27,16 +29,19 @@ import com.facebook.rendercore.BinderObserver
 /**
  * A [BinderObserver] responsible for tracking and recording state reads in binders.
  *
- * The [UiStateReadRecords] keeps a live [records] of which state is read and by who. This may be
- * used to determine which binders need to be re-bound when a state that they read is updated, which
- * can help to significantly optimize binder execution in the mount phase.
+ * The [UiStateReadRecords] keeps a live record of which state is read and by who. These records are
+ * exclusively those that are read during the UI phases (namely Mount and Draw). These records may
+ * be used to determine which binders or draw hosts need to be re-bound or invalidated respectively,
+ * when a state that they read is updated, which can help to significantly optimize execution in the
+ * respective phases.
  *
  * @param config Supplies the read tracking status and current treeId to the observer.
  */
 @MainThread
 internal class UiStateReadRecords(private val config: Config) : BinderObserver() {
 
-  private val records = mutableScatterMapOf<StateId, MutableScatterSet<BinderId>>()
+  private val onBindRecords = mutableScatterMapOf<StateId, MutableScatterSet<BinderId>>()
+  private val onDrawRecords = mutableScatterMapOf<StateId, MutableLongSet>()
 
   override fun onBind(binderId: BinderId, func: () -> Unit) {
     if (!config.isReadTrackingEnabled) return func()
@@ -45,7 +50,7 @@ internal class UiStateReadRecords(private val config: Config) : BinderObserver()
     val reads = StateReadRecorder.record(treeId, func)
     synchronized(this) {
       reads.forEach { stateId ->
-        val binderIds = records.getOrPut(stateId) { mutableScatterSetOf() }
+        val binderIds = onBindRecords.getOrPut(stateId) { mutableScatterSetOf() }
         binderIds.add(binderId)
       }
     }
@@ -57,25 +62,49 @@ internal class UiStateReadRecords(private val config: Config) : BinderObserver()
 
     synchronized(this) {
       val stateIdsToRemove = mutableSetOf<StateId>()
-      records.forEach { stateId, binderIds ->
+      onBindRecords.forEach { stateId, binderIds ->
         val removed = binderIds.remove(binderId)
         if (removed && binderIds.isEmpty()) stateIdsToRemove.add(stateId)
       }
-      stateIdsToRemove.forEach { records.remove(it) }
+      stateIdsToRemove.forEach { onBindRecords.remove(it) }
+    }
+  }
+
+  fun recordOnDraw(hostId: Long, func: () -> Unit) {
+    if (!config.isReadTrackingEnabled) return func()
+
+    val treeId = config.currentTreeId()
+    val reads = StateReadRecorder.record(treeId, func)
+    reads.forEach { stateId ->
+      val hostIds = onDrawRecords.getOrPut(stateId) { mutableLongSetOf() }
+      hostIds.add(hostId)
+    }
+  }
+
+  fun removeDrawScope(hostId: Long) {
+    if (!config.isReadTrackingEnabled) return
+
+    synchronized(this) {
+      val stateIdsToRemove = mutableSetOf<StateId>()
+      onDrawRecords.forEach { stateId, hostIds ->
+        val removed = hostIds.remove(hostId)
+        if (removed && hostIds.isEmpty()) stateIdsToRemove.add(stateId)
+      }
+      stateIdsToRemove.forEach { onDrawRecords.remove(it) }
     }
   }
 
   fun hasBindersForState(dirtyStates: Set<StateId>): Boolean {
     return synchronized(this) {
-      dirtyStates.any { stateId -> records[stateId]?.isNotEmpty() == true }
+      dirtyStates.any { stateId -> onBindRecords[stateId]?.isNotEmpty() == true }
     }
   }
 
-  fun takeSnapshotForState(dirtyStates: Set<StateId>): ScatterSet<BinderId> {
+  fun takeBinderSnapshotForState(dirtyStates: Set<StateId>): ScatterSet<BinderId> {
     val result = mutableScatterSetOf<BinderId>()
     synchronized(this) {
       dirtyStates.forEach { stateId ->
-        val binderIds = records[stateId]
+        val binderIds = onBindRecords[stateId]
         if (binderIds != null) result.addAll(binderIds)
       }
     }

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

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.facebook.litho.state
+
+/**
+ * Internal abstraction used to provide access to the [UiStateReadRecords] directly from the
+ * [LithoTree].
+ */
+internal fun interface UiStateReadRecordsProvider {
+  fun getUiStateReadRecords(): UiStateReadRecords
+}