feat: support source locations of heap

Author: zjffun

packages/vscode-js-profile-core/src/common/types.ts

@@ -5,6 +5,34 @@
 import { Protocol as Cdp } from 'devtools-protocol';
 import { ISourceLocation } from '../location-mapping';
 
+export interface IAnnotationLocation {
+  callFrame: Cdp.Runtime.CallFrame;
+  locations: ISourceLocation[];
+}
+
+/**
+ * Extra annotations added by js-debug.
+ */
+export interface IJsDebugAnnotations {
+  /**
+   * Workspace root path, if set.
+   */
+  rootPath?: string;
+
+  /**
+   * For each node in the profile, the list of locations in corresponds to
+   * in the workspace (if any).
+   */
+  locations: ReadonlyArray<IAnnotationLocation>;
+
+  /**
+   * Optional cell data saved from previously opening the profile as a notebook.
+   */
+  cellData?: {
+    version: number;
+  };
+}
+
 /**
  * Request from the webview to open a document
  */

packages/vscode-js-profile-core/src/cpu/model.ts

@@ -4,9 +4,11 @@
 
 import { Protocol as Cdp } from 'devtools-protocol';
 import { categorize, INode } from '../common/model';
-import { addRelativeDiskPath, ISourceLocation } from '../location-mapping';
+import { IAnnotationLocation } from '../common/types';
+import { getBestLocation } from '../getBestLocation';
+import { ISourceLocation } from '../location-mapping';
 import { maybeFileUrlToPath } from '../path';
-import { IAnnotationLocation, ICpuProfileRaw } from './types';
+import { ICpuProfileRaw } from './types';
 
 /**
  * One measured node in the call stack. Contains the time it spent in itself,
@@ -71,24 +73,6 @@ const computeAggregateTime = (index: number, nodes: IComputedNode[]): number =>
   return (row.aggregateTime = total);
 };
 
-const getBestLocation = (
-  profile: ICpuProfileRaw,
-  candidates: ReadonlyArray<ISourceLocation> = [],
-) => {
-  if (!profile.$vscode?.rootPath) {
-    return candidates[0];
-  }
-
-  for (const candidate of candidates) {
-    const mapped = addRelativeDiskPath(profile.$vscode.rootPath, candidate);
-    if (mapped.relativePath) {
-      return mapped;
-    }
-  }
-
-  return candidates[0];
-};
-
 /**
  * Ensures that all profile nodes have a location ID, setting them if they
  * aren't provided by default.

packages/vscode-js-profile-core/src/cpu/types.ts

@@ -3,40 +3,12 @@
  *--------------------------------------------------------*/
 
 import { Protocol as Cdp } from 'devtools-protocol';
-import { ISourceLocation } from '../location-mapping';
+import { IJsDebugAnnotations } from '../common/types';
 
 export const enum Constants {
   CurrentDataVersion = 1,
 }
 
