@woosh/meep-engine/src/engine/graphics/material/optimization/MaterialOptimizationContext.js

Pure JavaScript game engine. Fully featured and production ready.

import {
    AdditiveBlending,
    ClampToEdgeWrapping,
    DataTexture,
    Matrix3,
    MeshStandardMaterial,
    MultiplyBlending,
    NormalBlending
} from "three";
import { BinaryDataType } from "../../../../core/binary/type/BinaryDataType.js";
import {
    compute_typed_array_constructor_from_data_type
} from "../../../../core/binary/type/DataType2TypedArrayConstructorMapping.js";
import { array_push_if_unique } from "../../../../core/collection/array/array_push_if_unique.js";
import {
    compute_binary_data_type_from_typed_array
} from "../../../../core/collection/array/typed/compute_binary_data_type_from_typed_array.js";
import { collectIteratorValueToArray } from "../../../../core/collection/collectIteratorValueToArray.js";
import { HashMap } from "../../../../core/collection/map/HashMap.js";
import AABB2 from "../../../../core/geom/AABB2.js";
import { MaxRectanglesPacker } from "../../../../core/geom/packing/max-rect/MaxRectanglesPacker.js";
import { computeMaterialEquality } from "../../../asset/loaders/material/computeMaterialEquality.js";
import { computeMaterialHash } from "../../../asset/loaders/material/computeMaterialHash.js";
import { TextureAttachmentsByMaterialType } from "../../../asset/loaders/material/TextureAttachmensByMaterialType.js";
import { traverseThreeObject } from "../../ecs/highlight/renderer/traverseThreeObject.js";
import { channelCountToThreeTextureFormat } from "../../texture/channelCountToThreeTextureFormat.js";
import { computeThreeTextureTypeFromDataType } from "../../texture/computeThreeTextureTypeFromDataType.js";
import { normalized_internal_format } from "../../texture/normalized_internal_format.js";
import { convertTexture2Sampler2D } from "../../texture/sampler/convertTexture2Sampler2D.js";
import { Sampler2D } from "../../texture/sampler/Sampler2D.js";
import { sampler2d_copy_channel_data } from "../../texture/sampler/sampler2d_copy_channel_data.js";
import { sampler2d_sub_copy_same_item_size } from "../../texture/sampler/sampler2d_sub_copy_same_item_size.js";
import { sampler2d_transfer_data } from "../../texture/sampler/sampler2d_transfer_data.js";
import { BUFFER_GEOMETRY_UVS } from "./BUFFER_GEOMETRY_UVS.js";
import { is_compliant_mesh } from "./is_compliant_mesh.js";
import { MaterialDescriptor } from "./MaterialDescriptor.js";
import { prepare_atlas_texture } from "./prepare_atlas_texture.js";

const ATLAS_PATCH_BORDER = 0;


/**
 *
 * @param {THREE.Material} a
 * @param {THREE.Material} b
 * @returns {boolean}
 */
function materialsSimilarForCharting(a, b) {

    return a.visible === b.visible
        && a.transparent === b.transparent
        && a.blending === b.blending
        && a.alphaTest === b.alphaTest
        && a.depthTest === b.depthTest
        && a.depthWrite === b.depthWrite
        ;
}

/**
 *
 * @param {THREE.Material} a
 * @returns {number}
 */
function materialChartingHash(a) {
    const flags = (a.visible ? 1 : 0)
        | ((a.transparent ? 1 : 0) << 1)
        | ((a.depthTest ? 1 : 0) << 2)
        | ((a.depthWrite ? 1 : 0) << 3)
        | (([NormalBlending, AdditiveBlending, MultiplyBlending].indexOf(a.blending) & 0b11) << 4)
        | ((Math.floor(a.alphaTest * 15)) << 6) // 4 bit
    ;
    return flags;
}

/**
 *
 * @param {THREE.DataTexture} metalness
 * @param {THREE.DataTexture} roughness
 * @returns {THREE.DataTexture}
 */
