tjs-lang 0.7.8 → 0.8.0

package/src/lang/parser-transforms.ts

@@ -223,6 +223,342 @@ export function extractWasmBlocks(source: string): {
   return { source: result, blocks }
 }
 
+/**
+ * Extract top-level `wasm function NAME(params): RetType { body }` declarations.
+ *
+ * Unlike `extractWasmBlocks` (which finds `wasm { ... }` blocks nested inside
+ * regular tjs functions with auto-captured variables), this extractor finds
+ * top-level wasm function declarations with explicit parameters and an
+ * optional return type. The body is the wasm-subset source.
+ *
+ * Each declaration becomes a `WasmBlock` whose `captures` array holds the
+ * function's parameters with their type annotations (e.g. `['a: Float32Array',
+ * 'b: Float32Array', 'n: i32']`). The block ID is derived from the function
+ * name (`__tjs_wasm_<name>`), so the JS-side wrapper can reference it.
+ *
+ * The declaration is replaced in source with a regular JS function that
+ * forwards its args to the wasm export, preserving the `export` modifier
+ * if present:
+ *
+ *   (export)? wasm function dot(a: Float32Array, b: Float32Array, n: f64): f64 {
+ *     <wasm-source>
+ *   }
+ *
+ *   becomes:
+ *
+ *   (export)? function dot(a, b, n) { return globalThis.__tjs_wasm_dot(a, b, n) }
+ *
+ * This runs BEFORE `extractWasmBlocks` so its output (a regular JS function
+ * wrapper) isn't disturbed by the inline-block scanner.
+ *
+ * Return-type note: the underlying wasm bytecode builder currently emits f64
+ * or void return types only. The declared `: RetType` annotation is parsed
+ * and preserved on the block, but not yet validated against the backend's
+ * capabilities — `: f64` and omitted-return work today; other types (`: i32`,
+ * `: f32`, `: v128`) will be supported when the bytecode builder grows
+ * per-function return-type encoding.
+ */
+export function extractWasmFunctions(source: string): {
+  source: string
+  blocks: WasmBlock[]
+} {
+  const blocks: WasmBlock[] = []
+  let result = ''
+  let i = 0
+
+  while (i < source.length) {
+    // Match: `(export )?wasm function NAME(` (with leading whitespace allowed)
+    // The leading `\b` ensures we don't match inside identifiers like `mywasm`.
+    const declRe = /^\b(export\s+)?wasm\s+function\s+(\w+)\s*\(/
+    const m = source.slice(i).match(declRe)
+    if (!m) {
+      result += source[i]
+      i++
+      continue
+    }
+
+    const hasExport = !!m[1]
+    const name = m[2]
+    const matchStart = i
+
+    // After the opening `(`: scan balanced parens for the params block
+    const parensStart = i + m[0].length
+    let parenDepth = 1
+    let j = parensStart
+    while (j < source.length && parenDepth > 0) {
+      if (source[j] === '(') parenDepth++
+      else if (source[j] === ')') parenDepth--
+      j++
+    }
+    if (parenDepth !== 0) {
+      // Unmatched parens — leave source alone and move on
+      result += source[i]
+      i++
+      continue
+    }
+    let paramsSource = source.slice(parensStart, j - 1)
+    // j now points just past the closing `)`
+
+    // Phase 2: detect `(!` — reserved syntax for unsafe wasm functions
+    // (will allow host-import calls, side-effecting globals, etc.).
+    // Implementation deferred; for now we reject with a clear message so
+    // users either remove the bang or know to wait.
+    const unsafeMatch = paramsSource.match(/^\s*!/)
+    if (unsafeMatch) {
+      throw new SyntaxError(
+        `Unsafe wasm functions (with \`!\` marker) are reserved for a ` +
+          `future phase. Remove the bang from \`wasm function ${name}\` ` +
+          `to declare it as a regular (pure) wasm function, or wait until ` +
+          `the unsafe variant is implemented.`,
+        locAt(source, matchStart)
+      )
+    }
+
+    // Optional return-type annotation: `: TYPE` (TYPE is a single identifier).
+    // Pointer-style annotations like `Ptr<f32>` are reserved for a follow-up.
+    let returnType: string | undefined
+    let afterReturnType = j
+    const retMatch = source.slice(j).match(/^\s*:\s*(\w+)/)
+    if (retMatch) {
+      returnType = retMatch[1]
+      afterReturnType = j + retMatch[0].length
+    }
+
+    // Expect `{` next (with leading whitespace)
+    const braceMatch = source.slice(afterReturnType).match(/^\s*\{/)
+    if (!braceMatch) {
+      // Not a wasm function decl after all — pass through
+      result += source[i]
+      i++
+      continue
+    }
+
+    // Find body via balanced braces
+    const bodyStart = afterReturnType + braceMatch[0].length
+    let braceDepth = 1
+    let k = bodyStart
+    while (k < source.length && braceDepth > 0) {
+      if (source[k] === '{') braceDepth++
+      else if (source[k] === '}') braceDepth--
+      k++
+    }
+    if (braceDepth !== 0) {
+      result += source[i]
+      i++
+      continue
+    }
+    const body = source.slice(bodyStart, k - 1)
+    // k now points just past the closing `}`
+
+    // Parse params into the existing `captures` shape used by the wasm
+    // compiler: each entry is either `name` or `name: type`.
+    const captures = parseWasmFunctionParams(paramsSource)
+
+    const id = `__tjs_wasm_${name}`
+    const block: WasmBlock = {
+      id,
+      name, // enables Phase 3 cross-file matching against imported symbols
+      returnType, // undefined when the user wrote no `: T` annotation
+      body,
+      captures,
+      start: matchStart,
+      end: k,
+    }
+    blocks.push(block)
+
+    // Generate the JS wrapper function. The wasm runtime sets globalThis[id]
+    // (with type-aware marshalling already baked into the wrapper); we just
+    // forward args. Preserve `export` if present.
+    const argNames = captures.map((c) => c.split(':')[0].trim())
+    const exportKeyword = hasExport ? 'export ' : ''
+    const wrapper = `${exportKeyword}function ${name}(${argNames.join(
+      ', '
+    )}) { return globalThis.${id}(${argNames.join(', ')}) }`
+
+    result += wrapper
+    i = k
+  }
+
+  return { source: result, blocks }
+}
+
+/**
+ * Phase 3: cross-file wasm-function composition.
+ *
+ * Scans the consumer's source for `import { name1, name2, ... } from 'spec'`
+ * statements. For each, resolves the spec via the supplied ModuleLoader; if
+ * the imported file is tjs/ts and one of the imported names corresponds to a
+ * `wasm function` declared there, the function's WasmBlock is pulled into the
+ * consumer's compilation and replaced in source by a local JS wrapper. The
+ * import statement is rewritten to remove the satisfied names (or removed
+ * entirely if every imported name was wasm-composed).
+ *
+ * This is the heart of the cross-file composition story: imported wasm
+ * functions become local functions in the consumer's single WebAssembly.Module
+ * (via the Phase 0.5 consolidated-module path). The library's transpiled .js
+ * is NOT involved — the source is consumed at transpile time.
+ *
+ * Caller is responsible for providing a configured ModuleLoader. When no
+ * loader is supplied (the common case before Phase 3 is fully wired up), this
+ * function returns the source unchanged with an empty blocks array.
+ *
+ * @param source the consumer's source (after extractWasmFunctions on its own
+ *               wasm functions, before transformParenExpressions)
+ * @param options.loader the ModuleLoader to resolve imports through
+ * @param options.importerPath the path of the file being transpiled (used as
+ *               the resolver's importer context); optional
+ * @returns updated source + the list of imported wasm function blocks
+ */
+export function composeImportedWasmFunctions(
+  source: string,
+  options: {
+    loader?: {
+      load(specifier: string, importerPath?: string): {
+        parseResult: { wasmBlocks: WasmBlock[] }
+      } | null
+    }
+    importerPath?: string
+  }
+): { source: string; blocks: WasmBlock[] } {
+  const { loader, importerPath } = options
+  if (!loader) return { source, blocks: [] }
+
+  const composedBlocks: WasmBlock[] = []
+  // Track imported names that have been composed, so we don't pull the same
+  // wasm function in twice if it's imported from multiple specifiers.
+  const composedNames = new Set<string>()
+
+  /**
+   * Pull a wasm function into the consumer's composition and recursively
+   * pull in any other wasm functions it calls from the same source module.
+   *
+   * Transitive walking is required for correctness: if `outer` calls
+   * `inner` and the consumer only imports `outer`, we still need `inner`
+   * in the consumer's wasm module — otherwise outer's body's
+   * `call <inner-index>` instruction would reference a function index
+   * that doesn't exist. The wasm compiler would (correctly) reject the
+   * call at body-compile time.
+   *
+   * The body scan uses word-boundary regex against the known wasm-function
+   * names in the source module. False positives (matching a method call
+   * `obj.inner(x)` or a shadowed local) only cause bloat — not correctness
+   * problems — and TJS wasm bodies don't use either of those forms anyway.
+   */
+  function pullInTransitively(
+    block: WasmBlock,
+    sourceModuleFns: Map<string, WasmBlock>
+  ): void {
+    if (composedNames.has(block.id)) return
+    composedBlocks.push(block)
+    composedNames.add(block.id)
+    // Scan the body for calls to other wasm functions in this source module
+    for (const [name, target] of sourceModuleFns) {
+      if (name === block.name) continue // self isn't transitive
+      if (composedNames.has(target.id)) continue
+      const callRe = new RegExp(`\\b${name}\\s*\\(`)
+      if (callRe.test(block.body)) {
+        pullInTransitively(target, sourceModuleFns)
+      }
+    }
+  }
+
+  // Match `import { ... } from 'spec'` or `import { ... } from "spec"`.
+  // Default imports, namespace imports, and side-effect imports are left
+  // alone — only named-bindings are candidates for wasm composition.
+  // Multiline imports are supported via the [^}]*? non-greedy match.
+  const importRe = /^(\s*)import\s*\{([^}]*?)\}\s*from\s*(['"])([^'"]+)\3\s*;?\s*$/gm
+
+  const replaced = source.replace(
+    importRe,
+    (match, indent: string, bindings: string, _quote, spec: string) => {
+      const mod = loader.load(spec, importerPath)
+      if (!mod) return match // not a loadable tjs/ts module — leave as-is
+
+      const importedWasmFunctions = new Map<string, WasmBlock>()
+      for (const b of mod.parseResult.wasmBlocks) {
+        if (b.name) importedWasmFunctions.set(b.name, b)
+      }
+      if (importedWasmFunctions.size === 0) return match // no wasm here
+
+      // Parse the bindings list: `{ a, b as c, d }`
+      // Each binding is `name` or `name as local` (we keep `local` as the
+      // local name; for wasm-composed names, `local` must equal the wasm
+      // function name since the local wrapper uses the original name).
+      const parts = bindings
+        .split(',')
+        .map((p) => p.trim())
+        .filter((p) => p.length > 0)
+
+      const wrappers: string[] = []
+      const remainingBindings: string[] = []
+
+      for (const part of parts) {
+        const m = part.match(/^(\w+)(?:\s+as\s+(\w+))?$/)
+        if (!m) {
+          // Couldn't parse — be conservative and keep it as-is
+          remainingBindings.push(part)
+          continue
+        }
+        const imported = m[1]
+        const local = m[2] ?? m[1]
+        const wasmBlock = importedWasmFunctions.get(imported)
+        if (!wasmBlock) {
+          // Not a wasm function in the source module — keep the import
+          remainingBindings.push(part)
+          continue
+        }
+        // Composed: pull the block in (with any transitive callees from
+        // the same source module) and emit a local wrapper that uses the
+        // LOCAL name (in case of `as` renames) but forwards to the wasm
+        // export's original id. Transitive walking is what makes
+        // `import { outer } from 'lib'` work when outer internally calls
+        // inner — both end up in the consumer's composed module.
+        pullInTransitively(wasmBlock, importedWasmFunctions)
+        const argNames = wasmBlock.captures.map((c) =>
+          c.split(':')[0].trim()
+        )
+        wrappers.push(
+          `function ${local}(${argNames.join(
+            ', '
+          )}) { return globalThis.${wasmBlock.id}(${argNames.join(', ')}) }`
+        )
+      }
+
+      const wrapperBlock = wrappers.join('\n')
+
+      if (remainingBindings.length === 0) {
+        // All names were composed — drop the import statement entirely
+        return wrapperBlock ? `${indent}${wrapperBlock}` : `${indent}`
+      }
+      // Some names remain — keep them as a smaller import
+      const trimmedImport = `${indent}import { ${remainingBindings.join(
+        ', '
+      )} } from '${spec}'`
+      return wrapperBlock
+        ? `${trimmedImport}\n${indent}${wrapperBlock}`
+        : trimmedImport
+    }
+  )
+
+  return { source: replaced, blocks: composedBlocks }
+}
+
+/**
+ * Split a parameter list source into individual param strings.
+ * Returns entries like `name` (no type) or `name: type` (with type).
+ * Type values are single identifiers (`i32`, `Float32Array`, etc.) — generic
+ * forms like `Ptr<f32>` are reserved for a follow-up phase.
+ */
+function parseWasmFunctionParams(paramsSource: string): string[] {
+  const trimmed = paramsSource.trim()
+  if (!trimmed) return []
+  // Simple comma split is fine for the current type syntax (no generics yet).
+  return trimmed
+    .split(',')
+    .map((s) => s.trim())
+    .filter((s) => s.length > 0)
+}
+
 /** Check if an identifier is a WASM SIMD intrinsic (not a captured variable) */
 function isWasmIntrinsic(name: string): boolean {
   return name.startsWith('f32x4_') || name.startsWith('v128_')

package/src/lang/parser-types.ts

@@ -15,6 +15,17 @@ export interface ParseOptions {
    * When true, skips == to Is() transformation since the VM handles == correctly.
    */
   vmTarget?: boolean
+  /**
+   * Optional ModuleLoader for cross-file `wasm function` composition (Phase 3).
+   * When provided, imports are resolved at transpile time and matching wasm
+   * functions are composed into the consumer's WebAssembly.Module. When
+   * omitted, imports are preserved verbatim (the default — runtime resolves
+   * them as before).
+   *
+   * Type is left as `any` here to avoid a circular import with module-loader.ts;
+   * callers should pass a `ModuleLoader` instance.
+   */
+  moduleLoader?: any
 }
 
 /**
@@ -37,6 +48,21 @@ export interface ParseOptions {
 export interface WasmBlock {
   /** Unique ID for this block */
   id: string
+  /**
+   * Declared function name (only set for top-level `wasm function NAME(...)`
+   * declarations — Phase 1+). Used by Phase 3 cross-file composition to
+   * match an imported symbol against a wasm function declaration. Inline
+   * `wasm {}` blocks have no name and don't participate in composition.
+   */
+  name?: string
+  /**
+   * Declared return-type annotation, e.g. `'f64'`. Only set for top-level
+   * `wasm function NAME(...): RetType` declarations; presence/absence is
+   * used to determine `hasReturn` BEFORE the body is compiled, so the
+   * function index map can be built up-front for wasm-to-wasm calls.
+   * Inline blocks have no declared return type.
+   */
+  returnType?: string
   /** The body (JS subset that compiles to WASM, also used as fallback) */
   body: string
   /** Explicit fallback body (only if different from body) */
@@ -83,6 +109,13 @@ export interface PreprocessOptions {
    * Default: false (transform == to Is() for TJS code running in regular JS)
    */
   vmTarget?: boolean
+  /**
+   * Optional ModuleLoader for cross-file `wasm function` composition (Phase 3).
+   * See ParseOptions.moduleLoader for details.
+   */
+  moduleLoader?: any
+  /** Path of the file being preprocessed (used as importer context). */
+  filename?: string
 }
 
 /**

package/src/lang/parser.ts

@@ -31,6 +31,8 @@ import { transformParenExpressions } from './parser-params'
 import {
   transformTryWithoutCatch,
   extractWasmBlocks,
+  extractWasmFunctions,
+  composeImportedWasmFunctions,
   transformIsOperators,
   insertAsiProtection,
   transformEqualityToStructural,
@@ -262,6 +264,14 @@ export function preprocess(
   source = letAnnoResult.source
   const letAnnotations = letAnnoResult.annotations
 
+  // Extract `wasm function NAME(...) { ... }` declarations EARLY, before
+  // any source-level transforms that would mangle wasm-body text. In
+  // particular, the equality transforms below rewrite `==` to `Eq()` and
+  // `Is`/`IsNot` to function calls — wasm bodies use literal operators
+  // and shouldn't be affected.
+  const wasmFunctions = extractWasmFunctions(source)
+  source = wasmFunctions.source
+
   // Transform Is/IsNot infix operators to function calls
   // a Is b -> Is(a, b)
   // a IsNot b -> IsNot(a, b)
@@ -290,6 +300,18 @@ export function preprocess(
   // Foo = ... -> const Foo = ...
   source = transformBareAssignments(source)
 
+  // Phase 3: cross-file wasm-function composition. When a ModuleLoader is
+  // supplied, resolve `import { ... } from '<spec>'` statements at transpile
+  // time. Any imported names that correspond to `wasm function` declarations
+  // in the source module get pulled into the consumer's wasm module, with
+  // the import statement rewritten to a local JS wrapper. No loader supplied
+  // = no behavior change (imports stay verbatim, runtime resolves them).
+  const importedWasm = composeImportedWasmFunctions(source, {
+    loader: options.moduleLoader,
+    importerPath: options.filename,
+  })
+  source = importedWasm.source
+
   // Unified paren expression transformer
   // Handles: function params, arrow params, return types, safe/unsafe markers
   // Model: open paren can be ( or (? or (!, close can be ) or )-> or )-? or )-!
@@ -324,9 +346,24 @@ export function preprocess(
   source = polyResult.source
 
   // Extract WASM blocks: wasm(args) { ... } fallback { ... }
+  // `wasm function` declarations are already extracted earlier in the pipeline
+  // (see above, before transformParenExpressions). This finds the remaining
+  // inline `wasm { ... }` blocks inside regular tjs functions.
   const wasmBlocks = extractWasmBlocks(source)
   source = wasmBlocks.source
 
+  // Combine all flavors of wasm blocks for the downstream emitter.
+  // They're indistinguishable from the compiler's perspective — all have
+  // an id, body, captures, and need the same module composition treatment.
+  //   - wasmFunctions: top-level `wasm function NAME(...)` decls in this file
+  //   - importedWasm:  cross-file `wasm function`s pulled in via Phase 3
+  //   - wasmBlocks:    inline `wasm { ... }` blocks nested in tjs functions
+  const allWasmBlocks = [
+    ...wasmFunctions.blocks,
+    ...importedWasm.blocks,
+    ...wasmBlocks.blocks,
+  ]
+
   // Extract and run test blocks: test 'desc'? { body }
   // Tests run at transpile time and are stripped from output
   const testResult = extractAndRunTests(source, options.dangerouslySkipTests)
@@ -381,7 +418,7 @@ export function preprocess(
     requiredParams,
     unsafeFunctions,
     safeFunctions,
-    wasmBlocks: wasmBlocks.blocks,
+    wasmBlocks: allWasmBlocks,
     tests: testResult.tests,
     testErrors: testResult.errors,
     polymorphicNames: polyResult.polymorphicNames,
@@ -433,7 +470,11 @@ export function parse(
     letAnnotations,
     tjsModes,
   } = colonShorthand
-    ? preprocess(source, { vmTarget })
+    ? preprocess(source, {
+        vmTarget,
+        moduleLoader: options.moduleLoader,
+        filename: options.filename,
+      })
     : {
         source,
         returnType: undefined,