@smoothbricks/lmao 0.2.2 → 0.2.3

package/src/lib/wasm/wasmTraceRoot.ts

@@ -10,12 +10,25 @@
  */
 
 import { Nanoseconds } from '@smoothbricks/arrow-builder';
+import type { LogSchema } from '../schema/LogSchema.js';
 import { ENTRY_TYPE_SPAN_EXCEPTION, ENTRY_TYPE_SPAN_START } from '../schema/systemSchema.js';
 import { createTraceId, type TraceId } from '../traceId.js';
-import type { ITraceRoot, TracerLifecycleHooks } from '../traceRoot.js';
+import {
+  consumeSpanStartedAtAllocation,
+  type ITraceRoot,
+  type SpanEndPrimitive,
+  type SpanStartPrimitive,
+  type TimestampAppendPrimitive,
+  type TimestampNowPrimitive,
+  type TraceRootFactory,
+  type TracerLifecycleHooks,
+} from '../traceRoot.js';
+import { TraceTopology } from '../traceTopology.js';
 import type { AnySpanBuffer } from '../types.js';
 import type { WasmAllocator } from './wasmAllocator.js';
 
+const traceIdEncoder = new TextEncoder();
+
 // Spec link (88): realizes specs/lmao/01q_wasm_memory_architecture.md#smoo/lmao!n/wasm-mem (span lifecycle writer).
 //#region smoo/lmao!n/wasm-mem.trace-root
 
@@ -30,12 +43,17 @@ import type { WasmAllocator } from './wasmAllocator.js';
 export interface WasmSpanBufferLike {
   /** Byte offset into WASM memory for system columns (timestamp + entry_type) */
   readonly _systemPtr: number;
+  /** Concrete per-op capacity; may differ from the allocator default. */
+  readonly _capacity: number;
   /** Byte offset into WASM memory for identity block (writeIndex, span_id, trace_id) */
   readonly _identityPtr: number;
   /** Message column - string array for span names and log messages */
   readonly _message: string[];
   /** Current write position (getter/setter reads from WASM identity block) */
   _writeIndex: number;
+  readonly _identityOwner: boolean;
+  readonly timestamp: BigInt64Array;
+  readonly entry_type: Uint8Array;
 }
 
 /**
@@ -62,18 +80,96 @@ export function isWasmSpanBuffer(buffer: unknown): buffer is WasmSpanBufferLike
  * The WASM allocator's initTraceRoot() captures these values at trace start,
  * and span lifecycle methods use them to calculate timestamps.
  */
