@woosh/meep-engine/src/engine/graphics/particles/node-based/simulation/TransformFeedback.js

Pure JavaScript game engine. Fully featured and production ready.

import { getTypeByteSize } from "../codegen/glsl/getTypeByteSize.js";
import { ParticleDataTypes } from "../nodes/ParticleDataTypes.js";
import { genAttributeInputName } from "../codegen/glsl/genAttributeInputName.js";
import { genAttributeOutputName } from "../codegen/glsl/genAttributeOutputName.js";
import { getShaderErrors, WRITE_WHITE_FRAGMENT_SHADER } from "./GLSLSimulationShader.js";

export class TransformFeedback {
    constructor() {

        /**
         *
         * @type {String}
         * @private
         */
        this.__source = null;

        /**
         *
         * @type {WebGLProgram}
         * @private
         */
        this.__program = null;

        /**
         * Uniform location pointers from GL context. Order matches order of uniform specification array
         * @type {number[]}
         * @private
         */
        this.__uniform_locations = [];

        /**
         *
         * @type {ParticleAttributeSpecification[]}
         */
        this.attributes = [];

        /**
         * How much space does a single vertex require
         * @type {number}
         * @private
         */
        this.__attribute_byte_size = 0;

        /**
         *
         * @type {ParticleAttributeSpecification[]}
         */
        this.uniforms = [];

        /**
         *
         * @type {FunctionModuleRegistry}
         * @private
         */
        this.__function_registry = null;

        /**
         *
         * @type {WebGLTransformFeedback}
         * @private
         */
        this.__transform_feedback = null;

        /**
         *
         * @type {WebGL2RenderingContext}
         * @private
         */
        this.__context = null;


        /**
         *
         * @type {number}
         * @private
         */
        this.__output_buffer_handle = null;
        /**
         *
         * @type {number}
         * @private
         */
        this.__output_buffer_size = 0;
    }

    /**
     *
     * @return {String}
     */
    getSourceCode() {
        return this.__source;
    }

    /**
     * @param {string} source
     */
    setSourceCode(source) {
        if (this.isCompiled()) {
            throw new Error("Can't assign new source to a compiled program");
        }

        this.__source = source;
    }

    /**
     *
     * @param {ParticleAttributeSpecification[]} attributes
     */
    setAttributes(attributes) {
        this.attributes = attributes;
    }

    /**
     *
     * @param {ParticleAttributeSpecification[]} attributes
     */
    setUniforms(attributes) {
        this.uniforms = attributes;
    }

    __updateAttributeByteSize() {

        let offset = 0;

        const attributes = this.attributes;
        const attribute_count = attributes.length;

        for (let i = 0; i < attribute_count; i++) {
            const attribute = attributes[i];

            offset += getTypeByteSize(attribute.type);
        }

        this.__attribute_byte_size = offset;
    }

    /**
     *
     * @param {EmitterAttributeData} attribute_source
     * @private
     */
    __prepareOutputBuffer(attribute_source) {
        const min_size = attribute_source.data.getSize();

        if (this.__output_buffer_size < min_size) {
            this.__output_buffer_size = min_size;
            /**
             *
             * @type {WebGLRenderingContext}
             */
            const gl = this.__context;

            gl.bindBuffer(gl.ARRAY_BUFFER, this.__output_buffer_handle);
            gl.bufferData(gl.ARRAY_BUFFER, min_size * 4, gl.DYNAMIC_DRAW);

            gl.bindBuffer(gl.ARRAY_BUFFER, null);
        }

    }

    /**
     *
     * @param {Array} uniform_values
     * @param {EmitterAttributeData} attributeSource
     */
    execute(uniform_values, attributeSource) {
        const gl = this.__context;

        gl.useProgram(this.__program);

        // write uniforms
        const uniform_locations = this.__uniform_locations;
        const uniform_count = uniform_locations.length;

        for (let i = 0; i < uniform_count; i++) {
            const location = uniform_locations[i];

            const uniform_spec = this.uniforms[i];

            const uniformValue = uniform_values[i];

            switch (uniform_spec.type) {
                case ParticleDataTypes.Float:

                    gl.uniform1f(location, uniformValue);
                    break;
                case ParticleDataTypes.Matrix4:
                    gl.uniformMatrix4fv(location, false, uniformValue)
                    break;
                default:
                    throw new Error(`Unsupported uniform type`);
            }
        }

        // bind attributes
        const attributes = this.attributes;
        const attribute_count = attributes.length;

        let offset = 0;


        gl.bindBuffer(gl.ARRAY_BUFFER, attributeSource.data.gl_buffer_f32);
        //flush data to buffer
        gl.bufferData(gl.ARRAY_BUFFER, attributeSource.data.data_f32, gl.DYNAMIC_DRAW);

        for (let i = 0; i < attribute_count; i++) {
            const attribute = attributes[i];

            const size = getTypeByteSize(attribute.type);

            const componentCount = attribute.computeComponentCount();

            gl.enableVertexAttribArray(i);
            gl.vertexAttribPointer(i, componentCount, gl.FLOAT, false, this.__attribute_byte_size, offset);

            offset += size;
        }

        gl.bindTransformFeedback(gl.TRANSFORM_FEEDBACK, this.__transform_feedback);

        this.__prepareOutputBuffer(attributeSource);

        // bind target buffer
        gl.bindBufferRange(gl.TRANSFORM_FEEDBACK_BUFFER, 0, this.__output_buffer_handle, 0, attributeSource.count * this.__attribute_byte_size);

        //
        gl.enable(gl.RASTERIZER_DISCARD);
        gl.beginTransformFeedback(gl.POINTS);

        gl.drawArrays(gl.POINTS, 0, attributeSource.count);

        gl.endTransformFeedback();
        gl.disable(gl.RASTERIZER_DISCARD);

        // Unbind the transform feedback buffer so subsequent attempts
        gl.bindBufferBase(gl.TRANSFORM_FEEDBACK_BUFFER, 0, null);

        // unbind transform feedback
        gl.bindTransformFeedback(gl.TRANSFORM_FEEDBACK, null);

    }

