Implement compacting for the shared source table

Author: canova

src/profile-logic/profile-compacting.ts

@@ -18,34 +18,40 @@ import type {
 export type CompactedProfileWithTranslationMaps = {
   profile: Profile;
   oldStringToNewStringPlusOne: Int32Array;
+  oldSourceToNewSourcePlusOne: Int32Array;
 };
 
 /**
- * Returns a new profile with all unreferenced strings removed.
+ * Returns a new profile with all unreferenced strings and sources removed.
  *
- * Since the string table is shared between all threads, if the user asks for a
- * thread to be removed during sanitization, by default we'd keep the strings
- * from the removed threads in the profile.
+ * Since the string table and source table are shared between all threads, if
+ * the user asks for a thread to be removed during sanitization, by default
+ * we'd keep the strings and sources from the removed threads in the profile.
  *
- * By calling this function, you can get a profile with an adjusted string table
- * where those unused strings from the removed threads have been removed.
+ * By calling this function, you can get a profile with adjusted string and
+ * source tables where those unused strings and sources from the removed
+ * threads have been removed.
  */
 export function computeCompactedProfile(
   profile: Profile
 ): CompactedProfileWithTranslationMaps {
   const stringIndexMarkerFieldsByDataType =
     computeStringIndexMarkerFieldsByDataType(profile.meta.markerSchema);
 
-  // Step 1: Gather all references.
+  // Step 1: Gather all references of strings.
   const referencedStrings = _gatherStringReferencesInProfile(
     profile,
     stringIndexMarkerFieldsByDataType
   );
 
-  // Step 2: Adjust all tables to use new string indexes.
-  return _createProfileWithTranslatedStringIndexes(
+  // Step 2: Gather all references of sources.
+  const referencedSources = _gatherSourceReferencesInProfile(profile);
+
+  // Step 3: Adjust all tables to use new string and source indexes.
+  return _createProfileWithTranslatedIndexes(
     profile,
     referencedStrings,
+    referencedSources,
     stringIndexMarkerFieldsByDataType
   );
 }
@@ -68,36 +74,53 @@ function _gatherStringReferencesInProfile(
   return referencedStrings;
 }
 
-function _createProfileWithTranslatedStringIndexes(
+function _gatherSourceReferencesInProfile(profile: Profile): Uint8Array {
+  const referencedSources = new Uint8Array(profile.shared.sources.length);
+
+  for (const thread of profile.threads) {
+    _gatherSourceReferencesInThread(thread, referencedSources);
+  }
+
+  return referencedSources;
+}
+
+function _gatherSourceReferencesInThread(
+  thread: RawThread,
+  referencedSources: Uint8Array
+) {
+  for (let i = 0; i < thread.funcTable.length; i++) {
+    const sourceIndex = thread.funcTable.source[i];
+    if (sourceIndex !== null) {
+      referencedSources[sourceIndex] = 1;
+    }
+  }
+}
+
+function _createProfileWithTranslatedIndexes(
   profile: Profile,
   referencedStrings: Uint8Array,
+  referencedSources: Uint8Array,
   stringIndexMarkerFieldsByDataType: Map<string, string[]>
 ): CompactedProfileWithTranslationMaps {
   const { newStringArray, oldStringToNewStringPlusOne } =
     _createCompactedStringArray(profile.shared.stringArray, referencedStrings);
 
+  const { newSources, oldSourceToNewSourcePlusOne } =
+    _createCompactedSourceTable(
+      profile.shared.sources,
+      referencedSources,
+      oldStringToNewStringPlusOne
+    );
+
   const newThreads = profile.threads.map((thread) =>
-    _createThreadWithTranslatedStringIndexes(
+    _createThreadWithTranslatedIndexes(
       thread,
       oldStringToNewStringPlusOne,
+      oldSourceToNewSourcePlusOne,
       stringIndexMarkerFieldsByDataType
     )
   );
 
-  // Update sources table with translated string indexes
-  const newSources = {
-    ...profile.shared.sources,
-    filename: profile.shared.sources.filename.map((oldUrlIndex) => {
-      const newIndexPlusOne = oldStringToNewStringPlusOne[oldUrlIndex];
-      if (newIndexPlusOne === 0) {
-        throw new Error(
-          `String index ${oldUrlIndex} was not found in the translation map`
-        );
-      }
-      return newIndexPlusOne - 1;
-    }),
-  };
-
   const newShared: RawProfileSharedData = {
     stringArray: newStringArray,
     sources: newSources,
@@ -112,6 +135,7 @@ function _createProfileWithTranslatedStringIndexes(
   return {
     profile: newProfile,
     oldStringToNewStringPlusOne,
+    oldSourceToNewSourcePlusOne,
   };
 }
 
@@ -132,9 +156,10 @@ function _gatherStringReferencesInThread(
   _gatherReferencesInNativeSymbols(thread.nativeSymbols, referencedStrings);
 }
 
-function _createThreadWithTranslatedStringIndexes(
+function _createThreadWithTranslatedIndexes(
   thread: RawThread,
   oldStringToNewStringPlusOne: Int32Array,
+  oldSourceToNewSourcePlusOne: Int32Array,
   stringIndexMarkerFieldsByDataType: Map<string, string[]>
 ): RawThread {
   const newNativeSymbols = _createNativeSymbolsWithTranslatedStringIndexes(
@@ -145,9 +170,10 @@ function _createThreadWithTranslatedStringIndexes(
     thread.resourceTable,
     oldStringToNewStringPlusOne
   );
-  const newFuncTable = _createFuncTableWithTranslatedStringIndexes(
+  const newFuncTable = _createFuncTableWithTranslatedIndexes(
     thread.funcTable,
-    oldStringToNewStringPlusOne
+    oldStringToNewStringPlusOne,
+    oldSourceToNewSourcePlusOne
   );
   const newMarkers = _createMarkersWithTranslatedStringIndexes(
     thread.markers,
@@ -253,18 +279,23 @@ function _gatherReferencesInFuncTable(
   }
 }
 
-function _createFuncTableWithTranslatedStringIndexes(
+function _createFuncTableWithTranslatedIndexes(
   funcTable: FuncTable,
-  oldStringToNewStringPlusOne: Int32Array
+  oldStringToNewStringPlusOne: Int32Array,
+  oldSourceToNewSourcePlusOne: Int32Array
 ): FuncTable {
   const newFuncTableNameCol = funcTable.name.slice();
   const newFuncTableSourceCol = funcTable.source.slice();
   for (let i = 0; i < funcTable.length; i++) {
     const name = funcTable.name[i];
     newFuncTableNameCol[i] = oldStringToNewStringPlusOne[name] - 1;
 
-    // Note: source indexes don't need translation as they point to sources table, not strings
-    // Source table will be handled separately in _createSourcesTableWithTranslatedStringIndexes
+    // Translate source indexes to new compacted source table.
+    const sourceIndex = funcTable.source[i];
+    if (sourceIndex !== null) {
+      const newSourceIndexPlusOne = oldSourceToNewSourcePlusOne[sourceIndex];
+      newFuncTableSourceCol[i] = newSourceIndexPlusOne - 1;
+    }
   }
 
   const newFuncTable = {
@@ -357,3 +388,44 @@ function _createCompactedStringArray(
 
   return { newStringArray, oldStringToNewStringPlusOne };
 }
+
+function _createCompactedSourceTable(
+  sourceTable: SourceTable,
+  referencedSources: Uint8Array,
+  oldStringToNewStringPlusOne: Int32Array
+): { newSources: SourceTable; oldSourceToNewSourcePlusOne: Int32Array } {
+  const oldSourceToNewSourcePlusOne = new Int32Array(sourceTable.length);
+  let nextIndex = 0;
+  const newUuid = [];
+  const newFilename = [];
+
+  for (let i = 0; i < sourceTable.length; i++) {
+    if (referencedSources[i] === 0) {
+      continue;
+    }
+
+    const newIndex = nextIndex++;
+    newUuid[newIndex] = sourceTable.uuid[i];
+
+    // Translate the filename string index
+    const oldFilenameIndex = sourceTable.filename[i];
+    const newFilenameIndexPlusOne =
+      oldStringToNewStringPlusOne[oldFilenameIndex];
+    if (newFilenameIndexPlusOne === 0) {
+      throw new Error(
+        `String index ${oldFilenameIndex} was not found in the translation map`
+      );
+    }
+    newFilename[newIndex] = newFilenameIndexPlusOne - 1;
+
+    oldSourceToNewSourcePlusOne[i] = newIndex + 1;
+  }
+
+  const newSources: SourceTable = {
+    length: nextIndex,
+    uuid: newUuid,
+    filename: newFilename,
+  };
+
+  return { newSources, oldSourceToNewSourcePlusOne };
+}

src/test/unit/sanitize.test.ts

@@ -1206,4 +1206,57 @@ describe('sanitizePII', function () {
       ]);
     });
   });
+
+  it('should compact the source table when threads are removed', function () {
+    // Create a profile with multiple threads that reference different sources
+    const { profile } = getProfileFromTextSamples(
+      `A[file:file1.js]`,
+      `B[file:file2.js]`,
+      `C[file:file3.js]`
+    );
+
+    const originalSourcesLength = profile.shared.sources.length;
+
+    // Verify we have sources for each thread
+    expect(originalSourcesLength).toEqual(3);
+
+    // Verify that different threads reference different sources
+    const thread0SourceIndex = profile.threads[0].funcTable.source[0];
+    const thread1SourceIndex = profile.threads[1].funcTable.source[0];
+    const thread2SourceIndex = profile.threads[2].funcTable.source[0];
+
+    expect(thread0SourceIndex).not.toBe(thread1SourceIndex);
+    expect(thread1SourceIndex).not.toBe(thread2SourceIndex);
+
+    // Remove threads 0 and 2, keeping only thread 1.
+    const { sanitizedProfile } = setup(
+      {
+        shouldRemoveThreads: new Set([0, 2]),
+      },
+      profile
+    );
+
+    // The source table should be compacted to only contain sources referenced
+    // by remaining threads
+    expect(sanitizedProfile.shared.sources.length).toBeLessThan(
+      originalSourcesLength
+    );
+    expect(sanitizedProfile.shared.sources.length).toEqual(1);
+
+    // The remaining thread should still have a valid source reference
+    const remainingSourceIndex =
+      sanitizedProfile.threads[0].funcTable.source[0];
+    expect(remainingSourceIndex).not.toBeNull();
+    expect(remainingSourceIndex).toBeLessThan(
+      sanitizedProfile.shared.sources.length
+    );
+
+    // Verify that the filename string is still accessible
+    expect(remainingSourceIndex).not.toBeNull();
+    const filenameStringIndex =
+      sanitizedProfile.shared.sources.filename[remainingSourceIndex!];
+    expect(sanitizedProfile.shared.stringArray[filenameStringIndex]).toContain(
+      'file2.js'
+    );
+  });
 });