ui/packages/shared/profile: color the arrow based flame graph

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraphArrow/ColorStackLegend.tsx

@@ -22,17 +22,17 @@ import {EVERYTHING_ELSE, FeaturesMap} from '@parca/store';
 import type {NavigateFunction} from '@parca/utilities';
 
 interface Props {
-  mappingsFeatures: FeaturesMap;
+  mappingColors: FeaturesMap;
   navigateTo?: NavigateFunction;
   compareMode?: boolean;
 }
 
 const ColorStackLegend = ({
-  mappingsFeatures,
+  mappingColors,
   navigateTo,
   compareMode = false,
 }: Props): React.JSX.Element => {
-  if (mappingsFeatures === undefined) {
+  if (mappingColors === undefined) {
     return <></>;
   }
 
@@ -41,12 +41,12 @@ const ColorStackLegend = ({
   );
   const [currentSearchString, setSearchString] = useURLState({param: 'search_string', navigateTo});
 
-  if (Object.entries(mappingsFeatures).length === 0) {
+  if (Object.entries(mappingColors).length === 0) {
     return <></>;
   }
 
   const stackColorArray = useMemo(() => {
-    return Object.entries(mappingsFeatures).sort(([featureA], [featureB]) => {
+    return Object.entries(mappingColors).sort(([featureA], [featureB]) => {
       if (featureA === EVERYTHING_ELSE) {
         return 1;
       }
@@ -55,7 +55,7 @@ const ColorStackLegend = ({
       }
       return featureA?.localeCompare(featureB ?? '') ?? 0;
     });
-  }, [mappingsFeatures]);
+  }, [mappingColors]);
 
   if (colorProfileName === 'default' || compareMode) {
     return <></>;

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraphArrow/IcicleGraphNodes.tsx

@@ -17,17 +17,30 @@ import {Table} from 'apache-arrow';
 import cx from 'classnames';
 
 import {useKeyDown} from '@parca/components';
-import {selectBinaries, setHoveringRow, useAppDispatch, useAppSelector} from '@parca/store';
-import {isSearchMatch, scaleLinear} from '@parca/utilities';
+import {
+  EVERYTHING_ELSE,
+  selectBinaries,
+  setHoveringRow,
+  useAppDispatch,
+  useAppSelector,
+} from '@parca/store';
+import {getLastItem, isSearchMatch, scaleLinear} from '@parca/utilities';
 
-import {FIELD_CHILDREN, FIELD_CUMULATIVE, FIELD_DIFF, FIELD_FUNCTION_NAME} from './index';
+import {
+  FIELD_CHILDREN,
+  FIELD_CUMULATIVE,
+  FIELD_DIFF,
+  FIELD_FUNCTION_NAME,
+  FIELD_MAPPING_FILE,
+} from './index';
 import {nodeLabel} from './utils';
 
 export const RowHeight = 26;
 
 interface IcicleGraphNodesProps {
   table: Table<any>;
   row: number;
+  mappingColors: mappingColors;
   children: number[];
   x: number;
   y: number;
@@ -46,6 +59,7 @@ interface IcicleGraphNodesProps {
 export const IcicleGraphNodes = React.memo(function IcicleGraphNodesNoMemo({
   table,
   children,
+  mappingColors,
   x,
   y,
   xScale,
@@ -82,6 +96,7 @@ export const IcicleGraphNodes = React.memo(function IcicleGraphNodesNoMemo({
         key={`node-${level}-${i}`}
         table={table}
         row={child}
+        mappingColors={mappingColors}
         x={xStart}
         y={0}
         totalWidth={totalWidth}
@@ -102,6 +117,10 @@ export const IcicleGraphNodes = React.memo(function IcicleGraphNodesNoMemo({
   return <g transform={`translate(${x}, ${y})`}>{childrenElements}</g>;
 });
 
+interface mappingColors {
+  [key: string]: string;
+}
+
 interface IcicleNodeProps {
   x: number;
   y: number;
@@ -111,6 +130,7 @@ interface IcicleNodeProps {
   level: number;
   table: Table<any>;
   row: number;
+  mappingColors: mappingColors;
   path: string[];
   total: bigint;
   setCurPath: (path: string[]) => void;
@@ -134,6 +154,7 @@ const fadedIcicleRectStyles = {
 export const IcicleNode = React.memo(function IcicleNodeNoMemo({
   table,
   row,
+  mappingColors,
   x,
   y,
   height,
@@ -152,10 +173,14 @@ export const IcicleNode = React.memo(function IcicleNodeNoMemo({
   const {isShiftDown} = useKeyDown();
   const dispatch = useAppDispatch();
 
+  // get the columns to read from
+  const mappingColumn = table.getChild(FIELD_MAPPING_FILE);
   const functionNameColumn = table.getChild(FIELD_FUNCTION_NAME);
   const cumulativeColumn = table.getChild(FIELD_CUMULATIVE);
   const diffColumn = table.getChild(FIELD_DIFF);
-
+  // get the actual values from the columns
+  const mappingFile: string | null = mappingColumn?.get(row);
+  const functionName: string | null = functionNameColumn?.get(row);
   const cumulative: bigint = cumulativeColumn?.get(row);
   const children: number[] = Array.from(table.getChild(FIELD_CHILDREN)?.get(row) ?? []);
 
@@ -238,6 +263,14 @@ export const IcicleNode = React.memo(function IcicleNodeNoMemo({
     dispatch(setHoveringRow(undefined));
   };
 
+  // To get the color we first check if the function name starts with 'runtime'.
+  // If it does, we color it as runtime. Otherwise, we check the mapping file.
+  // If there is no mapping file, we color it as 'everything else'.
+  const color =
+    functionName?.startsWith('runtime') === true
+      ? mappingColors.runtime
+      : mappingColors[getLastItem(mappingFile ?? '') ?? EVERYTHING_ELSE];
+
   return (
     <>
       <g
@@ -255,7 +288,7 @@ export const IcicleNode = React.memo(function IcicleNodeNoMemo({
           width={width}
           height={height}
           style={{
-            fill: '#929FEB', // TODO: Introduce color coding for binaries again
+            fill: color,
           }}
           className={cx('stroke-white dark:stroke-gray-700', {
             'opacity-50': isHighlightEnabled && !isHighlighted,
@@ -273,6 +306,7 @@ export const IcicleNode = React.memo(function IcicleNodeNoMemo({
         <IcicleGraphNodes
           table={table}
           row={row}
+          mappingColors={mappingColors}
           children={children}
           x={x}
           y={RowHeight}

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraphArrow/index.tsx

@@ -19,6 +19,7 @@ import {USER_PREFERENCES, useUserPreference} from '@parca/hooks';
 import {
   FEATURE_TYPES,
   FeaturesMap,
+  getColorForFeature,
   selectDarkMode,
   setHoveringNode,
   useAppDispatch,
@@ -35,6 +36,7 @@ import {
 import GraphTooltip from '../../GraphTooltip';
 import ColorStackLegend from './ColorStackLegend';
 import {IcicleNode, RowHeight} from './IcicleGraphNodes';
+import {extractFeature} from './utils';
 
 export const FIELD_MAPPING_FILE = 'mapping_file';
 export const FIELD_MAPPING_BUILD_ID = 'mapping_build_id';
@@ -82,8 +84,9 @@ export const IcicleGraphArrow = memo(function IcicleGraphArrow({
     selectQueryParam('compare_a') === 'true' && selectQueryParam('compare_b') === 'true';
 
   const mappings = useMemo(() => {
+    // Reading the mappings from the dictionary that contains all mapping strings.
     const mappingsDict: Vector<Dictionary> | null = table.getChild(FIELD_MAPPING_FILE);
-    return Array.from(mappingsDict?.data.values() ?? [])
+    const mappings = Array.from(mappingsDict?.data.values() ?? [])
       .map((mapping): string[] => {
         const dict = mapping.dictionary;
         const len = dict?.data.length ?? 0;
@@ -94,19 +97,41 @@ export const IcicleGraphArrow = memo(function IcicleGraphArrow({
         }
         return values;
       })
-      .flat()
-      .sort((a, b) => a.localeCompare(b));
+      .flat();
+
+    // We add a EVERYTHING ELSE mapping to the list.
+    mappings.push('');
+
+    // We look through the function names to find out if there's a runtime function.
+    const functionNamesDict: Vector<Dictionary> | null = table.getChild(FIELD_FUNCTION_NAME);
+    // TODO: There must be a better way to do this. Somehow read the function name dictionary rather than iterating over all rows.
+    for (let i = 0; i < table.numRows; i++) {
+      const fn: string | null = functionNamesDict?.get(i);
+      if (fn?.startsWith('runtime') === true) {
+        mappings.push('runtime');
+        break;
+      }
+    }
+
+    // We sort the mappings alphabetically to make sure that the order is always the same.
+    mappings.sort((a, b) => a.localeCompare(b));
+    return mappings;
   }, [table]);
 
+  // TODO: Somehow figure out how to add runtime to this, if stacks are present.
+  // Potentially read the function name dictionary and check if it contains strings starting with runtime.
   const mappingFeatures = useMemo(() => {
-    const features: FeaturesMap = {
-      runtime: FEATURE_TYPES.Runtime,
-    };
-    mappings.forEach(mapping => {
-      features[mapping] = FEATURE_TYPES.Binary;
+    return mappings.map(mapping => extractFeature(mapping));
+  }, [mappings]);
+
+  // TODO: Unify with mappingFeatures
+  const mappingColors = useMemo(() => {
+    const colors = {};
+    Object.entries(mappingFeatures).forEach(([_, feature]) => {
+      colors[feature.name] = getColorForFeature(feature.name, isDarkMode, colorProfile);
     });
-    return features;
-  }, [mappings, colorProfile]);
+    return colors;
+  }, [colorProfile, isDarkMode, mappingFeatures]);
 
   useEffect(() => {
     if (ref.current != null) {
@@ -136,7 +161,7 @@ export const IcicleGraphArrow = memo(function IcicleGraphArrow({
         <option value="diff">Diff</option>
       </select>
       <ColorStackLegend
-        mappingsFeatures={mappingFeatures}
+        mappingColors={mappingColors}
         navigateTo={navigateTo}
         compareMode={compareMode}
       />
@@ -159,6 +184,7 @@ export const IcicleGraphArrow = memo(function IcicleGraphArrow({
             <IcicleNode
               table={table}
               row={0} // root is always row 0 in the arrow record
+              mappingColors={mappingColors}
               x={0}
               y={0}
               totalWidth={width}

ui/packages/shared/profile/src/ProfileIcicleGraph/IcicleGraphArrow/utils.ts

@@ -78,14 +78,13 @@ export function nodeLabel(table: Table<any>, row: number, showBinaryName: boolea
   //return fallback === '' ? '<unknown>' : fallback;
 }
 
-export const extractFeature = (table: Table<any>, row: number): Feature => {
-  const name = nodeLabel(table, row, false).trim();
-  if (name.startsWith('runtime') || name === 'root') {
+export const extractFeature = (mapping: string): Feature => {
+  if (mapping.startsWith('runtime') || mapping === 'root') {
     return {name: 'runtime', type: FEATURE_TYPES.Runtime};
   }
 
-  if (binaryName != null) {
-    return {name: binaryName, type: FEATURE_TYPES.Binary};
+  if (mapping != null && mapping !== '') {
+    return {name: mapping, type: FEATURE_TYPES.Binary};
   }
 
   return {name: EVERYTHING_ELSE, type: FEATURE_TYPES.Misc};

ui/packages/shared/store/src/slices/colorsSlice.ts

@@ -87,7 +87,7 @@ const findAColor = (colorIndex: number, colors: ColorsDuo[]): ColorsDuo => {
   // TODO: add some logic to find unallocated colors if this index is already allocated to another feature for better color distribution.
 };
 
-const getColorForFeature = (
+export const getColorForFeature = (
   feature: string,
   isDarkMode: boolean,
   colorProfileName: ColorProfileName