static/js/wasm_processor.js

/**
 * WebAssembly Face Mesh Processor
 * Handles high-performance mesh processing with WebAssembly
 */

class WasmFaceProcessor {
  constructor() {
    this.wasmModule = null;
    this.wasmMemory = null;
    this.wasmInstance = null;
    this.isReady = false;
    this.pendingTasks = [];
    
    // Load and initialize the WASM module
    this.initialize();
  }
  
  /**
   * Initialize the WebAssembly module
   */
  async initialize() {
    try {
      // Create memory with initial 16MB and maximum 128MB
      this.wasmMemory = new WebAssembly.Memory({ initial: 256, maximum: 2048 });
      
      // Import object for WASM instantiation
      const importObject = {
        env: {
          memory: this.wasmMemory,
          emscripten_notify_memory_growth: () => {},
          // Add any other required imports here
          // (Would typically include Math functions for Emscripten)
          powf: Math.pow,
          sqrt: Math.sqrt,
          expf: Math.exp,
          fabs: Math.abs,
          sinf: Math.sin,
          cosf: Math.cos,
          tanf: Math.tan,
          acosf: Math.acos,
          asinf: Math.asin,
          atanf: Math.atan,
          atan2f: Math.atan2,
          ceilf: Math.ceil,
          floorf: Math.floor,
          log: Math.log,
          console_log: function(ptr) {
            // Function to log messages from WASM
            console.log("WASM Log:", UTF8ToString(ptr));
          }
        }
      };
      
      // Fetch and instantiate the WebAssembly module
      const response = await fetch('/static/js/wasm/face_mesh_processor.wasm');
      const wasmBytes = await response.arrayBuffer();
      const { instance, module } = await WebAssembly.instantiate(wasmBytes, importObject);
      
      this.wasmModule = module;
      this.wasmInstance = instance;
      this.isReady = true;
      
      console.log("WebAssembly face processor module loaded successfully");
      
      // Process any tasks that were queued while loading
      this.processPendingTasks();
    } catch (error) {
      console.error("Failed to initialize WebAssembly face processor:", error);
    }
  }
  
  /**
   * Process any tasks that were queued while the WASM module was loading
   */
  processPendingTasks() {
    while (this.pendingTasks.length > 0) {
      const task = this.pendingTasks.shift();
      task();
    }
  }
  
  /**
   * Execute a function when the WASM module is ready
   * @param {Function} callback - Function to execute
   */
  whenReady(callback) {
    if (this.isReady) {
      callback();
    } else {
      this.pendingTasks.push(callback);
    }
  }
  
