Visualize skeleton in Cocos Creator 3D

SkeletonDisplay.ts

import { _decorator, Component, Node, renderer, director, Material, GFXPrimitiveMode, Mesh, GFXBufferUsageBit, GFXMemoryUsageBit, Mat4, Vec3, IGFXAttribute, GFXFormat, GFXAttributeName, Scene, GFXBuffer, GFXComparisonFunc } from 'cc';

type RenderingSubMesh = Parameters<renderer.Model['initSubModel']>[1];

class VisualRelations {
    private _buffer: Float32Array;
    private _relationIndices: number[] = [];
    private _watchedNodes: Node[] = [];
    private _minPos: Vec3 = new Vec3();
    private _maxPos: Vec3 = new Vec3();

    constructor(relations: [Node, Node][]) {
        this._relationIndices = new Array(relations.length * 2).fill(0);
        this._watchedNodes = [];

        const watchNode = (node: Node) => {
            let index = this._watchedNodes.indexOf(node);
            if (index < 0) {
                index = this._watchedNodes.length;
                this._watchedNodes.push(node);
            }
            return index;
        };

        relations.forEach(([parent, child], relationIndex) => {
            this._relationIndices[2 * relationIndex] = watchNode(parent);
            this._relationIndices[2 * relationIndex + 1] = watchNode(child);
        });

        this._buffer = new Float32Array(3 * this._relationIndices.length);
        this.update();
    }

    public update() {
        Vec3.set(this._minPos, Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY);
        Vec3.set(this._maxPos, -Number.POSITIVE_INFINITY, -Number.POSITIVE_INFINITY, -Number.POSITIVE_INFINITY);
        const tempPos = new Vec3();
        for (let iRelationIndex = 0; iRelationIndex < this._relationIndices.length; ++iRelationIndex) {
            const iWatchedNode = this._relationIndices[iRelationIndex];
            const worldPosition = this._watchedNodes[iWatchedNode].getWorldPosition(tempPos);
            Vec3.toArray(this._buffer, worldPosition, 3 * iWatchedNode);
            Vec3.min(this._minPos, this._minPos, worldPosition);
            Vec3.max(this._maxPos, this._maxPos, worldPosition);
        }
    }

    get buffer() {
        return this._buffer;
    }

    get stride() {
        return this._buffer.BYTES_PER_ELEMENT * 3;
    }

    public get attributes(): Array<{ name: string; format: GFXFormat; }> {
        return [{
            name: GFXAttributeName.ATTR_POSITION,
            format: GFXFormat.RGB32F,
        }];
    }

    get minPosition() {
        return this._minPos;
    }

    get maxPosition() {
        return this._maxPos;
    }
}

class SkeletonDisplayModel {
    private _model: renderer.Model;
    private _vertexBuffer: GFXBuffer;
    private _visualRelations: VisualRelations;

    constructor(rootNode: Node, displayScene: Scene) {
        const relations: [Node, Node][] = [];
        getRelations(rootNode, relations);

        const hostNode = new Node(`SkeletonDisplay - ${rootNode.name}`);
        hostNode.setParent(displayScene as unknown as Node);

        const material = new Material();
        material.reset({
            effectName: 'builtin-unlit',
            states: {
                primitive: GFXPrimitiveMode.LINE_LIST,
                depthStencilState: {
                    depthFunc: GFXComparisonFunc.ALWAYS,
                },
            },
        });

        const visualRelations = new VisualRelations(relations);

        const vertexBuffer = director.root.device.createBuffer({
            usage: GFXBufferUsageBit.VERTEX | GFXBufferUsageBit.TRANSFER_DST,
            memUsage: GFXMemoryUsageBit.HOST | GFXMemoryUsageBit.DEVICE,
            size: visualRelations.buffer.byteLength,
            stride: visualRelations.stride,
        });
        vertexBuffer.update(visualRelations.buffer);

        const renderingSubMesh: RenderingSubMesh = {
            vertexBuffers: [ vertexBuffer ],
            attributes: visualRelations.attributes,
            primitiveMode: GFXPrimitiveMode.LINE_LIST,
            geometricInfo: undefined,
            flatBuffers: undefined,
            jointMappedBuffers: [],
            enableVertexIdChannel: undefined,
            destroy: undefined,
        };

        const model = director.root.createModel(renderer.Model);
        model.initialize(hostNode);
        model.initSubModel(0, renderingSubMesh, material);
        model.enabled = true;
        model.createBoundingShape(visualRelations.minPosition, visualRelations.maxPosition);

        const renderScene = displayScene.renderScene;
        renderScene.addModel(model);

        this._model = model;
        this._vertexBuffer = vertexBuffer;
        this._visualRelations = visualRelations;
    }

    public update() {
        this._visualRelations.update();
        this._vertexBuffer.update(this._visualRelations.buffer);
    }

    public destroy() {
        this._vertexBuffer.destroy();
        this._model.destroy();
    }
}

@_decorator.ccclass('SkeletonDisplay')
export class SkeletonDisplay extends Component {
    @_decorator.type(Node)
    public root: Node | null = null;

    private _skeletonDisplayModel: SkeletonDisplayModel | null = null;

    public start() {
        const rootNode = this.root;
        if (!rootNode) {
            return;
        }
        this._skeletonDisplayModel = new SkeletonDisplayModel(
            rootNode,
            this.node.scene,
        );
    }

    public update() {
        this._skeletonDisplayModel?.update();
    }
}

function getRelations(node: Node, relations: [Node, Node][]) {
    for (const child of node.children) {
        relations.push([node, child]);
        getRelations(child, relations);
    }
}