Ignore recomposition counts after a reset

Ignore the recomposition counts from requests made before the
last recomposition reset command was sent.

Without this the UI may show recomposition counts that make it
look like the reset only worked for very shortly.

Fixes: 223246122
Test: Added unit tests
Change-Id: I94baacbf5e2f8fba0e0e4c7bc66461e9f335bc78

6 files changed, 47 insertions, 22 deletions

diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/compose/ComposeLayoutInspectorClient.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/compose/ComposeLayoutInspectorClient.kt
index da4719902da..c513f3d4da0 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/compose/ComposeLayoutInspectorClient.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/compose/ComposeLayoutInspectorClient.kt
@@ -77,6 +77,17 @@ const val COMPOSE_INSPECTION_NOT_AVAILABLE = "Compose inspection is not availabl
 private const val PROGUARD_LEARN_MORE = "https://d.android.com/r/studio-ui/layout-inspector/code-shrinking"
 
 /**
+ * Result from [ComposeLayoutInspectorClient.getComposeables].
+ */
+class GetComposablesResult(
+  /** The response received from the agent */
+  val response: GetComposablesResponse,
+
+  /** This is true, if a recomposition count reset command was sent after the GetComposables command was sent. */
+  val pendingRecompositionCountReset: Boolean
+)
+
+/**
  * The client responsible for interacting with the compose layout inspector running on the target
  * device.
  *
@@ -150,9 +161,19 @@ class ComposeLayoutInspectorClient(
   }
 
   val parametersCache = ComposeParametersCache(this, model)
-  var lastGeneration = 0
 
-  suspend fun getComposeables(rootViewId: Long, newGeneration: Int): GetComposablesResponse {
+  /**
+   * The caller will supply a running (increasing) number, that can be used to coordinate the responses from
+   * varies commands.
+   */
+  private var lastGeneration = 0
+
+  /**
+   * The value of [lastGeneration] when the last recomposition reset command was sent.
+   */
+  private var lastGenerationReset = 0
+
+  suspend fun getComposeables(rootViewId: Long, newGeneration: Int): GetComposablesResult {
     lastGeneration = newGeneration
     launchMonitor.updateProgress(AttachErrorState.COMPOSE_REQUEST_SENT)
     val response = messenger.sendCommand {
@@ -162,7 +183,7 @@ class ComposeLayoutInspectorClient(
       }.build()
     }
     launchMonitor.updateProgress(AttachErrorState.COMPOSE_RESPONSE_RECEIVED)
-    return response.getComposablesResponse
+    return GetComposablesResult(response.getComposablesResponse, lastGenerationReset >= newGeneration)
   }
 
   suspend fun getParameters(rootViewId: Long, composableId: Long): GetParametersResponse {
@@ -210,6 +231,7 @@ class ComposeLayoutInspectorClient(
   }
 
   suspend fun updateSettings(): UpdateSettingsResponse {
+    lastGenerationReset = lastGeneration
     val response = messenger.sendCommand {
       updateSettingsCommand = UpdateSettingsCommand.newBuilder().apply {
         includeRecomposeCounts = treeSettings.showRecompositions
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/compose/ComposeViewNodeCreator.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/compose/ComposeViewNodeCreator.kt
index 5090fb265c2..0395e48e75d 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/compose/ComposeViewNodeCreator.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/compose/ComposeViewNodeCreator.kt
@@ -33,7 +33,9 @@ private val semanticsSupport = EnumSet.of(Capability.SUPPORTS_COMPOSE, Capabilit
  * Helper class which handles the logic of using data from a [LayoutInspectorComposeProtocol.GetComposablesResponse] in order to
  * create [ComposeViewNode]s.
  */
-class ComposeViewNodeCreator(response: GetComposablesResponse) {
+class ComposeViewNodeCreator(result: GetComposablesResult) {
+  private val response = result.response
+  private val pendingRecompositionCountReset = result.pendingRecompositionCountReset
   private val stringTable = StringTableImpl(response.stringsList)
   private val roots = response.rootsList.map { it.viewId to it.nodesList }.toMap()
   private var composeFlags = 0
@@ -100,8 +102,8 @@ class ComposeViewNodeCreator(response: GetComposablesResponse) {
       null,
       "",
       0,
-      recomposeCount,
-      recomposeSkips,
+      if (pendingRecompositionCountReset) 0 else recomposeCount,
+      if (pendingRecompositionCountReset) 0 else recomposeSkips,
       stringTable[filename],
       packageHash,
       offset,
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewInspectorTreeLoader.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewInspectorTreeLoader.kt
index ec2a2b9d6fe..9ef3acbc219 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewInspectorTreeLoader.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewInspectorTreeLoader.kt
@@ -19,12 +19,12 @@ import com.android.tools.idea.appinspection.inspector.api.process.ProcessDescrip
 import com.android.tools.idea.layoutinspector.model.AndroidWindow
 import com.android.tools.idea.layoutinspector.pipeline.InspectorClient
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.AppInspectionTreeLoader
+import com.android.tools.idea.layoutinspector.pipeline.appinspection.compose.GetComposablesResult
 import com.android.tools.idea.layoutinspector.resource.ResourceLookup
 import com.android.tools.idea.layoutinspector.skia.SkiaParser
 import com.google.wireless.android.sdk.stats.DynamicLayoutInspectorEvent.DynamicLayoutInspectorEventType
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.LowMemoryWatcher
-import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol
 import layoutinspector.view.inspection.LayoutInspectorViewProtocol
 
 /**
@@ -36,7 +36,7 @@ class ViewInspectorTreeLoader(
   private val viewEvent: LayoutInspectorViewProtocol.LayoutEvent,
   private val resourceLookup: ResourceLookup,
   private val process: ProcessDescriptor,
-  composeEvent: LayoutInspectorComposeProtocol.GetComposablesResponse?,
+  composeResult: GetComposablesResult?,
   private val logEvent: (DynamicLayoutInspectorEventType) -> Unit,
 ) {
   private var folderConfig = LayoutInspectorViewProtocol.Configuration.getDefaultInstance().convert(1)
@@ -44,7 +44,7 @@ class ViewInspectorTreeLoader(
   // if true, exit immediately and return null
   private var isInterrupted = false
 
-  private val viewNodeCreator = ViewNodeCreator(viewEvent, composeEvent)
+  private val viewNodeCreator = ViewNodeCreator(viewEvent, composeResult)
 
   val dynamicCapabilities: Set<InspectorClient.Capability>
     get() = viewNodeCreator.dynamicCapabilities
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewLayoutInspectorClient.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewLayoutInspectorClient.kt
index c707d0bf99e..0a56c82fdf1 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewLayoutInspectorClient.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewLayoutInspectorClient.kt
@@ -26,6 +26,7 @@ import com.android.tools.idea.layoutinspector.model.InspectorModel
 import com.android.tools.idea.layoutinspector.pipeline.ConnectionFailedException
 import com.android.tools.idea.layoutinspector.pipeline.InspectorClientLaunchMonitor
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.compose.ComposeLayoutInspectorClient
+import com.android.tools.idea.layoutinspector.pipeline.appinspection.compose.GetComposablesResult
 import com.android.tools.idea.layoutinspector.snapshots.APP_INSPECTION_SNAPSHOT_VERSION
 import com.android.tools.idea.layoutinspector.snapshots.SnapshotMetadata
 import com.android.tools.idea.layoutinspector.snapshots.saveAppInspectorSnapshot
@@ -47,7 +48,6 @@ import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.future.asCompletableFuture
 import kotlinx.coroutines.launch
 import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetAllParametersResponse
-import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetComposablesResponse
 import layoutinspector.snapshots.Metadata
 import layoutinspector.view.inspection.LayoutInspectorViewProtocol
 import layoutinspector.view.inspection.LayoutInspectorViewProtocol.CaptureSnapshotCommand
@@ -105,7 +105,7 @@ class ViewLayoutInspectorClient(
     val generation: Int,
     val rootIds: List<Long>,
     val viewEvent: LayoutEvent,
-    val composeEvent: GetComposablesResponse?
+    val composeEvent: GetComposablesResult?
   )
 
   companion object {
@@ -278,13 +278,14 @@ class ViewLayoutInspectorClient(
     propertiesCache.clearFor(layoutEvent.rootView.id)
     composeInspector?.parametersCache?.clearFor(layoutEvent.rootView.id)
 
-    val composablesResponse = composeInspector?.getComposeables(layoutEvent.rootView.id, generation)
+    val composablesResult = composeInspector?.getComposeables(layoutEvent.rootView.id, generation)
 
     val data = Data(
       generation,
       currRoots,
       layoutEvent,
-      composablesResponse)
+      composablesResult
+    )
     if (!isFetchingContinuously) {
       lastData[layoutEvent.rootView.id] = data
     }
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewNodeCreator.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewNodeCreator.kt
index 9ffa2c71cd5..ec5188367ab 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewNodeCreator.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/view/ViewNodeCreator.kt
@@ -20,7 +20,7 @@ import com.android.tools.idea.layoutinspector.model.ComposeViewNode
 import com.android.tools.idea.layoutinspector.model.ViewNode
 import com.android.tools.idea.layoutinspector.pipeline.InspectorClient.Capability
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.compose.ComposeViewNodeCreator
-import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol
+import com.android.tools.idea.layoutinspector.pipeline.appinspection.compose.GetComposablesResult
 import layoutinspector.view.inspection.LayoutInspectorViewProtocol
 
 private const val ANDROID_VIEWS_HANDLER = "androidx.compose.ui.platform.AndroidViewsHandler"
@@ -29,16 +29,16 @@ private const val ANDROID_VIEWS_HANDLER = "androidx.compose.ui.platform.AndroidV
  * Helper class which handles the logic of converting a [LayoutInspectorViewProtocol.LayoutEvent] into
  * its corresponding [ViewNode]s.
  *
- * @param composeEvent If passed in, and the current view being processed is an AndroidComposeView, this will
+ * @param composeResult If passed in, and the current view being processed is an AndroidComposeView, this will
  *   be used to generate [ComposeViewNode] children.
  */
 class ViewNodeCreator(
   layoutEvent: LayoutInspectorViewProtocol.LayoutEvent,
-  composeEvent: LayoutInspectorComposeProtocol.GetComposablesResponse?
+  composeResult: GetComposablesResult?
 ) {
   val strings: StringTable = StringTableImpl(layoutEvent.stringsList)
   private val rootView = layoutEvent.rootView
-  private val composeNodeCreator = composeEvent?.let { ComposeViewNodeCreator(it) }
+  private val composeNodeCreator = composeResult?.let { ComposeViewNodeCreator(it) }
 
   /**
    * The collected capabilities based on the loaded data.
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/snapshots/AppInspectionSnapshotSupport.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/snapshots/AppInspectionSnapshotSupport.kt
index 16a7260761e..0ccda1a5d8d 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/snapshots/AppInspectionSnapshotSupport.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/snapshots/AppInspectionSnapshotSupport.kt
@@ -21,6 +21,7 @@ import com.android.tools.idea.layoutinspector.pipeline.InspectorClient
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.AppInspectionPropertiesProvider
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.AppInspectionTreeLoader
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.compose.ComposeParametersCache
+import com.android.tools.idea.layoutinspector.pipeline.appinspection.compose.GetComposablesResult
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.view.DisconnectedViewPropertiesCache
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.view.ViewLayoutInspectorClient
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.view.convert
@@ -28,7 +29,6 @@ import com.android.tools.idea.layoutinspector.skia.SkiaParserImpl
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.util.io.write
 import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetAllParametersResponse
-import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetComposablesResponse
 import layoutinspector.snapshots.Metadata
 import layoutinspector.snapshots.Snapshot
 import layoutinspector.view.inspection.LayoutInspectorViewProtocol
@@ -75,8 +75,8 @@ class AppInspectionSnapshotLoader : SnapshotLoader {
         // should always be true
         if (windowInfo != null) {
           val composeInfo = allComposeInfo[windowInfo.layout.rootView.id]
-          val data = ViewLayoutInspectorClient.Data(0, rootIds, windowInfo.layout,
-                                                    composeInfo?.composables)
+          val composeResult = composeInfo?.let { GetComposablesResult(it.composables, false) }
+          val data = ViewLayoutInspectorClient.Data(0, rootIds, windowInfo.layout, composeResult)
 
           val treeLoader = AppInspectionTreeLoader(model.project, metrics::logEvent, SkiaParserImpl({}))
           val treeData = treeLoader.loadComponentTree(data, model.resourceLookup, processDescriptor) ?: throw Exception()
@@ -123,7 +123,7 @@ fun saveAppInspectorSnapshot(
 fun saveAppInspectorSnapshot(
   path: Path,
   data: LayoutInspectorViewProtocol.CaptureSnapshotResponse,
-  composeInfo: Map<Long, Pair<GetComposablesResponse?, GetAllParametersResponse>>,
+  composeInfo: Map<Long, Pair<GetComposablesResult?, GetAllParametersResponse>>,
   snapshotMetadata: SnapshotMetadata,
   foldInfo: InspectorModel.FoldInfo?
 ) {
@@ -134,7 +134,7 @@ fun saveAppInspectorSnapshot(
       val (composables, composeParameters) = composableAndParameters
       Snapshot.ComposeInfo.newBuilder().apply {
         this.viewId = viewId
-        this.composables = composables
+        this.composables = composables?.response
         this.composeParameters = composeParameters
       }.build()
     })