  /**
   * Simplify a 3D facial mesh while preserving important features
   * @param {Object} meshData - Object containing mesh data
   * @param {Float32Array} meshData.vertices - Vertex positions and normals (x,y,z,nx,ny,nz,...)
   * @param {Uint32Array} meshData.indices - Triangle indices
   * @param {Array} meshData.landmarks - Facial landmarks
   * @param {number} targetRatio - Target ratio of vertices to keep (0-1)
   * @returns {Promise<Object>} - Simplified mesh data
   */
  async simplifyMesh(meshData, targetRatio) {
    return new Promise((resolve, reject) => {
      this.whenReady(() => {
        try {
          // Get WASM exports
          const { 
            initMesh, 
            setLandmarks, 
            calculateVertexImportance, 
            simplifyMesh, 
            freeMesh 
          } = this.wasmInstance.exports;
          
          // Calculate target vertex count
          const vertexCount = meshData.vertices.length / 6; // Each vertex has x,y,z,nx,ny,nz
          const targetVertexCount = Math.max(100, Math.floor(vertexCount * targetRatio));
          
          console.log(`Simplifying mesh from ${vertexCount} to approximately ${targetVertexCount} vertices`);
          
          // Allocate memory for vertex data in WASM
          const vertexDataPtr = this._allocateFloat32Array(meshData.vertices);
          
          // Allocate memory for triangle indices in WASM
          const indexDataPtr = this._allocateInt32Array(meshData.indices);
          
          // Initialize mesh in WASM
          const initResult = initMesh(
            vertexDataPtr, 
            vertexCount, 
            indexDataPtr, 
            meshData.indices.length / 3
          );
          
          if (!initResult) {
            freeMesh();
            this._freeMemory(vertexDataPtr);
            this._freeMemory(indexDataPtr);
            reject(new Error("Failed to initialize mesh in WebAssembly module"));
            return;
          }
          
          // If landmarks are provided, set them
          if (meshData.landmarks && meshData.landmarks.length > 0) {
            // Convert landmarks to flat array in format expected by WASM: x,y,z,type
            const landmarksFlat = new Float32Array(meshData.landmarks.length * 4);
            for (let i = 0; i < meshData.landmarks.length; i++) {
              const landmark = meshData.landmarks[i];
              landmarksFlat[i * 4] = landmark.position.x;
              landmarksFlat[i * 4 + 1] = landmark.position.y;
              landmarksFlat[i * 4 + 2] = landmark.position.z;
              landmarksFlat[i * 4 + 3] = landmark.type || 0;
            }
            
            // Allocate memory for landmarks in WASM
            const landmarksPtr = this._allocateFloat32Array(landmarksFlat);
            
            // Set landmarks in WASM
            const landmarkResult = setLandmarks(landmarksPtr, meshData.landmarks.length);
            
            // Free landmarks memory
            this._freeMemory(landmarksPtr);
            
            if (!landmarkResult) {
              freeMesh();
              this._freeMemory(vertexDataPtr);
              this._freeMemory(indexDataPtr);
              reject(new Error("Failed to set landmarks in WebAssembly module"));
              return;
            }
            
            // Calculate vertex importance based on landmarks
            calculateVertexImportance();
          }
          
          // Allocate memory for output parameters
          const resultVertexCountPtr = this._allocateInt32(0);
          const resultTriangleCountPtr = this._allocateInt32(0);
          
          // Simplify the mesh
          const resultDataPtr = simplifyMesh(targetVertexCount, resultVertexCountPtr, resultTriangleCountPtr);
          
          // Read output parameters
          const resultVertexCount = this._readInt32(resultVertexCountPtr);
          const resultTriangleCount = this._readInt32(resultTriangleCountPtr);
          
          console.log(`Mesh simplified to ${resultVertexCount} vertices and ${resultTriangleCount} triangles`);
          
          // Read result data
          let simplifiedMesh = null;
          
          if (resultDataPtr) {
            // Result format: [v1x, v1y, v1z, v1nx, v1ny, v1nz, ..., t1v1, t1v2, t1v3, ...]
            const resultDataSize = 2 + (resultVertexCount * 6) + (resultTriangleCount * 3);
            const resultData = this._readFloat32Array(resultDataPtr, resultDataSize);
            
            // Extract vertices and indices
            const vertices = new Float32Array(resultVertexCount * 6);
            const indices = new Uint32Array(resultTriangleCount * 3);
            
            // Copy vertex data
            for (let i = 0; i < resultVertexCount; i++) {
              for (let j = 0; j < 6; j++) {
                vertices[i * 6 + j] = resultData[2 + i * 6 + j];
              }
            }
            
            // Copy index data
            for (let i = 0; i < resultTriangleCount * 3; i++) {
              indices[i] = resultData[2 + resultVertexCount * 6 + i];
            }
            
            simplifiedMesh = {
              vertices,
              indices,
              vertexCount: resultVertexCount,
              triangleCount: resultTriangleCount
            };
            
            // Free result data memory
            this._freeMemory(resultDataPtr);
          }
          
          // Clean up
          freeMesh();
          this._freeMemory(vertexDataPtr);
          this._freeMemory(indexDataPtr);
          this._freeMemory(resultVertexCountPtr);
          this._freeMemory(resultTriangleCountPtr);
          
          if (simplifiedMesh) {
            resolve(simplifiedMesh);
          } else {
            reject(new Error("Mesh simplification failed"));
          }
          
        } catch (error) {
          console.error("Error in WebAssembly mesh simplification:", error);
          reject(error);
        }
      });
    });
  }
  
