@needle-tools/engine/src/engine/engine_assetdatabase.ts

Needle Engine is a web-based runtime for 3D apps. It runs on your machine for development with great integrations into editors like Unity or Blender - and can be deployed onto any device! It is flexible, extensible and networking and XR are built-in.

import { Bone, BufferAttribute, BufferGeometry, InterleavedBuffer, InterleavedBufferAttribute, Material, Mesh, NeverCompare, Object3D, Scene, Skeleton, SkinnedMesh, Source, Texture, Uniform, WebGLRenderer } from "three";

import { addPatch } from "./engine_patcher.js";
import { Physics } from "./engine_physics.js";
import { getParam } from "./engine_utils.js";
import { InternalUsageTrackerPlugin } from "./extensions/usage_tracker.js";


export class AssetDatabase {
    constructor() {
        window.addEventListener('unhandledrejection', (event: PromiseRejectionEvent) => {
            if (event.defaultPrevented) return;
            const pathArray = event?.reason?.path;
            if (pathArray) {
                const source = pathArray[0];
                if (source && source.tagName === "IMG") {
                    console.warn("Could not load image:\n" + source.src);
                    event.preventDefault();
                }
            }
        });
    }
}

const trackUsageParam = getParam("trackresources");

function autoDispose() {
    return trackUsageParam === "dispose";
}

let allowUsageTracking = true;// trackUsageParam === true || trackUsageParam === 1;
if (trackUsageParam === 0) allowUsageTracking = false;

/** 
 * Disable usage tracking
 */
export function setResourceTrackingEnabled(enabled: boolean) {
    allowUsageTracking = enabled;
}
export function isResourceTrackingEnabled() {
    return allowUsageTracking;
}

const $disposable = Symbol("disposable");
export function setDisposable(obj: object | null | undefined, disposable: boolean) {
    if (!obj) return;
    obj[$disposable] = disposable;
    if (debug) console.warn("Set disposable", disposable, obj);
}

const $disposed = Symbol("disposed");
export function isDisposed(obj: object) {
    return obj[$disposed] === true;
}

/** Recursive disposes all referenced resources by this object. Does not traverse children */
export function disposeObjectResources(obj: object | null | undefined) {
    if (!obj) return;
    if (obj[$disposable] === false) {
        if (debug) console.warn("Object is marked as not disposable", obj);
        return;
    }

    if(typeof obj === "object") {
        obj[$disposed] = true;
    }

    if (obj instanceof Scene) {
        disposeObjectResources(obj.environment);
        disposeObjectResources(obj.background);
        disposeObjectResources(obj.customDepthMaterial);
        disposeObjectResources(obj.customDistanceMaterial);
    }
    else if (obj instanceof SkinnedMesh) {
        disposeObjectResources(obj.geometry);
        disposeObjectResources(obj.material);
        disposeObjectResources(obj.skeleton);
        disposeObjectResources(obj.bindMatrix);
        disposeObjectResources(obj.bindMatrixInverse);
        disposeObjectResources(obj.customDepthMaterial);
        disposeObjectResources(obj.customDistanceMaterial);
        obj.geometry = null;
        obj.material = null;
        obj.visible = false;
    }
    else if (obj instanceof Mesh) {
        disposeObjectResources(obj.geometry);
        disposeObjectResources(obj.material);
        disposeObjectResources(obj.customDepthMaterial);
        disposeObjectResources(obj.customDistanceMaterial);
        obj.geometry = null;
        obj.material = null;
        obj.visible = false;
    }
    else if (obj instanceof BufferGeometry) {
        free(obj);
        for (const key of Object.keys(obj.attributes)) {
            const value = obj.attributes[key];
            disposeObjectResources(value);
            // deleting the attribute might lead to errors when raycasting
            // obj.deleteAttribute(key);
        }
    }
    else if (obj instanceof BufferAttribute || obj instanceof InterleavedBufferAttribute) {
        // Currently not supported by three 
        // https://github.com/mrdoob/three.js/issues/15261
        // https://github.com/mrdoob/three.js/pull/17063#issuecomment-737993363
        if (debug)
            console.warn("BufferAttribute dispose not supported", obj.count);
    }
    else if (obj instanceof Array) {
        for (const entry of obj) {
            if (entry instanceof Material)
                disposeObjectResources(entry);
        }
    }
    else if (obj instanceof Material) {
        free(obj);
        for (const key of Object.keys(obj)) {
            const value = obj[key];
            if (value instanceof Texture) {
                disposeObjectResources(value);
                obj[key] = null;
            }
        }
        const uniforms = obj["uniforms"];
        if (uniforms) {
            for (const key of Object.keys(uniforms)) {
                const value = uniforms[key];
                if (value instanceof Texture) {
                    disposeObjectResources(value);
                    uniforms[key] = null;
                }
                else if (value instanceof Uniform) {
                    disposeObjectResources(value.value);
                    value.value = null;
                }
            }
        }
    }
    else if (obj instanceof Texture) {
        free(obj);
        free(obj.source);
        if (obj.source?.data instanceof ImageBitmap) {
            free(obj.source.data);
        }
    }
    else if (obj instanceof Skeleton) {
        free(obj.boneTexture);
        obj.boneTexture = null;
    }
    else if (obj instanceof Bone) {
    }
    else {
        if (!(obj instanceof Object3D) && debug)
            console.warn("Unknown object type", obj);
    }
}