function merge_textures_metalness_rougness({
                                               metalness,
                                               roughness
                                           }) {
    const ref = metalness !== undefined ? metalness : roughness;

    const pixel_count = ref.image.width * ref.image.height;
    const result_data = new Uint8Array(pixel_count * 4);

    const result = new DataTexture(result_data, ref.image.width, ref.image.height);

    result.encoding = ref.encoding;

    if (metalness !== undefined) {
        for (let i = 0; i < pixel_count; i++) {
            result_data[i * 4 + 2] = metalness.image.data[i];
        }
    }

    if (roughness !== undefined) {
        for (let i = 0; i < pixel_count; i++) {
            result_data[i * 4 + 1] = roughness.image.data[i];
        }
    }

    result.needsUpdate = true;

    return result;
}

/**
 *
 * @param {THREE.Material[]} source
 * @param {Map<string, THREE.Texture>} textures
 * @returns {Map<THREE.Material, THREE.Material>}
 */
function merge_materials(source, textures) {
    // collect unique materials
    const unique_materials = new HashMap({
        keyEqualityFunction: materialsSimilarForCharting,
        keyHashFunction: materialChartingHash
    });

    for (let i = 0; i < source.length; i++) {
        const material = source[i];

        let bucket = unique_materials.get(material);
        if (bucket === undefined) {
            bucket = [];
            unique_materials.set(material, bucket);
        }

        bucket.push(material);
    }

    // build new materials
    const source_unique_keys = [];
    collectIteratorValueToArray(source_unique_keys, unique_materials.keys());

    /**
     *
     * @type {Map<THREE.Material, THREE.Material>}
     */
    const result = new Map();

    for (let i = 0; i < source_unique_keys.length; i++) {
        const source_material = source_unique_keys[i];

        const result_material = new MeshStandardMaterial();

        result_material.name = `Generated charted material ${i}`;

        result_material.visible = source_material.visible;
        result_material.transparent = source_material.transparent;
        result_material.blending = source_material.blending;
        result_material.alphaTest = source_material.alphaTest;
        result_material.depthTest = source_material.depthTest;
        result_material.depthWrite = source_material.depthWrite;

        // write atlased texture
        const texture_attachments = TextureAttachmentsByMaterialType[result_material.type];

        for (let j = 0; j < texture_attachments.length; j++) {
            const textureAttachment = texture_attachments[j];

            const existing_texture = textures.get(textureAttachment.name);

            if (existing_texture !== undefined) {
                textureAttachment.write(result_material, existing_texture);
            }
        }

        // write result mapping
        const destination_materials = unique_materials.get(source_material);
        for (let j = 0; j < destination_materials.length; j++) {
            result.set(destination_materials[j], result_material);
        }
    }

    return result;
}

/**
 *
 * @param {Sampler2D} destination
 * @param {Sampler2D} source
 * @param {string} texture_name
 */
function extract_texture_data(destination, source, texture_name) {
    switch (texture_name) {
        case 'ao':
            sampler2d_copy_channel_data(source, 0, destination, 0);
            break;
        case 'roughness':
            sampler2d_copy_channel_data(source, 1, destination, 0);
            break;
        case 'metalness':
            sampler2d_copy_channel_data(source, 2, destination, 0);
            break;
        default:
            sampler2d_transfer_data(source, destination);
    }
}

/**
 *
 * @param {number[]} destination
 * @param {THREE.Material|THREE.MeshStandardMaterial} source_material
 * @param {string} texture_name
 */
function compute_default_texture_texel_value(destination, source_material, texture_name) {
    switch (texture_name) {
        case 'diffuse':
            destination[0] = 255;
            destination[1] = 255;
            destination[2] = 255;

            if (destination.length > 3) {
                destination[3] = 255;
            }
            break;
        case 'metalness':
            for (let i = 0; i < destination.length; i++) {
                destination[i] = Math.round(source_material.metalness * 255);
            }
            if (destination.length > 3) {
                destination[3] = 255;
            }
            break;
        case 'roughness':
            for (let i = 0; i < destination.length; i++) {
                destination[i] = Math.round(source_material.roughness * 255);
            }
            if (destination.length > 3) {
                destination[3] = 255;
            }
            break;
        case 'ao':
            for (let i = 0; i < destination.length; i++) {
                destination[i] = 255;
            }
            break;
        case 'normal':
            destination[0] = 127;
            destination[1] = 127;
            destination[2] = 255;
            break;
    }
}

