tjs-lang 0.7.7 → 0.8.0

package/src/lang/wasm.ts

@@ -596,6 +596,21 @@ interface TypedParam {
 // Compilation Context
 // ============================================================================
 
+/**
+ * Signature of a wasm function in the current module. Used to resolve
+ * cross-function calls (wasm-to-wasm `call <index>` instructions) without
+ * routing through JS. Populated by compileBlocksToModule before compiling
+ * any individual body, so forward references and mutual recursion work.
+ */
+export interface ModuleFunctionSig {
+  /** Function index in the composed module */
+  index: number
+  /** Parameter types (used for arg-type checking + auto-conversion) */
+  params: TypedParam[]
+  /** Whether this function returns a value (f64) or is void */
+  hasReturn: boolean
+}
+
 interface CompileContext {
   /** Parameter definitions */
   params: TypedParam[]
@@ -621,9 +636,19 @@ interface CompileContext {
   wat: string[]
   /** Current indentation level for WAT */
   watIndent: number
+  /**
+   * Other wasm functions in the same composed module that this body can
+   * call directly via `call <index>` instructions (no JS↔wasm boundary).
+   * Populated when compiling via compileBlocksToModule; empty for the
+   * single-block compileToWasm path.
+   */
+  moduleFunctions: Map<string, ModuleFunctionSig>
 }
 
-function createContext(params: TypedParam[]): CompileContext {
+function createContext(
+  params: TypedParam[],
+  moduleFunctions: Map<string, ModuleFunctionSig> = new Map()
+): CompileContext {
   const ctx: CompileContext = {
     params,
     locals: new Map(),
@@ -637,6 +662,7 @@ function createContext(params: TypedParam[]): CompileContext {
     hasReturn: false,
     wat: [],
     watIndent: 1,
+    moduleFunctions,
   }
 
   // Add params to locals map
@@ -939,6 +965,10 @@ function inferExprType(
         const name = (call.callee as acorn.Identifier).name
         if (name === 'f32x4_extract_lane') return 'f32'
         if (name.startsWith('f32x4_') || name === 'v128_load') return 'v128'
+        // Wasm-to-wasm call: returns f64 if the called function has a return,
+        // or i32 (the dummy 0 pushed for void calls — see compileWasmFunctionCall)
+        const fn = ctx.moduleFunctions.get(name)
+        if (fn) return fn.hasReturn ? 'f64' : 'i32'
       }
       return 'f64'
     }
@@ -1579,9 +1609,27 @@ function compileCall(
 
   // Handle SIMD intrinsics: f32x4_xxx(...), v128_load(...)
   if (node.callee.type === 'Identifier') {
-    const name = (node.callee as acorn.Identifier).name
-    if (name.startsWith('f32x4_') || name.startsWith('v128_')) {
-      return compileSIMDCall(name, node.arguments as acorn.Expression[], ctx)
+    const calleeName = (node.callee as acorn.Identifier).name
+    if (calleeName.startsWith('f32x4_') || calleeName.startsWith('v128_')) {
+      return compileSIMDCall(
+        calleeName,
+        node.arguments as acorn.Expression[],
+        ctx
+      )
+    }
+
+    // Wasm-to-wasm call: the callee is another `wasm function` in the same
+    // composed module. This is the cross-function `call <index>` path —
+    // the consumer's wasm body calls a library kernel directly, with no
+    // JS↔wasm boundary crossing.
+    const fn = ctx.moduleFunctions.get(calleeName)
+    if (fn) {
+      return compileWasmFunctionCall(
+        fn,
+        calleeName,
+        node.arguments as acorn.Expression[],
+        ctx
+      )
     }
   }
 
@@ -1589,6 +1637,92 @@ function compileCall(
   return [Op.f64_const, ...encodeF64(0)]
 }
 
+/**
+ * Emit a `call <index>` to another wasm function in the same module.
+ * Each argument is compiled and, if its inferred type doesn't match the
+ * called function's parameter type, an explicit conversion instruction is
+ * inserted (truncate for f→i, promote/demote for f32↔f64, etc.).
+ * Void-returning calls push a dummy `i32.const 0` so an enclosing
+ * ExpressionStatement's automatic drop has something to discard.
+ */
+function compileWasmFunctionCall(
+  fn: ModuleFunctionSig,
+  name: string,
+  args: acorn.Expression[],
+  ctx: CompileContext
+): number[] {
+  if (args.length !== fn.params.length) {
+    ctx.errors.push(
+      `wasm function ${name} expects ${fn.params.length} arguments, got ${args.length}`
+    )
+    return [Op.f64_const, ...encodeF64(0)]
+  }
+
+  const code: number[] = []
+  for (let i = 0; i < args.length; i++) {
+    const arg = args[i]
+    const paramType = fn.params[i].type
+    const argType = inferExprType(arg, ctx)
+    code.push(...compileExpression(arg, ctx))
+    // Insert conversion when arg type doesn't match param type.
+    // Source target: argType → paramType.
+    if (argType !== paramType) {
+      const conv = convertOp(argType, paramType)
+      if (conv === undefined) {
+        ctx.errors.push(
+          `wasm function ${name} param ${i} expects ${paramType}, got ${argType} (no conversion available)`
+        )
+      } else if (conv !== null) {
+        code.push(conv)
+      }
+    }
+  }
+
+  code.push(Op.call, ...encodeULEB128(fn.index))
+
+  // Void calls don't leave a value on the stack — push a dummy i32 0 so
+  // ExpressionStatement's drop has something to remove. Same pattern used
+  // by f32x4_store (line ~1632) for stores that have no return value.
+  if (!fn.hasReturn) {
+    code.push(Op.i32_const, 0)
+  }
+
+  return code
+}
+
+/**
+ * Return the wasm opcode that converts `from` to `to`, or:
+ *   - `null` when no conversion is needed (types match — but caller should
+ *     have already checked this; included for safety)
+ *   - `undefined` when no conversion is available (e.g. anything ↔ v128)
+ */
+function convertOp(
+  from: WasmValueType,
+  to: WasmValueType
+): number | null | undefined {
+  if (from === to) return null
+  // v128 isn't convertible to/from scalar types
+  if (from === 'v128' || to === 'v128') return undefined
+  // i64 conversions aren't in the current opcode table — skip for now
+  if (from === 'i64' || to === 'i64') return undefined
+  switch (`${from}->${to}`) {
+    case 'f64->i32':
+      return Op.i32_trunc_f64_s
+    case 'f32->i32':
+      return Op.i32_trunc_f32_s
+    case 'i32->f64':
+      return Op.f64_convert_i32_s
+    case 'i32->f32':
+      return Op.f32_convert_i32_s
+    case 'f32->f64':
+      return Op.f64_promote_f32
+    case 'f64->f32':
+      return Op.f32_demote_f64
+    default:
+      return undefined
+  }
+}
+
 /** Compile SIMD intrinsic calls */
 function compileSIMDCall(
   name: string,
@@ -1802,63 +1936,35 @@ function parseTypeAnnotation(capture: string): TypedParam {
   return { name, type: typeMap[typeStr] ?? 'f64' }
 }
 
-/** Build a complete WASM module */
-function buildModule(
-  params: TypedParam[],
-  bodyCode: number[],
-  localTypes: WasmValueType[],
-  needsMemory: boolean,
+/**
+ * Per-function intermediate: everything needed to embed one function inside
+ * a (single- or multi-function) WASM module.
+ */
+interface CompiledFunction {
+  params: TypedParam[]
+  bodyCode: number[]
+  localTypes: WasmValueType[]
+  needsMemory: boolean
   hasReturn: boolean
-): number[] {
-  // Magic number and version
-  const header = [0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00]
+}
 
-  // Type section: function signature
+/** Encode a single function signature as a Type-section entry body */
+function encodeFuncType(params: TypedParam[], hasReturn: boolean): number[] {
   const paramWasmTypes = params.map((p) => Type[p.type])
   const returnSpec = hasReturn ? [0x01, Type.f64] : [0x00] // one f64 return OR void
-  const typeSection = encodeSection(Section.type, [
-    0x01, // one type
+  return [
     0x60, // func type
     ...encodeULEB128(params.length),
     ...paramWasmTypes,
     ...returnSpec,
-  ])
-
-  // Memory section (if needed)
-  const memorySection: number[] = []
-  if (needsMemory) {
-    // Import memory from JS instead of declaring it
-    // This lets us share memory with typed arrays
-  }
-
-  // Import section for memory
-  let importSection: number[] = []
-  if (needsMemory) {
-    importSection = encodeSection(Section.import, [
-      0x01, // one import
-      ...encodeString('env'),
-      ...encodeString('memory'),
-      0x02, // memory
-      0x00, // flags: no max
-      0x01, // initial: 1 page (64KB)
-    ])
-  }
-
-  // Function section: function 0 has type 0
-  const funcSection = encodeSection(Section.function, [
-    0x01, // one function
-    0x00, // type index 0
-  ])
-
-  // Export section: export function as "compute"
-  const exportSection = encodeSection(Section.export, [
-    0x01, // one export
-    ...encodeString('compute'),
-    0x00, // export kind: function
-    0x00, // function index 0
-  ])
+  ]
+}
 
-  // Code section
+/** Encode the locals + body bytes for one function as a Code-section entry */
+function encodeFuncBody(
+  bodyCode: number[],
+  localTypes: WasmValueType[]
+): number[] {
   // Encode locals: group by type
   const localGroups: number[][] = []
   if (localTypes.length > 0) {
@@ -1883,24 +1989,98 @@ function buildModule(
 
   const funcBody = [...localsEncoded, ...bodyCode, Op.end]
 
+  // The Code-section entry is a length-prefixed sequence of (locals + body)
+  return [...encodeULEB128(funcBody.length), ...funcBody]
+}
+
+/** Encode the memory-import section (used when any function needs memory) */
+function encodeMemoryImport(): number[] {
+  return encodeSection(Section.import, [
+    0x01, // one import
+    ...encodeString('env'),
+    ...encodeString('memory'),
+    0x02, // memory
+    0x00, // flags: no max
+    0x01, // initial: 1 page (64KB)
+  ])
+}
+
+/** Build a complete WASM module containing N functions. */
+function buildMultiFunctionModule(
+  functions: CompiledFunction[],
+  exportNames: string[]
+): number[] {
+  if (functions.length !== exportNames.length) {
+    throw new Error('functions and exportNames length mismatch')
+  }
+
+  // Magic number and version
+  const header = [0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00]
+
+  // Type section: one entry per function (no dedup; modules are tiny)
+  const typeEntries = functions.map((f) =>
+    encodeFuncType(f.params, f.hasReturn)
+  )
+  const typeSection = encodeSection(Section.type, [
+    ...encodeULEB128(typeEntries.length),
+    ...typeEntries.flat(),
+  ])
+
+  // Import section: shared memory if any function needs it
+  const anyNeedsMemory = functions.some((f) => f.needsMemory)
+  const importSection = anyNeedsMemory ? encodeMemoryImport() : []
+
+  // Function section: each function references its own type by index
+  const funcSection = encodeSection(Section.function, [
+    ...encodeULEB128(functions.length),
+    ...functions.map((_, i) => encodeULEB128(i)).flat(),
+  ])
+
+  // Export section: each function exported under its given name
+  const exportSection = encodeSection(Section.export, [
+    ...encodeULEB128(functions.length),
+    ...exportNames
+      .map((name, i) => [
+        ...encodeString(name),
+        0x00, // export kind: function
+        ...encodeULEB128(i), // function index
+      ])
+      .flat(),
+  ])
+
+  // Code section: one body per function
+  const codeBodies = functions.map((f) =>
+    encodeFuncBody(f.bodyCode, f.localTypes)
+  )
   const codeSection = encodeSection(Section.code, [
-    0x01, // one function
-    ...encodeULEB128(funcBody.length),
-    ...funcBody,
+    ...encodeULEB128(codeBodies.length),
+    ...codeBodies.flat(),
   ])
 
   // Assemble module
   const sections = [...header, ...typeSection]
-
-  if (importSection.length > 0) {
-    sections.push(...importSection)
-  }
-
+  if (importSection.length > 0) sections.push(...importSection)
   sections.push(...funcSection, ...exportSection, ...codeSection)
-
   return sections
 }
 
+/**
+ * Build a single-function module. Kept for the per-block path (legacy
+ * `compileToWasm`, `createWasmFunction`). Always exports `compute`.
+ */
+function buildModule(
+  params: TypedParam[],
+  bodyCode: number[],
+  localTypes: WasmValueType[],
+  needsMemory: boolean,
+  hasReturn: boolean
+): number[] {
+  return buildMultiFunctionModule(
+    [{ params, bodyCode, localTypes, needsMemory, hasReturn }],
+    ['compute']
+  )
+}
+
 // ============================================================================
 // Public API
 // ============================================================================
@@ -1922,9 +2102,36 @@ export interface WasmCompileResult {
 }
 
 /**
- * Compile a WASM block to WebAssembly
+ * Per-block compile result: the intermediate before module wrapping.
+ * Used to compose multiple blocks into a single multi-function module.
  */
-export function compileToWasm(block: WasmBlock): WasmCompileResult {
+interface BlockCompileResult {
+  success: boolean
+  /** Compiled function pieces (when success === true) */
+  fn?: CompiledFunction
+  /** WAT disassembly (when success === true) */
+  wat?: string
+  /** Warnings (always present) */
+  warnings: string[]
+  /** Error message (when success === false) */
+  error?: string
+}
+
+/**
+ * Compile a WASM block to its function-level intermediate (params, bytecode,
+ * locals, etc.) WITHOUT wrapping in a module. Used by both the single-block
+ * path (`compileToWasm`) and the multi-block path (`compileBlocksToModule`).
+ *
+ * @param moduleFunctions Map of other wasm functions in the same composed
+ *   module. When the body calls one of these by name, the compiler emits a
+ *   `call <index>` instruction instead of treating it as an unknown
+ *   identifier. Empty by default (single-block compilation has nothing to
+ *   call into).
+ */
+function compileBlockToFunction(
+  block: WasmBlock,
+  moduleFunctions: Map<string, ModuleFunctionSig> = new Map()
+): BlockCompileResult {
   try {
     // Parse type annotations from captures
     const params = block.captures.map(parseTypeAnnotation)
@@ -1939,9 +2146,8 @@ export function compileToWasm(block: WasmBlock): WasmCompileResult {
       ast = acorn.parse(wrapped, { ecmaVersion: 2022 }) as acorn.Program
     } catch (e: any) {
       return {
-        bytes: new Uint8Array(),
-        warnings: [],
         success: false,
+        warnings: [],
         error: `Parse error: ${e.message}`,
       }
     }
@@ -1951,7 +2157,7 @@ export function compileToWasm(block: WasmBlock): WasmCompileResult {
     const body = funcDecl.body.body
 
     // Create compilation context
-    const ctx = createContext(params)
+    const ctx = createContext(params, moduleFunctions)
 
     // Compile statements
     const code: number[] = []
@@ -1959,45 +2165,222 @@ export function compileToWasm(block: WasmBlock): WasmCompileResult {
       code.push(...compileStatement(stmt, ctx))
     }
 
-    // Check for errors
     if (ctx.errors.length > 0) {
       return {
-        bytes: new Uint8Array(),
-        warnings: ctx.warnings,
         success: false,
+        warnings: ctx.warnings,
         error: ctx.errors.join('; '),
       }
     }
 
-    // Build the module
-    const moduleBytes = buildModule(
-      params,
-      code,
-      ctx.localTypes,
-      ctx.needsMemory,
-      ctx.hasReturn
-    )
-
-    // Generate WAT disassembly for debugging
-    const watText = disassemble(code, params, ctx.localTypes)
-
     return {
-      bytes: new Uint8Array(moduleBytes),
-      warnings: ctx.warnings,
       success: true,
-      needsMemory: ctx.needsMemory,
-      wat: watText,
+      fn: {
+        params,
+        bodyCode: code,
+        localTypes: ctx.localTypes,
+        needsMemory: ctx.needsMemory,
+        hasReturn: ctx.hasReturn,
+      },
+      wat: disassemble(code, params, ctx.localTypes),
+      warnings: ctx.warnings,
     }
   } catch (e: any) {
     return {
-      bytes: new Uint8Array(),
-      warnings: [],
       success: false,
+      warnings: [],
       error: e.message,
     }
   }
 }
 
+/**
+ * Compile a single WASM block to a complete WebAssembly module.
+ * The module exports a single function named `compute`.
+ */
+export function compileToWasm(block: WasmBlock): WasmCompileResult {
+  const r = compileBlockToFunction(block)
+  if (!r.success || !r.fn) {
+    return {
+      bytes: new Uint8Array(),
+      warnings: r.warnings,
+      success: false,
+      error: r.error,
+    }
+  }
+  const moduleBytes = buildModule(
+    r.fn.params,
+    r.fn.bodyCode,
+    r.fn.localTypes,
+    r.fn.needsMemory,
+    r.fn.hasReturn
+  )
+  return {
+    bytes: new Uint8Array(moduleBytes),
+    warnings: r.warnings,
+    success: true,
+    needsMemory: r.fn.needsMemory,
+    wat: r.wat,
+  }
+}
+
+// ============================================================================
+// Multi-block module composition (one module, N exports)
+// ============================================================================
+
+/** Per-export metadata produced by compileBlocksToModule */
+export interface BlockExport {
+  /** Original block ID (assigned by the parser) */
+  id: string
+  /** Export name in the composed module (e.g. 'compute_0') */
+  exportName: string
+  /** Capture annotations (preserved for runtime wrapper) */
+  captures: string[]
+  /** Whether this function reads/writes memory */
+  needsMemory: boolean
+  /** WAT disassembly */
+  wat: string
+}
+
+/** Result of composing multiple blocks into one module */
+export interface MultiBlockCompileResult {
+  /** The composed module bytes, or empty if all blocks failed */
+  bytes: Uint8Array
+  /** Per-block compile status (preserves input order) */
+  results: {
+    id: string
+    success: boolean
+    error?: string
+    /** Index into `exports` (only when success === true) */
+    exportIndex?: number
+  }[]
+  /** Successfully-compiled exports (in module-index order) */
+  exports: BlockExport[]
+  /** True if any included function needs memory */
+  needsMemory: boolean
+  /** Aggregated warnings from all blocks */
+  warnings: string[]
+}
+
+/**
+ * Compile N WASM blocks into a single WebAssembly module with N exports.
+ * Failed blocks are skipped (their slot in `results` records the error)
+ * but do not abort compilation of the rest.
+ *
+ * Exports are named `compute_0`, `compute_1`, ... in input order, skipping
+ * indices that correspond to failed blocks.
+ */
+export function compileBlocksToModule(
+  blocks: WasmBlock[]
+): MultiBlockCompileResult {
+  // Pre-pass: assign a function index to every block and build the
+  // moduleFunctions map for NAMED wasm functions. Done before compiling
+  // any body, so wasm-to-wasm `call <index>` instructions resolve
+  // regardless of declaration order (forward references and mutual
+  // recursion both work).
+  //
+  // To keep call-instruction indices valid even when individual blocks
+  // fail compilation, indices are assigned densely from the input list
+  // and any failed block is replaced by a stub function with the same
+  // signature (returns 0 for value-returning, no-op for void). Callers
+  // never observe failed indices being skipped or renumbered.
+  const moduleFunctions = new Map<string, ModuleFunctionSig>()
+  const preSignatures: { params: TypedParam[]; hasReturn: boolean }[] = []
+  for (let i = 0; i < blocks.length; i++) {
+    const b = blocks[i]
+    const params = b.captures.map(parseTypeAnnotation)
+    const hasReturn = b.returnType !== undefined
+    preSignatures.push({ params, hasReturn })
+    if (b.name) {
+      moduleFunctions.set(b.name, { index: i, params, hasReturn })
+    }
+  }
+
+  // Pass 2: compile each block with the full map in scope. Failed
+  // blocks are replaced by stubs (same signature, trivial body).
+  const results: MultiBlockCompileResult['results'] = []
+  const compiledFns: CompiledFunction[] = []
+  const exports: BlockExport[] = []
+  const warnings: string[] = []
+
+  for (let i = 0; i < blocks.length; i++) {
+    const block = blocks[i]
+    const r = compileBlockToFunction(block, moduleFunctions)
+    warnings.push(...r.warnings)
+    if (!r.success || !r.fn) {
+      // Emit a stub so the function index stays valid (callers' encoded
+      // `call <i>` instructions still target a real wasm function — it
+      // just returns a default value or does nothing).
+      results.push({ id: block.id, success: false, error: r.error })
+      compiledFns.push(stubFunction(preSignatures[i]))
+      exports.push({
+        id: block.id,
+        exportName: `compute_${i}`,
+        captures: block.captures,
+        needsMemory: false,
+        wat: `(failed: ${r.error ?? 'unknown error'})`,
+      })
+      continue
+    }
+    compiledFns.push(r.fn)
+    exports.push({
+      id: block.id,
+      exportName: `compute_${i}`,
+      captures: block.captures,
+      needsMemory: r.fn.needsMemory,
+      wat: r.wat ?? '',
+    })
+    results.push({ id: block.id, success: true, exportIndex: i })
+  }
+
+  if (compiledFns.length === 0) {
+    return {
+      bytes: new Uint8Array(),
+      results,
+      exports: [],
+      needsMemory: false,
+      warnings,
+    }
+  }
+
+  const moduleBytes = buildMultiFunctionModule(
+    compiledFns,
+    exports.map((e) => e.exportName)
+  )
+
+  return {
+    bytes: new Uint8Array(moduleBytes),
+    results,
+    exports,
+    needsMemory: exports.some((e) => e.needsMemory),
+    warnings,
+  }
+}
+
+/**
+ * Build a stub `CompiledFunction` matching the given signature. Used in
+ * place of failed-compilation results so function indices remain valid
+ * for any wasm-to-wasm calls that target this slot. The stub's behavior
+ * is intentionally bland: return 0.0 for f64-returning, do nothing for
+ * void. (Callers shouldn't end up here if their own compilation
+ * succeeded — failed targets should be reported via the results array.)
+ */
+function stubFunction(sig: {
+  params: TypedParam[]
+  hasReturn: boolean
+}): CompiledFunction {
+  const bodyCode: number[] = sig.hasReturn
+    ? [Op.f64_const, ...encodeF64(0)]
+    : []
+  return {
+    params: sig.params,
+    bodyCode,
+    localTypes: [],
+    needsMemory: false,
+    hasReturn: sig.hasReturn,
+  }
+}
+
 /**
  * Instantiate a compiled WASM module
  */