GPUSkinningCompute.js

GPUSkinningCompute.js

import { StorageInstancedBufferAttribute, StorageTexture, FloatType, NearestFilter, Matrix4, Vector3, Quaternion } from 'three/webgpu';
import { Fn, If, storage, storageTexture, textureStore, instanceIndex, attribute, uniform, float, int, vec4, mat4, mod, floor, clamp, add, ivec2, NodeAccess, positionLocal, normalLocal } from 'three/tsl';
/**
 * Sorts bones hierarchically so parents always come before children.
 * This is required for proper bone matrix propagation in GPU compute.
 */
function sortBonesHierarchically(skeleton) {
    const bones = skeleton.bones;
    const result = [];
    const visited = new Set();
    const originalToSorted = new Map();
    function visit(boneIndex) {
        if (visited.has(boneIndex))
            return;
        const bone = bones[boneIndex];
        // Visit parent first if it exists and hasn't been visited
        if (bone.parent && bone.parent.isBone) {
            const parentIndex = bones.indexOf(bone.parent);
            if (parentIndex >= 0 && !visited.has(parentIndex)) {
                visit(parentIndex);
            }
        }
        visited.add(boneIndex);
        const sortedIndex = result.length;
        originalToSorted.set(boneIndex, sortedIndex);
        let parentSortedIndex = -1;
        if (bone.parent && bone.parent.isBone) {
            const parentOriginalIndex = bones.indexOf(bone.parent);
            if (parentOriginalIndex >= 0) {
                parentSortedIndex = originalToSorted.get(parentOriginalIndex);
            }
        }
        result.push({
            originalIndex: boneIndex,
            parentSortedIndex
        });
    }
    for (let i = 0; i < bones.length; i++) {
        visit(i);
    }
    return result;
}
/**
 * Parses animation track name to extract bone name and property.
 */
function parseTrackName(name) {
    // Try bracket notation: ".bones[BoneName].property"
    const bracketMatch = name.match(/\.bones\[(.+?)\]\.(\w+)/);
    if (bracketMatch) {
        return {
            bone: bracketMatch[1],
            property: bracketMatch[2]
        };
    }
    // Try dot notation: "BoneName.property"
    const dotMatch = name.match(/^(.+)\.(\w+)$/);
    if (dotMatch) {
        return {
            bone: dotMatch[1],
            property: dotMatch[2]
        };
    }
    return null;
}
/**
 * Packs animation data into a single GPU-friendly typed array.
 * Layout: [metadata] [posTimes] [quatTimes] [scaleTimes] [posValues] [quatValues] [scaleValues]
 * Metadata per bone stores GLOBAL offsets into the packed buffer.
 */