/**
 *
 * @param {Sampler2D} sampler
 * @param {THREE.MeshStandardMaterial} material
 * @param {string} texture_name
 */
function apply_material_shading_to_texture(sampler, material, texture_name) {
    switch (texture_name) {
        case 'diffuse': {
            const r = material.color.r;
            const g = material.color.g;
            const b = material.color.b;
            const opacity = material.opacity;

            const itemSize = sampler.itemSize;

            const data = sampler.data;
            const data_size = data.length;

            if (r < 1 || g < 1 || b < 1) {
                // apply color tint
                for (let i = 0; i < data_size; i += itemSize) {
                    data[i] *= r;
                    data[i + 1] *= g;
                    data[i + 2] *= b;
                }
            }

            if (itemSize > 3 && opacity < 1) {
                // apply opacity
                for (let i = 3; i < data_size; i += itemSize) {
                    data[i] *= opacity;
                }

            }
        }
            break;
        default:
        // do nothing
    }
}

/**
 *
 * @param {Sampler2D} destination
 * @param {THREE.Material} source_material
 * @param {string} texture_name
 */
function compute_default_texture_value(destination, source_material, texture_name) {
    const itemSize = destination.itemSize;

    const texel = new Array(itemSize);

    compute_default_texture_texel_value(texel, source_material, texture_name);

    const width = destination.width;
    const height = destination.height;

    const destination_data = destination.data;

    if (itemSize === 1) {
        // special fast case
        destination_data.fill(texel[0]);
    } else {

        const texel_count = width * height;

        for (let i = 0; i < texel_count; i++) {
            const destination_offset = i * itemSize;

            destination_data.set(texel, destination_offset);
        }
    }
}

/**
 *
 * @param {string} texture_name
 * @returns {{type:BinaryDataType,itemSize:number, default_value:[]}}
 */
function build_atlas_sampler_options_texture_name(texture_name) {
    let type;
    let itemSize;
    let default_value;

    switch (texture_name) {
        case 'normal':
            type = BinaryDataType.Uint8;
            itemSize = 3;
            default_value = [127, 127, 255];
            break;
        case 'ao':
            type = BinaryDataType.Uint8;
            itemSize = 1;
            default_value = [255];
            break;
        case "metalness":
            type = BinaryDataType.Uint8;
            itemSize = 1;
            default_value = [0];
            break;
        case "roughness":
            type = BinaryDataType.Uint8;
            itemSize = 1;
            default_value = [255];
            break;
        default:
            type = BinaryDataType.Uint8;
            itemSize = 4;
            default_value = [255, 255, 255, 255];
            break;
    }

    return {
        type,
        itemSize,
        default_value
    };
}

export class MaterialOptimizationContext {
    constructor() {
        /**
         *
         * @type {THREE.Texture[]}
         */
        this.textures = [];

        /**
         *
         * @type {THREE.Mesh[]}
         */
        this.meshes = [];


        /**
         *
         * @type {Map<string,Sampler2D>}
         */
        this.atlases = new Map();

        /**
         *
         * @type {Vector2[]}
         */
        this.atlas_size = [];
    }

