static/js/draco_gpu_decompressor.js

/**
 * GPU-Accelerated Draco Mesh Decompression
 * Uses WebGL compute shaders to efficiently decompress Draco-encoded 3D meshes
 */

class DracoGPUDecompressor {
  constructor() {
    this.gl = null;
    this.extensions = null;
    this.programs = {};
    this.buffers = {};
    this.textures = {};
    this.dracoDecoder = null;
    this.isInitialized = false;
    this.maxVertices = 100000; // Maximum vertices to process in one batch
    this.maxIndices = 300000;  // Maximum indices to process in one batch
    
    // Initialize WebGL and shaders
    this.initialize();
  }
  
  /**
   * Initialize WebGL context and required extensions
   */
  async initialize() {
    try {
      // Create WebGL 2.0 context
      const canvas = document.createElement('canvas');
      canvas.width = 1;
      canvas.height = 1;
      this.gl = canvas.getContext('webgl2', { 
        premultipliedAlpha: false,
        antialias: false,
        depth: false,
        stencil: false
      });
      
      if (!this.gl) {
        throw new Error('WebGL 2.0 not supported');
      }
      
      // Check for required extensions
      this.extensions = {
        computeShader: this.gl.getExtension('OES_texture_float_linear'),
        floatTextures: this.gl.getExtension('OES_texture_float'),
        drawBuffers: this.gl.getExtension('WEBGL_draw_buffers'),
        instancedArrays: this.gl.getExtension('ANGLE_instanced_arrays')
      };
      
      // Make sure we have the required extensions
      if (!this.extensions.floatTextures) {
        throw new Error('OES_texture_float extension not supported');
      }
      
      // Initialize Draco decoder (web worker)
      await this.initDracoDecoder();
      
      // Compile shaders
      await this.initShaders();
      
      console.log('GPU Draco decompressor initialized successfully');
      this.isInitialized = true;
      
    } catch (error) {
      console.error('Failed to initialize GPU Draco decompressor:', error);
      // Fall back to CPU-based decompression
      this.useFallback = true;
    }
  }
  
  /**
   * Initialize Draco decoder (CPU-side preprocessing)
   */
  async initDracoDecoder() {
    return new Promise((resolve, reject) => {
      // Create a web worker for initial Draco decoding steps
      this.dracoWorker = new Worker('/static/js/workers/draco_worker.js');
      
      // Handle worker messages
      this.dracoWorker.onmessage = (e) => {
        const { type, result } = e.data;
        
        if (type === 'initialized') {
          resolve();
        } else if (type === 'error') {
          reject(new Error(result.message));
        }
      };
      
      // Initialize worker
      this.dracoWorker.postMessage({ type: 'initialize' });
    });
  }
  
  /**
   * Initialize WebGL shader programs for decompression steps
   */
  async initShaders() {
    // Shader to decompress vertex positions
    const positionVertexShader = `
      #version 300 es
      in vec2 position;
      in vec2 texCoord;
      out vec2 vTexCoord;
      
      void main() {
        vTexCoord = texCoord;
        gl_Position = vec4(position, 0.0, 1.0);
      }
    `;
    
    // Shader to decompress vertex positions
    const positionFragmentShader = `
      #version 300 es
      precision highp float;
      precision highp sampler2D;
      
      in vec2 vTexCoord;
      out vec4 fragColor;
      
      uniform sampler2D uOctCoordsTex;
      uniform sampler2D uQuantizedPosTex;
      uniform vec3 uDequantizationFactors;
      uniform vec3 uPositionOffset;
      
      void main() {
        vec4 quantized = texture(uQuantizedPosTex, vTexCoord);
        vec3 position = quantized.xyz * uDequantizationFactors + uPositionOffset;
        fragColor = vec4(position, 1.0);
      }
    `;
    
    // Shader to decompress normals
    const normalFragmentShader = `
      #version 300 es
      precision highp float;
      precision highp sampler2D;
      
      in vec2 vTexCoord;
      out vec4 fragColor;
      
      uniform sampler2D uOctCoordsTex;
      
      // Octahedral normal unpacking
      vec3 decodeOctNormal(vec2 oct) {
        oct = oct * 2.0 - 1.0; // from [0,1] to [-1,1]
        
        // Decode the octahedral normal encoding
        vec3 normal;
        normal.z = 1.0 - abs(oct.x) - abs(oct.y);
        
        if (normal.z < 0.0) {
          normal.xy = (1.0 - abs(oct.yx)) * vec2(
            oct.x >= 0.0 ? 1.0 : -1.0,
            oct.y >= 0.0 ? 1.0 : -1.0
          );
        } else {
          normal.xy = oct.xy;
        }
        
        return normalize(normal);
      }
      
      void main() {
        vec2 octCoords = texture(uOctCoordsTex, vTexCoord).xy;
        vec3 normal = decodeOctNormal(octCoords);
        fragColor = vec4(normal * 0.5 + 0.5, 1.0);
      }
    `;
    
    // Shader to decompress texture coordinates
    const uvFragmentShader = `
      #version 300 es
      precision highp float;
      precision highp sampler2D;
      
      in vec2 vTexCoord;
      out vec4 fragColor;
      
      uniform sampler2D uQuantizedUVsTex;
      uniform vec2 uUVFactors;
      
      void main() {
        vec2 quantizedUV = texture(uQuantizedUVsTex, vTexCoord).xy;
        vec2 uv = quantizedUV * uUVFactors;
        fragColor = vec4(uv, 0.0, 1.0);
      }
    `;
    
    // Create shader programs
    this.programs.position = this.createProgram(positionVertexShader, positionFragmentShader);
    this.programs.normal = this.createProgram(positionVertexShader, normalFragmentShader);
    this.programs.uv = this.createProgram(positionVertexShader, uvFragmentShader);
    
    // Create a full-screen quad for rendering to textures
    this.createFullScreenQuad();
  }
  