    /**
     *
     * @param {string} name
     * @return {number}
     */
    getUniformIndexByName(name) {
        for (let i = 0; i < this.uniforms.length; i++) {
            const u = this.uniforms[i];

            if (u.name === name) {
                return i;
            }
        }

        return -1;
    }

    /**
     *
     * @private
     */
    __buildUniformPointers() {
        const gl = this.__context;

        const program = this.__program;

        // obtain uniform pointers
        const active_uniform_count = gl.getProgramParameter(program, gl.ACTIVE_UNIFORMS);

        for (let i = 0; i < active_uniform_count; i++) {
            /**
             *
             * @type {WebGLActiveInfo}
             */
            const uniform = gl.getActiveUniform(program, i);

            const name = uniform.name;

            const uniform_index = this.getUniformIndexByName(name);

            if (uniform_index === -1) {
                throw new Error(`Uniform '${name}' not found`);
            }

            this.__uniform_locations[uniform_index] = gl.getUniformLocation(program, name);
        }
    }

    /**
     * GL context uses it's own internal buffers that are not shared with JS context, because of this - it is necessary to copy data out
     * @param {EmitterAttributeData} target
     */
    readOutput(target) {
        /**
         *
         * @type {WebGL2RenderingContext}
         */
        const gl = this.__context;

        gl.bindBuffer(gl.ARRAY_BUFFER, this.__output_buffer_handle);
        gl.getBufferSubData(gl.ARRAY_BUFFER, 0, target.data.data_f32);
    }

    isCompiled() {
        return this.__program !== null;
    }

    dispose() {
        if (this.isCompiled()) {


            const gl = this.__context;
            gl.deleteBuffer(this.__output_buffer_handle);
            gl.deleteProgram(this.__program);

            this.__program = null;
            this.__context = null;
            this.__outpub_buffer_handle = null;
        }
    }

    /**
     *
     * @param {WebGL2RenderingContext} gl
     */
    compile(gl) {

        if (this.isCompiled()) {
            throw new Error('Already compiled, must call .dispose first before reusing this transform');
        }

        if (this.__source === null) {
            throw new Error('No shader source code, shader must be built before being compiled');
        }

        this.__updateAttributeByteSize();

        const vertex_shader = gl.createShader(gl.VERTEX_SHADER);
        const fragment_shader = gl.createShader(gl.FRAGMENT_SHADER);


        gl.shaderSource(vertex_shader, this.__source);
        gl.shaderSource(fragment_shader, WRITE_WHITE_FRAGMENT_SHADER);

        // Compile vertex shader
        gl.compileShader(vertex_shader);

        if (!gl.getShaderParameter(vertex_shader, gl.COMPILE_STATUS)) {
            const error_message = getShaderErrors(gl, vertex_shader, 'VERTEX');

            // cleanup
            gl.deleteShader(vertex_shader);

            throw new Error(error_message);
        }

        // Compile fragment shader
        gl.compileShader(fragment_shader);

        if (!gl.getShaderParameter(fragment_shader, gl.COMPILE_STATUS)) {
            const error_message = getShaderErrors(gl, fragment_shader, 'FRAGMENT');

            // cleanup
            gl.deleteShader(vertex_shader);
            gl.deleteShader(fragment_shader);

            throw new Error(error_message);
        }

        const program = gl.createProgram();


        gl.attachShader(program, vertex_shader);
        gl.attachShader(program, fragment_shader);

        // clean up allocated shaders, they are now part of the program and can be removed
        gl.deleteShader(vertex_shader);
        gl.deleteShader(fragment_shader);

        // bind attribute inputs
        for (let i = 0; i < this.attributes.length; i++) {
            const attribute = this.attributes[i];

            const input_name = genAttributeInputName(attribute);

            gl.bindAttribLocation(program, i, input_name);
        }

        // bind outputs
        const output_names = this.attributes.map(genAttributeOutputName);

        // We use interleaved attributes for cache coherence
        gl.transformFeedbackVaryings(program, output_names, gl.INTERLEAVED_ATTRIBS);

        gl.linkProgram(program);

        if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {

            const error_message = "Shader program failed to link" + gl.getProgramInfoLog(program);

            gl.deleteProgram(program);

            throw new Error(error_message);
        }

        this.__context = gl;
        this.__program = program;


        this.__buildUniformPointers(gl);

        this.__output_buffer_handle = gl.createBuffer();

        this.__transform_feedback = gl.createTransformFeedback();
    }
}