function packAnimationData(clip, skeleton, sortedBones) {
    const bones = skeleton.bones;
    const boneNameToIndex = new Map();
    const boneToIndex = new Map();
    bones.forEach((bone, i) => {
        boneNameToIndex.set(bone.name, i);
        boneToIndex.set(bone, i);
    });
    // Initialize per-bone track data with bind pose defaults
    const boneTrackData = [];
    // Prefer bind pose derived from skeleton.boneInverses to avoid using a potentially animated
    // runtime pose as the "default" for bones without tracks (can cause popping on some models).
    const canDeriveBindPose = Array.isArray(skeleton.boneInverses) && skeleton.boneInverses.length === bones.length;
    const bindPositions = new Array(bones.length);
    const bindQuaternions = new Array(bones.length);
    const bindScales = new Array(bones.length);
    if (canDeriveBindPose) {
        const worldBind = new Matrix4();
        const localBind = new Matrix4();
        // Compute inverse of skeleton root to factor it out from root bone bind poses
        // This prevents double-application (bind pose + compute shader both applying skeleton root)
        let skeletonRootInverse = null;
        for (let i = 0; i < bones.length; i++) {
            const parent = bones[i].parent;
            if (parent && !parent.isBone && parent.matrixWorld) {
                skeletonRootInverse = parent.matrixWorld.clone().invert();
                break;
            }
        }
        for (let i = 0; i < bones.length; i++) {
            // worldBind = inverse( boneInverse ) = bone's world matrix at bind time
            worldBind.copy(skeleton.boneInverses[i]).invert();
            const parent = bones[i].parent;
            const parentIndex = (parent && parent.isBone) ? boneToIndex.get(parent) : undefined;
            if (parentIndex !== undefined) {
                // Child bone: localBind = inverse(parentWorld) * childWorld
                localBind.copy(skeleton.boneInverses[parentIndex]).multiply(worldBind);
            }
            else {
                // Root bone: factor out skeleton root to get bone-local transform
                // worldBind = skeletonRoot * boneLocal, so boneLocal = inv(skeletonRoot) * worldBind
                if (skeletonRootInverse) {
                    localBind.copy(skeletonRootInverse).multiply(worldBind);
                }
                else {
                    localBind.copy(worldBind);
                }
            }
            const pos = new Vector3();
            const quat = new Quaternion();
            const scale = new Vector3();
            localBind.decompose(pos, quat, scale);
            bindPositions[i] = pos;
            bindQuaternions[i] = quat;
            bindScales[i] = scale;
        }
    }
    for (let i = 0; i < bones.length; i++) {
        const bone = bones[i];
        const position = canDeriveBindPose ? bindPositions[i] : bone.position;
        const quaternion = canDeriveBindPose ? bindQuaternions[i] : bone.quaternion;
        const scale = canDeriveBindPose ? bindScales[i] : bone.scale;
        boneTrackData.push({
            position: {
                times: [0],
                values: [position.x, position.y, position.z]
            },
            quaternion: {
                times: [0],
                values: [quaternion.x, quaternion.y, quaternion.z, quaternion.w]
            },
            scale: {
                times: [0],
                values: [scale.x, scale.y, scale.z]
            }
        });
    }
    // Parse animation tracks and populate bone data
    for (const track of clip.tracks) {
        const parsed = parseTrackName(track.name);
        if (!parsed)
            continue;
        const boneIndex = boneNameToIndex.get(parsed.bone);
        if (boneIndex === undefined)
            continue;
        if (parsed.property === 'position' || parsed.property === 'quaternion' || parsed.property === 'scale') {
            boneTrackData[boneIndex][parsed.property] = {
                times: Array.from(track.times),
                values: Array.from(track.values)
            };
        }
    }
    // Count total keyframes
    let totalPositionKeyframes = 0;
    let totalQuaternionKeyframes = 0;
    let totalScaleKeyframes = 0;
    for (const data of boneTrackData) {
        totalPositionKeyframes += data.position.times.length;
        totalQuaternionKeyframes += data.quaternion.times.length;
        totalScaleKeyframes += data.scale.times.length;
    }
    // Calculate section offsets in packed buffer
    const metadataSize = bones.length * 8;
    const metadataOffset = 0;
    const positionTimesOffset = metadataSize;
    const quaternionTimesOffset = positionTimesOffset + totalPositionKeyframes;
    const scaleTimesOffset = quaternionTimesOffset + totalQuaternionKeyframes;
    const positionValuesOffset = scaleTimesOffset + totalScaleKeyframes;
    const quaternionValuesOffset = positionValuesOffset + totalPositionKeyframes * 3;
    const scaleValuesOffset = quaternionValuesOffset + totalQuaternionKeyframes * 4;
    const totalSize = scaleValuesOffset + totalScaleKeyframes * 3;
    // Allocate single packed buffer
    const packedData = new Float32Array(totalSize);
    let posTimeIdx = 0;
    let quatTimeIdx = 0;
    let scaleTimeIdx = 0;
    let posValIdx = 0;
    let quatValIdx = 0;
    let scaleValIdx = 0;
    // Detect frame time from first track with multiple keyframes
    let detectedFrameTime = 1 / 30;
    // Pack data in hierarchically sorted order
    for (let sortedIdx = 0; sortedIdx < sortedBones.length; sortedIdx++) {
        const originalIdx = sortedBones[sortedIdx].originalIndex;
        const parentSortedIdx = sortedBones[sortedIdx].parentSortedIndex;
        const data = boneTrackData[originalIdx];
        const metaBase = metadataOffset + sortedIdx * 8;
        // Position track - store GLOBAL offset into packed buffer
        packedData[metaBase + 0] = positionTimesOffset + posTimeIdx; // global time offset
        packedData[metaBase + 1] = data.position.times.length;
        for (let i = 0; i < data.position.times.length; i++) {
            packedData[positionTimesOffset + posTimeIdx + i] = data.position.times[i];
        }
        const posValBase = positionValuesOffset + posValIdx * 3;
        for (let i = 0; i < data.position.times.length; i++) {
            packedData[posValBase + i * 3 + 0] = data.position.values[i * 3 + 0];
            packedData[posValBase + i * 3 + 1] = data.position.values[i * 3 + 1];
            packedData[posValBase + i * 3 + 2] = data.position.values[i * 3 + 2];
        }
        // Values offset computed by shader from time offset + section offsets
        posTimeIdx += data.position.times.length;
        posValIdx += data.position.times.length;
        // Quaternion track
        packedData[metaBase + 2] = quaternionTimesOffset + quatTimeIdx;
        packedData[metaBase + 3] = data.quaternion.times.length;
        for (let i = 0; i < data.quaternion.times.length; i++) {
            packedData[quaternionTimesOffset + quatTimeIdx + i] = data.quaternion.times[i];
            // Detect frame time
            if (i > 0 && detectedFrameTime === 1 / 30) {
                const delta = data.quaternion.times[i] - data.quaternion.times[i - 1];
                if (delta > 0.001) {
                    detectedFrameTime = delta;
                }
            }
        }
        const quatValBase = quaternionValuesOffset + quatValIdx * 4;
        for (let i = 0; i < data.quaternion.times.length; i++) {
            packedData[quatValBase + i * 4 + 0] = data.quaternion.values[i * 4 + 0];
            packedData[quatValBase + i * 4 + 1] = data.quaternion.values[i * 4 + 1];
            packedData[quatValBase + i * 4 + 2] = data.quaternion.values[i * 4 + 2];
            packedData[quatValBase + i * 4 + 3] = data.quaternion.values[i * 4 + 3];
        }
        quatTimeIdx += data.quaternion.times.length;
        quatValIdx += data.quaternion.times.length;
        // Scale track
        packedData[metaBase + 4] = scaleTimesOffset + scaleTimeIdx;
        packedData[metaBase + 5] = data.scale.times.length;
        for (let i = 0; i < data.scale.times.length; i++) {
            packedData[scaleTimesOffset + scaleTimeIdx + i] = data.scale.times[i];
        }
        const scaleValBase = scaleValuesOffset + scaleValIdx * 3;
        for (let i = 0; i < data.scale.times.length; i++) {
            packedData[scaleValBase + i * 3 + 0] = data.scale.values[i * 3 + 0];
            packedData[scaleValBase + i * 3 + 1] = data.scale.values[i * 3 + 1];
            packedData[scaleValBase + i * 3 + 2] = data.scale.values[i * 3 + 2];
        }
        scaleTimeIdx += data.scale.times.length;
        scaleValIdx += data.scale.times.length;
        // Parent and original indices
        packedData[metaBase + 6] = parentSortedIdx;
        packedData[metaBase + 7] = originalIdx;
    }
    return {
        packedData,
        metadataOffset,
        positionTimesOffset,
        quaternionTimesOffset,
        scaleTimesOffset,
        positionValuesOffset,
        quaternionValuesOffset,
        scaleValuesOffset,
        totalPositionKeyframes,
        totalQuaternionKeyframes,
        totalScaleKeyframes,
        detectedFrameTime
    };
}
/**
 * Bakes an animation clip into pre-computed final bone matrices at a fixed sample rate.
 * For each frame, evaluates all bones: interpolates keyframes on CPU, builds world
 * matrices via hierarchy traversal, applies boneInverse, and remaps to original order.
 *
 * Returns a Float32Array of size (boneCount * 4) * frameCount * 4, laid out as
 * RGBA texels: 4 texels per mat4, boneCount * 4 texels wide, frameCount rows tall.
 */