-export interface IAnnotationLocation {
-  callFrame: Cdp.Runtime.CallFrame;
-  locations: ISourceLocation[];
-}
-
-/**
- * Extra annotations added by js-debug.
- */
-export interface IJsDebugAnnotations {
-  /**
-   * Workspace root path, if set.
-   */
-  rootPath?: string;
-
-  /**
-   * For each node in the profile, the list of locations in corresponds to
-   * in the workspace (if any).
-   */
-  locations: ReadonlyArray<IAnnotationLocation>;
-
-  /**
-   * Optional cell data saved from previously opening the profile as a notebook.
-   */
-  cellData?: {
-    version: number;
-  };
-}
-
 export interface IProfileNode extends Cdp.Profiler.ProfileNode {
   locationId?: number;
   positionTicks?: (Cdp.Profiler.PositionTickInfo & {

packages/vscode-js-profile-core/src/getBestLocation.ts

@@ -0,0 +1,24 @@
+/*---------------------------------------------------------
+ * Copyright (C) Microsoft Corporation. All rights reserved.
+ *--------------------------------------------------------*/
+
+import { IJsDebugAnnotations } from './common/types';
+import { addRelativeDiskPath, ISourceLocation } from './location-mapping';
+
+export const getBestLocation = (
+  profile: { $vscode?: IJsDebugAnnotations },
+  candidates: ReadonlyArray<ISourceLocation> = [],
+) => {
+  if (!profile.$vscode?.rootPath) {
+    return candidates[0];
+  }
+
+  for (const candidate of candidates) {
+    const mapped = addRelativeDiskPath(profile.$vscode.rootPath, candidate);
+    if (mapped.relativePath) {
+      return mapped;
+    }
+  }
+
+  return candidates[0];
+};

packages/vscode-js-profile-core/src/heap/model.ts

@@ -4,6 +4,8 @@
 
 import { Protocol as Cdp } from 'devtools-protocol';
 import { INode } from '../common/model';
+import { IAnnotationLocation, IJsDebugAnnotations } from '../common/types';
+import { getBestLocation } from '../getBestLocation';
 import { ISourceLocation } from '../location-mapping';
 import { maybeFileUrlToPath } from '../path';
 
@@ -18,24 +20,16 @@ export interface ITreeNode extends IHeapProfileNode {
   parent?: ITreeNode;
 }
 
-/**
- * Extra annotations added by js-debug.
- */
-export interface IJsDebugAnnotations {
-  /**
-   * Workspace root path, if set.
-   */
-  rootPath?: string;
-}
-
 export interface IHeapProfileRaw extends Cdp.HeapProfiler.SamplingHeapProfile {
   $vscode?: IJsDebugAnnotations;
+  head: IProfileModelNode;
 }
 
 export interface IProfileModelNode
   extends Omit<Cdp.HeapProfiler.SamplingHeapProfileNode, 'children'> {
   src?: ISourceLocation;
   children: IProfileModelNode[];
+  locationId?: number;
 }
 
 /**
@@ -48,25 +42,84 @@ export type IProfileModel = {
 };
 
 /**
- * Computes the model for the given profile.
+ * Ensures that all profile nodes have a location ID, setting them if they
+ * aren't provided by default.
  */
-export const buildModel = (profile: IHeapProfileRaw): IProfileModel => {
-  let nodes = [profile.head];
+const ensureSourceLocations = (profile: IHeapProfileRaw): ReadonlyArray<IAnnotationLocation> => {
+  if (profile.$vscode) {
+    return profile.$vscode.locations; // profiles we generate are already good
+  }
 
-  while (nodes.length) {
-    const node = nodes.pop();
+  let locationIdCounter = 0;
+  const locationsByRef = new Map<
+    string,
+    { id: number; callFrame: Cdp.Runtime.CallFrame; location: ISourceLocation }
+  >();
 
-    if (node) {
-      const { callFrame } = node;
-      (node as unknown as IHeapProfileNode).src = {
+  const getLocationIdFor = (callFrame: Cdp.Runtime.CallFrame) => {
+    const ref = [
+      callFrame.functionName,
+      callFrame.url,
+      callFrame.scriptId,
+      callFrame.lineNumber,
+      callFrame.columnNumber,
+    ].join(':');
+
+    const existing = locationsByRef.get(ref);
+    if (existing) {
+      return existing.id;
+    }
+    const id = locationIdCounter++;
+    locationsByRef.set(ref, {
+      id,
+      callFrame,
+      location: {
         lineNumber: callFrame.lineNumber,
         columnNumber: callFrame.columnNumber,
         source: {
           name: maybeFileUrlToPath(callFrame.url),
           path: maybeFileUrlToPath(callFrame.url),
           sourceReference: 0,
         },
-      };
+      },
+    });
+
+    return id;
+  };
+
+  let nodes = [profile.head];
+
+  while (nodes.length) {
+    const node = nodes.pop();
+
+    if (node) {
+      const { callFrame } = node;
+      node.locationId = getLocationIdFor(callFrame);
+
+      nodes = nodes.concat(node.children);
+    }
+  }
+
+  return [...locationsByRef.values()]
+    .sort((a, b) => a.id - b.id)
+    .map(l => ({ locations: [l.location], callFrame: l.callFrame }));
+};
+
+/**
+ * Computes the model for the given profile.
+ */
+export const buildModel = (profile: IHeapProfileRaw): IProfileModel => {
+  let nodes = [profile.head];
+
+  const sourceLocations = ensureSourceLocations(profile);
+
+  while (nodes.length) {
+    const node = nodes.pop();
+
+    if (node) {
+      if (node.locationId) {
+        node.src = getBestLocation(profile, sourceLocations[node.locationId].locations);
+      }
       nodes = nodes.concat(node.children);
     }
   }