Construct new stack tables with the correct column lengths. (#5906)

Fixes #5905.

Author: mstange

src/profile-logic/data-structures.ts

@@ -13,7 +13,6 @@ import type {
   RawSamplesTable,
   FrameTable,
   RawStackTable,
-  StackTable,
   FuncTable,
   RawMarkerTable,
   JsAllocationsTable,
@@ -33,20 +32,6 @@ import type {
  * This module collects all of the creation of new empty profile data structures.
  */
 
-export function getEmptyStackTable(): StackTable {
-  return {
-    // Important!
-    // If modifying this structure, please update all callers of this function to ensure
-    // that they are pushing on correctly to the data structure. These pushes may not
-    // be caught by the type system.
-    frame: [],
-    prefix: [],
-    category: new Uint8Array(),
-    subcategory: new Uint8Array(),
-    length: 0,
-  };
-}
-
 export function getEmptySamplesTable(): RawSamplesTable {
   return {
     // Important!

src/profile-logic/profile-data.ts

@@ -2714,6 +2714,34 @@ export function createThreadFromDerivedTables(
   return thread;
 }
 
+/**
+ * Throws if the column lengths of a StackTable don't match stackTable.length.
+ * Call this after constructing a new StackTable to catch bugs early.
+ */
+export function validateStackTableShape(stackTable: StackTable): void {
+  const { length, frame, prefix, category, subcategory } = stackTable;
+  if (frame.length !== length) {
+    throw new Error(
+      `StackTable frame column length ${frame.length} does not match stackTable.length ${length}`
+    );
+  }
+  if (prefix.length !== length) {
+    throw new Error(
+      `StackTable prefix column length ${prefix.length} does not match stackTable.length ${length}`
+    );
+  }
+  if (category.length !== length) {
+    throw new Error(
+      `StackTable category column length ${category.length} does not match stackTable.length ${length}`
+    );
+  }
+  if (subcategory.length !== length) {
+    throw new Error(
+      `StackTable subcategory column length ${subcategory.length} does not match stackTable.length ${length}`
+    );
+  }
+}
+
 /**
  * Sometimes we want to update the stacks for a thread, for instance while searching
  * for a text string, or doing a call tree transformation. This function abstracts
@@ -2738,6 +2766,7 @@ export function updateThreadStacksByGeneratingNewStackColumns(
     markerData: Array<MarkerPayload | null>
   ) => Array<MarkerPayload | null>
 ): Thread {
+  validateStackTableShape(newStackTable);
   const { jsAllocations, nativeAllocations, samples, markers } = thread;
 
   const newSamples = {

src/profile-logic/transforms.ts

@@ -20,7 +20,6 @@ import {
   getSearchFilteredMarkerIndexes,
   stringsToMarkerRegExps,
 } from '../profile-logic/marker-data';
-import { getEmptyStackTable } from './data-structures';
 import { getFunctionName } from './function-info';
 import { splitSearchString } from '../utils/string';
 
@@ -766,7 +765,11 @@ export function mergeCallNode(
     // A root stack's prefix will be null. Maintain that relationship from old to new
     // stacks by mapping from null to null.
     oldStackToNewStack.set(null, null);
-    const newStackTable = getEmptyStackTable();
+    const newFrameCol = [];
+    const newPrefixCol = [];
+    const newCategoryCol = [];
+    const newSubcategoryCol = [];
+    let newLength = 0;
     // Provide two arrays to efficiently cache values for the algorithm. This probably
     // could be refactored to use only one array here.
     const stackDepths = [];
@@ -822,17 +825,28 @@ export function mergeCallNode(
           newStackPrefix === undefined ? null : newStackPrefix
         );
       } else {
-        const newStackIndex = newStackTable.length++;
+        const newStackIndex = newLength++;
         const newStackPrefix = oldStackToNewStack.get(prefix);
-        newStackTable.prefix[newStackIndex] =
+        newPrefixCol[newStackIndex] =
           newStackPrefix === undefined ? null : newStackPrefix;
-        newStackTable.frame[newStackIndex] = frameIndex;
-        newStackTable.category[newStackIndex] = category;
-        newStackTable.subcategory[newStackIndex] = subcategory;
+        newFrameCol[newStackIndex] = frameIndex;
+        newCategoryCol[newStackIndex] = category;
+        newSubcategoryCol[newStackIndex] = subcategory;
         oldStackToNewStack.set(stackIndex, newStackIndex);
       }
     }
 
+    const newStackTable = {
+      frame: newFrameCol,
+      prefix: newPrefixCol,
+      category: new Uint8Array(newCategoryCol),
+      subcategory:
+        stackTable.subcategory instanceof Uint8Array
+          ? new Uint8Array(newSubcategoryCol)
+          : new Uint16Array(newSubcategoryCol),
+      length: newLength,
+    };
+
     return updateThreadStacks(
       thread,
       newStackTable,
@@ -1247,7 +1261,11 @@ export function focusSubtree(
     // A root stack's prefix will be null. Maintain that relationship from old to new
     // stacks by mapping from null to null.
     oldStackToNewStack.set(null, null);
-    const newStackTable = getEmptyStackTable();
+    const newFrameCol = [];
+    const newPrefixCol = [];
+    const newCategoryCol = [];
+    const newSubcategoryCol = [];
+    let newLength = 0;
     for (let stackIndex = 0; stackIndex < stackTable.length; stackIndex++) {
       const prefix = stackTable.prefix[stackIndex];
       const prefixMatchesUpTo = prefix !== null ? stackMatches[prefix] : 0;
@@ -1267,18 +1285,29 @@ export function focusSubtree(
           }
         }
         if (stackMatchesUpTo === prefixDepth) {
-          const newStackIndex = newStackTable.length++;
+          const newStackIndex = newLength++;
           const newStackPrefix = oldStackToNewStack.get(prefix);
-          newStackTable.prefix[newStackIndex] = newStackPrefix ?? null;
-          newStackTable.frame[newStackIndex] = frame;
-          newStackTable.category[newStackIndex] = category;
-          newStackTable.subcategory[newStackIndex] = subcategory;
+          newPrefixCol[newStackIndex] = newStackPrefix ?? null;
+          newFrameCol[newStackIndex] = frame;
+          newCategoryCol[newStackIndex] = category;
+          newSubcategoryCol[newStackIndex] = subcategory;
           oldStackToNewStack.set(stackIndex, newStackIndex);
         }
       }
       stackMatches[stackIndex] = stackMatchesUpTo;
     }
 
+    const newStackTable = {
+      frame: newFrameCol,
+      prefix: newPrefixCol,
+      category: new Uint8Array(newCategoryCol),
+      subcategory:
+        stackTable.subcategory instanceof Uint8Array
+          ? new Uint8Array(newSubcategoryCol)
+          : new Uint16Array(newSubcategoryCol),
+      length: newLength,
+    };
+
     return updateThreadStacks(thread, newStackTable, (oldStack) => {
       if (oldStack === null || stackMatches[oldStack] !== prefixDepth) {
         return null;
@@ -1359,7 +1388,12 @@ export function focusFunction(
     // Typed arrays are initialized to zero, which we interpret as null.
     const oldStackToNewStackPlusOne = new Uint32Array(stackTable.length);
 
-    const newStackTable = getEmptyStackTable();
+    const newFrameCol = [];
+    const newPrefixCol = [];
+    const newCategoryCol = [];
+    const newSubcategoryCol = [];
+    let newLength = 0;
+
     for (let stackIndex = 0; stackIndex < stackTable.length; stackIndex++) {
       const prefix = stackTable.prefix[stackIndex];
       const frameIndex = stackTable.frame[stackIndex];
@@ -1369,16 +1403,26 @@ export function focusFunction(
         prefix === null ? 0 : oldStackToNewStackPlusOne[prefix];
       const newPrefix = newPrefixPlusOne === 0 ? null : newPrefixPlusOne - 1;
       if (newPrefix !== null || funcIndex === funcIndexToFocus) {
-        const newStackIndex = newStackTable.length++;
-        newStackTable.prefix[newStackIndex] = newPrefix;
-        newStackTable.frame[newStackIndex] = frameIndex;
-        newStackTable.category[newStackIndex] = stackTable.category[stackIndex];
-        newStackTable.subcategory[newStackIndex] =
-          stackTable.subcategory[stackIndex];
+        const newStackIndex = newLength++;
+        newPrefixCol[newStackIndex] = newPrefix;
+        newFrameCol[newStackIndex] = frameIndex;
+        newCategoryCol[newStackIndex] = stackTable.category[stackIndex];
+        newSubcategoryCol[newStackIndex] = stackTable.subcategory[stackIndex];
         oldStackToNewStackPlusOne[stackIndex] = newStackIndex + 1;
       }
     }
 
+    const newStackTable = {
+      frame: newFrameCol,
+      prefix: newPrefixCol,
+      category: new Uint8Array(newCategoryCol),
+      subcategory:
+        stackTable.subcategory instanceof Uint8Array
+          ? new Uint8Array(newSubcategoryCol)
+          : new Uint16Array(newSubcategoryCol),
+      length: newLength,
+    };
+
     return updateThreadStacks(thread, newStackTable, (oldStack) => {
       if (oldStack === null) {
         return null;
@@ -1445,7 +1489,11 @@ export function focusCategory(thread: Thread, category: IndexIntoCategoryList) {
     > = new Map();
     oldStackToNewStack.set(null, null);
 
-    const newStackTable = getEmptyStackTable();
+    const newFrameCol = [];
+    const newPrefixCol = [];
+    const newCategoryCol = [];
+    const newSubcategoryCol = [];
+    let newLength = 0;
 
     // fill the new stack table with the kept frames
     for (let stackIndex = 0; stackIndex < stackTable.length; stackIndex++) {
@@ -1460,14 +1508,25 @@ export function focusCategory(thread: Thread, category: IndexIntoCategoryList) {
         continue;
       }
 
-      const newStackIndex = newStackTable.length++;
-      newStackTable.prefix[newStackIndex] = newPrefix;
-      newStackTable.frame[newStackIndex] = stackTable.frame[stackIndex];
-      newStackTable.category[newStackIndex] = stackTable.category[stackIndex];
-      newStackTable.subcategory[newStackIndex] =
-        stackTable.subcategory[stackIndex];
+      const newStackIndex = newLength++;
+      newPrefixCol[newStackIndex] = newPrefix;
+      newFrameCol[newStackIndex] = stackTable.frame[stackIndex];
+      newCategoryCol[newStackIndex] = stackTable.category[stackIndex];
+      newSubcategoryCol[newStackIndex] = stackTable.subcategory[stackIndex];
       oldStackToNewStack.set(stackIndex, newStackIndex);
     }
+
+    const newStackTable = {
+      frame: newFrameCol,
+      prefix: newPrefixCol,
+      category: new Uint8Array(newCategoryCol),
+      subcategory:
+        stackTable.subcategory instanceof Uint8Array
+          ? new Uint8Array(newSubcategoryCol)
+          : new Uint16Array(newSubcategoryCol),
+      length: newLength,
+    };
+
     const updated = updateThreadStacks(
       thread,
       newStackTable,