    /**
     *
     * @param {MaterialDescriptor[]} materials
     */
    pack_atlas(materials) {
        const uv_indices = [];
        // figure out UV sets
        materials.forEach(m => {
            for (let i = 0; i < m.uvs.length; i++) {
                if (m.uvs[i] !== undefined) {
                    array_push_if_unique(uv_indices, i);
                }
            }
        });

        for (let i = 0; i < uv_indices.length; i++) {
            const uv_index = uv_indices[i];

            const packer = new MaxRectanglesPacker(1, 1);

            let pow_x = 0;
            let pow_y = 0;

            const uv_relevant_materials = materials.filter(d => d.uvs[uv_index] !== undefined);

            const rectangles = uv_relevant_materials.map(d => {
                const uv_context = d.uvs[uv_index];

                const resolution = uv_context.resolution;

                return new AABB2(0, 0, resolution.x + ATLAS_PATCH_BORDER * 2, resolution.y + ATLAS_PATCH_BORDER * 2);

            });

            let cycle_index = 0;

            while (!packer.addMany(rectangles)) {
                cycle_index++;

                if (cycle_index % 2 !== 0) {
                    pow_x++;
                } else {
                    pow_y++;
                }

                const size_x = Math.pow(2, pow_x);
                const size_y = Math.pow(2, pow_y);

                const max_texture_size = MaterialOptimizationContext.MAX_TEXTURE_SIZE;

                if (size_x > max_texture_size || size_y > max_texture_size) {
                    throw new Error(`Maximum texture size exceeded '${max_texture_size}'`);
                }

                packer.resize(size_x, size_y);
            }

            this.atlas_size[uv_index] = packer.size.clone();

            // all packed, write back bounds
            rectangles.forEach((r, i) => {
                const materialDescriptor = uv_relevant_materials[i];

                const p = materialDescriptor.uvs[uv_index].atlas_patch;

                p.set(r.x0 + ATLAS_PATCH_BORDER, r.y0 + ATLAS_PATCH_BORDER, r.x1 - ATLAS_PATCH_BORDER, r.y1 - ATLAS_PATCH_BORDER);
            });
        }
    }

    /**
     *
     * @param {MaterialDescriptor[]} materials
     */
    build_atlas_samplers(materials) {

        let i;

        this.atlases.clear();

        const names = [];
        const uv_indices = [];

        const material_count = materials.length;

        for (i = 0; i < material_count; i++) {
            const material = materials[i];

            const source_material = material.source_material;
            const materialType = source_material.type;

            /**
             *
             * @type {TextureAttachment[]}
             */
            const attachments = TextureAttachmentsByMaterialType[materialType];

            for (let i = 0; i < attachments.length; i++) {
                const attachment = attachments[i];

                const attachment_texture = attachment.read(source_material);

                if (!attachment_texture) {
                    // no attachment on the material
                    continue;
                }

                if (array_push_if_unique(names, attachment.name)) {
                    // added
                    uv_indices[names.length - 1] = attachment.uv_index;
                }

            }
        }

        for (i = 0; i < names.length; i++) {
            const name = names[i];
            const uv_index = uv_indices[i];

            const size = this.atlas_size[uv_index];

            const width = size.x;
            const height = size.y;

            const sampler_options = build_atlas_sampler_options_texture_name(name);

            const TA = compute_typed_array_constructor_from_data_type(sampler_options.type);

            const itemSize = sampler_options.itemSize;
            const sampler_data_size = itemSize * width * height;
            const sampler_data = new TA(sampler_data_size);

            for (let j = 0; j < sampler_data_size; j += itemSize) {
                sampler_data.set(sampler_options.default_value, j);
            }

            const atlas_sampler = new Sampler2D(sampler_data, itemSize, width, height);
            this.atlases.set(name, atlas_sampler);
        }
    }

