Add functionality to filter flame graphs by search query

src/app-state/index.ts

@@ -14,6 +14,9 @@ export const flattenRecursionAtom = new Atom<boolean>(false, 'flattenRecursion')
 export const searchIsActiveAtom = new Atom<boolean>(false, 'searchIsActive')
 export const searchQueryAtom = new Atom<string>('', 'searchQueryAtom')
 
+// True if the flamechart should only show the matching frames when searching.
+export const onlyMatchesAtom = new Atom<boolean>(false, 'onlyMatches')
+
 // Which top-level view should be displayed
 export const viewModeAtom = new Atom<ViewMode>(ViewMode.CHRONO_FLAME_CHART, 'viewMode')
 

src/lib/profile-search.ts

@@ -42,6 +42,7 @@ export class ProfileSearchResults {
   constructor(
     readonly profile: Profile,
     readonly searchQuery: string,
+    readonly onlyMatches: boolean,
   ) {}
 
   private matches: Map<Frame, [number, number][] | null> | null = null

src/lib/profile.ts

@@ -1,6 +1,7 @@
 import {lastOf, KeyedSet} from './utils'
 import {ValueFormatter, RawValueFormatter} from './value-formatters'
 import {FileFormat} from './file-format-spec'
+import {ProfileSearchResults} from './profile-search'
 
 export interface FrameInfo {
   key: string | number
@@ -196,6 +197,48 @@ export class Profile {
     visit(this.groupedCalltreeRoot)
   }
 
+  // Filter the calls of a traversal function (e.g. `forEachCall` or `forEachCallGrouped`)
+  // based on the results of a search.
+  filteredTraversal(
+    searchResults: ProfileSearchResults | null,
+    traversalFn: (
+      openFrame: (node: CallTreeNode, value: number) => void,
+      closeFrame: (node: CallTreeNode, value: number) => void,
+    ) => void,
+  ) {
+    if (searchResults == null || !searchResults.onlyMatches) {
+      return traversalFn
+    }
+    const traversal = (
+      openFrame: (node: CallTreeNode, value: number) => void,
+      closeFrame: (node: CallTreeNode, value: number) => void,
+    ) => {
+      const matchingFrames = new Set<Frame>()
+      this.forEachFrame(frame => {
+        if (searchResults.getMatchForFrame(frame)) {
+          matchingFrames.add(frame)
+        }
+      })
+      if (matchingFrames.size === 0) {
+        return
+      }
+      // Use the original traversal function but filter callbacks
+      traversalFn(
+        (node, value) => {
+          if (matchingFrames.has(node.frame)) {
+            openFrame(node, value)
+          }
+        },
+        (node, value) => {
+          if (matchingFrames.has(node.frame)) {
+            closeFrame(node, value)
+          }
+        },
+      )
+    }
+    return traversal
+  }
+
   forEachCallGrouped(
     openFrame: (node: CallTreeNode, value: number) => void,
     closeFrame: (node: CallTreeNode, value: number) => void,

src/views/flamechart-search-view.tsx

@@ -10,6 +10,8 @@ import {Rect, Vec2} from '../lib/math'
 import {h, createContext, ComponentChildren} from 'preact'
 import {Flamechart} from '../lib/flamechart'
 import {CallTreeNode} from '../lib/profile'
+import {useAtom} from '../lib/atom'
+import {onlyMatchesAtom} from '../app-state'
 
 export const FlamechartSearchContext = createContext<FlamechartSearchData | null>(null)
 
@@ -65,6 +67,7 @@ export const FlamechartSearchContextProvider = ({
 
 export const FlamechartSearchView = memo(() => {
   const flamechartData = useContext(FlamechartSearchContext)
+  const onlyMatches = useAtom(onlyMatchesAtom)
 
   // TODO(jlfwong): This pattern is pretty gross, but I really don't want values
   // that can be undefined or null.
@@ -84,6 +87,10 @@ export const FlamechartSearchView = memo(() => {
     return searchResults.indexOf(selectedNode)
   }, [searchResults, selectedNode])
 
+  const toggleOnlyMatches = useCallback(() => {
+    onlyMatchesAtom.set(!onlyMatches)
+  }, [onlyMatches])
+
   const selectAndZoomToMatch = useCallback(
     (match: FlamechartSearchMatch) => {
       if (!setSelectedNode) return
@@ -146,6 +153,7 @@ export const FlamechartSearchView = memo(() => {
       numResults={numResults}
       selectPrev={selectPrev}
       selectNext={selectNext}
+      toggleOnlyMatches={toggleOnlyMatches}
     />
   )
 })

src/views/flamechart-view-container.tsx

@@ -13,12 +13,14 @@ import {
   getFrameToColorBucket,
 } from '../app-state/getters'
 import {Vec2, Rect} from '../lib/math'
-import {memo, useCallback} from 'preact/compat'
+import {memo, useCallback, useContext} from 'preact/compat'
 import {ActiveProfileState} from '../app-state/active-profile-state'
 import {FlamechartSearchContextProvider} from './flamechart-search-view'
 import {Theme, useTheme} from './themes/theme'
 import {FlamechartID, FlamechartViewState} from '../app-state/profile-group'
 import {profileGroupAtom} from '../app-state'
+import {ProfileSearchContext} from './search-view'
+import {ProfileSearchResults} from '../lib/profile-search'
 
 interface FlamechartSetters {
   setLogicalSpaceViewportSize: (logicalSpaceViewportSize: Vec2) => void
@@ -70,15 +72,19 @@ export const getChronoViewFlamechart = memoizeByShallowEquality(
   ({
     profile,
     getColorBucketForFrame,
+    searchResults,
   }: {
     profile: Profile
     getColorBucketForFrame: (frame: Frame) => number
+    searchResults: ProfileSearchResults | null
   }): Flamechart => {
     return new Flamechart({
       getTotalWeight: profile.getTotalWeight.bind(profile),
-      forEachCall: profile.forEachCall.bind(profile),
       formatValue: profile.formatValue.bind(profile),
       getColorBucketForFrame,
+      forEachCall: profile
+        .filteredTraversal(searchResults, profile.forEachCall.bind(profile))
+        .bind(profile),
     })
   },
 )
@@ -115,13 +121,18 @@ export const ChronoFlamechartView = memo((props: FlamechartViewContainerProps) =
   const {profile, chronoViewState} = activeProfileState
 
   const theme = useTheme()
+  const profileSearchResults = useContext(ProfileSearchContext)
 
   const canvasContext = getCanvasContext({theme, canvas: glCanvas})
   const frameToColorBucket = getFrameToColorBucket(profile)
   const getColorBucketForFrame = createGetColorBucketForFrame(frameToColorBucket)
   const getCSSColorForFrame = createGetCSSColorForFrame({theme, frameToColorBucket})
 
-  const flamechart = getChronoViewFlamechart({profile, getColorBucketForFrame})
+  const flamechart = getChronoViewFlamechart({
+    profile,
+    getColorBucketForFrame,
+    searchResults: profileSearchResults,
+  })
   const flamechartRenderer = getChronoViewFlamechartRenderer({
     canvasContext,
     flamechart,
@@ -155,15 +166,19 @@ export const getLeftHeavyFlamechart = memoizeByShallowEquality(
   ({
     profile,
     getColorBucketForFrame,
+    searchResults,
   }: {
     profile: Profile
     getColorBucketForFrame: (frame: Frame) => number
+    searchResults: ProfileSearchResults | null
   }): Flamechart => {
     return new Flamechart({
       getTotalWeight: profile.getTotalNonIdleWeight.bind(profile),
-      forEachCall: profile.forEachCallGrouped.bind(profile),
       formatValue: profile.formatValue.bind(profile),
       getColorBucketForFrame,
+      forEachCall: profile
+        .filteredTraversal(searchResults, profile.forEachCallGrouped.bind(profile))
+        .bind(profile),
     })
   },
 )
@@ -176,7 +191,7 @@ export const LeftHeavyFlamechartView = memo((ownProps: FlamechartViewContainerPr
   const {profile, leftHeavyViewState} = activeProfileState
 
   const theme = useTheme()
-
+  const profileSearchResults = useContext(ProfileSearchContext)
   const canvasContext = getCanvasContext({theme, canvas: glCanvas})
   const frameToColorBucket = getFrameToColorBucket(profile)
   const getColorBucketForFrame = createGetColorBucketForFrame(frameToColorBucket)
@@ -185,6 +200,7 @@ export const LeftHeavyFlamechartView = memo((ownProps: FlamechartViewContainerPr
   const flamechart = getLeftHeavyFlamechart({
     profile,
     getColorBucketForFrame,
+    searchResults: profileSearchResults,
   })
   const flamechartRenderer = getLeftHeavyFlamechartRenderer({
     canvasContext,

src/views/sandwich-search-view.tsx

@@ -39,6 +39,7 @@ export const SandwichSearchView = memo(() => {
       numResults={numResults}
       selectPrev={selectPrev}
       selectNext={selectNext}
+      toggleOnlyMatches={null}
     />
   )
 })

src/views/search-view.tsx

@@ -7,7 +7,7 @@ import {ProfileSearchResults} from '../lib/profile-search'
 import {Profile} from '../lib/profile'
 import {useActiveProfileState} from '../app-state/active-profile-state'
 import {useTheme, withTheme} from './themes/theme'
-import {searchIsActiveAtom, searchQueryAtom} from '../app-state'
+import {searchIsActiveAtom, onlyMatchesAtom, searchQueryAtom} from '../app-state'
 import {useAtom} from '../lib/atom'
 
 function stopPropagation(ev: Event) {
@@ -21,13 +21,14 @@ export const ProfileSearchContextProvider = ({children}: {children: ComponentChi
   const profile: Profile | null = activeProfileState ? activeProfileState.profile : null
   const searchIsActive = useAtom(searchIsActiveAtom)
   const searchQuery = useAtom(searchQueryAtom)
+  const onlyMatches = useAtom(onlyMatchesAtom)
 
   const searchResults = useMemo(() => {
     if (!profile || !searchIsActive || searchQuery.length === 0) {
       return null
     }
-    return new ProfileSearchResults(profile, searchQuery)
-  }, [searchIsActive, searchQuery, profile])
+    return new ProfileSearchResults(profile, searchQuery, onlyMatches)
+  }, [searchIsActive, searchQuery, profile, onlyMatches])
 
   return (
     <ProfileSearchContext.Provider value={searchResults}>{children}</ProfileSearchContext.Provider>
@@ -39,10 +40,11 @@ interface SearchViewProps {
   numResults: number | null
   selectNext: () => void
   selectPrev: () => void
+  toggleOnlyMatches: (() => void) | null
 }
 
 export const SearchView = memo(
-  ({numResults, resultIndex, selectNext, selectPrev}: SearchViewProps) => {
+  ({numResults, resultIndex, selectNext, selectPrev, toggleOnlyMatches}: SearchViewProps) => {
     const theme = useTheme()
     const style = getStyle(theme)
     const searchIsActive = useAtom(searchIsActiveAtom)
@@ -157,6 +159,11 @@ export const SearchView = memo(
             <button className={css(style.icon, style.button)} onClick={selectNext}>
               ➡️
             </button>
+            {toggleOnlyMatches && (
+              <button className={css(style.icon, style.button)} onClick={toggleOnlyMatches}>
+                🎯
+              </button>
+            )}
           </Fragment>
         )}
         <svg