/**
* ndwasm.js
* Responsibility: Loading WASM, memory buffer management, NDArray bridge extensions.
*/
import { NDArray, DTYPE_MAP } from './ndarray_core.js';
/**
* WasmBuffer
*/
class WasmBuffer {
/**
* @param {Object} exports - The exports object from Go WASM.
* @param {number} size - Number of elements.
* @param {string} dtype - Data type (float64, float32, etc.).
*/
constructor(exports, size, dtype) {
this.exports = exports;
this.size = size;
this.dtype = dtype;
const Constructor = DTYPE_MAP[dtype];
if (!Constructor) throw new Error(`Unsupported dtype: ${dtype}`);
this.byteLength = size * Constructor.BYTES_PER_ELEMENT;
// 1. Allocate memory in the WASM heap.
if(this.byteLength === 0){
this.ptr = 0;
}else{
this.ptr = this.exports.malloc(this.byteLength);
if (!this.ptr) throw new Error(`WASM malloc failed to allocate ${this.byteLength} bytes`);
}
this.refresh();
}
/**
* Refreshes the view into WASM memory.
*/
refresh(){
const Constructor = DTYPE_MAP[this.dtype];
// Create a direct view into WASM memory (sharing the same ArrayBuffer).
this.view = new Constructor(this.exports.mem.buffer, this.ptr, this.size);
return this;
}
/** Synchronizes JS data into the WASM buffer. */
push(typedArray) {
this.refresh();
this.view.set(typedArray);
}
/** Pulls data from the WASM buffer back to JS (returns a copy).
* @returns {NDArray}
*/
pull() {
this.refresh();
return this.view.slice();
}
/** Disposes of the temporary buffer. */
dispose() {
if (this.ptr) {
this.exports.free(this.ptr);
this.ptr = null;
this.view = null;
}
}
}
/**
* WasmRuntime
*/
class WasmRuntime {
constructor() {
this.instance = null;
this.exports = null;
this.isLoaded = false;
}
/**
* Initializes the Go WASM environment.
* @param {object} [options] - Optional configuration.
* @param {string} [options.wasmUrl='./ndarray_plugin.wasm'] - Path or URL to the wasm file.
* @param {string} [options.execUrl='./wasm_exec.js'] - Path to the wasm_exec.js file (Node.js only).
* @param {number} [options.initialMemoryPages] - Initial memory size in 64KiB pages.
* @param {number} [options.maximumMemoryPages] - Maximum memory size in 64KiB pages.
* @param {string} [options.baseDir='.']
*/
async init(options = {}) {
options = {
initialMemoryPages : 1024*1024*64 / (64*1024),
maximumMemoryPages : 1024*1024*512 / (64*1024),
...options};
const wasmUrl = options.wasmUrl || `${options.baseDir || '.'}/ndarray_plugin.wasm`;
const execUrl = options.execUrl || `${options.baseDir || '.'}/wasm_exec.js`;
// Create the WebAssembly.Memory object
const wasmMemory = new WebAssembly.Memory({
initial: options.initialMemoryPages,
maximum: options.maximumMemoryPages
});
// Environment-specific loading
const isNode = typeof window === 'undefined';
if (isNode) {
// Node.js environment
const fs = require('fs');
const path = require('path');
// The wasm_exec.js from Go requires these globals.
global.fs = fs;
global.util = require('util');
// Load and execute the wasm_exec.js script to define the Go global.
require(path.resolve(process.cwd(), execUrl));
const wasmBytes = fs.readFileSync(path.resolve(process.cwd(), wasmUrl));
const go = new Go();
// Add the memory object to the import object
go.env['mem'] = wasmMemory;
const { instance } = await WebAssembly.instantiate(wasmBytes, go.importObject);
this.instance = instance;
go.run(this.instance);
} else {
// Browser environment
try{
await importScripts(execUrl);
}catch(e){
console.log(e);
}
if (typeof Go === 'undefined') {
console.error("Go's wasm_exec.js must be loaded with a <script> tag before initializing WasmRuntime in a browser.");
throw new Error("Missing Go global object.");
}
const go = new Go();
// Add the memory object to the import object
go.env['mem'] = wasmMemory;
const result = await WebAssembly.instantiateStreaming(fetch(wasmUrl), go.importObject);
this.instance = result.instance;
// The go.run promise resolves when the WASM exits, which in a reactor model is never. We don't await it.
go.run(this.instance);
}
this.exports = this.instance.exports;
this.isLoaded = true;
}
/**
* Helper method: Gets the suffix for Go export function names based on type.
* @private
*/
_getSuffix(dtype) {
if (dtype === 'float64') return '_F64';
if (dtype === 'float32') return '_F32';
throw new Error(`Go-Wasm compute does not support dtype: ${dtype}. Use float64 or float32.`);
}
/**
* Quickly allocates a buffer.
* @returns {WasmBuffer}
*/
createBuffer(size, dtype) {
return new WasmBuffer(this.exports, size, dtype);
}
}
/**
* Static factory: Creates a new NDArray directly from WASM computation results.
* @param {WasmBuffer} bridge - The WasmBuffer instance containing the result from WASM.
* @param {Array<number>} shape - The shape of the new NDArray.
* @param {string} [dtype] - The data type of the new NDArray. If not provided, uses bridge.dtype.
* @returns {NDArray} The new NDArray.
*/
export function fromWasm(bridge, shape, dtype) {
const data = bridge.pull(); // Data is a flat TypedArray pulled from WASM.
// Call the constructor directly as we already have the flat data array and the correct shape.
return new NDArray(data, { shape, dtype: dtype || bridge.dtype });
};
// --- Main WASM Bridge Object ---
import { NDWasmBlas } from './ndwasm_blas.js';
import { NDWasmDecomp } from './ndwasm_decomp.js';
import { NDWasmSignal } from './ndwasm_signal.js';
import { NDWasmAnalysis } from './ndwasm_analysis.js';
// --- Main WASM Bridge Object ---
/**
* @namespace NDWasm
*/
const NDWasm = {
runtime: null,
/** Binds a loaded WASM runtime to the bridge.
* @param {*} runtime
*/
bind(runtime) {
this.runtime = runtime;
},
/**
* Init the NDWasm
* @param {string} [baseDir='.']
*/
async init(baseDir = '.'){
const runtime = new WasmRuntime();
await runtime.init({baseDir});
NDWasm.bind(runtime);
},
/** Internal helper: executes a computation in WASM and manages memory. */
_compute(inputs, outShape, outDtype, computeFn) {
if (!this.runtime || !this.runtime.isLoaded) throw new Error("WasmRuntime not initialized");
// 1. Convert all input NDArrays to WASM buffer bridges.
const bridges = inputs.map(arr => arr.toWasm(this.runtime));
// 2. Prepare the output buffer bridge.
const size = outShape.reduce((a, b) => a * b, 1);
const outBridge = this.runtime.createBuffer(size, outDtype);
try {
// 3. Execute the Go function, passing pointers.
const status = computeFn(...bridges.map(b => b.ptr), outBridge.ptr);
if (status !== undefined && status !== 0) {
throw new Error(`WASM compute failed with status code: ${status}`);
}
// 4. Create a new JS NDArray by copying the data back from the output bridge.
return fromWasm(outBridge, outShape, outDtype);
} finally {
// 5. Free all temporary WASM memory.
bridges.forEach(b => b.dispose());
outBridge.dispose();
}
},
};
export const blas=NDWasmBlas
export const decomp=NDWasmDecomp;
export const signal=NDWasmSignal;
export const analysis=NDWasmAnalysis;
export { NDWasm, WasmRuntime, WasmBuffer };
export default NDWasm;