    /**
     *
     * @param {MaterialDescriptor} material
     */
    write_material_to_atlas(material) {
        //patch textures
        const source_material = material.source_material;
        const materialType = source_material.type;

        /**
         *
         * @type {TextureAttachment[]}
         */
        const attachments = TextureAttachmentsByMaterialType[materialType];

        const attachment_names = Array.from(this.atlases.keys());

        for (let i = 0; i < attachments.length; i++) {
            const attachment = attachments[i];

            const texture_name = attachment.name;

            if (!attachment_names.includes(texture_name)) {
                continue;
            }

            const uv_context = material.uvs[attachment.uv_index];

            if (uv_context === undefined) {
                // material does not participate in this UV set
                continue;
            }

            const atlas_sampler = this.atlases.get(texture_name);

            let source_sampler;

            const texture = attachment.read(source_material);

            const resolution = uv_context.resolution;

            const texture_width = resolution.x;
            const texture_height = resolution.y;

            if (texture !== undefined && texture !== null) {
                // source_sampler = Sampler2D.uint8(5, texture_width, texture_height);
                //
                // source_sampler.fill(0, 0, texture_width, texture_height, [255, 255, 255, 255]);

                source_sampler = convertTexture2Sampler2D(texture, texture_width, texture_height, false);

                if (source_sampler.itemSize !== atlas_sampler.itemSize) {
                    const destination = Sampler2D.uint8(atlas_sampler.itemSize, texture_width, texture_height);

                    extract_texture_data(destination, source_sampler, texture_name);

                    source_sampler = destination;
                }

            } else {
                // texture for a given slot is missing, create a default
                source_sampler = Sampler2D.uint8(atlas_sampler.itemSize, texture_width, texture_height);

                compute_default_texture_value(source_sampler, source_material, texture_name);
            }

            // apply modification to texture to remove dependency on source material
            apply_material_shading_to_texture(source_sampler, source_material, texture_name);

            const patch = uv_context.atlas_patch;

            sampler2d_sub_copy_same_item_size(
                atlas_sampler,
                source_sampler,
                0, 0,
                patch.x0, patch.y0, patch.getWidth(), patch.getHeight()
            );
        }
    }

    /**
     *
     * @param {THREE.Mesh} mesh
     */
    addMesh(mesh) {
        const material = mesh.material;

        if (material.isMeshStandardMaterial) {
            this.meshes.push(mesh);
        } else {
            // throw new Error(`Unsupported material type '${material.type}'`);
        }
    }

    /**
     *
     * @param {THREE.Object} o
     */
    addObject(o) {
        traverseThreeObject(o, (o) => {
            if (o.isMesh) {
                this.addMesh(o);
            }
        });
    }

    /**
     * Update materials and geometries
     */
    update() {
        // collect unique materials and their geometries

        /**
         *
         * @type {Map<MaterialDescriptor, THREE.BufferGeometry[]>}
         */
        const material_to_geometry_map = new Map();

        /**
         *
         * @type {Map<MaterialDescriptor, THREE.Mesh[]>}
         */
        const material_to_mesh_map = new Map();

        const unique_materials = new HashMap({
            keyHashFunction: computeMaterialHash,
            keyEqualityFunction: computeMaterialEquality
        });

        const input_meshes = this.meshes;
        const input_mesh_count = input_meshes.length;

        for (let i = 0; i < input_mesh_count; i++) {
            const mesh = input_meshes[i];

            if (!is_compliant_mesh(mesh)) {
                continue;
            }

            const material = mesh.material;

            let material_descriptor = unique_materials.get(material);

            if (material_descriptor === undefined) {
                material_descriptor = new MaterialDescriptor();
                material_descriptor.source_material = material;

                material_descriptor.build_uvs();

                unique_materials.set(material, material_descriptor);

                material_to_geometry_map.set(material_descriptor, []);
                material_to_mesh_map.set(material_descriptor, []);

            }

            const geometries = material_to_geometry_map.get(material_descriptor);

            array_push_if_unique(geometries, mesh.geometry);

            const mesh_bucket = material_to_mesh_map.get(material_descriptor);
            mesh_bucket.push(mesh);
        }

        // collect materials
        /**
         *
         * @type {MaterialDescriptor[]}
         */
        const unique_material_array = [];
        collectIteratorValueToArray(unique_material_array, unique_materials.values());

        const unique_material_count = unique_material_array.length;

        if (unique_material_count <= 4) {
            // too few unique materials, no point in atlasing
            return;
        }

        this.pack_atlas(unique_material_array);

        this.build_atlas_samplers(unique_material_array);

        const textures = this.initialize_atlas_textures();


        for (let i = 0; i < unique_material_count; i++) {
            const mat = unique_material_array[i];

            this.write_material_to_atlas(mat);

            // get all associated geometries
            const geometries = material_to_geometry_map.get(mat);

            for (let j = 0; j < geometries.length; j++) {
                const bufferGeometry = geometries[j];

                this.apply_geometry_uv_changes(bufferGeometry, mat);
            }
        }

        /*
         merge roughness and metalness. Three.js shaders sample roughness from G channel (index = 1) and metalness from B channel (index 2)
         Because of this, we can't pass B/W textures with only Red channel set, so we are forced to create a merged texture from metalness and roughnes
         */
        if (textures.get('roughness') !== undefined || textures.get('metalness') !== undefined) {
            const result_texture = merge_textures_metalness_rougness({
                roughness: textures.get('roughness'),
                metalness: textures.get('metalness')
            });

            if (textures.get('roughness') !== undefined) {
                textures.set('roughness', result_texture);
            }
            if (textures.get('metalness') !== undefined) {
                textures.set('roughness', result_texture);
            }
        }

        // update textures
        for (const [name, texture] of textures) {
            texture.needsUpdate = true;
        }

        // compute merged materials
        const unique_source_materials = [];
        collectIteratorValueToArray(unique_source_materials, unique_materials.keys());

        const merged_materials = merge_materials(unique_source_materials, textures);

        for (let i = 0; i < unique_material_count; i++) {
            const mat = unique_material_array[i];

            // bind material
            const new_material = merged_materials.get(mat.source_material);

            // set new material to all associated meshes
            const meshes = material_to_mesh_map.get(mat);

            for (let j = 0; j < meshes.length; j++) {
                const mesh = meshes[j];
                mesh.material = new_material;
            }
        }

    }

