BoneRendererUtil.cs

[InitializeOnLoad]
public static class BoneRendererUtil {
    class BatchRenderer {
        const int KMaxDrawMeshInstanceCount = 1023;

        public enum SubMeshType {
            BoneFaces,
            BoneWire,
            Count
        }

        public Mesh Mesh;
        public Material Material;

        List<Matrix4x4> _mMatrices = new List<Matrix4x4>();
        List<Vector4> _mColors = new List<Vector4>();
        List<Vector4> _mHighlights = new List<Vector4>();

        public void AddInstance(Matrix4x4 matrix, Color color, Color highlight) {
            _mMatrices.Add(matrix);
            _mColors.Add(color);
            _mHighlights.Add(highlight);
        }

        public void Clear() {
            _mMatrices.Clear();
            _mColors.Clear();
            _mHighlights.Clear();
        }

        static int RenderChunkCount(int totalCount) {
            return Mathf.CeilToInt((totalCount / (float) KMaxDrawMeshInstanceCount));
        }

        static T[] GetRenderChunk<T>(List<T> array, int chunkIndex) {
            var rangeCount = (chunkIndex < (RenderChunkCount(array.Count) - 1)) ? KMaxDrawMeshInstanceCount : array.Count - (chunkIndex * KMaxDrawMeshInstanceCount);

            return array.GetRange(chunkIndex * KMaxDrawMeshInstanceCount, rangeCount).ToArray();
        }

        public void Render() {
            if (_mMatrices.Count == 0 || _mColors.Count == 0 || _mHighlights.Count == 0)
                return;

            var count = System.Math.Min(_mMatrices.Count, System.Math.Min(_mColors.Count, _mHighlights.Count));

            var mat = Material;
            mat.SetPass(0);

            var propertyBlock = new MaterialPropertyBlock();
            var cb = new CommandBuffer();

            var chunkCount = RenderChunkCount(count);
            for (var i = 0; i < chunkCount; ++i) {
                cb.Clear();
                var matrices = GetRenderChunk(_mMatrices, i);
                propertyBlock.SetVectorArray("_Color", GetRenderChunk(_mColors, i));

                Material.DisableKeyword("WIRE_ON");
                cb.DrawMeshInstanced(Mesh, (int) SubMeshType.BoneFaces, Material, 0, matrices, matrices.Length, propertyBlock);
                Graphics.ExecuteCommandBuffer(cb);

                cb.Clear();
                propertyBlock.SetVectorArray("_Color", GetRenderChunk(_mHighlights, i));

                Material.EnableKeyword("WIRE_ON");
                cb.DrawMeshInstanced(Mesh, (int) SubMeshType.BoneWire, Material, 0, matrices, matrices.Length, propertyBlock);
                Graphics.ExecuteCommandBuffer(cb);
            }
        }
    }

    static List<BoneRenderer> _sBoneRendererComponents = new List<BoneRenderer>();

    static BatchRenderer _sPyramidMeshRenderer;
    static BatchRenderer _sBoxMeshRenderer;

    static Material _sMaterial;

    const float KEpsilon = 1e-5f;

    const float KBoneBaseSize = 2f;
    const float KBoneTipSize = 0.5f;

    static int _sButtonHash = "BoneHandle".GetHashCode();

    static BoneRendererUtil() {
        BoneRenderer.OnAddBoneRenderer += OnAddBoneRenderer;
        BoneRenderer.OnRemoveBoneRenderer += OnRemoveBoneRenderer;
        SceneVisibilityManager.visibilityChanged += OnVisibilityChanged;

        SceneView.duringSceneGui += DrawSkeletons;
    }

    static Material Material {
        get {
            if (_sMaterial) return _sMaterial;
            var shader = (Shader) EditorGUIUtility.LoadRequired("BoneHandles.shader");
            _sMaterial = new Material(shader) {hideFlags = HideFlags.DontSaveInEditor, enableInstancing = true};

            return _sMaterial;
        }
    }

    static BatchRenderer PyramidMeshRenderer {
        get {
            if (_sPyramidMeshRenderer == null) {
                var mesh = new Mesh {name = "BoneRendererPyramidMesh", subMeshCount = (int) BatchRenderer.SubMeshType.Count, hideFlags = HideFlags.DontSave};

                // Bone vertices
                var vertices = new Vector3[] {
                    new Vector3(0.0f, 1.0f, 0.0f),
                    new Vector3(0.0f, 0.0f, -1.0f),
                    new Vector3(-0.9f, 0.0f, 0.5f),
                    new Vector3(0.9f, 0.0f, 0.5f),
                };

                mesh.vertices = vertices;

                // Build indices for different sub meshes
                var boneFaceIndices = new int[] {
                    0, 2, 1,
                    0, 1, 3,
                    0, 3, 2,
                    1, 2, 3
                };
                mesh.SetIndices(boneFaceIndices, MeshTopology.Triangles, (int) BatchRenderer.SubMeshType.BoneFaces);

                var boneWireIndices = new int[] {
                    0, 1, 0, 2, 0, 3, 1, 2, 2, 3, 3, 1
                };
                mesh.SetIndices(boneWireIndices, MeshTopology.Lines, (int) BatchRenderer.SubMeshType.BoneWire);

                _sPyramidMeshRenderer = new BatchRenderer() {
                    Mesh = mesh,
                    Material = Material
                };
            }

            return _sPyramidMeshRenderer;
        }
    }

