Improve scene serialization efficiency

src/viser/client/src/FilePlayback.tsx

@@ -1,14 +1,8 @@
-import * as msgpack from "@msgpack/msgpack";
-import { Message } from "./WebsocketMessages";
-import { ZSTDDecoder } from "zstddec";
 import {
-  replaceBinaryPlaceholders,
-  computeBinaryOffsets,
-} from "./BinaryMessageDecode";
-
-// Initialize zstd decoder at module load.
-const zstdDecoder = new ZSTDDecoder();
-const zstdReady = zstdDecoder.init();
+  SerializedMessages,
+  deserializeEmbeddedData,
+  deserializeZstdMsgpackFile,
+} from "./PlaybackDecode";
 
 import {
   Dispatch,
@@ -59,122 +53,6 @@ import {
   IconPlayerPlayFilled,
 } from "@tabler/icons-react";
 
-/**
- * Decompress and decode a hybrid-format payload.
- *
- * Decompressed layout:
- *   [8 bytes] msgpack length (little-endian uint64)
- *   [N bytes] msgpack payload (with binary placeholders)
- *   [P bytes] padding + aligned binary buffers
- *
- * Binary placeholders are replaced with properly typed array views.
- */
-function decodeHybridPayload<T>(decompressed: Uint8Array): T {
-  const buf = decompressed.buffer as ArrayBuffer;
-  const base = decompressed.byteOffset;
-
-  // Read msgpack length from inner header.
-  const msgpackLength = Number(
-    new DataView(buf, base, 8).getBigUint64(0, true),
-  );
-
-  // Decode msgpack.
-  const msgpackData = new Uint8Array(buf, base + 8, msgpackLength);
-  const data = msgpack.decode(msgpackData) as T & {
-    binaryBufferLengths?: number[];
-  };
-
-  // Replace binary placeholders with typed array views.
-  const bufferLengths = data.binaryBufferLengths;
-  if (bufferLengths && bufferLengths.length > 0) {
-    const binaryOffsets = computeBinaryOffsets(
-      bufferLengths,
-      base + 8 + msgpackLength,
-    );
-    replaceBinaryPlaceholders(data, buf, binaryOffsets, bufferLengths);
-  }
-
-  return data;
-}
-
-/** Download, decompress, and deserialize a .viser recording file. */
-async function deserializeZstdMsgpackFile<T>(
-  fileUrl: string,
-  setStatus: (status: { downloaded: number; total: number }) => void,
-): Promise<T> {
-  const response = await fetch(fileUrl);
-  if (!response.ok) {
-    throw new Error(`Failed to fetch the file: ${response.statusText}`);
-  }
-
-  const totalLength = parseInt(response.headers.get("Content-Length")!);
-  setStatus({ downloaded: 0, total: totalLength });
-
-  // Stream the download to track progress.
-  const reader = response.body!.getReader();
-  const chunks: Uint8Array[] = [];
-  let downloadedLength = 0;
-
-  while (true) {
-    const { done, value } = await reader.read();
-    if (done) break;
-    chunks.push(value);
-    downloadedLength += value.length;
-    setStatus({ downloaded: downloadedLength, total: totalLength });
-  }
-
-  // Concatenate chunks into a single buffer.
-  const bytes = new Uint8Array(downloadedLength);
-  let offset = 0;
-  for (const chunk of chunks) {
-    bytes.set(chunk, offset);
-    offset += chunk.length;
-  }
-
-  // Read decompressed size from 8-byte little-endian header.
-  const view = new DataView(bytes.buffer);
-  const decompressedSize = Number(view.getBigUint64(0, true));
-  const compressedData = bytes.slice(8);
-
-  // Decompress and decode using shared hybrid format logic.
-  await zstdReady;
-  const decompressed = zstdDecoder.decode(compressedData, decompressedSize);
-  return decodeHybridPayload<T>(decompressed);
-}
-
-/** Deserialize embedded base64-encoded zstd-compressed data.
- * Used for static embedding in HTML pages (e.g., myst-nb documentation). */
-async function deserializeEmbeddedData<T>(
-  base64Data: string,
-  setStatus: (status: { downloaded: number; total: number }) => void,
-): Promise<T> {
-  // Decode base64 to Uint8Array.
-  const binaryString = atob(base64Data);
-  const bytes = new Uint8Array(binaryString.length);
-  for (let i = 0; i < binaryString.length; i++) {
-    bytes[i] = binaryString.charCodeAt(i);
-  }
-
-  // Data is already embedded, so mark download as complete.
-  setStatus({ downloaded: 1.0, total: 1.0 });
-
-  // Read decompressed size from 8-byte little-endian header.
-  const view = new DataView(bytes.buffer);
-  const decompressedSize = Number(view.getBigUint64(0, true));
-  const compressedData = bytes.slice(8);
-
-  // Decompress and decode using shared hybrid format logic.
-  await zstdReady;
-  const decompressed = zstdDecoder.decode(compressedData, decompressedSize);
-  return decodeHybridPayload<T>(decompressed);
-}
-
-export interface SerializedMessages {
-  durationSeconds: number;
-  messages: [number, Message][]; // (time in seconds, message).
-  viserVersion: string;
-}
-
 /** Shared playback UI and timing logic for recorded scenes.
  *
  * The two entry points -- downloading a `.viser` file vs. decoding embedded

src/viser/client/src/PlaybackDecode.test.ts

@@ -0,0 +1,111 @@
+import { describe, expect, it } from "vitest";
+import {
+  base64ToBytes,
+  decompressAndDecodeHybridPayload,
+  SerializedMessages,
+} from "./PlaybackDecode";
+import { MeshMessage } from "./WebsocketMessages";
+
+// A real `.viser` payload produced by the Python server, exercising the full
+// encode -> decode contract (zstd framing, msgpack header, aligned binary
+// buffers, placeholder replacement, and buffer deduplication). Regenerate
+// with:
+//
+//   import base64, numpy as np, viser
+//   server = viser.ViserServer(port=0, verbose=False)
+//   verts = np.array(
+//       [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32)
+//   faces = np.array([[0, 1, 2], [0, 2, 3]], dtype=np.uint32)
+//   server.scene.add_mesh_simple("/fixture_a", verts.copy(), faces.copy())
+//   server.scene.add_mesh_simple("/fixture_b", verts.copy(), faces.copy())
+//   s = server.get_scene_serializer()
+//   s.insert_sleep(0.25)
+//   print(base64.b64encode(s.serialize()).decode())
+const FIXTURE_B64 = [
+  "mAQAAAAAAAAotS/9YJgDRQ8AZtlaPkBtmwPGh6dv5gVOC6ZVnmdenq0fF58wZ0TEPd8MLrdo2xI5",
+  "AxHZCLmJTCKnWYvDeds3o/35VANWFRZlrmplSgBKAEwAGrOndtphg4aFp3us4/qJp6PqNLmKMpzn",
+  "bKMLpd6jwl/Ur+05QxArpw4IfLC1E0PhdRkanqf9Z4t7JZ7u46SQoUm7rmPRA8YTUAmWgjgx6hME",
+  "uJUQVL04foRs3RhJk3PTvDbqWk4lB4ZWu+oLjpQBT4ZQZUjEhiupvy0dFmytJ9DwpVUMiqfjauFv",
+  "T6Qj0sHgtKJQEx/VraQmRKsN96IdG8neF4+RDonOg60Ck6zcKuSBjbi9jOPsSPOdLUURyIRn2NXZ",
+  "wl3YiB2rWcs6ub1BZQhEihETIzpJhghadsFwmDbvvQAAJI+EN4CS8QucpcQqHygGJpRJAY5J+50v",
+  "nL8tGkniUM8WGbY0Q3YQzoX+wkz0jXPrc9W2c12Lq5O63Dj2YzbT7LhUMyBggkFFVTcM/LoCPBZ2",
+  "j2FUavK5ph8Dh3nrHFGsFDEzfO3XyuPKzw1DxHcHUrtqyIb1fB5T4hKkaWtwihwO/YJOHWrVl1nI",
+  "iCgH2HAeEiQMyUBJF4t/yArn7TGVI1xJAZyEk8iDnQ+DqVVEoXEj7BzEFJ7nX2MCHwI=",
+].join("");
+
+const EXPECTED_VERTICES = [0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1];
+const EXPECTED_FACES = [0, 1, 2, 0, 2, 3];
+
+async function decodeFixture(): Promise<SerializedMessages> {
+  const bytes = base64ToBytes(FIXTURE_B64);
+  return await decompressAndDecodeHybridPayload<SerializedMessages>(bytes);
+}
+
+function meshesOf(recording: SerializedMessages): MeshMessage[] {
+  return recording.messages
+    .map(([, message]) => message)
+    .filter((message): message is MeshMessage => {
+      return message.type === "MeshMessage";
+    });
+}
+
+describe("PlaybackDecode", () => {
+  it("decodes a Python-serialized recording", async () => {
+    const recording = await decodeFixture();
+
+    expect(recording.durationSeconds).toBe(0.25);
+    expect(typeof recording.viserVersion).toBe("string");
+
+    const meshes = meshesOf(recording);
+    expect(meshes.map((m) => m.name).sort()).toEqual([
+      "/fixture_a",
+      "/fixture_b",
+    ]);
+
+    for (const mesh of meshes) {
+      expect(mesh.props.vertices).toBeInstanceOf(Float32Array);
+      expect(Array.from(mesh.props.vertices)).toEqual(EXPECTED_VERTICES);
+      expect(mesh.props.faces).toBeInstanceOf(Uint32Array);
+      expect(Array.from(mesh.props.faces)).toEqual(EXPECTED_FACES);
+    }
+  });
+
+  it("resolves deduplicated buffers to shared byte ranges", async () => {
+    const recording = await decodeFixture();
+    const meshes = meshesOf(recording);
+    expect(meshes.length).toBe(2);
+
+    // The server stores byte-identical arrays once; both messages' views must
+    // alias the same region of the shared ArrayBuffer (while remaining
+    // distinct view objects).
+    const [a, b] = meshes;
+    expect(a.props.vertices).not.toBe(b.props.vertices);
+    expect(a.props.vertices.buffer).toBe(b.props.vertices.buffer);
+    expect(a.props.vertices.byteOffset).toBe(b.props.vertices.byteOffset);
+    expect(a.props.faces.byteOffset).toBe(b.props.faces.byteOffset);
+    // Vertices and faces are different content, so they must NOT alias.
+    expect(a.props.vertices.byteOffset).not.toBe(a.props.faces.byteOffset);
+  });
+
+  it("decodes base64 identically with and without native fromBase64", () => {
+    const uint8ArrayStatics = Uint8Array as unknown as {
+      fromBase64?: (data: string) => Uint8Array;
+    };
+    const original = uint8ArrayStatics.fromBase64;
+    try {
+      // Exercise the native branch even on runtimes without fromBase64, by
+      // shimming it (Buffer is the Node-native reference decoder).
+      uint8ArrayStatics.fromBase64 ??= (data) =>
+        new Uint8Array(Buffer.from(data, "base64"));
+      const native = base64ToBytes(FIXTURE_B64);
+
+      // Exercise the real atob fallback branch by hiding the native method.
+      uint8ArrayStatics.fromBase64 = undefined;
+      const fallback = base64ToBytes(FIXTURE_B64);
+
+      expect(Array.from(fallback)).toEqual(Array.from(native));
+    } finally {
+      uint8ArrayStatics.fromBase64 = original;
+    }
+  });
+});