    /**
     *
     * @return {Map<string, THREE.DataTexture>}
     */
    initialize_atlas_textures() {
        const textures = new Map();
        //build new textures
        for (const [name, sampler] of this.atlases) {

            const sampler_data_type = compute_binary_data_type_from_typed_array(sampler.data);

            const format = channelCountToThreeTextureFormat(sampler.itemSize);
            const type = computeThreeTextureTypeFromDataType(sampler_data_type);

            const texture = new DataTexture(sampler.data, sampler.width, sampler.height, format, type);

            texture.wrapS = ClampToEdgeWrapping;
            texture.wrapT = ClampToEdgeWrapping;
            texture.name = name;

            texture.internalFormat = normalized_internal_format(sampler_data_type, sampler.itemSize);

            prepare_atlas_texture(texture, name);

            textures.set(name, texture);
        }
        return textures;
    }

    /**
     *
     * @param {THREE.BufferGeometry} bufferGeometry
     * @param {MaterialDescriptor} mat
     */
    apply_geometry_uv_changes(bufferGeometry, mat) {
        const uv_count = BUFFER_GEOMETRY_UVS.length;

        for (let k = 0; k < uv_count; k++) {
            const uv_metadata = BUFFER_GEOMETRY_UVS[k];

            // update UVs
            const uv_attribute = bufferGeometry.attributes[uv_metadata.name];

            if (uv_attribute === undefined) {
                // no UVs, nothing to update
                continue;
            }


            const uv_index = uv_metadata.index;
            const uv_context = mat.uvs[uv_index];

            if (uv_context === undefined) {
                // no UV context present
                // TODO need to force context to exist in this case or remove the attribute
                continue;
            }

            const atlas_size = this.atlas_size[uv_index];

            const atlas_width = atlas_size.x;
            const atlas_height = atlas_size.y;

            const patch = uv_context.atlas_patch;

            const matrix3 = new Matrix3();
            matrix3.setUvTransform(
                patch.x0 / atlas_width,
                patch.y0 / atlas_height,
                patch.getWidth() / atlas_width,
                patch.getHeight() / atlas_height,
                0,
                0, 0
            );

            uv_attribute.applyMatrix3(matrix3);

            uv_attribute.needsUpdate = true;
        }
    }
}

/**
 *
 * @type {number}
 */
MaterialOptimizationContext.MAX_TEXTURE_SIZE = 4096;