static/js/wasm_loader.js

/**
 * WebAssembly Module Loader for MorphGuard
 * Handles loading and initialization of WebAssembly modules for high-performance tasks
 */

class WasmLoader {
    constructor(options = {}) {
        this.config = Object.assign({
            wasmPath: '/static/wasm/',
            simdSupported: null,
            threadingSupported: null,
            fallbackToJS: true
        }, options);

        // Track loaded modules
        this.modules = {};
        
        // Check for WebAssembly support and features
        this.detectFeatures();
    }
    
    /**
     * Detect WebAssembly features
     */
    async detectFeatures() {
        // Check basic WebAssembly support
        if (typeof WebAssembly === 'undefined') {
            this.config.simdSupported = false;
            this.config.threadingSupported = false;
            console.warn('WebAssembly not supported in this browser. Using JavaScript fallbacks.');
            return;
        }
        
        // Check for SIMD support
        if (this.config.simdSupported === null) {
            try {
                // Test for SIMD support
                const simdTest = await WebAssembly.instantiate(
                    new Uint8Array([0, 97, 115, 109, 1, 0, 0, 0, 1, 5, 1, 96, 0, 1, 123, 3, 2, 1, 0, 10, 10, 1, 8, 0, 65, 0, 253, 15, 253, 98, 11]),
                    {}
                );
                this.config.simdSupported = !!simdTest.instance.exports;
            } catch (e) {
                this.config.simdSupported = false;
            }
        }
        
        // Check for threading support
        if (this.config.threadingSupported === null) {
            this.config.threadingSupported = typeof SharedArrayBuffer !== 'undefined' && 
                                             typeof Atomics !== 'undefined';
        }
        
        console.log(`WebAssembly features: SIMD: ${this.config.simdSupported}, Threading: ${this.config.threadingSupported}`);
    }
    
    /**
     * Load a WebAssembly module
     * @param {string} moduleName - Name of the module to load
     * @param {object} importObject - Import object for WebAssembly instantiation
     * @returns {Promise<object>} - The loaded WebAssembly module
     */
    async loadModule(moduleName, importObject = {}) {
        // Check if already loaded
        if (this.modules[moduleName]) {
            return this.modules[moduleName];
        }
        
        try {
            // Determine file name based on available features
            let fileName = moduleName;
            
            if (this.config.simdSupported && this.config.threadingSupported) {
                fileName += '_simd_threaded.wasm';
            } else if (this.config.simdSupported) {
                fileName += '_simd.wasm';
            } else {
                fileName += '.wasm';
            }
            
            // Load the WebAssembly module
            const response = await fetch(`${this.config.wasmPath}${fileName}`);
            
            if (!response.ok) {
                throw new Error(`Failed to load WebAssembly module: ${response.statusText}`);
            }
            
            const wasmBinary = await response.arrayBuffer();
            
            // Set up default import object with memory and console logging
            const defaultImports = {
                env: {
                    memory: new WebAssembly.Memory({ 
                        initial: 256, 
                        maximum: 512,
                        shared: this.config.threadingSupported
                    }),
                    abort: (msg, file, line, column) => {
                        console.error(`Abort called from ${file}:${line}:${column}: ${msg}`);
                    },
                    log: (ptr, len) => {
                        // In a real implementation, this would extract string from memory
                        console.log('WASM log message');
                    }
                },
                wasi_snapshot_preview1: {
                    proc_exit: () => {},
                    fd_close: () => {},
                    fd_write: () => {},
                    fd_seek: () => {},
                    fd_read: () => {}
                }
            };
            
            // Merge with provided import object
            const mergedImports = this.mergeImportObjects(defaultImports, importObject);
            
            // Instantiate the WebAssembly module
            const result = await WebAssembly.instantiate(wasmBinary, mergedImports);
            
            // Extract exports and create wrapper
            const instance = result.instance;
            const exports = instance.exports;
            
            // Create module wrapper with exports and memory helpers
            const moduleWrapper = {
                instance,
                exports,
                memory: instance.exports.memory || mergedImports.env.memory,
                getMemoryBuffer: () => new Uint8Array(
                    (instance.exports.memory || mergedImports.env.memory).buffer
                ),
                allocateMemory: (size) => {
                    if (typeof exports.malloc !== 'function') {
                        throw new Error('Memory allocation function not available');
                    }
                    return exports.malloc(size);
                },
                freeMemory: (ptr) => {
                    if (typeof exports.free !== 'function') {
                        throw new Error('Memory deallocation function not available');
                    }
                    exports.free(ptr);
                }
            };
            
            // Cache the module
            this.modules[moduleName] = moduleWrapper;
            
            return moduleWrapper;
        } catch (error) {
            console.error(`Error loading WebAssembly module '${moduleName}':`, error);
            
            if (this.config.fallbackToJS) {
                console.warn(`Falling back to JavaScript implementation for '${moduleName}'`);
                return this.loadJSFallback(moduleName);
            }
            
            throw error;
        }
    }
    
