Add search box for available nodes in the specified frame

wenli-cai · wenli-cai · commit 325b23938650 · 2025-07-29T13:15:48.000-04:00
The goal is to have the viewport snap to the node when it's being searched for, but current use of Offset seems to give some positional errors.
diff --git a/workflow-trace-viewer/build.gradle.kts b/workflow-trace-viewer/build.gradle.kts
@@ -15,6 +15,7 @@ kotlin {
         implementation(compose.runtime)
         implementation(compose.foundation)
         implementation(compose.material)
+        implementation(compose.material3)
         implementation(compose.ui)
         implementation(compose.components.resources)
         implementation(compose.components.uiToolingPreview)
diff --git a/workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/App.kt b/workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/App.kt
@@ -1,23 +1,28 @@
 package com.squareup.workflow1.traceviewer
 
 import androidx.compose.foundation.layout.Box
+import androidx.compose.foundation.layout.Column
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.mutableFloatStateOf
 import androidx.compose.runtime.mutableIntStateOf
+import androidx.compose.runtime.mutableStateMapOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
 import androidx.compose.runtime.snapshotFlow
+import androidx.compose.runtime.snapshots.SnapshotStateMap
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.geometry.Offset
+import com.squareup.workflow1.traceviewer.model.Node
 import com.squareup.workflow1.traceviewer.model.NodeUpdate
 import com.squareup.workflow1.traceviewer.ui.ColorLegend
 import com.squareup.workflow1.traceviewer.ui.DisplayDevices
 import com.squareup.workflow1.traceviewer.ui.FrameSelectTab
 import com.squareup.workflow1.traceviewer.ui.RightInfoPanel
+import com.squareup.workflow1.traceviewer.ui.SearchBox
 import com.squareup.workflow1.traceviewer.ui.TraceModeToggleSwitch
 import com.squareup.workflow1.traceviewer.util.FileDump
 import com.squareup.workflow1.traceviewer.util.RenderTrace
@@ -37,11 +42,14 @@ internal fun App(
   var rawRenderPass by remember { mutableStateOf("") }
   var frameIndex by remember { mutableIntStateOf(0) }
   val sandboxState = remember { SandboxState() }
+  val nodeLocations = remember { mutableListOf<SnapshotStateMap<Node, Offset>>() }
 
   // Default to File mode, and can be toggled to be in Live mode.
   var active by remember { mutableStateOf(false) }
   var traceMode by remember { mutableStateOf<TraceMode>(TraceMode.File(null)) }
   var selectedTraceFile by remember { mutableStateOf<PlatformFile?>(null) }
+  // frameIndex is set to -1 when app is in Live Mode, so we increment it by one to avoid off-by-one errors
+  val frameInd = if (traceMode is TraceMode.Live) frameIndex + 1 else frameIndex
 
   LaunchedEffect(sandboxState) {
     snapshotFlow { frameIndex }.collect {
@@ -57,6 +65,8 @@ internal fun App(
       selectedNode = null
       frameIndex = 0
       frameSize = 0
+      active = false
+      nodeLocations.clear()
     }
 
     // Main content
@@ -76,16 +86,40 @@ internal fun App(
           onNodeSelect = { selectedNode = it },
           onNewFrame = { frameIndex += 1 },
           onNewData = { rawRenderPass += "$it," },
+          storeNodeLocation = { node, loc -> nodeLocations[frameInd] += (node to loc) }
         )
       }
     }
 
-    FrameSelectTab(
-      size = frameSize,
-      currentIndex = frameIndex,
-      onIndexChange = { frameIndex = it },
+    Column(
       modifier = Modifier.align(Alignment.TopCenter)
-    )
+    ) {
+      if (active) {
+        FrameSelectTab(
+          size = frameSize,
+          currentIndex = frameIndex,
+          onIndexChange = { frameIndex = it },
+        )
+
+        // Since we can jump from frame to frame, we fill in the map during each recomposition
+        if (nodeLocations.getOrNull(frameInd) == null) {
+          // frameSize has not been updated yet, so on the first frame, frameSize = nodeLocations.size = 0,
+          // and it will append a new map
+          while (nodeLocations.size <= frameSize) {
+            nodeLocations.add(mutableStateMapOf())
+          }
+        }
+
+        SearchBox(
+          nodes = nodeLocations[frameInd].keys.toList(),
+          onSearch = { name ->
+            val node = nodeLocations[frameInd].keys.firstOrNull { it.name == name }
+            sandboxState.offset = nodeLocations[frameInd][node] ?: sandboxState.offset
+          },
+          modifier = Modifier.align(Alignment.CenterHorizontally)
+        )
+      }
+    }
 
     TraceModeToggleSwitch(
       onToggle = {
diff --git a/workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/ui/SearchBox.kt b/workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/ui/SearchBox.kt
@@ -0,0 +1,78 @@
+package com.squareup.workflow1.traceviewer.ui
+
+import androidx.compose.foundation.clickable
+import androidx.compose.foundation.layout.Column
+import androidx.compose.material.Icon
+import androidx.compose.material.IconButton
+import androidx.compose.material.Text
+import androidx.compose.material.icons.Icons
+import androidx.compose.material.icons.filled.Close
+import androidx.compose.material3.DockedSearchBar
+import androidx.compose.material3.ExperimentalMaterial3Api
+import androidx.compose.material3.ListItem
+import androidx.compose.material3.SearchBarColors
+import androidx.compose.material3.SearchBarDefaults
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.setValue
+import androidx.compose.ui.Modifier
+import androidx.compose.ui.graphics.Color
+import com.squareup.workflow1.traceviewer.model.Node
+
+@OptIn(ExperimentalMaterial3Api::class)
+@Composable
+internal fun SearchBox(
+  nodes: List<Node>,
+  onSearch: (String) -> Unit,
+  modifier: Modifier = Modifier
+) {
+  var searchText by remember { mutableStateOf("") }
+  var expanded by remember { mutableStateOf(false) }
+
+  DockedSearchBar(
+    modifier = modifier,
+    inputField = {
+      SearchBarDefaults.InputField(
+        query = searchText,
+        onQueryChange = { searchText = it },
+        onSearch = {
+          expanded = false
+        },
+        expanded = expanded,
+        onExpandedChange = { expanded = it },
+        placeholder = { Text("search for a node...") },
+        trailingIcon = {
+          IconButton(
+            onClick = {
+              expanded = false
+            }
+          ) {
+            Icon(
+              imageVector = Icons.Default.Close,
+              contentDescription = "Clear search"
+            )
+          }
+        }
+      )
+    },
+    colors = SearchBarColors(Color.White, Color.Black),
+    expanded = expanded,
+    onExpandedChange = { expanded = it },
+  ) {
+    val relevantNodes = nodes.filter { it.name.contains(searchText, ignoreCase = true) }
+    Column {
+      relevantNodes.take(5).forEach { node ->
+        ListItem(
+          headlineContent = { Text(node.name) },
+          modifier = Modifier
+            .clickable {
+              onSearch(node.name)
+              expanded = false
+            }
+        )
+      }
+    }
+  }
+}
diff --git a/workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/ui/WorkflowTree.kt b/workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/ui/WorkflowTree.kt
@@ -17,10 +17,13 @@ import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.ExperimentalComposeUiApi
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.geometry.Offset
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.input.pointer.PointerEventType
 import androidx.compose.ui.input.pointer.isSecondaryPressed
 import androidx.compose.ui.input.pointer.onPointerEvent
+import androidx.compose.ui.layout.onGloballyPositioned
+import androidx.compose.ui.layout.positionInParent
 import androidx.compose.ui.unit.dp
 import androidx.compose.ui.unit.sp
 import com.squareup.workflow1.traceviewer.model.Node
@@ -40,6 +43,7 @@ internal fun DrawTree(
   affectedNodes: Set<Node>,
   expandedNodes: MutableMap<String, Boolean>,
   onNodeSelect: (NodeUpdate) -> Unit,
+  storeNodeLocation: (Node, Offset) -> Unit,
   modifier: Modifier = Modifier,
 ) {
   Column(
@@ -64,7 +68,8 @@ internal fun DrawTree(
       isAffected = isAffected,
       isExpanded = isExpanded,
       onNodeSelect = onNodeSelect,
-      onExpandToggle = { expandedNodes[node.id] = !expandedNodes[node.id]!! }
+      onExpandToggle = { expandedNodes[node.id] = !expandedNodes[node.id]!! },
+      storeNodeLocation = storeNodeLocation
     )
 
     if (isExpanded) {
@@ -77,15 +82,17 @@ internal fun DrawTree(
         previousFrameNode = previousFrameNode,
         affectedNodes = affectedNodes,
         expandedNodes = expandedNodes,
-        onNodeSelect = onNodeSelect
+        onNodeSelect = onNodeSelect,
+        storeNodeLocation = storeNodeLocation
       )
 
       AffectedChildrenGroup(
         children = affectedChildren,
         previousFrameNode = previousFrameNode,
         affectedNodes = affectedNodes,
         expandedNodes = expandedNodes,
-        onNodeSelect = onNodeSelect
+        onNodeSelect = onNodeSelect,
+        storeNodeLocation = storeNodeLocation
       )
     }
   }
@@ -106,7 +113,8 @@ private fun UnaffectedChildrenGroup(
   previousFrameNode: Node?,
   affectedNodes: Set<Node>,
   expandedNodes: MutableMap<String, Boolean>,
-  onNodeSelect: (NodeUpdate) -> Unit
+  onNodeSelect: (NodeUpdate) -> Unit,
+  storeNodeLocation: (Node, Offset) -> Unit
 ) {
   if (children.isEmpty()) return
 
@@ -152,7 +160,8 @@ private fun UnaffectedChildrenGroup(
           affectedNodes = affectedNodes,
           expandedNodes = expandedNodes,
           unaffected = true,
-          onNodeSelect = onNodeSelect
+          onNodeSelect = onNodeSelect,
+          storeNodeLocation = storeNodeLocation
         )
       }
     }
@@ -168,7 +177,8 @@ private fun AffectedChildrenGroup(
   previousFrameNode: Node?,
   affectedNodes: Set<Node>,
   expandedNodes: MutableMap<String, Boolean>,
-  onNodeSelect: (NodeUpdate) -> Unit
+  onNodeSelect: (NodeUpdate) -> Unit,
+  storeNodeLocation: (Node, Offset) -> Unit
 ) {
   if (children.isEmpty()) return
 
@@ -177,7 +187,8 @@ private fun AffectedChildrenGroup(
     previousFrameNode = previousFrameNode,
     affectedNodes = affectedNodes,
     expandedNodes = expandedNodes,
-    onNodeSelect = onNodeSelect
+    onNodeSelect = onNodeSelect,
+    storeNodeLocation = storeNodeLocation
   )
 }
 
@@ -193,8 +204,9 @@ private fun DrawChildrenInGroups(
   previousFrameNode: Node?,
   affectedNodes: Set<Node>,
   expandedNodes: MutableMap<String, Boolean>,
+  onNodeSelect: (NodeUpdate) -> Unit,
+  storeNodeLocation: (Node, Offset) -> Unit,
   unaffected: Boolean = false,
-  onNodeSelect: (NodeUpdate) -> Unit
 ) {
   // Split children into those with children (nested) and those without
   var (nestedChildren, simpleChildren) = children.partition { it.children.isNotEmpty() }
@@ -227,7 +239,8 @@ private fun DrawChildrenInGroups(
               previousFrameNode = previousFrameNode?.children?.get(childNode.id),
               affectedNodes = affectedNodes,
               expandedNodes = expandedNodes,
-              onNodeSelect = onNodeSelect
+              onNodeSelect = onNodeSelect,
+              storeNodeLocation = storeNodeLocation
             )
           }
         }
@@ -250,7 +263,8 @@ private fun DrawChildrenInGroups(
             previousFrameNode = previousFrameNode?.children?.get(childNode.id),
             affectedNodes = affectedNodes,
             expandedNodes = expandedNodes,
-            onNodeSelect = onNodeSelect
+            onNodeSelect = onNodeSelect,
+            storeNodeLocation = storeNodeLocation
           )
         }
       }
@@ -270,6 +284,7 @@ private fun DrawNode(
   isExpanded: Boolean,
   onNodeSelect: (NodeUpdate) -> Unit,
   onExpandToggle: (Node) -> Unit,
+  storeNodeLocation: (Node, Offset) -> Unit
 ) {
   val nodeUpdate = NodeUpdate.create(
     current = node,
@@ -289,6 +304,10 @@ private fun DrawNode(
         }
       }
       .padding(16.dp)
+      .onGloballyPositioned { coords ->
+        val offset = coords.positionInParent()
+        storeNodeLocation(node, offset)
+      }
   ) {
     Column(horizontalAlignment = Alignment.CenterHorizontally) {
       Row(
diff --git a/workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/util/TraceParser.kt b/workflow-trace-viewer/src/jvmMain/kotlin/com/squareup/workflow1/traceviewer/util/TraceParser.kt
@@ -10,6 +10,7 @@ import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.setValue
 import androidx.compose.ui.Modifier
+import androidx.compose.ui.geometry.Offset
 import com.squareup.moshi.JsonAdapter
 import com.squareup.workflow1.traceviewer.TraceMode
 import com.squareup.workflow1.traceviewer.model.Node
@@ -30,6 +31,7 @@ internal fun RenderTrace(
   onNodeSelect: (NodeUpdate) -> Unit,
   onNewFrame: () -> Unit,
   onNewData: (String) -> Unit,
+  storeNodeLocation: (Node, Offset) -> Unit,
   modifier: Modifier = Modifier
 ) {
   var isLoading by remember(traceSource) { mutableStateOf(true) }
@@ -110,6 +112,7 @@ internal fun RenderTrace(
       affectedNodes = affectedNodes[frameInd],
       expandedNodes = remember(frameInd) { mutableStateMapOf() },
       onNodeSelect = onNodeSelect,
+      storeNodeLocation = storeNodeLocation
     )
   }
 }