-export class WasmTraceRoot implements ITraceRoot {
+const timestampNow: TimestampNowPrimitive = (traceRoot) => {
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion -- WASM primitive table is only installed on WasmTraceRoot instances.
+  const root = traceRoot as WasmTraceRoot;
+  root._assertLive();
+  const elapsedNanos = BigInt(Math.round((performance.now() - root.anchorPerfNow) * 1_000_000));
+  return Nanoseconds.unsafe(root.anchorEpochNanos + elapsedNanos);
+};
+
+const appendLogEntry: TimestampAppendPrimitive = (traceRoot, buffer, entryType) => {
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion -- WASM primitive table is only installed on WasmTraceRoot instances.
+  const root = traceRoot as WasmTraceRoot;
+  root._assertLive();
+  if (isWasmSpanBuffer(buffer) && buffer._messagePhysicalLayout !== 'packed' && buffer._identityOwner) {
+    return root.allocator.writeLogEntry(
+      buffer._systemPtr,
+      buffer._identityPtr,
+      root._traceRootPtr,
+      entryType,
+      buffer._capacity,
+    );
+  }
+  const idx = buffer._writeIndex;
+  const entryTypes = buffer.entry_type;
+  if (entryTypes === undefined) throw new TypeError('Split WASM appender requires entry_type storage');
+  buffer.timestamp[idx] = timestampNow(root);
+  entryTypes[idx] = entryType;
+  buffer._writeIndex = idx + 1;
+  return idx;
+};
+
+const writeSpanStartPrimitive: SpanStartPrimitive = (traceRoot, buffer, spanName) => {
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion -- WASM primitive table is only installed on WasmTraceRoot instances.
+  const root = traceRoot as WasmTraceRoot;
+  root._assertLive();
+  if (isWasmSpanBuffer(buffer) && buffer._messagePhysicalLayout !== 'packed') {
+    if (!consumeSpanStartedAtAllocation(buffer)) {
+      root.allocator.spanStart(buffer._systemPtr, buffer._identityPtr, root._traceRootPtr, buffer._capacity);
+    }
+    buffer._message[0] = spanName;
+    return;
+  }
+  const entryTypes = buffer.entry_type;
+  if (entryTypes === undefined) throw new TypeError('Split WASM span-start appender requires entry_type storage');
+  buffer.timestamp[0] = timestampNow(root);
+  entryTypes[0] = ENTRY_TYPE_SPAN_START;
+  buffer.message(0, spanName);
+  entryTypes[1] = ENTRY_TYPE_SPAN_EXCEPTION;
+  buffer.timestamp[1] = 0n;
+  buffer._writeIndex = 2;
+};
+
+const writeSpanEndPrimitive: SpanEndPrimitive = (traceRoot, buffer, entryType) => {
+  // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion -- WASM primitive table is only installed on WasmTraceRoot instances.
+  const root = traceRoot as WasmTraceRoot;
+  root._assertLive();
+  if (isWasmSpanBuffer(buffer) && buffer._messagePhysicalLayout !== 'packed') {
+    if (entryType === 2) root.allocator.spanEndOk(buffer._systemPtr, root._traceRootPtr, buffer._capacity);
+    else root.allocator.spanEndErr(buffer._systemPtr, root._traceRootPtr, buffer._capacity);
+    if (entryType !== 2 && entryType !== 3) {
+      const splitEntryTypes = buffer.entry_type;
+      if (splitEntryTypes === undefined) throw new TypeError('Split WASM buffer is missing entry_type storage');
+      splitEntryTypes[1] = entryType;
+    }
+    return;
+  }
+  const entryTypes = buffer.entry_type;
+  if (entryTypes === undefined) throw new TypeError('Split WASM span-end appender requires entry_type storage');
+  buffer.timestamp[1] = timestampNow(root);
+  entryTypes[1] = entryType;
+};
+
+export class WasmTraceRoot<T extends LogSchema = LogSchema> implements ITraceRoot<T> {
   /** Byte offset in WASM memory where TraceRoot data lives */
   readonly _traceRootPtr: number;
+  private _state: 'live' | 'freed' = 'live';
 
   /** The WASM allocator instance */
   readonly allocator: WasmAllocator;
 
   /** Trace ID string */
   readonly trace_id: TraceId;
+  readonly _traceIdBytes: Uint8Array;
 
   /** Tracer reference for lifecycle hooks */
-  readonly tracer: TracerLifecycleHooks;
+  readonly tracer: TracerLifecycleHooks<T>;
+  readonly _topology: TraceTopology;
+  readonly _timestampNow = timestampNow;
+  readonly _appendLogEntry = appendLogEntry;
+  readonly _writeSpanStart = writeSpanStartPrimitive;
+  readonly _writeSpanEnd = writeSpanEndPrimitive;
 
   /**
    * Raw backing buffer - not used for WASM, but required by ITraceRoot.
@@ -81,28 +177,32 @@ export class WasmTraceRoot implements ITraceRoot {
    */
   readonly _system: ArrayBuffer;
 
-  constructor(allocator: WasmAllocator, trace_id: TraceId, tracer: TracerLifecycleHooks) {
+  constructor(allocator: WasmAllocator, trace_id: TraceId, tracer: TracerLifecycleHooks<T>) {
     this.allocator = allocator;
     this.trace_id = trace_id;
+    this._traceIdBytes = traceIdEncoder.encode(trace_id);
     this.tracer = tracer;
+    this._topology = new TraceTopology();
 
-    // Create minimal _system buffer (not actually used, but required by interface)
     this._system = new ArrayBuffer(16);
-
-    // Allocate an 8B column block for TraceRoot data (16 bytes needed, 8B block is large enough)
-    // The 8B column block size is: ceil(capacity/8) + capacity*8 = 8 + 512 = 520 bytes for capacity 64
-    // We only need 16 bytes, so this works fine
-    this._traceRootPtr = allocator.alloc8B();
-
-    // Initialize wall clock and monotonic start times in WASM memory
+    // Native TraceRoot is exactly two aligned 8-byte fields.
+    const traceRootPtr = allocator.allocExact(16, 8);
+    if (traceRootPtr === 0) {
+      throw new Error('WASM trace-root allocation failed');
+    }
+    this._traceRootPtr = traceRootPtr;
     allocator.initTraceRoot(this._traceRootPtr);
   }
+  _assertLive(): void {
+    if (this._state !== 'live') throw new Error('WASM trace root has been released');
+  }
 
   /**
    * Epoch time in nanoseconds when trace was created.
    * Read from WASM memory at _traceRootPtr offset 0.
    */
   get anchorEpochNanos(): bigint {
+    this._assertLive();
     // i64 view is indexed by 8-byte chunks
     return this.allocator.i64[this._traceRootPtr / 8];
   }
@@ -112,6 +212,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * Read from WASM memory at _traceRootPtr offset 8.
    */
   get anchorPerfNow(): number {
+    this._assertLive();
     // f64 view is indexed by 8-byte chunks
     return this.allocator.f64[(this._traceRootPtr + 8) / 8];
   }
@@ -126,9 +227,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * We convert to nanoseconds for consistency with the trace format.
    */
   getTimestampNanos(): Nanoseconds {
-    const elapsedMs = performance.now() - this.anchorPerfNow;
-    const elapsedNanos = BigInt(Math.round(elapsedMs * 1_000_000));
-    return Nanoseconds.unsafe(this.anchorEpochNanos + elapsedNanos);
+    return this._timestampNow(this);
   }
 
   // ===========================================================================
@@ -147,24 +246,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * @param spanName - Name for this span
    */
   writeSpanStart(buffer: AnySpanBuffer, spanName: string): void {
-    if (isWasmSpanBuffer(buffer)) {
-      // WASM path: delegate to WASM allocator
-      // spanStart writes row 0 (span-start) and row 1 (span-exception placeholder)
-      // Also sets writeIndex to 2 in the identity block
-      this.allocator.spanStart(buffer._systemPtr, buffer._identityPtr, this._traceRootPtr);
-      // Write span name to message column (JS string array)
-      buffer._message[0] = spanName;
-      // Note: _writeIndex is set by WASM spanStart to 2, but the getter reads from WASM
-      // so we don't need to set it here explicitly (the setter would write to WASM anyway)
-    } else {
-      // JS path: write directly to TypedArrays (fallback for mixed usage)
-      buffer.timestamp[0] = this.getTimestampNanos();
-      buffer.entry_type[0] = ENTRY_TYPE_SPAN_START;
-      buffer.message(0, spanName);
-      buffer.entry_type[1] = ENTRY_TYPE_SPAN_EXCEPTION;
-      buffer.timestamp[1] = 0n;
-      buffer._writeIndex = 2;
-    }
+    this._writeSpanStart(this, buffer, spanName);
   }
 
   /**
@@ -175,29 +257,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * @param entryType - Entry type (SPAN_OK, SPAN_ERR, or SPAN_EXCEPTION)
    */
   writeSpanEnd(buffer: AnySpanBuffer, entryType: number): void {
-    if (isWasmSpanBuffer(buffer)) {
-      // WASM path: use appropriate end function based on entry type
-      // Note: WASM spanEndOk/spanEndErr write to row 1
-      if (entryType === 2) {
-        // SPAN_OK
-        this.allocator.spanEndOk(buffer._systemPtr, this._traceRootPtr);
-      } else if (entryType === 3) {
-        // SPAN_ERR
-        this.allocator.spanEndErr(buffer._systemPtr, this._traceRootPtr);
-      } else {
-        // SPAN_EXCEPTION or other - write via generic log entry
-        // For now, use spanEndErr as fallback (entry_type will be overwritten)
-        this.allocator.spanEndErr(buffer._systemPtr, this._traceRootPtr);
-        // Overwrite entry_type if needed (for EXCEPTION)
-        if (entryType !== 3) {
-          buffer.entry_type[1] = entryType;
-        }
-      }
-    } else {
-      // JS path
-      buffer.timestamp[1] = this.getTimestampNanos();
-      buffer.entry_type[1] = entryType;
-    }
+    this._writeSpanEnd(this, buffer, entryType);
   }
 
   /**
@@ -209,16 +269,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * @returns The row index where entry was written
    */
   writeLogEntry(buffer: AnySpanBuffer, entryType: number): number {
-    if (isWasmSpanBuffer(buffer)) {
-      // WASM path: bump writeIndex, write timestamp + entry_type, return idx
-      return this.allocator.writeLogEntry(buffer._systemPtr, buffer._identityPtr, this._traceRootPtr, entryType);
-    }
-    // JS path: bump writeIndex, write directly, return idx
-    const idx = buffer._writeIndex;
-    buffer.timestamp[idx] = this.getTimestampNanos();
-    buffer.entry_type[idx] = entryType;
-    buffer._writeIndex = idx + 1;
-    return idx;
+    return this._appendLogEntry(this, buffer, entryType);
   }
 
   // ===========================================================================
@@ -235,6 +286,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * @param identityPtr - Byte offset of buffer's identity block in WASM memory
    */
   writeSpanStartPtr(systemPtr: number, identityPtr: number): void {
+    this._assertLive();
     this.allocator.spanStart(systemPtr, identityPtr, this._traceRootPtr);
   }
 
@@ -243,6 +295,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * @param systemPtr - Byte offset of buffer's system block in WASM memory
    */
   writeSpanEndOkPtr(systemPtr: number): void {
+    this._assertLive();
     this.allocator.spanEndOk(systemPtr, this._traceRootPtr);
   }
 
@@ -251,6 +304,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * @param systemPtr - Byte offset of buffer's system block in WASM memory
    */
   writeSpanEndErrPtr(systemPtr: number): void {
+    this._assertLive();
     this.allocator.spanEndErr(systemPtr, this._traceRootPtr);
   }
 
@@ -264,6 +318,7 @@ export class WasmTraceRoot implements ITraceRoot {
    * @returns Row index where entry was written
    */
   writeLogEntryPtr(systemPtr: number, identityPtr: number, entryType: number): number {
+    this._assertLive();
     return this.allocator.writeLogEntry(systemPtr, identityPtr, this._traceRootPtr, entryType);
   }
 
@@ -272,14 +327,20 @@ export class WasmTraceRoot implements ITraceRoot {
    * @param identityPtr - Byte offset of buffer's identity block in WASM memory
    */
   readWriteIndex(identityPtr: number): number {
+    this._assertLive();
     return this.allocator.readWriteIndex(identityPtr);
   }
 
-  /**
-   * Free the TraceRoot's memory block when trace completes.
-   */
+  /** Mark the handle stale when its allocator is reset wholesale. */
+  invalidate(): void {
+    this._state = 'freed';
+  }
+
+  /** Free the TraceRoot allocation exactly once. */
   free(): void {
-    this.allocator.free8B(this._traceRootPtr);
+    if (this._state === 'freed') return;
+    this.allocator.freeExact(this._traceRootPtr, 16, 8);
+    this._state = 'freed';
   }
 }
 
@@ -295,12 +356,12 @@ export class WasmTraceRoot implements ITraceRoot {
  * @param trace_id - Trace ID for this trace
  * @param tracer - Tracer lifecycle hooks
  */
-export function createWasmTraceRoot(
+export function createWasmTraceRoot<T extends LogSchema = LogSchema>(
   allocator: WasmAllocator,
   trace_id: string,
-  tracer: TracerLifecycleHooks,
-): WasmTraceRoot {
-  return new WasmTraceRoot(allocator, createTraceId(trace_id), tracer);
+  tracer: TracerLifecycleHooks<T>,
+): WasmTraceRoot<T> {
+  return new WasmTraceRoot<T>(allocator, createTraceId(trace_id), tracer);
 }
 
 /**
@@ -311,9 +372,10 @@ export function createWasmTraceRoot(
  * @param allocator - WASM allocator instance
  * @returns Factory function compatible with TraceRootFactory type
  */
-export function createWasmTraceRootFactory(allocator: WasmAllocator) {
-  return (trace_id: string, tracer: TracerLifecycleHooks): WasmTraceRoot => {
-    return new WasmTraceRoot(allocator, createTraceId(trace_id), tracer);
-  };
+export function createWasmTraceRootFactory<T extends LogSchema = LogSchema>(
+  allocator: WasmAllocator,
+): TraceRootFactory<T> {
+  return (trace_id: string, tracer: TracerLifecycleHooks<T>): WasmTraceRoot<T> =>
+    new WasmTraceRoot<T>(allocator, createTraceId(trace_id), tracer);
 }
 //#endregion smoo/lmao!n/wasm-mem.trace-root

package/src/vocabulary/register/v1.ts

@@ -0,0 +1 @@
+import '../../lib/vocabularyRegistry.js';

package/src/wasm.ts

@@ -0,0 +1,19 @@
+export { WasmBufferStrategy, type WasmBufferStrategyOptions } from './lib/wasm/WasmBufferStrategy.js';
+export {
+  createWasmAllocator,
+  createWasmAllocatorSync,
+  WASM_NO_LAYOUT_OFFSET,
+  WASM_SPAN_IDENTITY_CHILD,
+  WASM_SPAN_IDENTITY_ROOT,
+  type WasmAllocator,
+  type WasmAllocatorOptions,
+} from './lib/wasm/wasmAllocator.js';
+export {
+  isWasmSpanBufferInstance,
+  type WasmSpanBufferInstance,
+} from './lib/wasm/wasmSpanBuffer.js';
+export {
+  createWasmTraceRoot,
+  createWasmTraceRootFactory,
+  WasmTraceRoot,
+} from './lib/wasm/wasmTraceRoot.js';