ui/src/base/proto_utils_wasm.ts

// Copyright (C) 2024 The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

import protos from '../protos';
import {assetSrc} from './assets';
import {errResult, okResult, type Result} from './result';
import {utf8Decode, utf8Encode} from './string_utils';
import WasmModuleGen from '../wasm/proto_utils';
import {PerfettoWasmModule} from '../wasm/perfetto_wasm_module';

/**
 * This file is the TS-equivalent of src/proto_utils.
 * It exposes functions to convert protos to/from text proto format.
 * It guarantees to have the same behaviour of perfetto_cmd and trace_processor
 * by using precisely the same code via WebAssembly.
 */
interface WasmModule {
  module: PerfettoWasmModule;
  buf: Uint8Array;
}

let moduleInstance: WasmModule | undefined = undefined;

/**
 * Generic helper: encodes a proto message to Uint8Array (if needed), sends it
 * to WASM, and returns the text proto output.
 */
async function pbToText<T>(
  input: Uint8Array | T,
  encode: (val: T) => {finish(): Uint8Array},
  ccallName: string,
): Promise<string> {
  const wasm = await initWasmOnce();
  const inputU8: Uint8Array =
    input instanceof Uint8Array ? input : encode(input).finish();
  if (inputU8.length > wasm.buf.length) {
    throw new Error(
      `Proto input too large: ${inputU8.length} bytes exceeds buffer size ${wasm.buf.length}`,
    );
  }
  wasm.buf.set(inputU8);

  const txtSize =
    wasm.module.ccall(ccallName, 'number', ['number'], [inputU8.length]) >>> 0;
  return utf8Decode(wasm.buf.subarray(0, txtSize));
}

/**
 * Convert a binary-encoded protos.TraceConfig to pbtxt (i.e. the text format
 * that can be passed to perfetto --txt).
 */
export async function traceConfigToTxt(
  config: Uint8Array | protos.ITraceConfig,
): Promise<string> {
  return pbToText(config, protos.TraceConfig.encode, 'trace_config_pb_to_txt');
}

/**
 * Generic helper: sends a text proto string to WASM for conversion to binary
 * proto format. Mirrors pbToText but in the opposite direction.
 */
async function textToPb(
  input: string,
  ccallName: string,
): Promise<Result<Uint8Array>> {
  const wasm = await initWasmOnce();

  const inputUtf8 = utf8Encode(input);
  if (inputUtf8.length > wasm.buf.length) {
    return errResult(
      `Text too large: ${inputUtf8.length} bytes exceeds buffer size ${wasm.buf.length}`,
    );
  }
  wasm.buf.set(inputUtf8);

  const resSize =
    wasm.module.ccall(ccallName, 'number', ['number'], [inputUtf8.length]) >>>
    0;

  const success = wasm.buf.at(0) === 1;
  const payload = wasm.buf.slice(1, 1 + resSize);
  return success ? okResult(payload) : errResult(utf8Decode(payload));
}

/** Convert a pbtxt (text-proto) text to a proto-encoded TraceConfig. */
export async function traceConfigToPb(
  configTxt: string,
): Promise<Result<Uint8Array>> {
  return textToPb(configTxt, 'trace_config_txt_to_pb');
}

/**
 * Convert a binary-encoded TraceSummarySpec to text proto format.
 */
export async function traceSummarySpecToText(
  spec: Uint8Array | protos.ITraceSummarySpec,
): Promise<string> {
  return pbToText(
    spec,
    protos.TraceSummarySpec.encode,
    'trace_summary_spec_to_text',
  );
}

/**
 * Convert a text proto (pbtxt) string to a proto-encoded TraceSummarySpec.
 */
export async function traceSummarySpecToPb(
  specTxt: string,
): Promise<Result<Uint8Array>> {
  return textToPb(specTxt, 'trace_summary_spec_txt_to_pb');
}

async function initWasmOnce(): Promise<WasmModule> {
  if (moduleInstance === undefined) {
    // We have to fetch the .wasm file manually because the stub generated by
    // emscripten uses sync-loading, which works only in Workers.
    const resp = await fetch(assetSrc('proto_utils.wasm'));
    const wasmBinary = await resp.arrayBuffer();
    const instance = await WasmModuleGen({
      noInitialRun: true,
      locateFile: (s: string) => s,
      print: (s: string) => console.log(s),
      printErr: (s: string) => console.error(s),
      wasmBinary,
    });
    const bufAddr = instance.ccall('proto_utils_buf', 'number', [], []) >>> 0;
    const bufSize =
      instance.ccall('proto_utils_buf_size', 'number', [], []) >>> 0;
    moduleInstance = {
      module: instance,
      buf: instance.HEAPU8.subarray(bufAddr, bufAddr + bufSize),
    };
  }
  return moduleInstance;
}