Control the view of recomposition counts in the inspector

This change:
- Adds a StudioFlag for use of this feature
- Adds a tree view option for displaying recomposition counts
- Adds a TreeSettings for the viewing option
- Adds a capability in case the app is using an older compose API

Hiding the recompose counts will reset the recompose counts.

Bug: 172894679
Test: Unit tests added
Change-Id: Ic6687c33f464e01ae76250e2398114c446e8804b

9 files changed, 87 insertions, 8 deletions

diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/LayoutInspectorToolWindowFactory.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/LayoutInspectorToolWindowFactory.kt
index 7f3780b7ff2..2f721a132ac 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/LayoutInspectorToolWindowFactory.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/LayoutInspectorToolWindowFactory.kt
@@ -100,7 +100,8 @@ class LayoutInspectorToolWindowFactory : ToolWindowFactory {
         lateinit var launcher: InspectorClientLauncher
         val treeSettings = InspectorTreeSettings { launcher.activeClient }
         val stats = SessionStatistics(model, treeSettings)
-        launcher = InspectorClientLauncher.createDefaultLauncher(processes, model, LayoutInspectorMetrics(project, null, stats), workbench)
+        val metrics = LayoutInspectorMetrics(project, null, stats)
+        launcher = InspectorClientLauncher.createDefaultLauncher(processes, model, metrics, treeSettings, workbench)
         val layoutInspector = LayoutInspector(launcher, model, stats, treeSettings)
         val deviceViewPanel = DeviceViewPanel(processes, layoutInspector, viewSettings, workbench)
         DataManager.registerDataProvider(workbench, dataProviderForLayoutInspector(layoutInspector, deviceViewPanel))
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/model/InspectorModel.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/model/InspectorModel.kt
index 308701ee84d..3373aa4957a 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/model/InspectorModel.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/model/InspectorModel.kt
@@ -140,6 +140,13 @@ class InspectorModel(val project: Project) : ViewNodeAndResourceLookup {
   }
 
   /**
+   * In-place update of all nodes (no structural changes should be made).
+   */
+  fun updateAll(operation: (ViewNode) -> Unit) {
+    ViewNode.readAccess { root.flatten().forEach { operation(it) } }
+  }
+
+  /**
    * Update [root]'s bounds and children based on any updates to [windows]
    * Also adds a dark layer between windows if DIM_BEHIND is set.
    */
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/InspectorClient.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/InspectorClient.kt
index ce19845ba06..7fd8df33d85 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/InspectorClient.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/InspectorClient.kt
@@ -68,6 +68,11 @@ interface InspectorClient: Disposable {
      * Indicates that this client is able to inspect compose parts of the application.
      */
     SUPPORTS_COMPOSE,
+
+    /**
+     * Indicates that this client is able to inspect compose recomposition counts of the application.
+     */
+    SUPPORTS_COMPOSE_RECOMPOSITION_COUNTS,
   }
 
   /**
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/InspectorClientLauncher.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/InspectorClientLauncher.kt
index 6af6e854d18..4e4596924a3 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/InspectorClientLauncher.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/InspectorClientLauncher.kt
@@ -23,6 +23,7 @@ import com.android.tools.idea.layoutinspector.metrics.LayoutInspectorMetrics
 import com.android.tools.idea.layoutinspector.model.InspectorModel
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.AppInspectionInspectorClient
 import com.android.tools.idea.layoutinspector.pipeline.legacy.LegacyClient
+import com.android.tools.idea.layoutinspector.tree.TreeSettings
 import com.android.tools.idea.layoutinspector.ui.InspectorBannerService
 import com.google.common.annotations.VisibleForTesting
 import com.google.wireless.android.sdk.stats.DynamicLayoutInspectorEvent
@@ -64,6 +65,7 @@ class InspectorClientLauncher(
       processes: ProcessesModel,
       model: InspectorModel,
       metrics: LayoutInspectorMetrics,
+      treeSettings: TreeSettings,
       parentDisposable: Disposable
     ): InspectorClientLauncher {
       return InspectorClientLauncher(
@@ -71,7 +73,7 @@ class InspectorClientLauncher(
         listOf(
           { params ->
             if (params.process.device.apiLevel >= AndroidVersion.VersionCodes.Q) {
-              AppInspectionInspectorClient(params.process, params.isInstantlyAutoConnected, model, metrics, parentDisposable)
+              AppInspectionInspectorClient(params.process, params.isInstantlyAutoConnected, model, metrics, treeSettings, parentDisposable)
             }
             else {
               null
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/AppInspectionInspectorClient.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/AppInspectionInspectorClient.kt
index 5b3aeb35ec9..1f88ca058a9 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/AppInspectionInspectorClient.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/pipeline/appinspection/AppInspectionInspectorClient.kt
@@ -44,6 +44,7 @@ import com.android.tools.idea.layoutinspector.pipeline.appinspection.compose.Com
 import com.android.tools.idea.layoutinspector.pipeline.appinspection.view.ViewLayoutInspectorClient
 import com.android.tools.idea.layoutinspector.properties.PropertiesProvider
 import com.android.tools.idea.layoutinspector.skia.SkiaParserImpl
+import com.android.tools.idea.layoutinspector.tree.TreeSettings
 import com.android.tools.idea.layoutinspector.ui.InspectorBannerService
 import com.android.tools.idea.progress.StudioLoggerProgressIndicator
 import com.android.tools.idea.progress.StudioProgressRunner
@@ -92,6 +93,7 @@ class AppInspectionInspectorClient(
   isInstantlyAutoConnected: Boolean,
   private val model: InspectorModel,
   private val metrics: LayoutInspectorMetrics,
+  private val treeSettings: TreeSettings,
   parentDisposable: Disposable,
   @TestOnly private val apiServices: AppInspectionApiServices = AppInspectionDiscoveryService.instance.apiServices,
   @TestOnly private val sdkHandler: AndroidSdkHandler = AndroidSdks.getInstance().tryToChooseSdkHandler()
@@ -169,7 +171,7 @@ class AppInspectionInspectorClient(
       apiServices.attachToProcess(process, model.project.name)
       launchMonitor.updateProgress(DynamicLayoutInspectorErrorInfo.AttachErrorState.ATTACH_SUCCESS)
 
-      composeInspector = ComposeLayoutInspectorClient.launch(apiServices, process, model, launchMonitor)
+      composeInspector = ComposeLayoutInspectorClient.launch(apiServices, process, model, treeSettings, capabilities, launchMonitor)
       val viewIns = ViewLayoutInspectorClient.launch(apiServices, process, model, scope, composeInspector, ::fireError, ::fireTreeEvent,
                                                      launchMonitor)
       propertiesProvider = AppInspectionPropertiesProvider(viewIns.propertiesCache, composeInspector?.parametersCache, model)
@@ -259,6 +261,12 @@ class AppInspectionInspectorClient(
     capabilities.addAll(dynamicCapabilities)
   }
 
+  fun updateRecompositionCountSettings() {
+    scope.launch(loggingExceptionHandler) {
+      composeInspector?.updateSettings()
+    }
+  }
+
   @Slow
   override fun saveSnapshot(path: Path) {
     val startTime = System.currentTimeMillis()
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 986d6a744e0..27154951e68 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
@@ -29,12 +29,14 @@ import com.android.tools.idea.appinspection.inspector.api.launch.LaunchParameter
 import com.android.tools.idea.appinspection.inspector.api.process.ProcessDescriptor
 import com.android.tools.idea.flags.StudioFlags
 import com.android.tools.idea.layoutinspector.model.InspectorModel
+import com.android.tools.idea.layoutinspector.pipeline.InspectorClient
+import com.android.tools.idea.layoutinspector.pipeline.InspectorClient.Capability
 import com.android.tools.idea.layoutinspector.pipeline.InspectorClientLaunchMonitor
+import com.android.tools.idea.layoutinspector.tree.TreeSettings
 import com.android.tools.idea.layoutinspector.ui.InspectorBannerService
 import com.google.common.annotations.VisibleForTesting
 import com.google.wireless.android.sdk.stats.DynamicLayoutInspectorErrorInfo.AttachErrorState
 import kotlinx.coroutines.cancel
-import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol
 import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.Command
 import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetAllParametersCommand
 import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetAllParametersResponse
@@ -45,6 +47,9 @@ import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetPara
 import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetParametersCommand
 import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.GetParametersResponse
 import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.Response
+import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.UpdateSettingsCommand
+import layoutinspector.compose.inspection.LayoutInspectorComposeProtocol.UpdateSettingsResponse
+import java.util.EnumSet
 
 const val COMPOSE_LAYOUT_INSPECTOR_ID = "layoutinspector.compose.inspection"
 
@@ -68,10 +73,13 @@ private const val PROGUARD_LEARN_MORE = "https://d.android.com/r/studio-ui/layou
  * device.
  *
  * @param messenger The messenger that lets us communicate with the view inspector.
+ * @param capabilities Of the containing [InspectorClient]. Some capabilities may be added by this class.
  */
 class ComposeLayoutInspectorClient(
   model: InspectorModel,
+  private val treeSettings: TreeSettings,
   private val messenger: AppInspectorMessenger,
+  private val capabilities: EnumSet<Capability>,
   private val launchMonitor: InspectorClientLaunchMonitor
 ) {
 
@@ -84,6 +92,8 @@ class ComposeLayoutInspectorClient(
       apiServices: AppInspectionApiServices,
       process: ProcessDescriptor,
       model: InspectorModel,
+      treeSettings: TreeSettings,
+      capabilities: EnumSet<Capability>,
       launchMonitor: InspectorClientLaunchMonitor
     ): ComposeLayoutInspectorClient? {
       val jar = if (StudioFlags.APP_INSPECTION_USE_DEV_JAR.get()) {
@@ -109,7 +119,7 @@ class ComposeLayoutInspectorClient(
       val params = LaunchParameters(process, COMPOSE_LAYOUT_INSPECTOR_ID, jar, model.project.name, MINIMUM_COMPOSE_COORDINATE, force = true)
       return try {
         val messenger = apiServices.launchInspector(params)
-        ComposeLayoutInspectorClient(model, messenger, launchMonitor).apply { updateSettings() }
+        ComposeLayoutInspectorClient(model, treeSettings, messenger, capabilities, launchMonitor).apply { updateSettings() }
       }
       catch (ignored: AppInspectionVersionIncompatibleException) {
         InspectorBannerService.getInstance(model.project).setNotification(INCOMPATIBLE_LIBRARY_MESSAGE)
@@ -188,12 +198,15 @@ class ComposeLayoutInspectorClient(
     return response.getParameterDetailsResponse
   }
 
-  suspend fun updateSettings(): LayoutInspectorComposeProtocol.UpdateSettingsResponse {
+  suspend fun updateSettings(): UpdateSettingsResponse {
     val response = messenger.sendCommand {
-      updateSettingsCommand = LayoutInspectorComposeProtocol.UpdateSettingsCommand.newBuilder().apply {
-        includeRecomposeCounts = true
+      updateSettingsCommand = UpdateSettingsCommand.newBuilder().apply {
+        includeRecomposeCounts = treeSettings.showRecompositions
       }.build()
     }
+    if (response.hasUpdateSettingsResponse()) {
+      capabilities.add(Capability.SUPPORTS_COMPOSE_RECOMPOSITION_COUNTS)
+    }
     return response.updateSettingsResponse
   }
 
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/LayoutInspectorTreePanel.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/LayoutInspectorTreePanel.kt
index f1a5c867747..23284abe916 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/LayoutInspectorTreePanel.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/LayoutInspectorTreePanel.kt
@@ -70,6 +70,7 @@ class LayoutInspectorTreePanel(parentDisposable: Disposable) : ToolContent<Layou
   private val componentTreePanel: JComponent
   @VisibleForTesting
   val componentTreeModel: ComponentTreeModel
+  private val setColumnVisibility: (columnIndex: Int, visible: Boolean) -> Unit
   private val nodeType = InspectorViewNodeType()
   // synthetic node to hold the root of the tree.
   private var root: TreeViewNode = ViewNode("root").treeNode
@@ -115,6 +116,7 @@ class LayoutInspectorTreePanel(parentDisposable: Disposable) : ToolContent<Layou
     focusComponent = result.focusComponent
     componentTreeModel = result.model
     componentTreeSelectionModel = result.selectionModel
+    setColumnVisibility = result.setColumnVisibility
     ActionManager.getInstance()?.getAction(IdeActions.ACTION_GOTO_DECLARATION)?.shortcutSet
       ?.let { GotoDeclarationAction.registerCustomShortcutSet(it, componentTreePanel, parentDisposable) }
     componentTreeSelectionModel.addSelectionListener {
@@ -171,6 +173,9 @@ class LayoutInspectorTreePanel(parentDisposable: Disposable) : ToolContent<Layou
     GotoDeclarationAction.findNavigatable(model)?.navigate(true)
   }
 
+  fun showRecompositionColumn(show: Boolean) =
+    setColumnVisibility(1, show)
+
   // TODO: There probably can only be 1 layout inspector per project. Do we need to handle changes?
   override fun setToolContext(toolContext: LayoutInspector?) {
     layoutInspector?.layoutInspectorModel?.modificationListeners?.remove(modelModifiedListener)
@@ -181,6 +186,7 @@ class LayoutInspectorTreePanel(parentDisposable: Disposable) : ToolContent<Layou
     componentTreeModel.treeRoot = root
     layoutInspector?.layoutInspectorModel?.selectionListeners?.add(selectionChangedListener)
     layoutInspector?.layoutInspectorModel?.windows?.values?.forEach { modelModified(null, it, true) }
+    layoutInspector?.let { showRecompositionColumn(it.treeSettings.showRecompositions) }
   }
 
   override fun getAdditionalActions() = additionalActions
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/TreeSettings.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/TreeSettings.kt
index a098c20385f..c0e0a2c884b 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/TreeSettings.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/TreeSettings.kt
@@ -32,6 +32,9 @@ const val DEFAULT_COMPOSE_AS_CALLSTACK = true
 const val KEY_SUPPORT_LINES = "live.layout.inspector.tree.lines"
 const val DEFAULT_SUPPORT_LINES = true
 
+const val KEY_RECOMPOSITIONS = "live.layout.inspector.tree.recompositions"
+const val DEFAULT_RECOMPOSITIONS = false
+
 /**
  * Miscellaneous tree settings.
  */
@@ -41,6 +44,7 @@ interface TreeSettings {
   var composeAsCallstack: Boolean
   var highlightSemantics: Boolean
   var supportLines: Boolean
+  var showRecompositions: Boolean
 
   fun isInComponentTree(node: ViewNode): Boolean =
     !(hideSystemNodes && node.isSystemNode)
@@ -66,6 +70,13 @@ class InspectorTreeSettings(private val activeClient: () -> InspectorClient) : T
     get() = get(KEY_SUPPORT_LINES, DEFAULT_SUPPORT_LINES)
     set(value) = set(KEY_SUPPORT_LINES, value, DEFAULT_SUPPORT_LINES)
 
+  override var showRecompositions: Boolean
+    get() = hasCapability(Capability.SUPPORTS_COMPOSE_RECOMPOSITION_COUNTS) &&
+            get(KEY_RECOMPOSITIONS, DEFAULT_RECOMPOSITIONS) &&
+            StudioFlags.DYNAMIC_LAYOUT_INSPECTOR_ENABLE_RECOMPOSITION_COUNTS.get() &&
+            StudioFlags.USE_COMPONENT_TREE_TABLE.get()
+    set(value) = set(KEY_RECOMPOSITIONS, value, DEFAULT_RECOMPOSITIONS)
+
   @Suppress("SameParameterValue")
   private fun hasCapability(capability: Capability): Boolean {
     val client = activeClient()
@@ -91,4 +102,5 @@ class EditorTreeSettings(capabilities: Set<Capability>) : TreeSettings {
   override var composeAsCallstack: Boolean = DEFAULT_COMPOSE_AS_CALLSTACK
   override var highlightSemantics: Boolean = DEFAULT_HIGHLIGHT_SEMANTICS
   override var supportLines: Boolean = DEFAULT_SUPPORT_LINES
+  override var showRecompositions: Boolean = DEFAULT_RECOMPOSITIONS
 }
diff --git a/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/TreeSettingsActions.kt b/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/TreeSettingsActions.kt
index ed7e36ceed6..b1a871e5a58 100644
--- a/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/TreeSettingsActions.kt
+++ b/layout-inspector/src/com/android/tools/idea/layoutinspector/tree/TreeSettingsActions.kt
@@ -18,8 +18,10 @@ package com.android.tools.idea.layoutinspector.tree
 import com.android.tools.adtui.actions.DropDownAction
 import com.android.tools.idea.flags.StudioFlags
 import com.android.tools.idea.layoutinspector.LayoutInspector
+import com.android.tools.idea.layoutinspector.model.ComposeViewNode
 import com.android.tools.idea.layoutinspector.model.SelectionOrigin
 import com.android.tools.idea.layoutinspector.pipeline.InspectorClient.Capability
+import com.android.tools.idea.layoutinspector.pipeline.appinspection.AppInspectionInspectorClient
 import com.android.tools.idea.layoutinspector.ui.DEVICE_VIEW_MODEL_KEY
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.ToggleAction
@@ -37,6 +39,7 @@ object FilterGroupAction : DropDownAction("Filter", "View options for Component
     add(SystemNodeFilterAction)
     add(HighlightSemanticsAction)
     add(CallstackAction)
+    add(RecompositionCounts)
     add(SupportLines)
   }
 }
@@ -123,6 +126,28 @@ object SupportLines : ToggleAction("Show Support Lines", null, null) {
   }
 }
 
+object RecompositionCounts : ToggleAction("Show Recomposition Counts", null, null) {
+
+  override fun isSelected(event: AnActionEvent): Boolean =
+    LayoutInspector.get(event)?.treeSettings?.showRecompositions ?: DEFAULT_RECOMPOSITIONS
+
+  override fun setSelected(event: AnActionEvent, state: Boolean) {
+    val inspector = LayoutInspector.get(event) ?: return
+    val client = inspector.currentClient as? AppInspectionInspectorClient ?: return
+    inspector.treeSettings.showRecompositions = state
+    client.updateRecompositionCountSettings()
+    event.treePanel()?.showRecompositionColumn(state)
+    inspector.layoutInspectorModel.updateAll { (it as? ComposeViewNode)?.recomposeCount = 0 }
+  }
+
+  override fun update(event: AnActionEvent) {
+    super.update(event)
+    event.presentation.isVisible = isActionVisible(event, Capability.SUPPORTS_COMPOSE_RECOMPOSITION_COUNTS) &&
+                                   StudioFlags.DYNAMIC_LAYOUT_INSPECTOR_ENABLE_RECOMPOSITION_COUNTS.get() &&
+                                   StudioFlags.USE_COMPONENT_TREE_TABLE.get()
+  }
+}
+
 private fun isActionVisible(event: AnActionEvent, vararg capabilities: Capability): Boolean =
   LayoutInspector.get(event)?.currentClient?.let { client ->
     !client.isConnected || // If not running, default to visible so user can modify selection when next client is connected