function bakeAnimation( clip, skeleton, sortedBones, sampleRate = 30 ) {

	const boneCount = skeleton.bones.length;
	const duration = clip.duration;
	const frameCount = Math.max( 1, Math.ceil( duration * sampleRate ) + 1 );

	// Parse tracks into per-bone arrays for fast CPU interpolation
	const boneNameToIndex = new Map();
	skeleton.bones.forEach( ( bone, i ) => boneNameToIndex.set( bone.name, i ) );

	// Derive bind pose from boneInverses (same logic as packAnimationData)
	const canDeriveBindPose = Array.isArray( skeleton.boneInverses ) && skeleton.boneInverses.length === boneCount;
	const bindPos = new Array( boneCount );
	const bindQuat = new Array( boneCount );
	const bindScale = new Array( boneCount );
	const boneToIndex = new Map();
	skeleton.bones.forEach( ( bone, i ) => boneToIndex.set( bone, i ) );

	if ( canDeriveBindPose ) {

		const worldBind = new Matrix4();
		const localBind = new Matrix4();
		let skeletonRootInverse = null;

		for ( let i = 0; i < boneCount; i ++ ) {

			const parent = skeleton.bones[ i ].parent;
			if ( parent && ! parent.isBone && parent.matrixWorld ) {

				skeletonRootInverse = parent.matrixWorld.clone().invert();
				break;

			}

		}

		for ( let i = 0; i < boneCount; i ++ ) {

			worldBind.copy( skeleton.boneInverses[ i ] ).invert();
			const parent = skeleton.bones[ i ].parent;
			const parentIndex = ( parent && parent.isBone ) ? boneToIndex.get( parent ) : undefined;

			if ( parentIndex !== undefined ) {

				localBind.copy( skeleton.boneInverses[ parentIndex ] ).multiply( worldBind );

			} else {

				if ( skeletonRootInverse ) {

					localBind.copy( skeletonRootInverse ).multiply( worldBind );

				} else {

					localBind.copy( worldBind );

				}

			}

			const p = new Vector3();
			const q = new Quaternion();
			const s = new Vector3();
			localBind.decompose( p, q, s );
			bindPos[ i ] = [ p.x, p.y, p.z ];
			bindQuat[ i ] = [ q.x, q.y, q.z, q.w ];
			bindScale[ i ] = [ s.x, s.y, s.z ];

		}

	} else {

		for ( let i = 0; i < boneCount; i ++ ) {

			const bone = skeleton.bones[ i ];
			bindPos[ i ] = [ bone.position.x, bone.position.y, bone.position.z ];
			bindQuat[ i ] = [ bone.quaternion.x, bone.quaternion.y, bone.quaternion.z, bone.quaternion.w ];
			bindScale[ i ] = [ bone.scale.x, bone.scale.y, bone.scale.z ];

		}

	}

	// Per-bone track data: { posTimes, posValues, quatTimes, quatValues, scaleTimes, scaleValues }
	const tracks = [];

	for ( let i = 0; i < boneCount; i ++ ) {

		tracks.push( {
			posTimes: [ 0 ], posValues: bindPos[ i ].slice(),
			quatTimes: [ 0 ], quatValues: bindQuat[ i ].slice(),
			scaleTimes: [ 0 ], scaleValues: bindScale[ i ].slice()
		} );

	}

	for ( const track of clip.tracks ) {

		const parsed = parseTrackName( track.name );
		if ( ! parsed ) continue;
		const idx = boneNameToIndex.get( parsed.bone );
		if ( idx === undefined ) continue;

		if ( parsed.property === 'position' ) {

			tracks[ idx ].posTimes = Array.from( track.times );
			tracks[ idx ].posValues = Array.from( track.values );

		} else if ( parsed.property === 'quaternion' ) {

			tracks[ idx ].quatTimes = Array.from( track.times );
			tracks[ idx ].quatValues = Array.from( track.values );

		} else if ( parsed.property === 'scale' ) {

			tracks[ idx ].scaleTimes = Array.from( track.times );
			tracks[ idx ].scaleValues = Array.from( track.values );

		}

	}

	// CPU interpolation helpers
	function lerpScalar( a, b, t ) {

		return a + ( b - a ) * t;

	}

	function findKeyframeIndex( times, time ) {

		if ( times.length <= 1 ) return { i0: 0, i1: 0, t: 0 };
		if ( time <= times[ 0 ] ) return { i0: 0, i1: 0, t: 0 };
		if ( time >= times[ times.length - 1 ] ) return { i0: times.length - 1, i1: times.length - 1, t: 0 };

		for ( let i = 0; i < times.length - 1; i ++ ) {

			if ( time >= times[ i ] && time < times[ i + 1 ] ) {

				const span = times[ i + 1 ] - times[ i ];
				return { i0: i, i1: i + 1, t: span > 0 ? ( time - times[ i ] ) / span : 0 };

			}

		}

		return { i0: times.length - 1, i1: times.length - 1, t: 0 };

	}

	function sampleVec3( times, values, time ) {

		const { i0, i1, t } = findKeyframeIndex( times, time );
		return [
			lerpScalar( values[ i0 * 3 ], values[ i1 * 3 ], t ),
			lerpScalar( values[ i0 * 3 + 1 ], values[ i1 * 3 + 1 ], t ),
			lerpScalar( values[ i0 * 3 + 2 ], values[ i1 * 3 + 2 ], t )
		];

	}

	function sampleQuat( times, values, time ) {

		const { i0, i1, t } = findKeyframeIndex( times, time );
		let x0 = values[ i0 * 4 ], y0 = values[ i0 * 4 + 1 ], z0 = values[ i0 * 4 + 2 ], w0 = values[ i0 * 4 + 3 ];
		let x1 = values[ i1 * 4 ], y1 = values[ i1 * 4 + 1 ], z1 = values[ i1 * 4 + 2 ], w1 = values[ i1 * 4 + 3 ];

		// Shortest path
		if ( x0 * x1 + y0 * y1 + z0 * z1 + w0 * w1 < 0 ) {

			x1 = - x1; y1 = - y1; z1 = - z1; w1 = - w1;

		}

		let rx = lerpScalar( x0, x1, t );
		let ry = lerpScalar( y0, y1, t );
		let rz = lerpScalar( z0, z1, t );
		let rw = lerpScalar( w0, w1, t );

		// Normalize
		const len = Math.sqrt( rx * rx + ry * ry + rz * rz + rw * rw ) || 1;
		rx /= len; ry /= len; rz /= len; rw /= len;

		return [ rx, ry, rz, rw ];

	}

	// Compute skeleton root transform
	let skeletonRootMatrix = null;
	const roots = sortedBones.filter( s => s.parentSortedIndex === - 1 );

	if ( roots.length > 0 ) {

		const rootBone = skeleton.bones[ roots[ 0 ].originalIndex ];
		const parent = rootBone.parent;

		if ( parent && ! parent.isBone && parent.matrixWorld ) {

			skeletonRootMatrix = parent.matrixWorld.clone();

		}

	}

	// Allocate output: frameCount × boneCount mat4s (16 floats each)
	// Layout: [frame0_bone0, frame0_bone1, ..., frame1_bone0, ...]
	const output = new Float32Array( frameCount * boneCount * 16 );

	const localMat = new Matrix4();
	const worldMat = new Matrix4();
	const tmpPos = new Vector3();
	const tmpQuat = new Quaternion();
	const tmpScale = new Vector3();
	const worldMatrices = new Array( boneCount );

	for ( let i = 0; i < boneCount; i ++ ) {

		worldMatrices[ i ] = new Matrix4();

	}

	for ( let frame = 0; frame < frameCount; frame ++ ) {

		const time = Math.min( ( frame / sampleRate ), duration );

		// Pass 1: compute world matrices in hierarchical (sorted) order
		for ( let sortedIdx = 0; sortedIdx < sortedBones.length; sortedIdx ++ ) {

			const originalIdx = sortedBones[ sortedIdx ].originalIndex;
			const parentSortedIdx = sortedBones[ sortedIdx ].parentSortedIndex;
			const tr = tracks[ originalIdx ];

			const pos = sampleVec3( tr.posTimes, tr.posValues, time );
			const quat = sampleQuat( tr.quatTimes, tr.quatValues, time );
			const scale = sampleVec3( tr.scaleTimes, tr.scaleValues, time );

			tmpPos.set( pos[ 0 ], pos[ 1 ], pos[ 2 ] );
			tmpQuat.set( quat[ 0 ], quat[ 1 ], quat[ 2 ], quat[ 3 ] );
			tmpScale.set( scale[ 0 ], scale[ 1 ], scale[ 2 ] );
			localMat.compose( tmpPos, tmpQuat, tmpScale );

			if ( parentSortedIdx >= 0 ) {

				worldMatrices[ sortedIdx ].multiplyMatrices( worldMatrices[ parentSortedIdx ], localMat );

			} else if ( skeletonRootMatrix ) {

				worldMatrices[ sortedIdx ].multiplyMatrices( skeletonRootMatrix, localMat );

			} else {

				worldMatrices[ sortedIdx ].copy( localMat );

			}

		}

		// Pass 2: apply boneInverse, remap to original order, write as mat4
		for ( let sortedIdx = 0; sortedIdx < sortedBones.length; sortedIdx ++ ) {

			const originalIdx = sortedBones[ sortedIdx ].originalIndex;
			worldMat.multiplyMatrices( worldMatrices[ sortedIdx ], skeleton.boneInverses[ originalIdx ] );

			// Write mat4 at [frame * boneCount + originalIdx]
			worldMat.toArray( output, ( frame * boneCount + originalIdx ) * 16 );

		}

	}

	return { data: output, frameCount, sampleRate };

}
/**
 * GPU-driven skeletal animation class for WebGPU compute shaders.
 * Supports animation blending between two clips and instanced rendering.
 *
 * @example
 * const gpuAnim = new GPUSkinningCompute(skinnedMesh, clips, { instanceCount: 100 });
 * gpuAnim.setAnimationA(0);  // walk
 * gpuAnim.setAnimationB(1);  // run
 * gpuAnim.blendWeightUniform.value = 0.5;  // 50% blend
 *
 * const compute = gpuAnim.createComputeNode();
 * // In render loop:
 * renderer.compute(compute);
 */