  /**
   * Create a WebGL shader program from vertex and fragment shader sources
   */
  createProgram(vertexShaderSource, fragmentShaderSource) {
    const gl = this.gl;
    
    // Create and compile vertex shader
    const vertexShader = gl.createShader(gl.VERTEX_SHADER);
    gl.shaderSource(vertexShader, vertexShaderSource);
    gl.compileShader(vertexShader);
    
    if (!gl.getShaderParameter(vertexShader, gl.COMPILE_STATUS)) {
      const error = gl.getShaderInfoLog(vertexShader);
      gl.deleteShader(vertexShader);
      throw new Error(`Failed to compile vertex shader: ${error}`);
    }
    
    // Create and compile fragment shader
    const fragmentShader = gl.createShader(gl.FRAGMENT_SHADER);
    gl.shaderSource(fragmentShader, fragmentShaderSource);
    gl.compileShader(fragmentShader);
    
    if (!gl.getShaderParameter(fragmentShader, gl.COMPILE_STATUS)) {
      const error = gl.getShaderInfoLog(fragmentShader);
      gl.deleteShader(vertexShader);
      gl.deleteShader(fragmentShader);
      throw new Error(`Failed to compile fragment shader: ${error}`);
    }
    
    // Create and link program
    const program = gl.createProgram();
    gl.attachShader(program, vertexShader);
    gl.attachShader(program, fragmentShader);
    gl.linkProgram(program);
    
    if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
      const error = gl.getProgramInfoLog(program);
      gl.deleteProgram(program);
      gl.deleteShader(vertexShader);
      gl.deleteShader(fragmentShader);
      throw new Error(`Failed to link shader program: ${error}`);
    }
    
    // Clean up shaders, they're linked into the program now
    gl.deleteShader(vertexShader);
    gl.deleteShader(fragmentShader);
    
    // Get attribute and uniform locations
    const attributes = {};
    const uniforms = {};
    
    // Get attribute locations
    const numAttributes = gl.getProgramParameter(program, gl.ACTIVE_ATTRIBUTES);
    for (let i = 0; i < numAttributes; i++) {
      const info = gl.getActiveAttrib(program, i);
      attributes[info.name] = gl.getAttribLocation(program, info.name);
    }
    
    // Get uniform locations
    const numUniforms = gl.getProgramParameter(program, gl.ACTIVE_UNIFORMS);
    for (let i = 0; i < numUniforms; i++) {
      const info = gl.getActiveUniform(program, i);
      uniforms[info.name] = gl.getUniformLocation(program, info.name);
    }
    
