tjs-lang 0.7.7 → 0.8.0

package/src/lang/wasm.test.ts

@@ -5,6 +5,7 @@
 import { describe, it, expect } from 'bun:test'
 import {
   compileToWasm,
+  compileBlocksToModule,
   instantiateWasm,
   registerWasmBlock,
   createWasmFunction,
@@ -1617,8 +1618,10 @@ function inc(arr: Float32Array, len: 0) {
 }
 `
         const result = tjs(source)
-        // The wrapper should check buffer identity
-        expect(result.code).toContain('a.buffer===mem.buffer')
+        // The wrapper should check buffer identity against the shared memory.
+        // After Phase 0.5 consolidation there's one __wasmMem per file, so
+        // wrappers reference it directly (no per-block `mem` alias).
+        expect(result.code).toContain('a.buffer===__wasmMem.buffer')
       })
     })
 
@@ -1877,3 +1880,1291 @@ function simdScale(arr: Float32Array, len: 0, factor: 0.0) {
     })
   })
 })
+
+describe('module consolidation (Phase 0.5)', () => {
+  describe('compileBlocksToModule', () => {
+    it('emits one module with N exports for N blocks', async () => {
+      const blocks: WasmBlock[] = [
+        {
+          id: 'b0',
+          captures: ['a: f64', 'b: f64'],
+          body: 'return a + b',
+          hasReturn: true,
+        } as WasmBlock,
+        {
+          id: 'b1',
+          captures: ['x: f64', 'y: f64'],
+          body: 'return x * y',
+          hasReturn: true,
+        } as WasmBlock,
+      ]
+
+      const result = compileBlocksToModule(blocks)
+      expect(result.exports).toHaveLength(2)
+      expect(result.exports[0]).toMatchObject({ id: 'b0', exportName: 'compute_0' })
+      expect(result.exports[1]).toMatchObject({ id: 'b1', exportName: 'compute_1' })
+
+      // Instantiate and confirm both exports work
+      const instance = await instantiateWasm(result.bytes)
+      const add = instance.exports.compute_0 as (a: number, b: number) => number
+      const mul = instance.exports.compute_1 as (a: number, b: number) => number
+      expect(add(2, 3)).toBe(5)
+      expect(mul(4, 5)).toBe(20)
+    })
+
+    it('preserves input order in results, including failures', () => {
+      const blocks: WasmBlock[] = [
+        {
+          id: 'ok0',
+          captures: ['a: f64'],
+          body: 'return a + 1',
+          hasReturn: true,
+        } as WasmBlock,
+        {
+          id: 'bad',
+          captures: ['x: f64'],
+          // Syntax error — fails compilation
+          body: 'this is not valid js {{{',
+          hasReturn: true,
+        } as WasmBlock,
+        {
+          id: 'ok1',
+          captures: ['b: f64'],
+          body: 'return b * 2',
+          hasReturn: true,
+        } as WasmBlock,
+      ]
+
+      const result = compileBlocksToModule(blocks)
+      expect(result.results).toHaveLength(3)
+      expect(result.results[0]).toMatchObject({ id: 'ok0', success: true })
+      expect(result.results[1]).toMatchObject({ id: 'bad', success: false })
+      expect(result.results[2]).toMatchObject({ id: 'ok1', success: true })
+
+      // All three blocks appear in exports — failed blocks become stub
+      // functions so that wasm-to-wasm `call <index>` instructions targeting
+      // their slot stay valid. The `results` array distinguishes successes
+      // from failures.
+      expect(result.exports).toHaveLength(3)
+      expect(result.exports.map((e) => e.id)).toEqual(['ok0', 'bad', 'ok1'])
+      // Export indices stay dense and aligned with input order
+      expect(result.exports.map((e) => e.exportName)).toEqual([
+        'compute_0',
+        'compute_1',
+        'compute_2',
+      ])
+    })
+
+    it('produces a valid module even when all blocks fail (stubs only)', () => {
+      // All blocks fail compilation — module is still well-formed with
+      // stub functions in each slot. This preserves index stability for
+      // any callers that might reference these by index.
+      const blocks: WasmBlock[] = [
+        {
+          id: 'bad',
+          captures: ['x: f64'],
+          body: 'this is not valid js {{{',
+          hasReturn: true,
+        } as WasmBlock,
+      ]
+
+      const result = compileBlocksToModule(blocks)
+      expect(result.exports).toHaveLength(1)
+      expect(result.results[0].success).toBe(false)
+      // Module bytes are emitted; the stub is callable but returns 0
+      expect(result.bytes.length).toBeGreaterThan(0)
+    })
+
+    it('imports memory only when at least one function needs it', () => {
+      // Pure-arithmetic block: no memory needed
+      const noMemBlocks: WasmBlock[] = [
+        {
+          id: 'b0',
+          captures: ['a: f64'],
+          body: 'return a + 1',
+          hasReturn: true,
+        } as WasmBlock,
+      ]
+      const noMemResult = compileBlocksToModule(noMemBlocks)
+      expect(noMemResult.needsMemory).toBe(false)
+      // Should instantiate with no imports
+      expect(async () => {
+        await instantiateWasm(noMemResult.bytes)
+      }).not.toThrow()
+
+      // Mixed: one block uses Float32Array (needs memory), one doesn't.
+      // The whole module imports memory; the pure block coexists fine.
+      const mixedBlocks: WasmBlock[] = [
+        {
+          id: 'b0',
+          captures: ['a: f64'],
+          body: 'return a + 1',
+          hasReturn: true,
+        } as WasmBlock,
+        {
+          id: 'b1',
+          captures: ['arr: Float32Array', 'len: f64'],
+          body: 'for (let i = 0; i < len; i++) arr[i] = arr[i] + 1.0',
+          hasReturn: false,
+        } as WasmBlock,
+      ]
+      const mixedResult = compileBlocksToModule(mixedBlocks)
+      expect(mixedResult.needsMemory).toBe(true)
+      expect(mixedResult.exports).toHaveLength(2)
+    })
+
+    it('handles void and value-returning functions in the same module', async () => {
+      const blocks: WasmBlock[] = [
+        {
+          id: 'sum',
+          captures: ['a: f64', 'b: f64'],
+          body: 'return a + b',
+          hasReturn: true,
+        } as WasmBlock,
+        {
+          id: 'noop',
+          captures: ['x: f64'],
+          body: 'let y = x', // No return — void function
+          hasReturn: false,
+        } as WasmBlock,
+      ]
+
+      const result = compileBlocksToModule(blocks)
+      expect(result.exports).toHaveLength(2)
+
+      const instance = await instantiateWasm(result.bytes)
+      const sum = instance.exports.compute_0 as (a: number, b: number) => number
+      const noop = instance.exports.compute_1 as (x: number) => void
+      expect(sum(7, 8)).toBe(15)
+      expect(noop(42)).toBeUndefined()
+    })
+  })
+
+  describe('emitted bootstrap (single WebAssembly.compile per file)', () => {
+    it('two wasm blocks in one source produce ONE compile call', async () => {
+      const { tjs } = await import('./index')
+      const source = `
+function inc(arr: Float32Array, len: 0) {
+  wasm {
+    for (let i = 0; i < len; i++) { arr[i] = arr[i] + 1.0 }
+  } fallback {
+    for (let i = 0; i < len; i++) arr[i] += 1
+  }
+}
+
+function dbl(arr: Float32Array, len: 0) {
+  wasm {
+    for (let i = 0; i < len; i++) { arr[i] = arr[i] * 2.0 }
+  } fallback {
+    for (let i = 0; i < len; i++) arr[i] *= 2
+  }
+}
+`
+      const result = tjs(source)
+
+      // The hallmark of consolidation: exactly one compile call across all blocks
+      const compileCalls = (result.code.match(/WebAssembly\.compile\(/g) || [])
+        .length
+      expect(compileCalls).toBe(1)
+
+      // Both functions appear under their own export names in the module
+      expect(result.code).toContain('"compute_0"')
+      expect(result.code).toContain('"compute_1"')
+    })
+
+    it('two wasm functions actually run after consolidated bootstrap', async () => {
+      const { tjs } = await import('./index')
+      const { createRuntime } = await import('./runtime')
+      const source = `
+function inc(arr: Float32Array, len: 0) {
+  wasm {
+    for (let i = 0; i < len; i++) { arr[i] = arr[i] + 1.0 }
+  } fallback {
+    for (let i = 0; i < len; i++) arr[i] += 1
+  }
+}
+
+function dbl(arr: Float32Array, len: 0) {
+  wasm {
+    for (let i = 0; i < len; i++) { arr[i] = arr[i] * 2.0 }
+  } fallback {
+    for (let i = 0; i < len; i++) arr[i] *= 2
+  }
+}
+`
+      const result = tjs(source)
+      const savedTjs = globalThis.__tjs
+      try {
+        globalThis.__tjs = createRuntime()
+        // Wrap in IIFE so the emitted `const __tjs = ...` doesn't clash with
+        // the outer parameter; expose the user functions via globalThis.
+        await new Function(
+          '__tjs',
+          `return (async () => { ${result.code}\n` +
+            `globalThis.__test_inc = inc;\n` +
+            `globalThis.__test_dbl = dbl;\n` +
+            `})();`
+        )(globalThis.__tjs)
+
+        // Wait for the single async bootstrap (one instantiate) to complete
+        await new Promise((r) => setTimeout(r, 100))
+
+        const wasmBuffer = (globalThis as any).wasmBuffer
+        expect(typeof wasmBuffer).toBe('function')
+
+        const buf = wasmBuffer(Float32Array, 4)
+        buf[0] = 1
+        buf[1] = 2
+        buf[2] = 3
+        buf[3] = 4
+
+        ;(globalThis as any).__test_inc(buf, 4) // [2, 3, 4, 5]
+        expect(Array.from(buf)).toEqual([2, 3, 4, 5])
+
+        ;(globalThis as any).__test_dbl(buf, 4) // [4, 6, 8, 10]
+        expect(Array.from(buf)).toEqual([4, 6, 8, 10])
+      } finally {
+        globalThis.__tjs = savedTjs
+        delete (globalThis as any).wasmBuffer
+        delete (globalThis as any).__test_inc
+        delete (globalThis as any).__test_dbl
+      }
+    })
+  })
+})
+
+describe('wasm function declarations (Phase 1)', () => {
+  it('extracts a top-level wasm function as a WasmBlock', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+wasm function add(a: f64, b: f64): f64 {
+  return a + b
+}
+`
+    const result = tjs(source)
+    expect(result.wasmCompiled).toBeDefined()
+    expect(result.wasmCompiled).toHaveLength(1)
+    expect(result.wasmCompiled![0]).toMatchObject({
+      id: '__tjs_wasm_add',
+      success: true,
+    })
+  })
+
+  it('emits a regular JS wrapper that forwards to the wasm export', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+wasm function add(a: f64, b: f64): f64 {
+  return a + b
+}
+`
+    const result = tjs(source)
+    // The declaration is replaced with a wrapper function — that wrapper
+    // forwards to globalThis.__tjs_wasm_add, which the bootstrap sets up.
+    expect(result.code).toContain('function add(a, b)')
+    expect(result.code).toContain('globalThis.__tjs_wasm_add(a, b)')
+  })
+
+  it('preserves the export modifier on the wrapper', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+export wasm function mul(a: f64, b: f64): f64 {
+  return a * b
+}
+`
+    const result = tjs(source)
+    expect(result.code).toContain('export function mul(a, b)')
+    expect(result.code).toContain('globalThis.__tjs_wasm_mul(a, b)')
+  })
+
+  it('runs end-to-end: declare wasm function, call it, get correct result', async () => {
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    const source = `
+wasm function add(a: f64, b: f64): f64 {
+  return a + b
+}
+
+wasm function sub(a: f64, b: f64): f64 {
+  return a - b
+}
+`
+    const result = tjs(source)
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${result.code}\n` +
+          `globalThis.__test_add = add;\n` +
+          `globalThis.__test_sub = sub;\n` +
+          `})();`
+      )(globalThis.__tjs)
+
+      // Wait for the async bootstrap to complete
+      await new Promise((r) => setTimeout(r, 100))
+
+      expect((globalThis as any).__test_add(3, 4)).toBe(7)
+      expect((globalThis as any).__test_sub(10, 7)).toBe(3)
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).__test_add
+      delete (globalThis as any).__test_sub
+    }
+  })
+
+  it('works with Float32Array parameters (zero-copy via wasmBuffer)', async () => {
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    const source = `
+wasm function scaleArray(arr: Float32Array, len: f64, factor: f64) {
+  for (let i = 0; i < len; i++) {
+    arr[i] = arr[i] * factor
+  }
+}
+`
+    const result = tjs(source)
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${result.code}\n` +
+          `globalThis.__test_scale = scaleArray;\n` +
+          `})();`
+      )(globalThis.__tjs)
+
+      await new Promise((r) => setTimeout(r, 100))
+
+      const wasmBuffer = (globalThis as any).wasmBuffer
+      expect(typeof wasmBuffer).toBe('function')
+      const buf = wasmBuffer(Float32Array, 4)
+      buf[0] = 1
+      buf[1] = 2
+      buf[2] = 3
+      buf[3] = 4
+      ;(globalThis as any).__test_scale(buf, 4, 3.0)
+      expect(Array.from(buf)).toEqual([3, 6, 9, 12])
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).wasmBuffer
+      delete (globalThis as any).__test_scale
+    }
+  })
+
+  it('coexists with inline wasm {} blocks in the same file', async () => {
+    const { tjs } = await import('./index')
+    // No return-type annotation on `inline` — auto signature tests would
+    // call the function at transpile time before wasm is instantiated and
+    // fail because the dispatch returns undefined. The functional check is
+    // about block extraction + module consolidation, not the auto-test.
+    const source = `
+wasm function topLevel(a: f64, b: f64): f64 {
+  return a * b
+}
+
+function inline(x: 0, y: 0) {
+  return wasm {
+    return x + y
+  }
+}
+`
+    const result = tjs(source, { runTests: false })
+    expect(result.wasmCompiled).toHaveLength(2)
+    const ids = result.wasmCompiled!.map((b) => b.id).sort()
+    expect(ids[0]).toBe('__tjs_wasm_0') // inline block
+    expect(ids[1]).toBe('__tjs_wasm_topLevel') // named wasm function
+    // Both compile into the same consolidated module — verify exactly one
+    // WebAssembly.compile call in the output.
+    const compileCalls = (result.code.match(/WebAssembly\.compile\(/g) || [])
+      .length
+    expect(compileCalls).toBe(1)
+  })
+
+  it('handles wasm function with no params', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+wasm function answer(): f64 {
+  return 42
+}
+`
+    const result = tjs(source)
+    expect(result.wasmCompiled).toHaveLength(1)
+    expect(result.wasmCompiled![0].success).toBe(true)
+    expect(result.code).toContain('function answer()')
+    expect(result.code).toContain('globalThis.__tjs_wasm_answer()')
+  })
+
+  it('does not match identifiers that contain "wasm" (e.g. mywasm)', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+function mywasm(x: 0): 0 { return x }
+`
+    const result = tjs(source)
+    // No wasm blocks should be extracted from this — the source contains no
+    // actual `wasm function` declaration.
+    expect(result.wasmCompiled ?? []).toHaveLength(0)
+  })
+})
+
+describe('boundary distribution form (Phase 4)', () => {
+  // Write the transpiled library to a tmp .mjs file and dynamically import
+  // it. This is the most authentic test of the boundary form: real ESM
+  // resolution, real exports, real WebAssembly instantiation.
+  async function dynamicImportLibrary(transpiled: string): Promise<any> {
+    const { tmpdir } = await import('node:os')
+    const { join } = await import('node:path')
+    const { writeFileSync, unlinkSync } = await import('node:fs')
+    const path = join(
+      tmpdir(),
+      `tjs-lib-${Date.now()}-${Math.random().toString(36).slice(2, 8)}.mjs`
+    )
+    writeFileSync(path, transpiled)
+    try {
+      const mod = await import(path)
+      // Wait for the async wasm bootstrap inside the module to finish.
+      // The bootstrap runs as a top-level IIFE; instantiation is async.
+      await new Promise((r) => setTimeout(r, 100))
+      return mod
+    } finally {
+      try {
+        unlinkSync(path)
+      } catch {
+        /* ignore */
+      }
+    }
+  }
+
+  it('emits a self-contained ES module with exported wasm wrappers', async () => {
+    const { tjs } = await import('./index')
+    const librarySource = `
+export wasm function add(a: f64, b: f64): f64 { return a + b }
+export wasm function mul(a: f64, b: f64): f64 { return a * b }
+`
+    const result = tjs(librarySource, { runTests: false })
+
+    // Both wrappers exported
+    expect(result.code).toContain('export function add(a, b)')
+    expect(result.code).toContain('export function mul(a, b)')
+
+    // Each wrapper forwards to its globalThis-registered wasm function
+    expect(result.code).toContain('globalThis.__tjs_wasm_add(a, b)')
+    expect(result.code).toContain('globalThis.__tjs_wasm_mul(a, b)')
+
+    // The wasm module is base64-embedded and instantiated at the top
+    expect(result.code).toContain('__wasmModuleB64')
+    expect(result.code).toContain('WebAssembly.compile')
+
+    // No external runtime setup required — the inline __tjs fallback
+    // covers everything actually used. (Only the helpers this file needs
+    // are inlined, so simple wasm-wrapper libraries get a small fallback;
+    // libraries with type checks would also inline MonadicError etc.)
+    expect(result.code).toContain('globalThis.__tjs?.createRuntime?.()')
+  })
+
+  it('dynamic import of the boundary form gives a working module', async () => {
+    const { tjs } = await import('./index')
+    const librarySource = `
+export wasm function add(a: f64, b: f64): f64 { return a + b }
+export wasm function mul(a: f64, b: f64): f64 { return a * b }
+`
+    const result = tjs(librarySource, { runTests: false })
+
+    const lib = await dynamicImportLibrary(result.code)
+    expect(typeof lib.add).toBe('function')
+    expect(typeof lib.mul).toBe('function')
+    expect(lib.add(7, 5)).toBe(12)
+    expect(lib.mul(6, 7)).toBe(42)
+  })
+
+  it('boundary form and composed form return identical results', async () => {
+    // Same library source, consumed two different ways:
+    //  - boundary:  transpile → write to disk → dynamic import → call exports
+    //  - composed:  Phase 3 — moduleLoader pulls the wasm body into the
+    //               consumer's own module
+    // Both should produce the same numeric results.
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    const { ModuleLoader, inMemoryFileSystem } = await import(
+      './module-loader'
+    )
+
+    const librarySource = `
+export wasm function dot3(
+  ax: f64, ay: f64, az: f64,
+  bx: f64, by: f64, bz: f64
+): f64 {
+  return ax * bx + ay * by + az * bz
+}
+`
+    // Boundary form
+    const libCompiled = tjs(librarySource, { runTests: false })
+    const lib = await dynamicImportLibrary(libCompiled.code)
+    const boundary = lib.dot3(1, 2, 3, 4, 5, 6)
+
+    // Composed form (Phase 3 path)
+    const loader = new ModuleLoader({
+      fs: inMemoryFileSystem({ '/proj/linalg.tjs': librarySource }),
+      baseDir: '/proj',
+    })
+    const consumerSource = `
+import { dot3 } from './linalg.tjs'
+`
+    const consumerCompiled = tjs(consumerSource, {
+      moduleLoader: loader,
+      filename: '/proj/app.tjs',
+      runTests: false,
+    })
+
+    let composed: number
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${consumerCompiled.code}\n` +
+          `globalThis.__test_dot3 = dot3;\n` +
+          `})();`
+      )(globalThis.__tjs)
+      await new Promise((r) => setTimeout(r, 100))
+      composed = (globalThis as any).__test_dot3(1, 2, 3, 4, 5, 6)
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).__test_dot3
+    }
+
+    expect(boundary).toBe(composed)
+    expect(boundary).toBe(1 * 4 + 2 * 5 + 3 * 6) // 32 — pen-and-paper truth
+  })
+
+  it('boundary form library works for a plain JS consumer (no tjs involvement)', async () => {
+    // Build the library, write to disk, import it, then call from a
+    // plain JS function (simulating a consumer with no tjs in the chain).
+    // The library's wasm bootstrap should run, the wrapper should be
+    // callable, and the wasm function should produce correct results.
+    const { tjs } = await import('./index')
+    const librarySource = `
+export wasm function square(x: f64): f64 { return x * x }
+`
+    const result = tjs(librarySource, { runTests: false })
+    const lib = await dynamicImportLibrary(result.code)
+
+    // Use the import from a plain JS function — no tjs runtime involved
+    function jsConsumer(x: number): number {
+      return lib.square(x) + 1
+    }
+    expect(jsConsumer(5)).toBe(26)
+  })
+})
+
+describe('cross-file wasm composition (Phase 3)', () => {
+  // Build a minimal ModuleLoader backed by an in-memory FS for hermetic tests
+  async function buildLoader(files: Record<string, string>) {
+    const { ModuleLoader, inMemoryFileSystem } = await import('./module-loader')
+    return new ModuleLoader({
+      fs: inMemoryFileSystem(files),
+      baseDir: '/proj',
+    })
+  }
+
+  it('composes an imported wasm function into the consumer module', async () => {
+    const { tjs } = await import('./index')
+    const loader = await buildLoader({
+      '/proj/linalg.tjs': `
+wasm function dot(a: f64, b: f64): f64 {
+  return a * b
+}
+`,
+    })
+    const source = `
+import { dot } from './linalg.tjs'
+
+function compute(a: 0.0, b: 0.0): 0.0 {
+  return dot(a, b)
+}
+`
+    const result = tjs(source, {
+      moduleLoader: loader,
+      filename: '/proj/app.tjs',
+      runTests: false,
+    })
+
+    // The imported wasm function got composed in
+    expect(result.wasmCompiled).toBeDefined()
+    expect(result.wasmCompiled).toHaveLength(1)
+    expect(result.wasmCompiled![0]).toMatchObject({
+      id: '__tjs_wasm_dot',
+      success: true,
+    })
+
+    // The import statement was rewritten to a local wrapper
+    expect(result.code).not.toContain("import { dot } from './linalg.tjs'")
+    expect(result.code).toContain('function dot(a, b)')
+    expect(result.code).toContain('globalThis.__tjs_wasm_dot(a, b)')
+  })
+
+  it('keeps imports unchanged when no loader is supplied', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+import { dot } from './linalg.tjs'
+
+function compute(a: 0.0, b: 0.0): 0.0 {
+  return dot(a, b)
+}
+`
+    // No moduleLoader option — default behavior preserved
+    const result = tjs(source, { runTests: false })
+    expect(result.code).toContain("import { dot } from './linalg.tjs'")
+    expect(result.wasmCompiled ?? []).toHaveLength(0)
+  })
+
+  it('preserves imports that do NOT resolve to wasm functions', async () => {
+    const { tjs } = await import('./index')
+    const loader = await buildLoader({
+      '/proj/regular.tjs': `
+export function helper(x: 0.0): 0.0 { return x + 1.0 }
+`,
+    })
+    const source = `
+import { helper } from './regular.tjs'
+
+function compute(x: 0.0): 0.0 {
+  return helper(x)
+}
+`
+    const result = tjs(source, {
+      moduleLoader: loader,
+      filename: '/proj/app.tjs',
+      runTests: false,
+    })
+    // `helper` is a regular function, not a wasm function — import stays
+    expect(result.code).toContain("import { helper } from './regular.tjs'")
+    expect(result.wasmCompiled ?? []).toHaveLength(0)
+  })
+
+  it('handles mixed imports (some wasm, some regular) in one statement', async () => {
+    const { tjs } = await import('./index')
+    const loader = await buildLoader({
+      '/proj/lib.tjs': `
+wasm function fast(a: f64): f64 { return a * 2 }
+export function slow(x: 0.0): 0.0 { return x + 1.0 }
+`,
+    })
+    const source = `
+import { fast, slow } from './lib.tjs'
+
+function compute(x: 0.0): 0.0 {
+  return fast(x) + slow(x)
+}
+`
+    const result = tjs(source, {
+      moduleLoader: loader,
+      filename: '/proj/app.tjs',
+      runTests: false,
+    })
+    // fast was composed; slow remains imported
+    expect(result.wasmCompiled).toHaveLength(1)
+    expect(result.wasmCompiled![0].id).toBe('__tjs_wasm_fast')
+    expect(result.code).toContain('function fast(a)')
+    expect(result.code).toContain('globalThis.__tjs_wasm_fast(a)')
+    // The remaining import keeps `slow` only
+    expect(result.code).toMatch(/import\s*\{\s*slow\s*\}\s*from/)
+    expect(result.code).not.toMatch(
+      /import\s*\{\s*fast\s*,\s*slow\s*\}\s*from/
+    )
+  })
+
+  it('composes multiple wasm functions from one library', async () => {
+    const { tjs } = await import('./index')
+    const loader = await buildLoader({
+      '/proj/linalg.tjs': `
+wasm function dot(a: f64, b: f64): f64 { return a * b }
+wasm function add(a: f64, b: f64): f64 { return a + b }
+wasm function unused(x: f64): f64 { return x }
+`,
+    })
+    const source = `
+import { dot, add } from './linalg.tjs'
+
+function compute(a: 0.0, b: 0.0): 0.0 {
+  return add(dot(a, b), b)
+}
+`
+    const result = tjs(source, {
+      moduleLoader: loader,
+      filename: '/proj/app.tjs',
+      runTests: false,
+    })
+    expect(result.wasmCompiled).toHaveLength(2)
+    const ids = result.wasmCompiled!.map((b) => b.id).sort()
+    expect(ids).toEqual(['__tjs_wasm_add', '__tjs_wasm_dot'])
+    // One consolidated WebAssembly.Module per file (Phase 0.5 acceptance)
+    const compileCalls = (result.code.match(/WebAssembly\.compile\(/g) || [])
+      .length
+    expect(compileCalls).toBe(1)
+  })
+
+  it('module shape: composed functions are local (no extra imports beyond env.memory)', async () => {
+    // This is the Phase 3 acceptance criterion: imported wasm functions
+    // are LOCAL to the consumer's module, not imported at the wasm level.
+    const { tjs } = await import('./index')
+    const { compileBlocksToModule } = await import('./wasm')
+    const loader = await buildLoader({
+      '/proj/linalg.tjs': `
+wasm function dot(a: f64, b: f64): f64 { return a * b }
+`,
+    })
+    const source = `
+import { dot } from './linalg.tjs'
+function compute(a: 0.0, b: 0.0): 0.0 { return dot(a, b) }
+`
+    tjs(source, {
+      moduleLoader: loader,
+      filename: '/proj/app.tjs',
+      runTests: false,
+    })
+
+    // Now compile the loaded module's wasm blocks and inspect the bytes.
+    // The wasm binary format: after the magic + version, sections appear.
+    // Section 2 is "import" — we want to verify ONLY env.memory is imported
+    // (no host function imports).
+    const linalgBlocks = (await loader.load('./linalg.tjs', '/proj/app.tjs'))!
+      .parseResult.wasmBlocks
+    const composed = compileBlocksToModule(linalgBlocks)
+    expect(composed.exports).toHaveLength(1)
+    expect(composed.exports[0].id).toBe('__tjs_wasm_dot')
+
+    // The composed module needs no memory for a pure-scalar `dot(a, b) = a*b`
+    expect(composed.needsMemory).toBe(false)
+
+    // Confirm the bytecode parses as a valid WebAssembly.Module
+    const mod = new WebAssembly.Module(composed.bytes)
+    expect(mod).toBeInstanceOf(WebAssembly.Module)
+
+    // Inspect imports: should have NO function imports.
+    // WebAssembly.Module.imports returns [{ module, name, kind }, ...]
+    const imports = WebAssembly.Module.imports(mod)
+    const functionImports = imports.filter((i) => i.kind === 'function')
+    expect(functionImports).toHaveLength(0)
+  })
+
+  it('end-to-end: imported wasm function runs and produces correct results', async () => {
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    const loader = await buildLoader({
+      '/proj/linalg.tjs': `
+wasm function add(a: f64, b: f64): f64 { return a + b }
+wasm function mul(a: f64, b: f64): f64 { return a * b }
+`,
+    })
+    const source = `
+import { add, mul } from './linalg.tjs'
+`
+    const result = tjs(source, {
+      moduleLoader: loader,
+      filename: '/proj/app.tjs',
+      runTests: false,
+    })
+
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${result.code}\n` +
+          `globalThis.__test_add = add;\n` +
+          `globalThis.__test_mul = mul;\n` +
+          `})();`
+      )(globalThis.__tjs)
+
+      await new Promise((r) => setTimeout(r, 100))
+
+      expect((globalThis as any).__test_add(2, 3)).toBe(5)
+      expect((globalThis as any).__test_mul(4, 5)).toBe(20)
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).__test_add
+      delete (globalThis as any).__test_mul
+    }
+  })
+
+  describe('tree-shaking & transitive deps', () => {
+    async function loaderWith(files: Record<string, string>) {
+      const { ModuleLoader, inMemoryFileSystem } = await import(
+        './module-loader'
+      )
+      return new ModuleLoader({
+        fs: inMemoryFileSystem(files),
+        baseDir: '/proj',
+      })
+    }
+
+    it('only pulls in explicitly-imported wasm functions (tree-shaking)', async () => {
+      const { tjs } = await import('./index')
+      const loader = await loaderWith({
+        '/proj/lib.tjs': `
+wasm function wanted(x: f64): f64 { return x * 2 }
+wasm function unwanted(x: f64): f64 { return x + 999 }
+wasm function also_unwanted(x: f64): f64 { return x - 1 }
+`,
+      })
+      const source = `
+import { wanted } from './lib.tjs'
+function go() { return wanted(5) }
+`
+      const result = tjs(source, {
+        moduleLoader: loader,
+        filename: '/proj/app.tjs',
+        runTests: false,
+      })
+      // Only `wanted` should be in the composed module
+      expect(result.wasmCompiled).toHaveLength(1)
+      expect(result.wasmCompiled![0].id).toBe('__tjs_wasm_wanted')
+    })
+
+    it('pulls in transitive wasm-to-wasm callees automatically', async () => {
+      // `outer` calls `inner`; the consumer imports only `outer`. Without
+      // transitive walking, outer's body's `call <inner-index>` would fail
+      // (the inner function wouldn't be in the consumer's module). With
+      // transitive walking, both end up in the composed module.
+      const { tjs } = await import('./index')
+      const { createRuntime } = await import('./runtime')
+      const loader = await loaderWith({
+        '/proj/lib.tjs': `
+wasm function inner(x: f64): f64 { return x * 2 }
+wasm function outer(x: f64): f64 { return inner(x) + 1 }
+`,
+      })
+      const source = `
+import { outer } from './lib.tjs'
+`
+      const result = tjs(source, {
+        moduleLoader: loader,
+        filename: '/proj/app.tjs',
+        runTests: false,
+      })
+      // Both pulled in
+      expect(result.wasmCompiled).toHaveLength(2)
+      expect(result.wasmCompiled!.every((b) => b.success)).toBe(true)
+      const ids = result.wasmCompiled!.map((b) => b.id).sort()
+      expect(ids).toEqual(['__tjs_wasm_inner', '__tjs_wasm_outer'])
+
+      // Run end-to-end to confirm the call actually works
+      const savedTjs = globalThis.__tjs
+      try {
+        globalThis.__tjs = createRuntime()
+        await new Function(
+          '__tjs',
+          `return (async () => { ${result.code}\n` +
+            `globalThis.__test_outer = outer;\n` +
+            `})();`
+        )(globalThis.__tjs)
+        await new Promise((r) => setTimeout(r, 100))
+        // outer(5) = inner(5) + 1 = 10 + 1 = 11
+        expect((globalThis as any).__test_outer(5)).toBe(11)
+      } finally {
+        globalThis.__tjs = savedTjs
+        delete (globalThis as any).__test_outer
+        for (const key of Object.keys(globalThis)) {
+          if (key.startsWith('__tjs_wasm_')) {
+            delete (globalThis as any)[key]
+          }
+        }
+      }
+    })
+
+    it('handles multi-step transitive chains', async () => {
+      // a → b → c. Importing only `a` should pull in b and c.
+      const { tjs } = await import('./index')
+      const loader = await loaderWith({
+        '/proj/lib.tjs': `
+wasm function c(x: f64): f64 { return x + 1 }
+wasm function b(x: f64): f64 { return c(x) * 2 }
+wasm function a(x: f64): f64 { return b(x) + 3 }
+wasm function unrelated(x: f64): f64 { return x }
+`,
+      })
+      const source = `
+import { a } from './lib.tjs'
+`
+      const result = tjs(source, {
+        moduleLoader: loader,
+        filename: '/proj/app.tjs',
+        runTests: false,
+      })
+      // a, b, c pulled in; unrelated is not
+      expect(result.wasmCompiled).toHaveLength(3)
+      const ids = result.wasmCompiled!.map((b) => b.id).sort()
+      expect(ids).toEqual(['__tjs_wasm_a', '__tjs_wasm_b', '__tjs_wasm_c'])
+    })
+
+    it('handles mutual recursion across the import boundary', async () => {
+      // Two mutually-recursive library functions. Importing either one
+      // should pull in both.
+      const { tjs } = await import('./index')
+      const loader = await loaderWith({
+        '/proj/lib.tjs': `
+wasm function ping(n: i32): f64 {
+  if (n == 0) return 0.0
+  return pong(n - 1) + 1
+}
+wasm function pong(n: i32): f64 {
+  if (n == 0) return 0.0
+  return ping(n - 1) + 1
+}
+`,
+      })
+      const source = `
+import { ping } from './lib.tjs'
+`
+      const result = tjs(source, {
+        moduleLoader: loader,
+        filename: '/proj/app.tjs',
+        runTests: false,
+      })
+      expect(result.wasmCompiled).toHaveLength(2)
+      expect(result.wasmCompiled!.every((b) => b.success)).toBe(true)
+    })
+
+    it('does not over-pull when an identifier appears as a substring', async () => {
+      // The body's regex scan uses word-boundary anchors. A function named
+      // `add` shouldn't trigger if the body has `add2` or `_add` or similar.
+      const { tjs } = await import('./index')
+      const loader = await loaderWith({
+        '/proj/lib.tjs': `
+wasm function add(a: f64, b: f64): f64 { return a + b }
+wasm function user(x: f64): f64 {
+  // The local "_add" is just a variable name — not a call to add()
+  let _add = x * 2
+  return _add + 1
+}
+`,
+      })
+      const source = `
+import { user } from './lib.tjs'
+`
+      const result = tjs(source, {
+        moduleLoader: loader,
+        filename: '/proj/app.tjs',
+        runTests: false,
+      })
+      // user is pulled in; add is NOT (word-boundary regex doesn't match _add as add)
+      expect(result.wasmCompiled).toHaveLength(1)
+      expect(result.wasmCompiled![0].id).toBe('__tjs_wasm_user')
+    })
+  })
+})
+
+describe('wasm function purity & unsafe marker (Phase 2)', () => {
+  it('rejects `wasm function (! ...)` with a clear error', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+wasm function dangerous(! aPtr: i32, n: i32): f64 {
+  return 0
+}
+`
+    expect(() => tjs(source)).toThrow(/Unsafe wasm functions/)
+    expect(() => tjs(source)).toThrow(/dangerous/)
+    expect(() => tjs(source)).toThrow(/reserved for a future phase/)
+  })
+
+  it('accepts wasm function without the bang marker', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+wasm function safe(a: f64, b: f64): f64 {
+  return a + b
+}
+`
+    expect(() => tjs(source)).not.toThrow()
+  })
+
+  it('purity: host-import calls in a wasm function body fail with a clear error', async () => {
+    // Math.sin requires a host import; the wasm bytecode builder doesn't
+    // support that path today, so this fails at compile time. This test
+    // documents the property: wasm function bodies are pure compute — any
+    // host-import call is a compile error.
+    const { tjs } = await import('./index')
+    const source = `
+wasm function tryHostCall(x: f64): f64 {
+  return Math.sin(x)
+}
+`
+    const result = tjs(source)
+    // Block extraction succeeds but compilation fails
+    expect(result.wasmCompiled).toBeDefined()
+    expect(result.wasmCompiled).toHaveLength(1)
+    expect(result.wasmCompiled![0].success).toBe(false)
+    expect(result.wasmCompiled![0].error).toMatch(/Math\.sin/)
+    expect(result.wasmCompiled![0].error).toMatch(/import/i)
+  })
+
+  it('purity: inline Math ops that DO compile (sqrt, abs, etc.) work fine', async () => {
+    // sqrt, abs, floor, ceil, min, max compile to wasm intrinsics — no host
+    // imports needed. Confirms the constraint is "no host imports", not
+    // "no Math.* at all".
+    const { tjs } = await import('./index')
+    const source = `
+wasm function magnitude(x: f64, y: f64): f64 {
+  return Math.sqrt(x * x + y * y)
+}
+`
+    const result = tjs(source)
+    expect(result.wasmCompiled).toBeDefined()
+    expect(result.wasmCompiled![0].success).toBe(true)
+  })
+})
+
+describe('wasm-to-wasm calls (Phase 1.5)', () => {
+  it('a wasm function can call another wasm function in the same file', async () => {
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    const source = `
+wasm function square(x: f64): f64 {
+  return x * x
+}
+
+wasm function sumOfSquares(a: f64, b: f64): f64 {
+  return square(a) + square(b)
+}
+`
+    const result = tjs(source, { runTests: false })
+    expect(result.wasmCompiled).toHaveLength(2)
+    expect(result.wasmCompiled!.every((b) => b.success)).toBe(true)
+
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${result.code}\n` +
+          `globalThis.__test_sos = sumOfSquares;\n` +
+          `})();`
+      )(globalThis.__tjs)
+      await new Promise((r) => setTimeout(r, 100))
+
+      // sumOfSquares(3, 4) = 9 + 16 = 25
+      expect((globalThis as any).__test_sos(3, 4)).toBe(25)
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).__test_sos
+      for (const key of Object.keys(globalThis)) {
+        if (key.startsWith('__tjs_wasm_')) {
+          delete (globalThis as any)[key]
+        }
+      }
+    }
+  })
+
+  it('forward references work (caller declared before callee)', async () => {
+    // The pre-pass builds the function map before any body is compiled,
+    // so a wasm function can call one declared LATER in the file.
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    const source = `
+wasm function caller(x: f64): f64 {
+  return callee(x) + 1
+}
+
+wasm function callee(x: f64): f64 {
+  return x * 2
+}
+`
+    const result = tjs(source, { runTests: false })
+    expect(result.wasmCompiled!.every((b) => b.success)).toBe(true)
+
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${result.code}\n` +
+          `globalThis.__test_caller = caller;\n` +
+          `})();`
+      )(globalThis.__tjs)
+      await new Promise((r) => setTimeout(r, 100))
+      expect((globalThis as any).__test_caller(5)).toBe(11) // 5*2 + 1
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).__test_caller
+      for (const key of Object.keys(globalThis)) {
+        if (key.startsWith('__tjs_wasm_')) {
+          delete (globalThis as any)[key]
+        }
+      }
+    }
+  })
+
+  it('mutual recursion compiles and runs', async () => {
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    // Classic mutual-recursion shape: isEven(n) calls isOdd(n-1); isOdd(n)
+    // calls isEven(n-1). Returns 1.0 (true) or 0.0 (false).
+    const source = `
+wasm function isEven(n: i32): f64 {
+  if (n == 0) return 1.0
+  return isOdd(n - 1)
+}
+
+wasm function isOdd(n: i32): f64 {
+  if (n == 0) return 0.0
+  return isEven(n - 1)
+}
+`
+    const result = tjs(source, { runTests: false })
+    expect(result.wasmCompiled!.every((b) => b.success)).toBe(true)
+
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${result.code}\n` +
+          `globalThis.__test_isEven = isEven;\n` +
+          `globalThis.__test_isOdd = isOdd;\n` +
+          `})();`
+      )(globalThis.__tjs)
+      await new Promise((r) => setTimeout(r, 100))
+
+      expect((globalThis as any).__test_isEven(10)).toBe(1)
+      expect((globalThis as any).__test_isEven(7)).toBe(0)
+      expect((globalThis as any).__test_isOdd(7)).toBe(1)
+      expect((globalThis as any).__test_isOdd(10)).toBe(0)
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).__test_isEven
+      delete (globalThis as any).__test_isOdd
+      for (const key of Object.keys(globalThis)) {
+        if (key.startsWith('__tjs_wasm_')) {
+          delete (globalThis as any)[key]
+        }
+      }
+    }
+  })
+
+  it('cross-file: composed imports can be called from a wasm function', async () => {
+    // The big payoff: a consumer's wasm function calls imported library
+    // kernels via wasm `call` instructions — no JS↔wasm boundary in the
+    // inner loop.
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    const { ModuleLoader, inMemoryFileSystem } = await import(
+      './module-loader'
+    )
+
+    const loader = new ModuleLoader({
+      fs: inMemoryFileSystem({
+        '/proj/lib.tjs': `
+wasm function double(x: f64): f64 { return x * 2 }
+wasm function triple(x: f64): f64 { return x * 3 }
+`,
+      }),
+      baseDir: '/proj',
+    })
+    const consumerSource = `
+import { double, triple } from './lib.tjs'
+
+wasm function fancy(x: f64): f64 {
+  return double(x) + triple(x)
+}
+`
+    const result = tjs(consumerSource, {
+      moduleLoader: loader,
+      filename: '/proj/app.tjs',
+      runTests: false,
+    })
+
+    // All three wasm functions (double, triple, fancy) live in one module
+    expect(result.wasmCompiled).toHaveLength(3)
+    expect(result.wasmCompiled!.every((b) => b.success)).toBe(true)
+    const compileCalls = (result.code.match(/WebAssembly\.compile\(/g) || [])
+      .length
+    expect(compileCalls).toBe(1)
+
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${result.code}\n` +
+          `globalThis.__test_fancy = fancy;\n` +
+          `})();`
+      )(globalThis.__tjs)
+      await new Promise((r) => setTimeout(r, 100))
+
+      // fancy(5) = double(5) + triple(5) = 10 + 15 = 25
+      expect((globalThis as any).__test_fancy(5)).toBe(25)
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).__test_fancy
+      for (const key of Object.keys(globalThis)) {
+        if (key.startsWith('__tjs_wasm_')) {
+          delete (globalThis as any)[key]
+        }
+      }
+    }
+  })
+
+  it('arg type mismatch is detected with a clear error', async () => {
+    // A wasm function expects i32 but is called with an f64 expression
+    // that has no way to be converted safely without losing data. The
+    // compiler should still accept it (lossy conversion is allowed) —
+    // truncate is the standard wasm move for f64→i32.
+    //
+    // For an actually-incompatible case we'd need v128; testing that
+    // would require setting up SIMD types. Instead we verify the
+    // conversion path: pass an f64 expression to an i32 param and check
+    // it works (the truncation happens automatically).
+    const { tjs } = await import('./index')
+    const { createRuntime } = await import('./runtime')
+    const source = `
+wasm function takesInt(n: i32): f64 {
+  return n + 0.5
+}
+
+wasm function caller(): f64 {
+  // 3.7 is f64; takesInt expects i32 — the compiler inserts a truncate
+  return takesInt(3.7)
+}
+`
+    const result = tjs(source, { runTests: false })
+    expect(result.wasmCompiled!.every((b) => b.success)).toBe(true)
+
+    const savedTjs = globalThis.__tjs
+    try {
+      globalThis.__tjs = createRuntime()
+      await new Function(
+        '__tjs',
+        `return (async () => { ${result.code}\n` +
+          `globalThis.__test_caller = caller;\n` +
+          `})();`
+      )(globalThis.__tjs)
+      await new Promise((r) => setTimeout(r, 100))
+
+      // 3.7 truncates to 3 (i32), then 3 + 0.5 = 3.5
+      expect((globalThis as any).__test_caller()).toBe(3.5)
+    } finally {
+      globalThis.__tjs = savedTjs
+      delete (globalThis as any).__test_caller
+      for (const key of Object.keys(globalThis)) {
+        if (key.startsWith('__tjs_wasm_')) {
+          delete (globalThis as any)[key]
+        }
+      }
+    }
+  })
+
+  it('argument count mismatch produces a clear compile error', async () => {
+    const { tjs } = await import('./index')
+    const source = `
+wasm function takesTwo(a: f64, b: f64): f64 { return a + b }
+wasm function caller(): f64 {
+  return takesTwo(1.0)
+}
+`
+    const result = tjs(source, { runTests: false })
+    // caller fails because takesTwo gets one arg instead of two
+    const callerResult = result.wasmCompiled!.find((b) => b.id === '__tjs_wasm_caller')
+    expect(callerResult).toBeDefined()
+    expect(callerResult!.success).toBe(false)
+    expect(callerResult!.error).toMatch(/takesTwo expects 2 arguments, got 1/)
+  })
+})