@statsim/compiler 0.1.0

package/package.json

@@ -0,0 +1,23 @@
+{
+  "name": "@statsim/compiler",
+  "version": "0.1.0",
+  "description": "StatSim language compiler — parse .sm models and compile to stoch, PyMC, Stan, Z3",
+  "type": "module",
+  "main": "src/index.js",
+  "exports": {
+    ".": "./src/index.js"
+  },
+  "files": [
+    "src/",
+    "wasm/"
+  ],
+  "publishConfig": {
+    "access": "public"
+  },
+  "scripts": {
+    "test": "node --experimental-vm-modules test/test.js"
+  },
+  "keywords": ["statsim", "probabilistic-programming", "bayesian", "compiler"],
+  "author": "Anton Zemlyansky",
+  "license": "MIT"
+}

package/src/index.js

@@ -0,0 +1,83 @@
+import { loadWasm, getWasm } from './wasm.js'
+import { TfpjsRunner } from './runners/tfpjs.js'
+
+// Ensure WASM is loaded before any operations
+let initialized = false
+
+async function ensureInit () {
+  if (!initialized) {
+    await loadWasm()
+    initialized = true
+  }
+}
+
+/**
+ * Parse .sm source into an opaque block handle.
+ * @param {string} source - StatSim model source
+ * @returns {Promise<number>} block handle
+ */
+export async function parse (source) {
+  await ensureInit()
+  const wasm = getWasm()
+  const h1 = wasm.parse(source)
+  const h2 = wasm.expand(h1)
+  wasm.free(h1)
+  return h2
+}
+
+/**
+ * Compile blocks to target code (script mode).
+ * @param {number} handle - block handle from parse()
+ * @param {object} opts - { target: 'tfpjs'|'pymc'|'stan'|'z3' }
+ * @returns {string} generated code
+ */
+export function compileScript (handle, opts) {
+  const wasm = getWasm()
+  return wasm.compileScript(handle, opts.target)
+}
+
+/**
+ * Compile blocks to target code (module mode) with metadata.
+ * @param {number} handle - block handle from parse()
+ * @param {object} opts - { target: 'tfpjs'|'pymc'|'stan'|'z3' }
+ * @returns {{ code: string, dataNames: string[], paramNames: string[], inferenceMode: string }}
+ */
+export function compile (handle, opts) {
+  const wasm = getWasm()
+  return wasm.compileModule(handle, opts.target)
+}
+
+/**
+ * Free a block handle when no longer needed.
+ * @param {number} handle
+ */
+export function free (handle) {
+  const wasm = getWasm()
+  wasm.free(handle)
+}
+
+/**
+ * High-level API: parse + compile + return a runner.
+ * @param {string} source - StatSim model source
+ * @param {object} opts - { target: 'tfpjs'|'pymc'|'stan'|'z3' }
+ * @returns {Promise<TfpjsRunner>} runner with .run(), .build(), .code
+ */
+export async function model (source, opts = { target: 'tfpjs' }) {
+  const handle = await parse(source)
+  const result = compile(handle, opts)
+  free(handle)
+
+  if (opts.target === 'tfpjs') {
+    return new TfpjsRunner(result)
+  }
+
+  // Generic runner for other targets — code + metadata only
+  return {
+    code: result.code,
+    dataNames: result.dataNames,
+    paramNames: result.paramNames,
+    inferenceMode: result.inferenceMode
+  }
+}
+
+export { loadWasm }

package/src/runners/tfpjs.js

@@ -0,0 +1,99 @@
+/**
+ * TFP.js runner — executes compiled module code in the host environment.
+ *
+ * Supports three modes:
+ *   runner.run(data, options)   — run inference, return results
+ *   runner.build(data)          — return model object for custom inference
+ *   runner.code                 — raw code string
+ *   runner.exports              — lazy build() with no data
+ *
+ * Host functions are passed alongside data — they're auto-separated:
+ *   runner.run({ images, labels, encode: myEncoder })
+ */
+export class TfpjsRunner {
+  /**
+   * @param {{ code: string, dataNames: string[], paramNames: string[], inferenceMode: string }} result
+   */
+  constructor (result) {
+    this.code = result.code
+    this.dataNames = result.dataNames
+    this.paramNames = result.paramNames
+    this.inferenceMode = result.inferenceMode
+    this._module = null
+    this._exports = null
+  }
+
+  /**
+   * Lazily load the generated module code.
+   * @returns {Promise<object>} the evaluated module exports
+   */
+  async _loadModule () {
+    if (this._module) return this._module
+    const blob = new Blob([this.code], { type: 'text/javascript' })
+    const url = URL.createObjectURL(blob)
+    try {
+      this._module = await import(url)
+    } finally {
+      URL.revokeObjectURL(url)
+    }
+    return this._module
+  }
+
+  /**
+   * Separate data from host functions in a merged input object.
+   * Functions are extracted into ext, everything else stays as data.
+   */
+  _separate (input) {
+    const data = {}
+    const ext = {}
+    for (const [k, v] of Object.entries(input || {})) {
+      if (typeof v === 'function') {
+        ext[k] = v
+      } else {
+        data[k] = v
+      }
+    }
+    return { data, ext }
+  }
+
+  /**
+   * Run inference.
+   * @param {object} [input={}] - data + host functions (auto-separated)
+   * @param {object} [options={}] - inference options (numResults, numBurninSteps, etc.)
+   * @returns {Promise<object>} inference results keyed by parameter name
+   */
+  async run (input = {}, options = {}) {
+    const mod = await this._loadModule()
+    const { data, ext } = this._separate(input)
+    if (this.inferenceMode === 'optimize') {
+      return mod.runOptimization(data, ext, options)
+    }
+    return mod.runInference(data, ext, options)
+  }
+
+  /**
+   * Build model without running inference.
+   * Returns the model object for custom extension.
+   * @param {object} [input={}] - data + host functions (auto-separated)
+   * @returns {Promise<object>} model object with targetLogProb, bijectors, etc.
+   */
+  async build (input = {}) {
+    const mod = await this._loadModule()
+    const { data, ext } = this._separate(input)
+    if (this.inferenceMode === 'optimize') {
+      return mod.createOptimizer(data, ext)
+    }
+    return mod.createModel(data, ext)
+  }
+
+  /**
+   * Lazy exports — calls build() with empty data.
+   * @returns {Promise<object>} model object
+   */
+  get exports () {
+    if (!this._exports) {
+      this._exports = this.build()
+    }
+    return this._exports
+  }
+}

package/src/wasm.js

@@ -0,0 +1,16 @@
+// WASM loader — loads the statsim compiler WASM module
+let wasmModule = null
+
+export async function loadWasm () {
+  if (wasmModule) return wasmModule
+
+  // Dynamic import of the emscripten-generated glue code
+  const { default: createStatsim } = await import('../wasm/statsim.js')
+  wasmModule = await createStatsim()
+  return wasmModule
+}
+
+export function getWasm () {
+  if (!wasmModule) throw new Error('WASM not loaded — call loadWasm() first')
+  return wasmModule
+}