    return { program, attributes, uniforms };
  }
  
  /**
   * Create a full-screen quad for rendering to textures
   */
  createFullScreenQuad() {
    const gl = this.gl;
    
    // Vertex positions for a full-screen quad (2 triangles)
    const positions = new Float32Array([
      -1, -1,
       1, -1,
      -1,  1,
       1,  1
    ]);
    
    // Texture coordinates for the quad
    const texCoords = new Float32Array([
      0, 0,
      1, 0,
      0, 1,
      1, 1
    ]);
    
    // Create and bind vertex position buffer
    const positionBuffer = gl.createBuffer();
    gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer);
    gl.bufferData(gl.ARRAY_BUFFER, positions, gl.STATIC_DRAW);
    this.buffers.quadPosition = positionBuffer;
    
    // Create and bind texture coordinate buffer
    const texCoordBuffer = gl.createBuffer();
    gl.bindBuffer(gl.ARRAY_BUFFER, texCoordBuffer);
    gl.bufferData(gl.ARRAY_BUFFER, texCoords, gl.STATIC_DRAW);
    this.buffers.quadTexCoord = texCoordBuffer;
    
    // Vertex Array Object for the quad
    if (this.gl instanceof WebGL2RenderingContext) {
      const vao = gl.createVertexArray();
      gl.bindVertexArray(vao);
      
      // Set up position attribute
      gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer);
      gl.enableVertexAttribArray(0);
      gl.vertexAttribPointer(0, 2, gl.FLOAT, false, 0, 0);
      
      // Set up texture coordinate attribute
      gl.bindBuffer(gl.ARRAY_BUFFER, texCoordBuffer);
      gl.enableVertexAttribArray(1);
      gl.vertexAttribPointer(1, 2, gl.FLOAT, false, 0, 0);
      
      gl.bindVertexArray(null);
      this.buffers.quadVAO = vao;
    }
  }
  
  /**
   * Create a framebuffer with texture attachments for output
   * @param {number} width - Texture width
   * @param {number} height - Texture height 
   * @returns {Object} - Framebuffer and attached textures
   */
  createOutputFramebuffer(width, height) {
    const gl = this.gl;
    
    // Create framebuffer
    const framebuffer = gl.createFramebuffer();
    gl.bindFramebuffer(gl.FRAMEBUFFER, framebuffer);
    
    // Create output texture
    const texture = gl.createTexture();
    gl.bindTexture(gl.TEXTURE_2D, texture);
    gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA32F, width, height, 0, gl.RGBA, gl.FLOAT, null);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
    
    // Attach texture to framebuffer
    gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture, 0);
    
    // Check framebuffer status
    const status = gl.checkFramebufferStatus(gl.FRAMEBUFFER);
    if (status !== gl.FRAMEBUFFER_COMPLETE) {
      throw new Error(`Framebuffer not complete: ${status}`);
    }
    
    // Unbind framebuffer
    gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    
    return { framebuffer, texture };
  }
  
  /**
   * Create a texture from data
   * @param {TypedArray} data - Data to upload to the texture
   * @param {number} width - Texture width
   * @param {number} height - Texture height
   * @param {number} format - Texture format (e.g., gl.RGBA)
   * @param {number} type - Texture data type (e.g., gl.FLOAT)
   * @param {number} internalFormat - Texture internal format (e.g., gl.RGBA32F)
   * @returns {WebGLTexture} - Created texture
   */
  createDataTexture(data, width, height, format, type, internalFormat) {
    const gl = this.gl;
    
    const texture = gl.createTexture();
    gl.bindTexture(gl.TEXTURE_2D, texture);
    gl.texImage2D(gl.TEXTURE_2D, 0, internalFormat, width, height, 0, format, type, data);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.NEAREST);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
    gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
    
    return texture;
  }
  
  /**
   * Run a decompression shader
   * @param {string} programName - Name of the shader program to use
   * @param {Object} textures - Input textures
   * @param {Object} uniforms - Uniform values
   * @param {number} width - Output width
   * @param {number} height - Output height
   * @returns {Float32Array} - Decompressed data
   */
  runDecompressionShader(programName, textures, uniforms, width, height) {
    const gl = this.gl;
    const program = this.programs[programName];
    
    // Create output framebuffer
    const output = this.createOutputFramebuffer(width, height);
    
    // Bind framebuffer
    gl.bindFramebuffer(gl.FRAMEBUFFER, output.framebuffer);
    gl.viewport(0, 0, width, height);
    
    // Use shader program
    gl.useProgram(program.program);
    
    // Bind quad VAO
    if (this.gl instanceof WebGL2RenderingContext) {
      gl.bindVertexArray(this.buffers.quadVAO);
    } else {
      // For WebGL 1, set up attributes manually
      gl.bindBuffer(gl.ARRAY_BUFFER, this.buffers.quadPosition);
      gl.enableVertexAttribArray(program.attributes.position);
      gl.vertexAttribPointer(program.attributes.position, 2, gl.FLOAT, false, 0, 0);
      
      gl.bindBuffer(gl.ARRAY_BUFFER, this.buffers.quadTexCoord);
      gl.enableVertexAttribArray(program.attributes.texCoord);
      gl.vertexAttribPointer(program.attributes.texCoord, 2, gl.FLOAT, false, 0, 0);
    }
    
    // Bind input textures
    let textureUnit = 0;
    for (const [name, texture] of Object.entries(textures)) {
      gl.activeTexture(gl.TEXTURE0 + textureUnit);
      gl.bindTexture(gl.TEXTURE_2D, texture);
      gl.uniform1i(program.uniforms[name], textureUnit);
      textureUnit++;
    }
    
    // Set uniforms
    for (const [name, value] of Object.entries(uniforms)) {
      const location = program.uniforms[name];
      if (location) {
        if (Array.isArray(value) && value.length === 2) {
          gl.uniform2fv(location, value);
        } else if (Array.isArray(value) && value.length === 3) {
          gl.uniform3fv(location, value);
        } else if (Array.isArray(value) && value.length === 4) {
          gl.uniform4fv(location, value);
        } else if (Array.isArray(value) && value.length === 9) {
          gl.uniformMatrix3fv(location, false, value);
        } else if (Array.isArray(value) && value.length === 16) {
          gl.uniformMatrix4fv(location, false, value);
        } else {
          gl.uniform1f(location, value);
        }
      }
    }
    
    // Draw quad
    gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);
    
    // Read back data
    const data = new Float32Array(width * height * 4);
    gl.readPixels(0, 0, width, height, gl.RGBA, gl.FLOAT, data);
    
    // Clean up
    gl.bindFramebuffer(gl.FRAMEBUFFER, null);
    gl.deleteFramebuffer(output.framebuffer);
    gl.deleteTexture(output.texture);
    
    return data;
  }
  
  /**
   * Decompress a Draco-encoded mesh using GPU acceleration
   * @param {ArrayBuffer} dracoData - Draco-encoded mesh data
   * @returns {Promise<Object>} - Decompressed mesh data
   */
  async decompressMesh(dracoData) {
    if (!this.isInitialized) {
      await new Promise(resolve => {
        const checkInitialized = () => {
          if (this.isInitialized) resolve();
          else setTimeout(checkInitialized, 100);
        };
        checkInitialized();
      });
    }
    
    // If we had to fall back to CPU decompression
    if (this.useFallback) {
      return this.decompressMeshCPU(dracoData);
    }
    
    try {
      // First, use Draco decoder to extract the compressed data into a format
      // suitable for GPU processing
      const preparedData = await this.prepareDracoData(dracoData);
      
      // Extract geometry metadata
      const {
        vertexCount, 
        indexCount, 
        positionData, 
        normalData, 
        uvData,
        indices,
        positionDequantizationFactors,
        positionOffset,
        uvFactors
      } = preparedData;
      
      // Calculate texture dimensions
      const texWidth = Math.ceil(Math.sqrt(vertexCount));
      const texHeight = Math.ceil(vertexCount / texWidth);
      
      // Create input textures
      const textures = {};
      
      // Position data texture
      if (positionData) {
        textures.uQuantizedPosTex = this.createDataTexture(
          positionData, 
          texWidth, 
          texHeight, 
          this.gl.RGBA, 
          this.gl.FLOAT, 
          this.gl.RGBA32F
        );
      }
      
      // Normal data texture (octahedral encoded)
      if (normalData) {
        textures.uOctCoordsTex = this.createDataTexture(
          normalData, 
          texWidth, 
          texHeight, 
          this.gl.RGBA, 
          this.gl.FLOAT, 
          this.gl.RGBA32F
        );
      }
      
      // UV data texture
      if (uvData) {
        textures.uQuantizedUVsTex = this.createDataTexture(
          uvData, 
          texWidth, 
          texHeight, 
          this.gl.RGBA, 
          this.gl.FLOAT, 
          this.gl.RGBA32F
        );
      }
      
      // Decompress positions
      const positions = positionData ? this.runDecompressionShader(
        'position',
        {
          uQuantizedPosTex: textures.uQuantizedPosTex
        },
        {
          uDequantizationFactors: positionDequantizationFactors,
          uPositionOffset: positionOffset
        },
        texWidth,
        texHeight
      ) : null;
      
      // Decompress normals
      const normals = normalData ? this.runDecompressionShader(
        'normal',
        {
          uOctCoordsTex: textures.uOctCoordsTex
        },
        {},
        texWidth,
        texHeight
      ) : null;
      
      // Decompress UVs
      const uvs = uvData ? this.runDecompressionShader(
        'uv',
        {
          uQuantizedUVsTex: textures.uQuantizedUVsTex
        },
        {
          uUVFactors: uvFactors
        },
        texWidth,
        texHeight
      ) : null;
      
      // Clean up textures
      for (const texture of Object.values(textures)) {
        this.gl.deleteTexture(texture);
      }
      
      // Prepare result
      const result = {
        vertexCount,
        indexCount,
        indices
      };
      
      if (positions) {
        // Extract positions from RGBA texture data
        result.positions = new Float32Array(vertexCount * 3);
        for (let i = 0; i < vertexCount; i++) {
          result.positions[i * 3] = positions[i * 4];
          result.positions[i * 3 + 1] = positions[i * 4 + 1];
          result.positions[i * 3 + 2] = positions[i * 4 + 2];
        }
      }
      
      if (normals) {
        // Extract normals from RGBA texture data
        result.normals = new Float32Array(vertexCount * 3);
        for (let i = 0; i < vertexCount; i++) {
          // Convert from [0,1] to [-1,1] range
          result.normals[i * 3] = normals[i * 4] * 2 - 1;
          result.normals[i * 3 + 1] = normals[i * 4 + 1] * 2 - 1;
          result.normals[i * 3 + 2] = normals[i * 4 + 2] * 2 - 1;
        }
      }
      
      if (uvs) {
        // Extract UVs from RGBA texture data
        result.uvs = new Float32Array(vertexCount * 2);
        for (let i = 0; i < vertexCount; i++) {
          result.uvs[i * 2] = uvs[i * 4];
          result.uvs[i * 2 + 1] = uvs[i * 4 + 1];
        }
      }
      
      return result;
      
    } catch (error) {
      console.error('GPU decompression failed, falling back to CPU:', error);
      return this.decompressMeshCPU(dracoData);
    }
  }
  
  /**
   * Prepare Draco data for GPU processing
   * @param {ArrayBuffer} dracoData - Draco-encoded mesh data
   * @returns {Promise<Object>} - Prepared data for GPU processing
   */
  async prepareDracoData(dracoData) {
    return new Promise((resolve, reject) => {
      // Send data to worker for preparation
      this.dracoWorker.onmessage = (e) => {
        const { type, result } = e.data;
        
        if (type === 'prepared') {
          resolve(result);
        } else if (type === 'error') {
          reject(new Error(result.message));
        }
      };
      
      // Send Draco data to worker
      this.dracoWorker.postMessage(
        {
          type: 'prepare',
          buffer: dracoData
        },
        [dracoData]
      );
    });
  }
  
  /**
   * Fallback CPU-based Draco decompression
   * @param {ArrayBuffer} dracoData - Draco-encoded mesh data
   * @returns {Promise<Object>} - Decompressed mesh data
   */
  async decompressMeshCPU(dracoData) {
    return new Promise((resolve, reject) => {
      // Send data to worker for CPU decompression
      this.dracoWorker.onmessage = (e) => {
        const { type, result } = e.data;
        
        if (type === 'decompressed') {
          resolve(result);
        } else if (type === 'error') {
          reject(new Error(result.message));
        }
      };
      
      // Send Draco data to worker
      this.dracoWorker.postMessage(
        {
          type: 'decompress',
          buffer: dracoData
        },
        [dracoData]
      );
    });
  }
  
  /**
   * Dispose of resources
   */
  dispose() {
    const gl = this.gl;
    
    // Delete programs
    for (const { program } of Object.values(this.programs)) {
      gl.deleteProgram(program);
    }
    
    // Delete buffers
    for (const buffer of Object.values(this.buffers)) {
      gl.deleteBuffer(buffer);
    }
    
    // Delete textures
    for (const texture of Object.values(this.textures)) {
      gl.deleteTexture(texture);
    }
    
    // Terminate worker
    if (this.dracoWorker) {
      this.dracoWorker.terminate();
    }
    
    this.isInitialized = false;
  }
}

export { DracoGPUDecompressor };