  /**
   * Allocate memory for a Float32Array in the WASM module
   * @param {Float32Array} array - The array to allocate
   * @returns {number} - Pointer to the allocated memory
   */
  _allocateFloat32Array(array) {
    const bytesPerElement = Float32Array.BYTES_PER_ELEMENT;
    const numBytes = array.length * bytesPerElement;
    const ptr = this.wasmInstance.exports.malloc(numBytes);
    
    const heap = new Uint8Array(this.wasmMemory.buffer);
    const heapView = new Float32Array(this.wasmMemory.buffer);
    
    for (let i = 0; i < array.length; i++) {
      heapView[ptr / bytesPerElement + i] = array[i];
    }
    
    return ptr;
  }
  
  /**
   * Allocate memory for an Int32Array in the WASM module
   * @param {Int32Array | Uint32Array} array - The array to allocate
   * @returns {number} - Pointer to the allocated memory
   */
  _allocateInt32Array(array) {
    const bytesPerElement = Int32Array.BYTES_PER_ELEMENT;
    const numBytes = array.length * bytesPerElement;
    const ptr = this.wasmInstance.exports.malloc(numBytes);
    
    const heap = new Uint8Array(this.wasmMemory.buffer);
    const heapView = new Int32Array(this.wasmMemory.buffer);
    
    for (let i = 0; i < array.length; i++) {
      heapView[ptr / bytesPerElement + i] = array[i];
    }
    
    return ptr;
  }
  
  /**
   * Allocate memory for a single int in the WASM module
   * @param {number} value - The value to store
   * @returns {number} - Pointer to the allocated memory
   */
  _allocateInt32(value) {
    const bytesPerElement = Int32Array.BYTES_PER_ELEMENT;
    const ptr = this.wasmInstance.exports.malloc(bytesPerElement);
    
    const heapView = new Int32Array(this.wasmMemory.buffer);
    heapView[ptr / bytesPerElement] = value;
    
    return ptr;
  }
  
  /**
   * Read a Float32Array from WASM memory
   * @param {number} ptr - Pointer to the memory
   * @param {number} length - Number of elements to read
   * @returns {Float32Array} - The read array
   */
  _readFloat32Array(ptr, length) {
    const bytesPerElement = Float32Array.BYTES_PER_ELEMENT;
    return new Float32Array(
      this.wasmMemory.buffer.slice(ptr, ptr + length * bytesPerElement)
    );
  }
  
  /**
   * Read an int from WASM memory
   * @param {number} ptr - Pointer to the memory
   * @returns {number} - The read value
   */
  _readInt32(ptr) {
    const bytesPerElement = Int32Array.BYTES_PER_ELEMENT;
    const heapView = new Int32Array(this.wasmMemory.buffer);
    return heapView[ptr / bytesPerElement];
  }
  
  /**
   * Free allocated memory in the WASM module
   * @param {number} ptr - Pointer to the memory to free
   */
  _freeMemory(ptr) {
    if (ptr) {
      this.wasmInstance.exports.free(ptr);
    }
  }
}

/**
 * Worker wrapper for the WebAssembly face processor
 * This handles communication with the main thread
 */
class WasmProcessorWorker {
  constructor() {
    this.processor = new WasmFaceProcessor();
    
    // Set up message handling
    self.addEventListener('message', this.handleMessage.bind(this));
  }
  
  /**
   * Handle messages from the main thread
   * @param {MessageEvent} event - The message event
   */
  async handleMessage(event) {
    const { type, id, data } = event.data;
    
    switch (type) {
      case 'simplifyMesh':
        try {
          const result = await this.processor.simplifyMesh(data.mesh, data.targetRatio);
          self.postMessage({ id, type: 'result', result });
        } catch (error) {
          self.postMessage({ id, type: 'error', error: error.message });
        }
        break;
        
      default:
        self.postMessage({ id, type: 'error', error: `Unknown command: ${type}` });
    }
  }
}

// Initialize the worker if running in a worker context
if (typeof WorkerGlobalScope !== 'undefined' && self instanceof WorkerGlobalScope) {
  new WasmProcessorWorker();
} else {
  // For main thread use, export the processor class
  export { WasmFaceProcessor };
}