    /**
     * Load a JavaScript fallback implementation
     * @param {string} moduleName - Name of the module
     * @returns {Promise<object>} - JavaScript fallback module
     */
    async loadJSFallback(moduleName) {
        try {
            // Check if fallback script exists
            const scriptPath = `${this.config.wasmPath}fallbacks/${moduleName}.js`;
            
            // Create and append script element
            return new Promise((resolve, reject) => {
                const script = document.createElement('script');
                script.src = scriptPath;
                script.onload = () => {
                    // Access global object created by script
                    const fallbackModule = window[`${moduleName}Fallback`];
                    if (!fallbackModule) {
                        reject(new Error(`JavaScript fallback module '${moduleName}' not found`));
                        return;
                    }
                    
                    // Cache the fallback module
                    this.modules[moduleName] = {
                        isJSFallback: true,
                        ...fallbackModule
                    };
                    
                    resolve(this.modules[moduleName]);
                };
                script.onerror = () => {
                    reject(new Error(`Failed to load JavaScript fallback for '${moduleName}'`));
                };
                
                document.head.appendChild(script);
            });
        } catch (error) {
            console.error(`Error loading JavaScript fallback for '${moduleName}':`, error);
            throw error;
        }
    }
    
    /**
     * Merge two import objects
     * @param {object} obj1 - First import object
     * @param {object} obj2 - Second import object
     * @returns {object} - Merged import object
     */
    mergeImportObjects(obj1, obj2) {
        const result = { ...obj1 };
        
        for (const key in obj2) {
            if (obj2.hasOwnProperty(key)) {
                if (result[key] && typeof result[key] === 'object' && typeof obj2[key] === 'object') {
                    result[key] = { ...result[key], ...obj2[key] };
                } else {
                    result[key] = obj2[key];
                }
            }
        }
        
        return result;
    }
    
    /**
     * Check if a WebAssembly module is available
     * @param {string} moduleName - Name of the module
     * @returns {boolean} - True if module is loaded
     */
    isModuleLoaded(moduleName) {
        return !!this.modules[moduleName];
    }
    
    /**
     * Get a loaded WebAssembly module
     * @param {string} moduleName - Name of the module
     * @returns {object|null} - The loaded module or null if not loaded
     */
    getModule(moduleName) {
        return this.modules[moduleName] || null;
    }
    
    /**
     * Copy data to WebAssembly memory
     * @param {object} module - WebAssembly module object
     * @param {ArrayBuffer|TypedArray} data - Data to copy
     * @returns {number} - Pointer to the allocated memory
     */
    copyToWasmMemory(module, data) {
        if (!module || module.isJSFallback) {
            throw new Error('Cannot copy to memory of JavaScript fallback module');
        }
        
        // Allocate memory
        const size = data.byteLength || data.length;
        const ptr = module.allocateMemory(size);
        
        if (!ptr) {
            throw new Error('Memory allocation failed');
        }
        
        // Get memory buffer
        const memory = module.getMemoryBuffer();
        
        // Copy data
        if (data instanceof ArrayBuffer) {
            memory.set(new Uint8Array(data), ptr);
        } else if (ArrayBuffer.isView(data)) {
            memory.set(new Uint8Array(data.buffer, data.byteOffset, data.byteLength), ptr);
        } else {
            throw new Error('Unsupported data type for copying to WebAssembly memory');
        }
        
        return ptr;
    }
    
    /**
     * Copy data from WebAssembly memory
     * @param {object} module - WebAssembly module object
     * @param {number} ptr - Pointer to memory
     * @param {number} size - Size of data to copy
     * @returns {Uint8Array} - Copied data
     */
    copyFromWasmMemory(module, ptr, size) {
        if (!module || module.isJSFallback) {
            throw new Error('Cannot copy from memory of JavaScript fallback module');
        }
        
        // Get memory buffer
        const memory = module.getMemoryBuffer();
        
        // Copy data
        return memory.slice(ptr, ptr + size);
    }
    
    /**
     * Convert string to WebAssembly memory
     * @param {object} module - WebAssembly module object
     * @param {string} str - String to convert
     * @returns {object} - Pointer and size of string in memory
     */
    stringToWasmMemory(module, str) {
        const encoder = new TextEncoder();
        const encoded = encoder.encode(str);
        const ptr = this.copyToWasmMemory(module, encoded);
        
        return {
            ptr,
            size: encoded.length
        };
    }
    
    /**
     * Convert WebAssembly memory to string
     * @param {object} module - WebAssembly module object
     * @param {number} ptr - Pointer to memory
     * @param {number} size - Size of string data
     * @returns {string} - Decoded string
     */
    wasmMemoryToString(module, ptr, size) {
        const memory = this.copyFromWasmMemory(module, ptr, size);
        const decoder = new TextDecoder();
        return decoder.decode(memory);
    }
}

// Export for use in main script
window.WasmLoader = WasmLoader;