class GPUSkinningCompute {
    /**
     * Creates a GPU skinned animation controller.
     * Pre-bakes animation data into a storage buffer for fast GPU lookup.
     */
    constructor(skinnedMesh, clips, options = {}) {
        this.skinnedMesh = skinnedMesh;
        this.skeleton = skinnedMesh.skeleton;
        this.clips = clips;
        this.instanceCount = options.instanceCount || 1;
        this.boneCount = this.skeleton.bones.length;
        this.animationNames = clips.map((c, i) => c.name || `Animation ${i}`);
        // Sort bones hierarchically (needed by bakeAnimation)
        this._sortedBones = sortBonesHierarchically(this.skeleton);
        // Bake all clips into pre-computed bone matrices
        const bakedClips = [];
        let maxFrameCount = 0;
        for (const clip of clips) {
            const baked = bakeAnimation(clip, this.skeleton, this._sortedBones);
            bakedClips.push(baked);
            maxFrameCount = Math.max(maxFrameCount, baked.frameCount);
        }
        this._bakedClips = bakedClips;
        this._maxFrameCount = maxFrameCount;
        // Pack all clips into a single storage buffer
        // Layout: [clip0_frame0_bone0..boneN, clip0_frame1_bone0..boneN, ..., clip1_frame0_bone0..boneN, ...]
        const boneCount = this.boneCount;
        const totalMat4s = clips.length * maxFrameCount * boneCount;
        const bakedData = new Float32Array(totalMat4s * 16);
        for (let layer = 0; layer < clips.length; layer++) {
            const baked = bakedClips[layer];
            const layerOffset = layer * maxFrameCount * boneCount * 16;
            // Copy baked frames
            bakedData.set(baked.data, layerOffset);
            // If this clip has fewer frames than max, repeat last frame to fill
            if (baked.frameCount < maxFrameCount) {
                const lastFrameData = baked.data.subarray(
                    (baked.frameCount - 1) * boneCount * 16,
                    baked.frameCount * boneCount * 16
                );
                for (let f = baked.frameCount; f < maxFrameCount; f++) {
                    bakedData.set(lastFrameData, layerOffset + f * boneCount * 16);
                }
            }
        }
        this._bakedBuffer = new StorageInstancedBufferAttribute(bakedData, 16);
        this._totalBakedMat4s = totalMat4s;
        // Current animation indices
        this._animationAIndex = 0;
        this._animationBIndex = clips.length > 1 ? 1 : 0;
        // Uniforms
        this.timeUniform = uniform(0);
        this.blendWeightUniform = uniform(0);
        this._durationAUniform = uniform(clips[this._animationAIndex].duration);
        this._durationBUniform = uniform(clips[this._animationBIndex].duration);
        this._clipOffsetAUniform = uniform(this._animationAIndex * maxFrameCount * boneCount);
        this._clipOffsetBUniform = uniform(this._animationBIndex * maxFrameCount * boneCount);
        this._frameCountAUniform = uniform(bakedClips[this._animationAIndex].frameCount);
        this._frameCountBUniform = uniform(bakedClips[this._animationBIndex].frameCount);
        this._sampleRateUniform = uniform(bakedClips[0].sampleRate);
        // Storage buffers: instanceOffsets + output boneMatrices
        this._createStorageBuffers();
        this._computeNode = null;
        this._readBoneTex = null;
    }
    /**
     * Creates GPU storage buffers for instance offsets and output bone matrices.
     */
    _createStorageBuffers() {
        const boneCount = this.boneCount;
        const instanceCount = this.instanceCount;
        // Instance time offsets
        const instanceOffsetsArray = new Float32Array(instanceCount);
        const duration = this.clips[this._animationAIndex].duration;
        for (let i = 0; i < instanceCount; i++) {
            instanceOffsetsArray[i] = (i / instanceCount) * duration;
        }
        this._instanceOffsetsBuffer = new StorageInstancedBufferAttribute(instanceOffsetsArray, 1);
        // Output bone matrices as StorageTexture for texture cache benefits on tile-based GPUs.
        // Each mat4 = 4 vec4 columns = 4 RGBA texels.
        const texWidth = boneCount * 4;
        const texHeight = instanceCount;
        this._boneMatricesTex = new StorageTexture( texWidth, texHeight );
        this._boneMatricesTex.type = FloatType;
        this._boneMatricesTex.minFilter = NearestFilter;
        this._boneMatricesTex.magFilter = NearestFilter;
    }
    /**
     * Sets the first animation (A) by index.
     */
    setAnimationA(index) {
        this._setAnimation('A', index);
    }
    /**
     * Sets the second animation (B) by index for blending.
     */
    setAnimationB(index) {
        this._setAnimation('B', index);
    }
    _setAnimation(slot, index) {
        if (index < 0 || index >= this.clips.length)
            return;
        const isA = slot === 'A';
        const clipOffset = index * this._maxFrameCount * this.boneCount;
        if (isA) {
            this._animationAIndex = index;
            this._durationAUniform.value = this.clips[index].duration;
            this._clipOffsetAUniform.value = clipOffset;
            this._frameCountAUniform.value = this._bakedClips[index].frameCount;
        } else {
            this._animationBIndex = index;
            this._durationBUniform.value = this.clips[index].duration;
            this._clipOffsetBUniform.value = clipOffset;
            this._frameCountBUniform.value = this._bakedClips[index].frameCount;
        }
    }
    /**
     * Sets per-instance time offsets for animation staggering.
     */
    setInstanceOffsets(offsets) {
        if (offsets.length !== this.instanceCount) {
            console.warn('GPUSkinningCompute: Instance offsets length mismatch');
            return;
        }
        this._instanceOffsetsBuffer.array.set(offsets);
        this._instanceOffsetsBuffer.needsUpdate = true;
    }
    /**
     * Creates the compute node for GPU animation.
     * Uses baked DataArrayTexture for trivial frame lookup — no branching, no hierarchy traversal.
     * Dispatches instanceCount × boneCount threads for maximum GPU occupancy.
     */
    createComputeNode(instanceOffsetsBuffer) {
        const boneCount = this.boneCount;
        const instanceCount = this.instanceCount;
        const totalThreads = instanceCount * boneCount;
        // Storage nodes
        const instanceOffsets = storage(instanceOffsetsBuffer || this._instanceOffsetsBuffer, 'float', instanceCount);
        const bakedMatrices = storage(this._bakedBuffer, 'mat4', this._totalBakedMat4s);
        // StorageTexture nodes for bone matrix output (texture cache path)
        const writeBoneTex = storageTexture( this._boneMatricesTex ).setAccess( NodeAccess.WRITE_ONLY );
        this._readBoneTex = storageTexture( this._boneMatricesTex ).setAccess( NodeAccess.READ_ONLY );
        // Helper: store a mat4 as 4 RGBA texels (one per column)
        const storeMat4ToTex = ( mat, baseX, y ) => {
            const v = mat.toVar();
            textureStore( writeBoneTex, ivec2( baseX, y ), v.element( 0 ) );
            textureStore( writeBoneTex, ivec2( baseX.add( 1 ), y ), v.element( 1 ) );
            textureStore( writeBoneTex, ivec2( baseX.add( 2 ), y ), v.element( 2 ) );
            textureStore( writeBoneTex, ivec2( baseX.add( 3 ), y ), v.element( 3 ) );
        };
        // Main compute: one thread per instance × bone
        this._computeNode = Fn(() => {
            const threadId = instanceIndex;
            const instId = int(threadId).div(int(boneCount));
            const boneIdx = int(threadId).mod(int(boneCount));
            const baseX = boneIdx.mul( 4 );
            const instanceOffset = instanceOffsets.element(instId);
            // Compute time for animation A (looping)
            const timeBase = this.timeUniform.add(instanceOffset);
            const timeA = mod(timeBase, this._durationAUniform);
            // Convert time to fractional frame index
            const rawFrameA = timeA.mul(this._sampleRateUniform);
            const frame0A = clamp(int(floor(rawFrameA)), 0, int(this._frameCountAUniform).sub(1));
            const frame1A = clamp(frame0A.add(1), 0, int(this._frameCountAUniform).sub(1));
            const fracA = rawFrameA.sub(float(frame0A));
            // Read two adjacent frames from baked storage: clipOffset + frame * boneCount + boneIdx
            const idxA0 = int(this._clipOffsetAUniform).add(frame0A.mul(int(boneCount))).add(boneIdx);
            const idxA1 = int(this._clipOffsetAUniform).add(frame1A.mul(int(boneCount))).add(boneIdx);
            // Manual lerp: WGSL mix() doesn't support mat4, but mat4 arithmetic (+, -, * scalar) is valid
            const m0A = bakedMatrices.element(idxA0);
            const m1A = bakedMatrices.element(idxA1);
            const matA = m0A.add(m1A.sub(m0A).mul(fracA));
            // Early-out: skip animation B reads entirely when blendWeight is 0
            // Since blendWeight is uniform across all threads, no warp divergence occurs
            If(this.blendWeightUniform.greaterThan(0), () => {
                // Compute time for animation B (phase-aligned to A)
                const phase = timeA.div(this._durationAUniform.max(0.0001));
                const timeB = phase.mul(this._durationBUniform);
                const rawFrameB = timeB.mul(this._sampleRateUniform);
                const frame0B = clamp(int(floor(rawFrameB)), 0, int(this._frameCountBUniform).sub(1));
                const frame1B = clamp(frame0B.add(1), 0, int(this._frameCountBUniform).sub(1));
                const fracB = rawFrameB.sub(float(frame0B));
                const idxB0 = int(this._clipOffsetBUniform).add(frame0B.mul(int(boneCount))).add(boneIdx);
                const idxB1 = int(this._clipOffsetBUniform).add(frame1B.mul(int(boneCount))).add(boneIdx);
                const m0B = bakedMatrices.element(idxB0);
                const m1B = bakedMatrices.element(idxB1);
                const matB = m0B.add(m1B.sub(m0B).mul(fracB));
                // Blend A and B, store to texture
                storeMat4ToTex( matA.add( matB.sub( matA ).mul( this.blendWeightUniform ) ), baseX, instId );
            }).Else(() => {
                // No blending needed — use animation A directly
                storeMat4ToTex( matA, baseX, instId );
            });
        })().compute(totalThreads);
        return this._computeNode;
    }
    /**
     * Returns shared TSL skinning nodes (skinIndex, skinWeight, getBoneMatrix).
     * Used by applyToMaterial, getSkinnedPosition, and getSkinnedNormal.
     */
    _getSkinningContext() {
        const boneCount = this.boneCount;
        const readTex = this._readBoneTex;
        const skinIndexNode = attribute('skinIndex', 'uvec4');
        const skinWeightNode = attribute('skinWeight', 'vec4');
        const getBoneMatrix = ( boneIndex ) => {
            const idx = clamp( int( boneIndex ), 0, int( boneCount - 1 ) );
            const baseX = idx.mul( 4 );
            const y = instanceIndex;
            const col0 = readTex.load( ivec2( baseX, y ) );
            const col1 = readTex.load( ivec2( baseX.add( 1 ), y ) );
            const col2 = readTex.load( ivec2( baseX.add( 2 ), y ) );
            const col3 = readTex.load( ivec2( baseX.add( 3 ), y ) );
            return mat4( col0, col1, col2, col3 );
        };
        return { skinIndexNode, skinWeightNode, getBoneMatrix };
    }
    /**
     * Applies GPU skinning to a NodeMaterial.
     */
    applyToMaterial(material) {
        if (!this._readBoneTex) {
            console.warn('GPUSkinningCompute: Call createComputeNode() first');
            return;
        }
        // Compute bone matrices once, reuse for both position and normal skinning.
        // Single skinMatrix avoids redundant texture reads and mat4 arithmetic.
        const { skinIndexNode, skinWeightNode, getBoneMatrix } = this._getSkinningContext();
        const boneMatX = getBoneMatrix(skinIndexNode.x);
        const boneMatY = getBoneMatrix(skinIndexNode.y);
        const boneMatZ = getBoneMatrix(skinIndexNode.z);
        const boneMatW = getBoneMatrix(skinIndexNode.w);
        // Blend bone matrices into a single skinMatrix, then transform position and normal
        const skinMatrix = add(
            skinWeightNode.x.mul(boneMatX),
            skinWeightNode.y.mul(boneMatY),
            skinWeightNode.z.mul(boneMatZ),
            skinWeightNode.w.mul(boneMatW)
        );
        const skinnedPosition = skinMatrix.mul(vec4(positionLocal, 1)).xyz;
        const skinnedNormal = skinMatrix.transformDirection(normalLocal).xyz;
        const origSetupPosition = material.setupPosition;
        // Apply skinning BEFORE origSetupPosition so that the InstanceNode
        // (inside origSetupPosition) transforms already-skinned vertices.
        // This matches the standard SkinningNode order: skin → instance.
        material.setupPosition = function ( builder ) {
            positionLocal.assign( skinnedPosition );
            normalLocal.assign( skinnedNormal );
            return origSetupPosition.call( this, builder );
        };
    }
    /**
     * Returns the skinned position node for manual material setup.
     */
    getSkinnedPosition() {
        if (!this._readBoneTex) {
            console.warn('GPUSkinningCompute: Call createComputeNode() first');
            return positionLocal;
        }
        const { skinIndexNode, skinWeightNode, getBoneMatrix } = this._getSkinningContext();
        const boneMatX = getBoneMatrix(skinIndexNode.x);
        const boneMatY = getBoneMatrix(skinIndexNode.y);
        const boneMatZ = getBoneMatrix(skinIndexNode.z);
        const boneMatW = getBoneMatrix(skinIndexNode.w);
        const skinVertex = vec4(positionLocal, 1.0);
        const skinned = add(boneMatX.mul(skinWeightNode.x).mul(skinVertex), boneMatY.mul(skinWeightNode.y).mul(skinVertex), boneMatZ.mul(skinWeightNode.z).mul(skinVertex), boneMatW.mul(skinWeightNode.w).mul(skinVertex));
        return skinned.xyz;
    }
    /**
     * Returns the skinned normal node for manual material setup.
     */
    getSkinnedNormal() {
        if (!this._readBoneTex) {
            console.warn('GPUSkinningCompute: Call createComputeNode() first');
            return normalLocal;
        }
        const { skinIndexNode, skinWeightNode, getBoneMatrix } = this._getSkinningContext();
        const boneMatX = getBoneMatrix(skinIndexNode.x);
        const boneMatY = getBoneMatrix(skinIndexNode.y);
        const boneMatZ = getBoneMatrix(skinIndexNode.z);
        const boneMatW = getBoneMatrix(skinIndexNode.w);
        const skinMatrix = add(skinWeightNode.x.mul(boneMatX), skinWeightNode.y.mul(boneMatY), skinWeightNode.z.mul(boneMatZ), skinWeightNode.w.mul(boneMatW));
        return skinMatrix.transformDirection(normalLocal).xyz;
    }
    /**
     * Disposes of GPU resources.
     */
    dispose() {
        if ( this._computeNode ) {
            this._computeNode.dispose();
        }

        this._bakedBuffer = null;
        this._bakedClips = null;
        this._instanceOffsetsBuffer = null;
        if ( this._boneMatricesTex ) this._boneMatricesTex.dispose();
        this._boneMatricesTex = null;
        this._readBoneTex = null;
        this._computeNode = null;
    }
}
export { GPUSkinningCompute, sortBonesHierarchically, parseTrackName, packAnimationData, bakeAnimation };
export default GPUSkinningCompute;