    static BatchRenderer BoxMeshRenderer {
        get {
            if (_sBoxMeshRenderer != null) return _sBoxMeshRenderer;
            var mesh = new Mesh {name = "BoneRendererBoxMesh", subMeshCount = (int) BatchRenderer.SubMeshType.Count, hideFlags = HideFlags.DontSave};

            // Bone vertices
            var vertices = new Vector3[] {
                new Vector3(-0.5f, 0.0f, 0.5f),
                new Vector3(0.5f, 0.0f, 0.5f),
                new Vector3(0.5f, 0.0f, -0.5f),
                new Vector3(-0.5f, 0.0f, -0.5f),
                new Vector3(-0.5f, 1.0f, 0.5f),
                new Vector3(0.5f, 1.0f, 0.5f),
                new Vector3(0.5f, 1.0f, -0.5f),
                new Vector3(-0.5f, 1.0f, -0.5f)
            };

            mesh.vertices = vertices;

            // Build indices for different sub meshes
            var boneFaceIndices = new[] {
                0, 2, 1,
                0, 3, 2,

                0, 1, 5,
                0, 5, 4,

                1, 2, 6,
                1, 6, 5,

                2, 3, 7,
                2, 7, 6,

                3, 0, 4,
                3, 4, 7,

                4, 5, 6,
                4, 6, 7
            };
            mesh.SetIndices(boneFaceIndices, MeshTopology.Triangles, (int) BatchRenderer.SubMeshType.BoneFaces);

            var boneWireIndices = new int[] {
                0, 1, 1, 2, 2, 3, 3, 0,
                4, 5, 5, 6, 6, 7, 7, 4,
                0, 4, 1, 5, 2, 6, 3, 7
            };
            mesh.SetIndices(boneWireIndices, MeshTopology.Lines, (int) BatchRenderer.SubMeshType.BoneWire);

            _sBoxMeshRenderer = new BatchRenderer() {
                Mesh = mesh,
                Material = Material
            };

            return _sBoxMeshRenderer;
        }
    }

    static Matrix4x4 ComputeBoneMatrix(Vector3 start, Vector3 end, float length, float size) {
        var direction = (end - start) / length;
        var tangent = Vector3.Cross(direction, Vector3.up);
        if (Vector3.SqrMagnitude(tangent) < 0.1f)
            tangent = Vector3.Cross(direction, Vector3.right);
        tangent.Normalize();
        var bitangent = Vector3.Cross(direction, tangent);

        var scale = length * KBoneBaseSize * size;

        return new Matrix4x4(
            new Vector4(tangent.x * scale, tangent.y * scale, tangent.z * scale, 0f),
            new Vector4(direction.x * length, direction.y * length, direction.z * length, 0f),
            new Vector4(bitangent.x * scale, bitangent.y * scale, bitangent.z * scale, 0f),
            new Vector4(start.x, start.y, start.z, 1f));
    }

    static void DrawSkeletons(SceneView sceneview) {
        var gizmoColor = Gizmos.color;

        PyramidMeshRenderer.Clear();
        BoxMeshRenderer.Clear();

        foreach (var boneRenderer in _sBoneRendererComponents) {
            if (boneRenderer.Bones == null)
                continue;

            var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
            if (prefabStage != null) {
                var stageHandle = prefabStage.stageHandle;
                if (stageHandle.IsValid() && !stageHandle.Contains(boneRenderer.gameObject))
                    continue;
            }

            if (boneRenderer.drawBones) {
                var size = boneRenderer.boneSize * 0.025f;
                var shape = boneRenderer.boneShape;
                var color = boneRenderer.boneColor;
                var nubColor = new Color(color.r, color.g, color.b, color.a);
                var selectionColor = Color.white;

                foreach (var bone in boneRenderer.Bones) {
                    if (bone.First == null || bone.Second == null)
                        continue;

                    DoBoneRender(bone.First, bone.Second, shape, color, size);
                }

                foreach (var tip in boneRenderer.Tips) {
                    if (tip == null)
                        continue;

                    DoBoneRender(tip, null, shape, color, size);
                }
            }

            if (boneRenderer.drawTripods) {
                var size = boneRenderer.tripodSize * 0.025f;
                foreach (var t in boneRenderer.Transforms) {
                    const float tripodSize = 1f;
                    var transform = t;
                    if (transform == null)
                        continue;

                    var position = transform.position;
                    var sz = (size * tripodSize);
                    var rotation = transform.rotation;
                    var xAxis = position + rotation * Vector3.right * sz;
                    var yAxis = position + rotation * Vector3.up * sz;
                    var zAxis = position + rotation * Vector3.forward * sz;

                    Handles.color = Color.red;
                    Handles.DrawLine(position, xAxis);
                    Handles.color = Color.green;
                    Handles.DrawLine(position, yAxis);
                    Handles.color = Color.blue;
                    Handles.DrawLine(position, zAxis);
                }
            }
        }

        PyramidMeshRenderer.Render();
        BoxMeshRenderer.Render();

        Gizmos.color = gizmoColor;
    }


