backend/libs/three/renderers/webgl/WebGLShaderCache.js

let _id = 0;

class WebGLShaderCache {
	constructor() {
		this.shaderCache = new Map();
		this.materialCache = new Map();
	}

	update(material) {
		const vertexShader = material.vertexShader;
		const fragmentShader = material.fragmentShader;

		const vertexShaderStage = this._getShaderStage(vertexShader);
		const fragmentShaderStage = this._getShaderStage(fragmentShader);

		const materialShaders = this._getShaderCacheForMaterial(material);

		if (materialShaders.has(vertexShaderStage) === false) {
			materialShaders.add(vertexShaderStage);
			vertexShaderStage.usedTimes++;
		}

		if (materialShaders.has(fragmentShaderStage) === false) {
			materialShaders.add(fragmentShaderStage);
			fragmentShaderStage.usedTimes++;
		}

		return this;
	}

	remove(material) {
		const materialShaders = this.materialCache.get(material);

		for (const shaderStage of materialShaders) {
			shaderStage.usedTimes--;

			if (shaderStage.usedTimes === 0) this.shaderCache.delete(shaderStage);
		}

		this.materialCache.delete(material);

		return this;
	}

	getVertexShaderID(material) {
		return this._getShaderStage(material.vertexShader).id;
	}

	getFragmentShaderID(material) {
		return this._getShaderStage(material.fragmentShader).id;
	}

	dispose() {
		this.shaderCache.clear();
		this.materialCache.clear();
	}

	_getShaderCacheForMaterial(material) {
		const cache = this.materialCache;

		if (cache.has(material) === false) {
			cache.set(material, new Set());
		}

		return cache.get(material);
	}

	_getShaderStage(code) {
		const cache = this.shaderCache;

		if (cache.has(code) === false) {
			const stage = new WebGLShaderStage();
			cache.set(code, stage);
		}

		return cache.get(code);
	}
}

class WebGLShaderStage {
	constructor() {
		this.id = _id++;

		this.usedTimes = 0;
	}
}

export { WebGLShaderCache };