tjs-lang 0.7.7 → 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_')
@@ -3595,3 +3931,307 @@ function findMatchingOpen(
   }
   return pos
 }
+
+/**
+ * Transform `let x: <example>` and `let x: <example> = value` declarations.
+ *
+ * Strips the `: <example>` annotation so Acorn can parse, and records the
+ * variable name + example text so the linter and (later) type inference can
+ * use the annotation. Acorn rejects the colon since it is not valid JS.
+ *
+ *   let x: ''           → let x         (annotation: x → '')
+ *   let x: 0 = 5        → let x = 5     (annotation: x → 0)
+ *   let result: { ok: false } = ...     (annotation: result → { ok: false })
+ *
+ * Only `let` is processed. `const` always has an initializer, so the type
+ * is always inferable. `var` is rejected by TjsNoVar mode.
+ */
+export function transformLetTypeAnnotations(source: string): {
+  source: string
+  annotations: Map<string, string>
+} {
+  const annotations = new Map<string, string>()
+  if (!source.includes('let ')) return { source, annotations }
+
+  type Replacement = { start: number; end: number; replacement: string }
+  const replacements: Replacement[] = []
+
+  let i = 0
+  let state: TokenizerState = 'normal'
+  const templateStack: number[] = []
+
+  while (i < source.length) {
+    const char = source[i]
+    const nextChar = source[i + 1]
+
+    switch (state) {
+      case 'single-string':
+        if (char === '\\' && i + 1 < source.length) {
+          i += 2
+          continue
+        }
+        if (char === "'") state = 'normal'
+        i++
+        continue
+      case 'double-string':
+        if (char === '\\' && i + 1 < source.length) {
+          i += 2
+          continue
+        }
+        if (char === '"') state = 'normal'
+        i++
+        continue
+      case 'template-string':
+        if (char === '\\' && i + 1 < source.length) {
+          i += 2
+          continue
+        }
+        if (char === '$' && nextChar === '{') {
+          i += 2
+          templateStack.push(1)
+          state = 'normal'
+          continue
+        }
+        if (char === '`') state = 'normal'
+        i++
+        continue
+      case 'line-comment':
+        if (char === '\n') state = 'normal'
+        i++
+        continue
+      case 'block-comment':
+        if (char === '*' && nextChar === '/') {
+          i += 2
+          state = 'normal'
+          continue
+        }
+        i++
+        continue
+      case 'regex':
+        if (char === '\\' && i + 1 < source.length) {
+          i += 2
+          continue
+        }
+        if (char === '[') {
+          i++
+          while (i < source.length && source[i] !== ']') {
+            if (source[i] === '\\' && i + 1 < source.length) i += 2
+            else i++
+          }
+          if (i < source.length) i++
+          continue
+        }
+        if (char === '/') {
+          i++
+          while (i < source.length && /[gimsuy]/.test(source[i])) i++
+          state = 'normal'
+          continue
+        }
+        i++
+        continue
+      case 'normal':
+        if (templateStack.length > 0) {
+          if (char === '{') {
+            templateStack[templateStack.length - 1]++
+          } else if (char === '}') {
+            templateStack[templateStack.length - 1]--
+            if (templateStack[templateStack.length - 1] === 0) {
+              templateStack.pop()
+              i++
+              state = 'template-string'
+              continue
+            }
+          }
+        }
+        if (char === "'") {
+          i++
+          state = 'single-string'
+          continue
+        }
+        if (char === '"') {
+          i++
+          state = 'double-string'
+          continue
+        }
+        if (char === '`') {
+          i++
+          state = 'template-string'
+          continue
+        }
+        if (char === '/' && nextChar === '/') {
+          i += 2
+          state = 'line-comment'
+          continue
+        }
+        if (char === '/' && nextChar === '*') {
+          i += 2
+          state = 'block-comment'
+          continue
+        }
+        if (char === '/') {
+          let j = i - 1
+          while (j >= 0 && /\s/.test(source[j])) j--
+          const beforeChar = j >= 0 ? source[j] : ''
+          const isRegexContext =
+            !beforeChar ||
+            /[=(!,;:{[&|?+\-*%<>~^]/.test(beforeChar) ||
+            (j >= 5 &&
+              /\b(return|case|throw|in|of|typeof|instanceof|new|delete|void)$/.test(
+                source.slice(Math.max(0, j - 10), j + 1)
+              ))
+          if (isRegexContext) {
+            i++
+            state = 'regex'
+            continue
+          }
+        }
+
+        // Detect `let <ident> :` at top-level normal state
+        if (
+          char === 'l' &&
+          source.slice(i, i + 4) === 'let ' &&
+          (i === 0 || !/[\w$]/.test(source[i - 1]))
+        ) {
+          // Skip past `let` and whitespace
+          let j = i + 4
+          while (j < source.length && /\s/.test(source[j])) j++
+          // Match identifier
+          if (j < source.length && /[a-zA-Z_$]/.test(source[j])) {
+            const nameStart = j
+            while (j < source.length && /[\w$]/.test(source[j])) j++
+            const nameEnd = j
+            const varName = source.slice(nameStart, nameEnd)
+            // Skip whitespace; require a `:` (not `::` or part of `?:`)
+            let k = j
+            while (k < source.length && /\s/.test(source[k])) k++
+            if (
+              k < source.length &&
+              source[k] === ':' &&
+              source[k + 1] !== ':'
+            ) {
+              const colonPos = k
+              // Skip whitespace after colon
+              let exStart = colonPos + 1
+              while (exStart < source.length && /[ \t]/.test(source[exStart])) {
+                exStart++
+              }
+              // Scan example expression until `=`, `,`, `;`, or newline at depth 0
+              const exEnd = scanExampleEnd(source, exStart)
+              if (exEnd > exStart) {
+                const example = source.slice(exStart, exEnd).trim()
+                annotations.set(varName, example)
+                replacements.push({
+                  start: nameEnd,
+                  end: exEnd,
+                  replacement: '',
+                })
+                i = exEnd
+                continue
+              }
+            }
+          }
+        }
+        break
+    }
+    i++
+  }
+
+  if (replacements.length === 0) return { source, annotations }
+
+  // Apply right-to-left to preserve positions
+  let result = source
+  for (let k = replacements.length - 1; k >= 0; k--) {
+    const r = replacements[k]
+    result = result.slice(0, r.start) + r.replacement + result.slice(r.end)
+  }
+  return { source: result, annotations }
+}
+
+/**
+ * Scan forward from `start` and return the position where the example
+ * expression ends. Stops at `=`, `,`, `;`, or a newline when paren/brace/
+ * bracket depth is 0. Skips through nested brackets, strings, and templates.
+ */
+function scanExampleEnd(source: string, start: number): number {
+  let i = start
+  let parens = 0
+  let braces = 0
+  let brackets = 0
+  let state: 'normal' | 'sq' | 'dq' | 'tpl' = 'normal'
+  const templateStack: number[] = []
+  while (i < source.length) {
+    const c = source[i]
+    if (state === 'sq') {
+      if (c === '\\') {
+        i += 2
+        continue
+      }
+      if (c === "'") state = 'normal'
+      i++
+      continue
+    }
+    if (state === 'dq') {
+      if (c === '\\') {
+        i += 2
+        continue
+      }
+      if (c === '"') state = 'normal'
+      i++
+      continue
+    }
+    if (state === 'tpl') {
+      if (c === '\\') {
+        i += 2
+        continue
+      }
+      if (c === '$' && source[i + 1] === '{') {
+        templateStack.push(1)
+        state = 'normal'
+        i += 2
+        continue
+      }
+      if (c === '`') state = 'normal'
+      i++
+      continue
+    }
+    // normal
+    if (templateStack.length > 0) {
+      if (c === '{') templateStack[templateStack.length - 1]++
+      else if (c === '}') {
+        templateStack[templateStack.length - 1]--
+        if (templateStack[templateStack.length - 1] === 0) {
+          templateStack.pop()
+          state = 'tpl'
+          i++
+          continue
+        }
+      }
+    }
+    if (c === "'") {
+      state = 'sq'
+      i++
+      continue
+    }
+    if (c === '"') {
+      state = 'dq'
+      i++
+      continue
+    }
+    if (c === '`') {
+      state = 'tpl'
+      i++
+      continue
+    }
+    if (c === '(') parens++
+    else if (c === ')') parens--
+    else if (c === '{') braces++
+    else if (c === '}') braces--
+    else if (c === '[') brackets++
+    else if (c === ']') brackets--
+    if (parens === 0 && braces === 0 && brackets === 0) {
+      if (c === '=' || c === ',' || c === ';' || c === '\n') return i
+    }
+    i++
+  }
+  return i
+}