function free(obj: any) {
    if (!obj) {
        return;
    }
    if (debug || autoDispose() || trackUsageParam) console.warn("🧨 FREE", obj);
    if (obj instanceof ImageBitmap) {
        obj.close();
    }
    else if (obj instanceof Source) {
        obj.data = null;
    }
    else {
        obj.dispose();
    }
}

export function __internalNotifyObjectDestroyed(obj: Object3D) {
    if (obj instanceof Mesh || obj instanceof SkinnedMesh) {
        obj.material = null;
        obj.geometry = null;
    }
}

const usersBuffer = new Set<object>();

export type UserFilter = (user: object) => boolean;

/**
 * Find all users of an object
 * @param object Object to find users of
 * @param recursive Find users of users
 * @param predicate Filter users
 * @param set Set to add users to, a new one will be created if none is provided
 * @returns a set of users
 */
export function findResourceUsers(object: object, recursive: boolean, predicate: UserFilter | null | undefined = null, set?: Set<object>): Set<object> {
    if (!set) {
        set = usersBuffer;
        set.clear();
    }
    if (!object) return set;
    const users = object[$objectUsersKey] as Set<object>;
    if (users) {
        for (const user of users) {
            // Prevent infinite loop if recursive references
            if (set.has(user)) continue;
            //  Allow filtering
            if (predicate?.call(null, user) === false) continue;
            set.add(user);
            if (recursive)
                findResourceUsers(user, true, predicate, set);
        }
    }
    return set;
}

export function getResourceUserCount(object: object): number | undefined {
    return object[$objectUsersCountKey];
}



const debug = getParam("debugresourceusers") || getParam("debugmemory");

// Should we check if the type has the
const $objectUsersKey = Symbol("needle-resource-users");
const $objectUsersCountKey = Symbol("needle-resource-users-count");

function trackValueChange(prototype, fieldName) {
    addPatch(prototype, fieldName, function (this: object, oldValue, newValue) {
        if (allowUsageTracking && !Physics.raycasting) {
            updateUsers($objectUsersKey, this, oldValue, false);
            updateUsers($objectUsersKey, this, newValue, true);
        }
    });
}

// function stopTracking(prototype, fieldName) {
//     const $key = Symbol("needle-using-" + fieldName);
//     const currentValue = prototype[$key];
//     delete prototype[$key];
//     prototype[fieldName] = currentValue;
//     updateUsers($objectUsersKey, fieldName, prototype, currentValue, true);
// }


