add "return to sampler" button

src/viewer/sampler/components/Sampler.tsx

@@ -27,6 +27,7 @@ import AllView from './views/AllView';
 import FlatView from './views/FlatView';
 import SourcesView from './views/SourcesView';
 import { View, VIEW_ALL, VIEW_FLAT } from './views/types';
+import useExpanded from '../hooks/useExpanded';
 
 const Graph = dynamic(() => import('./graph/Graph'));
 
@@ -53,6 +54,7 @@ export default function Sampler({
 }: SamplerProps) {
     const searchQuery = useSearchQuery(data);
     const highlighted = useHighlight();
+    const expanded = useExpanded(searchQuery, highlighted);
     const [labelMode, setLabelMode] = useState(false);
     const timeSelector = useTimeSelector(
         data.timeWindows,
@@ -181,6 +183,16 @@ export default function Sampler({
                     mappings={mappings.mappingsResolver}
                     metadata={metadata}
                     timeSelector={timeSelector}
+                    onReturnToSampler={(node: VirtualNode) => {
+                        expanded.clearAll(); // fold all nodes
+                        // Expand selected node and its parents
+                        while (node != null) {
+                            expanded.set(node, true);
+                            node = node.getParents()[0];
+                        }
+                        // Return to sampler view
+                        setFlameData(undefined);
+                    }}
                 />
             )}
 
@@ -189,6 +201,7 @@ export default function Sampler({
                     mappings={mappings.mappingsResolver}
                     infoPoints={infoPoints}
                     highlighted={highlighted}
+                    expanded={expanded}
                     searchQuery={searchQuery}
                     labelMode={labelMode}
                     metadata={metadata}

src/viewer/sampler/components/SamplerContext.tsx

@@ -5,6 +5,7 @@ import { InfoPointsHook } from '../hooks/useInfoPoints';
 import { SearchQuery } from '../hooks/useSearchQuery';
 import { TimeSelector } from '../hooks/useTimeSelector';
 import { MappingsResolver } from '../mappings/resolver';
+import { Expanded } from '../hooks/useExpanded';
 
 export const MappingsContext = createContext<MappingsResolver | undefined>(
     undefined
@@ -25,11 +26,13 @@ export const LabelModeContext = createContext<boolean>(false);
 export const TimeSelectorContext = createContext<TimeSelector | undefined>(
     undefined
 );
+export const ExpandedContext = createContext<Expanded | undefined>(undefined);
 
 export default function SamplerContext({
     mappings,
     infoPoints,
     highlighted,
+    expanded,
     searchQuery,
     labelMode,
     metadata,
@@ -39,6 +42,7 @@ export default function SamplerContext({
     mappings: MappingsResolver;
     infoPoints: InfoPointsHook;
     highlighted: Highlight;
+    expanded: Expanded;
     searchQuery: SearchQuery;
     labelMode: boolean;
     metadata: SamplerMetadata;
@@ -51,15 +55,15 @@ export default function SamplerContext({
             <InfoPointsContext.Provider value={infoPoints}>
                 <HighlightedContext.Provider value={highlighted}>
                     <SearchQueryContext.Provider value={searchQuery}>
-                        <LabelModeContext.Provider value={labelMode}>
-                            <MetadataContext.Provider value={metadata}>
-                                <TimeSelectorContext.Provider
-                                    value={timeSelector}
-                                >
-                                    {children}
-                                </TimeSelectorContext.Provider>
-                            </MetadataContext.Provider>
-                        </LabelModeContext.Provider>
+                        <ExpandedContext.Provider value={expanded}>
+                            <LabelModeContext.Provider value={labelMode}>
+                                <MetadataContext.Provider value={metadata}>
+                                    <TimeSelectorContext.Provider value={timeSelector}>
+                                        {children}
+                                    </TimeSelectorContext.Provider>
+                                </MetadataContext.Provider>
+                            </LabelModeContext.Provider>
+                        </ExpandedContext.Provider>
                     </SearchQueryContext.Provider>
                 </HighlightedContext.Provider>
             </InfoPointsContext.Provider>

src/viewer/sampler/components/flamegraph/Flame.tsx

@@ -1,6 +1,6 @@
 // @ts-ignore
 import { FlameGraph } from '@lucko/react-flame-graph';
-import { useMemo } from 'react';
+import { useMemo, useRef, useState } from 'react';
 import { AutoSizer } from 'react-virtualized';
 import { formatBytesShort } from '../../../common/util/format';
 import {
@@ -10,21 +10,28 @@ import {
 import { TimeSelector } from '../../hooks/useTimeSelector';
 import { MappingsResolver } from '../../mappings/resolver';
 import VirtualNode from '../../node/VirtualNode';
+import { Menu, Item, ItemParams } from 'react-contexify';
+import { useContextMenu } from 'react-contexify';
+import 'react-contexify/dist/ReactContexify.css';
 
 export interface FlameProps {
     flameData: VirtualNode;
     mappings: MappingsResolver;
     metadata: SamplerMetadata;
     timeSelector: TimeSelector;
+    onReturnToSampler: (node: VirtualNode) => void;
 }
 
 export default function Flame({
     flameData,
     mappings,
     metadata,
     timeSelector,
+    onReturnToSampler,
 }: FlameProps) {
     const getTimeFunction = timeSelector.getTime;
+    const flameRef = useRef<HTMLDivElement>(null);
+    const { show } = useContextMenu({ id: 'flame-cm' });
 
     const isAlloc =
         metadata.samplerMode === SamplerMetadata_SamplerMode.ALLOCATION;
@@ -35,13 +42,45 @@ export default function Flame({
     );
     const calcHeight = Math.min(depth * 20, 5000);
 
+    const [selectedNode, setSelectedNode] = useState<VirtualNode | null>(null);
+
+    function handleDivContextMenu(event: React.MouseEvent) {
+        event.preventDefault();
+        if (selectedNode) {
+            show({ event, props: { node: selectedNode } });
+        }
+    }
+
+    function handleMouseOver(flameNode: any) {
+        if (flameNode && flameNode.virtualNode) {
+            setSelectedNode(flameNode.virtualNode);
+        } else {
+            setSelectedNode(null);
+        }
+    }
+
     return (
-        <div className="flame" style={{ height: `${calcHeight}px` }}>
+        <div
+            className="flame"
+            style={{ height: `${calcHeight}px` }}
+            ref={flameRef}
+            onContextMenu={handleDivContextMenu}
+        >
             <AutoSizer>
-                {({ width }) => (
-                    <FlameGraph data={data} height={calcHeight} width={width} />
+                {({ width }: { width: number }) => (
+                    <FlameGraph
+                        data={data}
+                        height={calcHeight}
+                        width={width}
+                        onMouseOver={(_e: any, node: any) => handleMouseOver(node)}
+                    />
                 )}
             </AutoSizer>
+            <Menu id="flame-cm" theme="dark">
+                <Item onClick={({ props }: ItemParams<{ node: VirtualNode }>) => props?.node && onReturnToSampler(props.node)}>
+                    View in sampler
+                </Item>
+            </Menu>
         </div>
     );
 }
@@ -51,6 +90,7 @@ interface FlameNode {
     tooltip?: string;
     value: number;
     children: FlameNode[];
+    virtualNode: VirtualNode;
 }
 
 function toFlameNode(
@@ -114,7 +154,7 @@ function toFlameNode(
         children.push(childData);
     }
 
-    return [{ name, tooltip, value, children }, depth];
+    return [{ name, tooltip, value, children, virtualNode: node }, depth];
 }
 
 function simplifyPackageName(

src/viewer/sampler/components/tree/BaseNode.tsx

@@ -1,9 +1,10 @@
 import classnames from 'classnames';
-import React, { useContext, useMemo, useState } from 'react';
+import React, { useContext, useMemo, useEffect } from 'react';
 import { useContextMenu } from 'react-contexify';
 import SourceThreadVirtualNode from '../../node/SourceThreadVirtualNode';
 import VirtualNode from '../../node/VirtualNode';
 import {
+    ExpandedContext,
     HighlightedContext,
     InfoPointsContext,
     MappingsContext,
@@ -31,27 +32,21 @@ const BaseNode = React.memo(({ parents, node, forcedTime }: BaseNodeProps) => {
     const highlighted = useContext(HighlightedContext)!;
     const searchQuery = useContext(SearchQueryContext)!;
     const timeSelector = useContext(TimeSelectorContext)!;
+    const expanded = useContext(ExpandedContext)!;
 
     const bottomUp = useContext(BottomUpContext) && parents.length !== 0;
 
     const directParent =
         parents.length !== 0 ? parents[parents.length - 1] : null;
 
-    const [expanded, setExpanded] = useState(() => {
-        if (highlighted.check(node)) {
-            return true;
+    // Get expanded state and check default expansion logic
+    const isExpanded = expanded.getOrDefault(node, directParent, bottomUp);
+    // schedule automatic fold on unmount
+    useEffect(() => {
+        return () => {
+            expanded.set(node, undefined);
         }
-        if (directParent == null) {
-            return false;
-        }
-
-        const nodes = bottomUp
-            ? directParent.getParents()
-            : directParent.getChildren();
-
-        const count = nodes.filter(n => searchQuery.matches(n)).length;
-        return count <= 1;
-    });
+    }, []);
 
     const parentsForChildren = useMemo(
         () => parents.concat([node]),
@@ -66,7 +61,7 @@ const BaseNode = React.memo(({ parents, node, forcedTime }: BaseNodeProps) => {
 
     const classNames = classnames({
         node: true,
-        collapsed: !expanded,
+        collapsed: !isExpanded,
         parent: parents.length === 0,
     });
     const nodeInfoClassNames = classnames({
@@ -82,7 +77,7 @@ const BaseNode = React.memo(({ parents, node, forcedTime }: BaseNodeProps) => {
         if (e.altKey) {
             highlighted.toggle(node);
         } else {
-            setExpanded(!expanded);
+            expanded.toggle(node);
         }
     }
 
@@ -150,7 +145,7 @@ const BaseNode = React.memo(({ parents, node, forcedTime }: BaseNodeProps) => {
                     <LineNumber node={node} parent={directParent} />
                 </NodeInfo>
             </div>
-            {expanded && (
+            {isExpanded && (
                 <ul className="children">
                     {(bottomUp ? node.getParents() : node.getChildren())
                         .sort(

src/viewer/sampler/hooks/useExpanded.ts

@@ -0,0 +1,75 @@
+import { useCallback, useState } from 'react';
+import VirtualNode from '../node/VirtualNode';
+import { SearchQuery } from './useSearchQuery';
+import { Highlight } from './useHighlight';
+
+export interface Expanded {
+    set: (node: VirtualNode, value: boolean | undefined) => void;
+    get: (node: VirtualNode) => boolean | undefined;
+    getOrDefault: (node: VirtualNode, directParent: VirtualNode | null, bottomUp: boolean) => boolean;
+    toggle: (node: VirtualNode) => void;
+    clearAll: () => void;
+}
+
+export default function useExpanded(searchQuery: SearchQuery, highlighted: Highlight): Expanded {
+    const [expanded, setExpanded] = useState<Map<string, boolean>>(() => new Map());
+
+    // Set expanded state for a node
+    const set: Expanded['set'] = useCallback(
+        (node, value) => {
+            setExpanded(prev => {
+                const map = new Map(prev);
+                if (value === undefined) {
+                    map.delete(String(node.getId()));
+                } else {
+                    map.set(String(node.getId()), value);
+                }
+                return map;
+            });
+        },
+        []
+    );
+
+    // Get expanded state for a node
+    const get: Expanded['get'] = useCallback(
+        node => {
+            return expanded.get(String(node.getId()));
+        },
+        [expanded]
+    );
+
+    // Get expanded state, or default (sometimes, nodes should be expanded by default)
+    const getOrDefault: Expanded['getOrDefault'] = (node, directParent, bottomUp) => {
+        // Try to get value from map first
+        const value = get(node);
+        if (value !== undefined) return value;
+
+        // auto-expand logic
+
+        if (highlighted.check(node)) return true;
+
+        if (directParent == null) return false;
+        let nodes;
+        if (bottomUp) nodes = directParent.getParents();
+        else nodes = directParent.getChildren();
+
+        nodes = nodes.filter(n => searchQuery.matches(n));
+        return nodes.length <= 1;
+    };
+
+    // Toggle expanded state for a node
+    const toggle: Expanded['toggle'] = useCallback(
+        node => {
+            set(node, !get(node));
+        },
+        [set, get]
+    );
+
+    // Clear all expanded nodes
+    const clearAll: Expanded['clearAll'] = useCallback(
+        () => setExpanded(new Map()),
+        []
+    );
+
+    return { set, get, getOrDefault, toggle, clearAll };
+}

src/viewer/sampler/node/VirtualNode.ts

@@ -1,5 +1,7 @@
 import { NodeDetails } from '../../proto/nodes';
 
+// represents the data of a node (frame)
+// One instance per ConcreteClass::method, meaning LivingEntity::tick() can have multiple virtual nodes, one per living entity
 export default interface VirtualNode {
     getId(): number | number[];