@needle-tools/engine/src/engine/engine_pmrem.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 { createLoaders } from "@needle-tools/gltf-progressive";
import { CubeUVReflectionMapping, SRGBColorSpace, Texture, TextureLoader, WebGLRenderer } from "three";
import { EXRLoader } from "three/examples/jsm/loaders/EXRLoader";
import { KTX2Loader } from "three/examples/jsm/loaders/KTX2Loader";
import { RGBELoader } from "three/examples/jsm/loaders/RGBELoader";

import { disposeObjectResources, setDisposable } from "./engine_assetdatabase.js";

const running: Map<string, Promise<Texture | null>> = new Map();

// #region api

export function loadPMREM(url: string, renderer: WebGLRenderer): Promise<Texture | null> {
    if (running.has(url)) {
        return running.get(url)!;
    }
    const actualUrl = new URL(url, window.location.href);
    const promise = internalLoadPMREM(actualUrl, renderer);
    running.set(url, promise);
    return promise;
}



// #region Cache

declare type SkyboxCacheEntry = { src: string, texture: Promise<Texture | null> };
function ensureGlobalCache() {
    if (!globalThis["NEEDLE_ENGINE_SKYBOX_TEXTURES"])
        globalThis["NEEDLE_ENGINE_SKYBOX_TEXTURES"] = new Array<SkyboxCacheEntry>();
    return globalThis["NEEDLE_ENGINE_SKYBOX_TEXTURES"] as Array<SkyboxCacheEntry>;
}

function tryGetPreviouslyLoadedTexture(src: string) {
    const cache = ensureGlobalCache();
    const found = cache.find(x => x.src === src);
    if (found) {
        return found.texture;
    }
    return null;
}
async function disposeCachedTexture(tex: Promise<Texture | null>) {
    const texture = await tex;
    if (!texture) return;
    setDisposable(texture, true);
    disposeObjectResources(texture);
}
function registerPromise(src: string, texture: Promise<Texture | null>) {
    const cache = ensureGlobalCache();
    // Make sure the cache doesnt get too big
    while (cache.length > 5) {
        const entry = cache.shift();
        if (entry) { disposeCachedTexture(entry.texture); }
    }
    texture.then(t => { return setDisposable(t, false) });
    cache.push({ src, texture });
}



// #region loading


async function internalLoadPMREM(url: URL, renderer: WebGLRenderer) {
    if (!url) return Promise.resolve(null);

    const pathname = url.pathname;
    const isPMREM_URL: boolean = url.toString().toLowerCase().includes("pmrem") || url.searchParams.get("pmrem") != null;

    const cached = tryGetPreviouslyLoadedTexture(pathname);
    if (cached) {
        const res = await cached;
        if (res?.source?.data?.length > 0 || res?.source?.data?.data?.length) return res;
    }
    const isEXR = pathname.endsWith(".exr");
    const isHdr = pathname.endsWith(".hdr");
    const isKtx2 = pathname.endsWith(".ktx2");

    let loader: RGBELoader | EXRLoader | TextureLoader | KTX2Loader;

    if (isEXR) {
        loader = new EXRLoader();
    }
    else if (isHdr) {
        loader = new RGBELoader();
    }
    else if (isKtx2) {
        const { ktx2Loader } = createLoaders(renderer);
        loader = ktx2Loader;
    }
    else {
        loader = new TextureLoader();
    }

    const str = url.toString();
    const promise = loader.loadAsync(str)
        .then(tex => {
            if (tex) {
                const nameIndex = pathname.lastIndexOf("/");
                tex.name = pathname.substring(nameIndex >= 0 ? nameIndex + 1 : 0);

                if (isPMREM_URL) {
                    tex.mapping = CubeUVReflectionMapping;
                }

                if (loader instanceof TextureLoader) {
                    tex.colorSpace = SRGBColorSpace;
                }
            }
            return tex;
        });

    registerPromise(str, promise);
    const texture = await promise.catch(_err => {
        console.warn("Failed to load texture from url:", url);
        return null;
    });
    return texture;
}