    static void DoBoneRender(Transform transform, Transform childTransform, BoneShape shape, Color color, float size) {
        var start = transform.position;
        var end = childTransform != null ? childTransform.position : start;

        var boneGo = transform.gameObject;

        var length = (end - start).magnitude;
        var tipBone = (length < KEpsilon);

        var id = GUIUtility.GetControlID(_sButtonHash, FocusType.Passive);
        var evt = Event.current;

        switch (evt.GetTypeForControl(id)) {
            case EventType.Layout: {
                HandleUtility.AddControl(id, tipBone ? HandleUtility.DistanceToCircle(start, KBoneTipSize * size * 0.5f) : HandleUtility.DistanceToLine(start, end));
                break;
            }
            case EventType.MouseMove:
                if (id == HandleUtility.nearestControl)
                    HandleUtility.Repaint();
                break;
            case EventType.MouseDown: {
                if (HandleUtility.nearestControl == id && evt.button == 0) {
#if UNITY_2019_3_OR_NEWER
                    if (!SceneVisibilityManager.instance.IsPickingDisabled(boneGo, false))
#endif
                    {
                        GUIUtility.hotControl = id; // Grab mouse focus
                        EditorHelper.HandleClickSelection(boneGo, evt);
                        evt.Use();
                    }
                }

                break;
            }
            case EventType.MouseDrag: {
                if (!evt.alt && GUIUtility.hotControl == id) {
#if UNITY_2019_3_OR_NEWER
                    if (!SceneVisibilityManager.instance.IsPickingDisabled(boneGo, false))
#endif
                    {
                        DragAndDrop.PrepareStartDrag();
                        DragAndDrop.objectReferences = new Object[] {transform};
                        DragAndDrop.StartDrag(ObjectNames.GetDragAndDropTitle(transform));

                        GUIUtility.hotControl = 0;

                        evt.Use();
                    }
                }

                break;
            }
            case EventType.MouseUp: {
                if (GUIUtility.hotControl == id && (evt.button == 0 || evt.button == 2)) {
                    GUIUtility.hotControl = 0;
                    evt.Use();
                }

                break;
            }
            case EventType.Repaint: {
                var highlight = color;

                var hoveringBone = GUIUtility.hotControl == 0 && HandleUtility.nearestControl == id;
#if UNITY_2019_3_OR_NEWER
                hoveringBone = hoveringBone && !SceneVisibilityManager.instance.IsPickingDisabled(transform.gameObject, false);
#endif

                if (hoveringBone) {
                    highlight = Handles.preselectionColor;
                }
                else if (Selection.Contains(boneGo) || Selection.activeObject == boneGo) {
                    highlight = Handles.selectedColor;
                }

                if (tipBone) {
                    Handles.color = highlight;
                    Handles.SphereHandleCap(0, start, Quaternion.identity, KBoneTipSize * size, EventType.Repaint);
                }
                else if (shape == BoneShape.Line) {
                    Handles.color = highlight;
                    Handles.DrawLine(start, end);
                }
                else {
                    if (shape == BoneShape.Pyramid)
                        PyramidMeshRenderer.AddInstance(ComputeBoneMatrix(start, end, length, size), color, highlight);
                    else // if (shape == BoneShape.Box)
                        BoxMeshRenderer.AddInstance(ComputeBoneMatrix(start, end, length, size), color, highlight);
                }
            }
                break;
        }
    }

    public static void OnAddBoneRenderer(BoneRenderer obj) {
        _sBoneRendererComponents.Add(obj);
    }

    public static void OnRemoveBoneRenderer(BoneRenderer obj) {
        _sBoneRendererComponents.Remove(obj);
    }

    public static void OnVisibilityChanged() {
        foreach (var boneRenderer in _sBoneRendererComponents) {
            boneRenderer.Invalidate();
        }
    }
}