if (allowUsageTracking) {
    trackValueChange(Mesh.prototype, "material");
    trackValueChange(Mesh.prototype, "geometry");
    trackValueChange(Material.prototype, "map");
    trackValueChange(Material.prototype, "bumpMap");
    trackValueChange(Material.prototype, "alphaMap");
    trackValueChange(Material.prototype, "normalMap");
    trackValueChange(Material.prototype, "displacementMap");
    trackValueChange(Material.prototype, "roughnessMap");
    trackValueChange(Material.prototype, "metalnessMap");
    trackValueChange(Material.prototype, "emissiveMap");
    trackValueChange(Material.prototype, "specularMap");
    trackValueChange(Material.prototype, "envMap");
    trackValueChange(Material.prototype, "lightMap");
    trackValueChange(Material.prototype, "aoMap");
    trackValueChange(Material.prototype, "gradientMap");
}



// TODO: patch dispose?


function onDispose(obj: object) {
    if (allowUsageTracking === false) return;
    const users = obj[$objectUsersKey] as Set<object>;
    if (users) {
        for (const user of users) {
            updateUsers($objectUsersKey, user, obj, false);
        }
    }
}

if (allowUsageTracking) {
    addPatch(Material.prototype, "dispose", function (this: object) { onDispose(this) });
}




// This variable is crucial for performance: 
// it is incremented during rendering to prevent usage updates during the three.js render loop 
// where materials and properties are updated every frame (e.g. the DepthMaterial)
// and we don't care about those
let noUpdateScope = 0;

// Main method called by wrapped fields/properties to update the users for an object
function updateUsers(symbol: symbol, user: object, object: object | object[], added: boolean) {

    // If we are rendering we dont want to update users
    if (noUpdateScope > 0) return;

    if (Array.isArray(object)) {
        for (const m of object) {
            updateUsers(symbol, user, m, added);
        }
        return;
    }
    if (!object) return;

    let users = object[symbol];
    if (!users) users = new Set();
    if (added) {
        if (user && !users.has(user)) {
            users.add(user);
            let count = object[$objectUsersCountKey] || 0;
            count += 1;
            object[$objectUsersCountKey] = count;
            if (debug) console.warn(`🟢 Added user of "${object["type"]}"`, user, object, count, "users:", users);
        }
    } else {
        if (user && users.has(user)) {
            users.delete(user);
            let count = object[$objectUsersCountKey] || 0;
            if (count > 0) {
                count -= 1;
                object[$objectUsersCountKey] = count;
            }
            if (debug) console.warn(`🔴 Removed user of "${object["type"]}"`, user, object, count, "users:", users);
            if (count <= 0) {
                if (!InternalUsageTrackerPlugin.isLoading(object)) {
                    if (trackUsageParam)
                        console.warn(`🔴 Removed all user of "${object["type"]}"`, object);
                    if (autoDispose())
                        disposeObjectResources(object);
                }
            }
        }
    }
    object[symbol] = users;
}


// We dont want to update users during rendering
try {
    addPatch(WebGLRenderer.prototype, "render",
        function () {
            noUpdateScope++;
        },
        function () {
            noUpdateScope--;
        }
    );
}
catch (e) {
    console.warn("Could not wrap WebGLRenderer.render", e);
}


// addGltfLoadEventListener(GltfLoadEventType.BeforeLoad, (_) => {
//     noUpdateScope++;
// });
// addGltfLoadEventListener(GltfLoadEventType.AfterLoaded, (_) => {
//     noUpdateScope--;
// });




// addPatch(Object3D.prototype, "add", (obj: Object3D) => {
// });


// addPatch(Object3D.prototype, "remove", (obj: Object3D) => {
//     if(obj instanceof Mesh) {
//     }
// });





// class MyObject {
//     myNumber: number = 1;
// }

// addPatch(MyObject.prototype, "myNumber", (obj, oldValue, newValue) => {
//     console.log("myNumber changed", oldValue, newValue);
// });

// const i = new MyObject();
// setInterval(() => {
//     console.log("RUN");
//     i.myNumber += 1;
// }, 1000);