as-labs 0.1.0

package/branch-hinting/transform/index.ts

@@ -0,0 +1,584 @@
+import { decode } from "@webassemblyjs/wasm-parser";
+import {
+  CallExpression,
+  FunctionDeclaration,
+  IdentifierExpression,
+  MethodDeclaration,
+  Node,
+  NodeKind,
+  ParenthesizedExpression,
+  SourceKind,
+  getFunctionName,
+} from "assemblyscript/dist/assemblyscript.js";
+import type {
+  Expression,
+  Program,
+  Source,
+} from "assemblyscript/dist/assemblyscript.js";
+import type binaryen from "assemblyscript/lib/binaryen.js";
+import { Transform } from "assemblyscript/dist/transform.js";
+
+const CUSTOM_SECTION_NAME = "metadata.code.branch_hint";
+
+type HintValue = 0 | 1;
+
+type HintRecord = {
+  value: HintValue;
+};
+
+type ResolvedHint = {
+  offset: number;
+  value: HintValue;
+};
+
+type FunctionHintGroup = {
+  functionIndex: number;
+  hints: ResolvedHint[];
+};
+
+type SectionInfo = {
+  id: number;
+  start: number;
+  payloadOffset: number;
+  end: number;
+};
+
+type U32 = {
+  value: number;
+  nextOffset: number;
+};
+
+type BranchHintCall = {
+  value: HintValue;
+  expression: Expression;
+};
+
+type BinaryWithSourceMap = {
+  binary: Uint8Array;
+  sourceMap?: Uint8Array | string | null;
+};
+
+type BinaryEmitter = (...args: unknown[]) => BinaryWithSourceMap;
+
+type ConditionalStatement = {
+  condition: Expression;
+};
+
+export default class BranchHintTransform extends Transform {
+  private functionHintsByDeclaration = new Map<
+    FunctionDeclaration | MethodDeclaration,
+    HintRecord[]
+  >();
+  private functionHintsByInternalName = new Map<string, HintRecord[]>();
+
+  afterParse(parser: { sources: Source[] }): void {
+    for (const source of parser.sources) {
+      if (source.sourceKind === SourceKind.Library) continue;
+      if (source.internalPath.startsWith("~lib/")) continue;
+      for (const statement of source.statements) {
+        this.visitNode(statement, null);
+      }
+    }
+  }
+
+  afterInitialize(program: Program): void {
+    for (const [declaration, hints] of this.functionHintsByDeclaration) {
+      const element = program.elementsByDeclaration.get(declaration);
+      if (!element || !element.internalName || hints.length === 0) continue;
+      this.functionHintsByInternalName.set(element.internalName, hints);
+    }
+  }
+
+  afterCompile(module: binaryen.Module): void {
+    if (this.functionHintsByInternalName.size === 0) return;
+
+    const originalEmitBinary = module.emitBinary.bind(module) as BinaryEmitter;
+    module.emitBinary = ((...args: unknown[]) => {
+      const result = originalEmitBinary(...args);
+      result.binary = injectBranchHints(
+        result.binary,
+        module,
+        this.functionHintsByInternalName,
+      );
+      return result;
+    }) as typeof module.emitBinary;
+  }
+
+  private visitNode(
+    node: Node | Node[] | null | undefined,
+    currentFunction: FunctionDeclaration | MethodDeclaration | null,
+  ): void {
+    if (!node) return;
+    if (Array.isArray(node)) {
+      for (const item of node) this.visitNode(item, currentFunction);
+      return;
+    }
+
+    switch (node.kind) {
+      case NodeKind.FunctionDeclaration:
+        this.visitFunctionLike(node as FunctionDeclaration, node as FunctionDeclaration);
+        return;
+      case NodeKind.MethodDeclaration:
+        this.visitFunctionLike(node as MethodDeclaration, node as MethodDeclaration);
+        return;
+      case NodeKind.If:
+        this.visitConditionalStatement(
+          node as unknown as ConditionalStatement,
+          currentFunction,
+          [(node as any).ifTrue, (node as any).ifFalse],
+        );
+        return;
+      case NodeKind.While:
+        this.visitConditionalStatement(
+          node as unknown as ConditionalStatement,
+          currentFunction,
+          [(node as any).body],
+        );
+        return;
+      case NodeKind.Do:
+        this.visitConditionalStatement(
+          node as unknown as ConditionalStatement,
+          currentFunction,
+          [(node as any).body],
+        );
+        return;
+      case NodeKind.For:
+        this.visitConditionalStatement(
+          node as unknown as ConditionalStatement,
+          currentFunction,
+          [(node as any).initializer, (node as any).incrementor, (node as any).body],
+        );
+        return;
+      case NodeKind.Block:
+        this.visitNode((node as any).statements, currentFunction);
+        return;
+      case NodeKind.Switch:
+        this.visitNode((node as any).condition, currentFunction);
+        this.visitNode((node as any).cases, currentFunction);
+        return;
+      case NodeKind.SwitchCase:
+        this.visitNode((node as any).label, currentFunction);
+        this.visitNode((node as any).statements, currentFunction);
+        return;
+      case NodeKind.Expression:
+        this.visitNode((node as any).expression, currentFunction);
+        return;
+      case NodeKind.Return:
+        this.visitNode((node as any).value, currentFunction);
+        return;
+      case NodeKind.Variable:
+        this.visitNode((node as any).declarations, currentFunction);
+        return;
+      case NodeKind.VariableDeclaration:
+        this.visitNode((node as any).initializer, currentFunction);
+        return;
+      case NodeKind.Call:
+        this.visitNode((node as any).expression, currentFunction);
+        this.visitNode((node as any).args, currentFunction);
+        return;
+      case NodeKind.Binary:
+        this.visitNode((node as any).left, currentFunction);
+        this.visitNode((node as any).right, currentFunction);
+        return;
+      case NodeKind.UnaryPrefix:
+      case NodeKind.UnaryPostfix:
+      case NodeKind.Parenthesized:
+      case NodeKind.Assertion:
+        this.visitNode((node as any).operand ?? (node as any).expression, currentFunction);
+        return;
+      case NodeKind.Ternary:
+        this.visitNode((node as any).condition, currentFunction);
+        this.visitNode((node as any).ifThen, currentFunction);
+        this.visitNode((node as any).ifElse, currentFunction);
+        return;
+      case NodeKind.ForOf:
+        this.visitNode((node as any).variable, currentFunction);
+        this.visitNode((node as any).iterable, currentFunction);
+        this.visitNode((node as any).body, currentFunction);
+        return;
+      case NodeKind.Try:
+        this.visitNode((node as any).bodyStatements, currentFunction);
+        this.visitNode((node as any).catchStatements, currentFunction);
+        this.visitNode((node as any).finallyStatements, currentFunction);
+        return;
+      case NodeKind.Throw:
+        this.visitNode((node as any).value, currentFunction);
+        return;
+      case NodeKind.Void:
+        this.visitNode((node as any).expression, currentFunction);
+        return;
+      default:
+        return;
+    }
+  }
+
+  private visitFunctionLike(
+    node: FunctionDeclaration | MethodDeclaration,
+    declaration: FunctionDeclaration | MethodDeclaration,
+  ): void {
+    if (!node.body) return;
+    this.visitNode(node.body, declaration);
+  }
+
+  private visitConditionalStatement(
+    node: ConditionalStatement,
+    currentFunction: FunctionDeclaration | MethodDeclaration | null,
+    extraChildren: Array<Node | Node[] | null | undefined>,
+  ): void {
+    if (!currentFunction) return;
+
+    const hint = unwrapHint(node.condition);
+    if (hint) {
+      node.condition = hint.expression;
+      let hints = this.functionHintsByDeclaration.get(currentFunction);
+      if (!hints) {
+        hints = [];
+        this.functionHintsByDeclaration.set(currentFunction, hints);
+      }
+      hints.push({ value: hint.value });
+    } else {
+      this.visitNode(node.condition, currentFunction);
+    }
+
+    for (const child of extraChildren) {
+      this.visitNode(child, currentFunction);
+    }
+  }
+}
+
+function unwrapHint(expression: Expression): BranchHintCall | null {
+  let current: Expression = expression;
+  while (current instanceof ParenthesizedExpression) {
+    current = current.expression;
+  }
+  if (!(current instanceof CallExpression)) return null;
+
+  const callee = current.expression;
+  if (!(callee instanceof IdentifierExpression)) return null;
+  if (current.args.length !== 1) return null;
+
+  if (callee.text === "likely") {
+    return { value: 1, expression: current.args[0] };
+  }
+  if (callee.text === "unlikely") {
+    return { value: 0, expression: current.args[0] };
+  }
+  return null;
+}
+
+function injectBranchHints(
+  binary: Uint8Array,
+  module: binaryen.Module,
+  functionHintsByInternalName: Map<string, HintRecord[]>,
+): Uint8Array {
+  const importedFunctionCount = countImportedFunctions(binary);
+  const functionOrder: string[] = [];
+  for (let i = 0, n = module.getNumFunctions(); i < n; i++) {
+    const name = getFunctionName(module.getFunctionByIndex(i));
+    functionOrder.push(name ?? "");
+  }
+
+  const conditionalOffsets = collectConditionalOffsets(binary);
+  const functionHints: FunctionHintGroup[] = [];
+
+  for (let definedIndex = 0; definedIndex < conditionalOffsets.length; definedIndex++) {
+    const internalName = functionOrder[importedFunctionCount + definedIndex];
+    const wanted = functionHintsByInternalName.get(internalName);
+    if (!wanted || wanted.length === 0) continue;
+
+    const available = conditionalOffsets[definedIndex] || [];
+    const matched: ResolvedHint[] = [];
+    const count = Math.min(wanted.length, available.length);
+    for (let i = 0; i < count; i++) {
+      matched.push({
+        offset: available[i],
+        value: wanted[i].value,
+      });
+    }
+    if (matched.length === 0) continue;
+
+    functionHints.push({
+      functionIndex: importedFunctionCount + definedIndex,
+      hints: matched,
+    });
+  }
+
+  if (functionHints.length === 0) return binary;
+
+  const payload = encodeBranchHintPayload(functionHints);
+  const customSection = encodeCustomSection(CUSTOM_SECTION_NAME, payload);
+  return insertCustomSectionBeforeCode(binary, customSection, CUSTOM_SECTION_NAME);
+}
+
+function countImportedFunctions(binary: Uint8Array): number {
+  const sections = parseSections(binary);
+  const importSection = sections.find((section) => section.id === 2);
+  if (!importSection) return 0;
+
+  let offset = importSection.payloadOffset;
+  const end = importSection.end;
+  const countInfo = readU32(binary, offset);
+  offset = countInfo.nextOffset;
+  let functionImports = 0;
+
+  for (let i = 0; i < countInfo.value && offset < end; i++) {
+    const moduleName = readName(binary, offset);
+    offset = moduleName.nextOffset;
+    const fieldName = readName(binary, offset);
+    offset = fieldName.nextOffset;
+    const kind = binary[offset++] as number;
+
+    switch (kind) {
+      case 0:
+        functionImports++;
+        offset = readU32(binary, offset).nextOffset;
+        break;
+      case 1:
+        offset = skipTableType(binary, offset);
+        break;
+      case 2:
+        offset = skipMemoryType(binary, offset);
+        break;
+      case 3:
+        offset = skipGlobalType(binary, offset);
+        break;
+      case 4:
+        offset = skipTagType(binary, offset);
+        break;
+      default:
+        throw new Error(`Unsupported import kind ${kind}`);
+    }
+  }
+
+  return functionImports;
+}
+
+function collectConditionalOffsets(binary: Uint8Array): number[][] {
+  const ast = decode(binary, {
+    dump: false,
+    ignoreCodeSection: false,
+    ignoreDataSection: true,
+  });
+
+  const functions: any[] = [];
+  walk(ast, (node: any) => {
+    if (node?.type === "Func") {
+      functions.push(node);
+    }
+  });
+
+  return functions.map((func) => {
+    const bodySize = readU32(binary, func.loc.start.column as number);
+    const localsOffset = bodySize.nextOffset;
+    const offsets: number[] = [];
+
+    walk(func.body, (node: any) => {
+      if (!node || !node.loc) return;
+      if (node.type === "IfInstruction") {
+        offsets.push((node.loc.start.column as number) - localsOffset - 1);
+      } else if (node.type === "Instr" && node.id === "br_if") {
+        offsets.push((node.loc.start.column as number) - localsOffset);
+      }
+    });
+
+    return offsets.filter((offset) => offset >= 0);
+  });
+}
+
+function walk(node: unknown, visit: (node: unknown) => void): void {
+  if (!node || typeof node !== "object") return;
+  if (Array.isArray(node)) {
+    for (const item of node) walk(item, visit);
+    return;
+  }
+
+  visit(node);
+  for (const [key, value] of Object.entries(node)) {
+    if (key === "loc") continue;
+    walk(value, visit);
+  }
+}
+
+function encodeBranchHintPayload(functionHints: FunctionHintGroup[]): Uint8Array {
+  const bytes: number[] = [];
+  writeU32(bytes, functionHints.length);
+
+  for (const { functionIndex, hints } of functionHints) {
+    writeU32(bytes, functionIndex);
+    writeU32(bytes, hints.length);
+    for (const hint of hints) {
+      writeU32(bytes, hint.offset);
+      writeU32(bytes, 1);
+      writeU32(bytes, hint.value);
+    }
+  }
+
+  return Uint8Array.from(bytes);
+}
+
+function encodeCustomSection(name: string, payload: Uint8Array): Uint8Array {
+  const nameBytes = new TextEncoder().encode(name);
+  const contents: number[] = [];
+  writeU32(contents, nameBytes.length);
+  contents.push(...nameBytes);
+  contents.push(...payload);
+
+  const section: number[] = [];
+  section.push(0);
+  writeU32(section, contents.length);
+  section.push(...contents);
+  return Uint8Array.from(section);
+}
+
+function insertCustomSectionBeforeCode(
+  binary: Uint8Array,
+  customSection: Uint8Array,
+  sectionName: string,
+): Uint8Array {
+  const sections = parseSections(binary);
+  const filtered = sections.filter((section) => {
+    if (section.id !== 0) return true;
+    return readCustomSectionName(binary, section) !== sectionName;
+  });
+
+  const codeSection = filtered.find((section) => section.id === 10);
+  if (!codeSection) {
+    throw new Error("WebAssembly module has no code section");
+  }
+
+  const output: number[] = [];
+  output.push(...binary.slice(0, 8));
+
+  for (const section of filtered) {
+    if (section.start === codeSection.start) {
+      output.push(...customSection);
+    }
+    output.push(...binary.slice(section.start, section.end));
+  }
+
+  return Uint8Array.from(output);
+}
+
+function parseSections(binary: Uint8Array): SectionInfo[] {
+  const sections: SectionInfo[] = [];
+  let offset = 8;
+
+  while (offset < binary.length) {
+    const start = offset;
+    const id = binary[offset++] as number;
+    const size = readU32(binary, offset);
+    offset = size.nextOffset;
+    const payloadOffset = offset;
+    const end = payloadOffset + size.value;
+    sections.push({
+      id,
+      start,
+      payloadOffset,
+      end,
+    });
+    offset = end;
+  }
+
+  return sections;
+}
+
+function readCustomSectionName(binary: Uint8Array, section: SectionInfo): string {
+  const name = readName(binary, section.payloadOffset);
+  return name.value;
+}
+
+function readName(binary: Uint8Array, offset: number): { value: string; nextOffset: number } {
+  const length = readU32(binary, offset);
+  const start = length.nextOffset;
+  const end = start + length.value;
+  return {
+    value: new TextDecoder().decode(binary.slice(start, end)),
+    nextOffset: end,
+  };
+}
+
+function readU32(binary: Uint8Array, offset: number): U32 {
+  let value = 0;
+  let shift = 0;
+  let current = offset;
+
+  while (true) {
+    const byte = binary[current++] as number;
+    value |= (byte & 0x7f) << shift;
+    if ((byte & 0x80) === 0) {
+      return { value: value >>> 0, nextOffset: current };
+    }
+    shift += 7;
+  }
+}
+
+function skipTableType(binary: Uint8Array, offset: number): number {
+  offset = skipReferenceType(binary, offset);
+  return skipLimits(binary, offset);
+}
+
+function skipMemoryType(binary: Uint8Array, offset: number): number {
+  return skipLimits(binary, offset);
+}
+
+function skipGlobalType(binary: Uint8Array, offset: number): number {
+  offset = skipValueType(binary, offset);
+  return offset + 1;
+}
+
+function skipTagType(binary: Uint8Array, offset: number): number {
+  offset += 1;
+  return readU32(binary, offset).nextOffset;
+}
+
+function skipLimits(binary: Uint8Array, offset: number): number {
+  const flags = binary[offset++] as number;
+  offset = readU32(binary, offset).nextOffset;
+  if (flags & 0x01) {
+    offset = readU32(binary, offset).nextOffset;
+  }
+  if (flags & 0x04) {
+    offset = readU32(binary, offset).nextOffset;
+  }
+  return offset;
+}
+
+function skipReferenceType(binary: Uint8Array, offset: number): number {
+  const type = binary[offset] as number;
+  if (type === 0x63 || type === 0x64) {
+    return offset + 1;
+  }
+  return skipHeapType(binary, offset + 1);
+}
+
+function skipValueType(binary: Uint8Array, offset: number): number {
+  const type = binary[offset] as number;
+  if (
+    type === 0x7f ||
+    type === 0x7e ||
+    type === 0x7d ||
+    type === 0x7c ||
+    type === 0x7b ||
+    type === 0x70 ||
+    type === 0x6f
+  ) {
+    return offset + 1;
+  }
+  return skipReferenceType(binary, offset);
+}
+
+function skipHeapType(binary: Uint8Array, offset: number): number {
+  const type = binary[offset] as number;
+  if (type <= 0x7f) return offset + 1;
+  return readU32(binary, offset).nextOffset;
+}
+
+function writeU32(target: number[], value: number): void {
+  let current = value >>> 0;
+  do {
+    let byte = current & 0x7f;
+    current >>>= 7;
+    if (current !== 0) byte |= 0x80;
+    target.push(byte);
+  } while (current !== 0);
+}

package/dist/branch-hinting/transform/index.d.ts

@@ -0,0 +1,15 @@
+import type { Program, Source } from "assemblyscript/dist/assemblyscript.js";
+import type binaryen from "assemblyscript/lib/binaryen.js";
+import { Transform } from "assemblyscript/dist/transform.js";
+export default class BranchHintTransform extends Transform {
+    private functionHintsByDeclaration;
+    private functionHintsByInternalName;
+    afterParse(parser: {
+        sources: Source[];
+    }): void;
+    afterInitialize(program: Program): void;
+    afterCompile(module: binaryen.Module): void;
+    private visitNode;
+    private visitFunctionLike;
+    private visitConditionalStatement;
+}