SteamVR_RenderModel.cs

//======= Copyright (c) Valve Corporation, All rights reserved. ===============
//
// Purpose: Render model of associated tracked object
//
//=============================================================================

using UnityEngine;
using System.Collections;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using Valve.VR;

[ExecuteInEditMode]
public class SteamVR_RenderModel : MonoBehaviour
{
    public SteamVR_TrackedObject.EIndex index = SteamVR_TrackedObject.EIndex.None;

    public const string modelOverrideWarning = "Model override is really only meant to be used in " +
        "the scene view for lining things up; using it at runtime is discouraged.  Use tracked device " +
        "index instead to ensure the correct model is displayed for all users.";

    [Tooltip(modelOverrideWarning)]
    public string modelOverride;

    [Tooltip("Shader to apply to model.")]
    public Shader shader;

    [Tooltip("Enable to print out when render models are loaded.")]
    public bool verbose = false;

    [Tooltip("If available, break down into separate components instead of loading as a single mesh.")]
    public bool createComponents = true;

    [Tooltip("Update transforms of components at runtime to reflect user action.")]
    public bool updateDynamically = true;

    // Additional controller settings for showing scrollwheel, etc.
    public RenderModel_ControllerMode_State_t controllerModeState;

    // Name of the sub-object which represents the "local" coordinate space for each component.
    public const string k_localTransformName = "attach";

    // Cached name of this render model for updating component transforms at runtime.
    public string renderModelName { get; private set; }

    // If someone knows how to keep these from getting cleaned up every time
    // you exit play mode, let me know.  I've tried marking the RenderModel
    // class below as [System.Serializable] and switching to normal public
    // variables for mesh and material to get them to serialize properly,
    // as well as tried marking the mesh and material objects as
    // DontUnloadUnusedAsset, but Unity was still unloading them.
    // The hashtable is preserving its entries, but the mesh and material
    // variables are going null.

    public class RenderModel
    {
        public RenderModel(Mesh mesh, Material material)
        {
            this.mesh = mesh;
            this.material = material;
        }
        public Mesh mesh { get; private set; }
        public Material material { get; private set; }
    }

    public static Hashtable models = new Hashtable();
    public static Hashtable materials = new Hashtable();

    // Helper class to load render models interface on demand and clean up when done.
    public sealed class RenderModelInterfaceHolder : System.IDisposable
    {
        private bool needsShutdown, failedLoadInterface;
        private CVRRenderModels _instance;
        public CVRRenderModels instance
        {
            get
            {
                if (_instance == null && !failedLoadInterface)
                {
                    if (!SteamVR.active && !SteamVR.usingNativeSupport)
                    {
                        var error = EVRInitError.None;
                        OpenVR.Init(ref error, EVRApplicationType.VRApplication_Utility);
                        needsShutdown = true;
                    }

                    _instance = OpenVR.RenderModels;
                    if (_instance == null)
                    {
                        Debug.LogError("Failed to load IVRRenderModels interface version " + OpenVR.IVRRenderModels_Version);
                        failedLoadInterface = true;
                    }
                }
                return _instance;
            }
        }
        public void Dispose()
        {
            if (needsShutdown)
                OpenVR.Shutdown();
        }
    }

    private void OnModelSkinSettingsHaveChanged(VREvent_t vrEvent)
    {
        if (!string.IsNullOrEmpty(renderModelName))
        {
            renderModelName = "";
            UpdateModel();
        }
    }

    private void OnHideRenderModels(bool hidden)
    {
        var meshRenderer = GetComponent<MeshRenderer>();
        if (meshRenderer != null)
            meshRenderer.enabled = !hidden;
        foreach (var child in transform.GetComponentsInChildren<MeshRenderer>())
            child.enabled = !hidden;
    }

    private void OnDeviceConnected(int i, bool connected)
    {
        if (i != (int)index)
            return;

        if (connected)
        {
            UpdateModel();
        }
    }

    public void UpdateModel()
    {
        var system = OpenVR.System;
        if (system == null)
            return;

        var error = ETrackedPropertyError.TrackedProp_Success;
        var capacity = system.GetStringTrackedDeviceProperty((uint)index, ETrackedDeviceProperty.Prop_RenderModelName_String, null, 0, ref error);
        if (capacity <= 1)
        {
            Debug.LogError("Failed to get render model name for tracked object " + index);
            return;
        }

        var buffer = new System.Text.StringBuilder((int)capacity);
        system.GetStringTrackedDeviceProperty((uint)index, ETrackedDeviceProperty.Prop_RenderModelName_String, buffer, capacity, ref error);

        var s = buffer.ToString();
        if (renderModelName != s)
        {
            renderModelName = s;
            StartCoroutine(SetModelAsync(s));
        }
    }

    IEnumerator SetModelAsync(string renderModelName)
    {
        if (string.IsNullOrEmpty(renderModelName))
            yield break;

        // Preload all render models before asking for the data to create meshes.
        using (var holder = new RenderModelInterfaceHolder())
        {
            var renderModels = holder.instance;
            if (renderModels == null)
                yield break;

            // Gather names of render models to preload.
            string[] renderModelNames;

            var count = renderModels.GetComponentCount(renderModelName);
            if (count > 0)
            {
                renderModelNames = new string[count];

                for (int i = 0; i < count; i++)
                {
                    var capacity = renderModels.GetComponentName(renderModelName, (uint)i, null, 0);
                    if (capacity == 0)
                        continue;

                    var componentName = new System.Text.StringBuilder((int)capacity);
                    if (renderModels.GetComponentName(renderModelName, (uint)i, componentName, capacity) == 0)
                        continue;

                    capacity = renderModels.GetComponentRenderModelName(renderModelName, componentName.ToString(), null, 0);
                    if (capacity == 0)
                        continue;

                    var name = new System.Text.StringBuilder((int)capacity);
                    if (renderModels.GetComponentRenderModelName(renderModelName, componentName.ToString(), name, capacity) == 0)
                        continue;

                    var s = name.ToString();

                    // Only need to preload if not already cached.
                    var model = models[s] as RenderModel;
                    if (model == null || model.mesh == null)
                    {
                        renderModelNames[i] = s;
                    }
                }
            }
            else
            {
                // Only need to preload if not already cached.
                var model = models[renderModelName] as RenderModel;
                if (model == null || model.mesh == null)
                {
                    renderModelNames = new string[] { renderModelName };
                }
                else
                {
                    renderModelNames = new string[0];
                }
            }

            // Keep trying every 100ms until all components finish loading.
            while (true)
            {
                var loading = false;
                foreach (var name in renderModelNames)
                {
                    if (string.IsNullOrEmpty(name))
                        continue;

                    var pRenderModel = System.IntPtr.Zero;

                    var error = renderModels.LoadRenderModel_Async(name, ref pRenderModel);
                    if (error == EVRRenderModelError.Loading)
                    {
                        loading = true;
                    }
                    else if (error == EVRRenderModelError.None)
                    {
                        // Preload textures as well.
                        var renderModel = MarshalRenderModel(pRenderModel);

                        // Check the cache first.
                        var material = materials[renderModel.diffuseTextureId] as Material;
                        if (material == null || material.mainTexture == null)
                        {
                            var pDiffuseTexture = System.IntPtr.Zero;

                            error = renderModels.LoadTexture_Async(renderModel.diffuseTextureId, ref pDiffuseTexture);
                            if (error == EVRRenderModelError.Loading)
                            {
                                loading = true;
                            }
                        }
                    }
                }

                if (loading)
                {
                    yield return new WaitForSecondsRealtime(0.1f);
                }
                else
                {
                    break;
                }
            }
        }

        bool success = SetModel(renderModelName);
        SteamVR_Events.RenderModelLoaded.Send(this, success);
    }

    private bool SetModel(string renderModelName)
    {
        StripMesh(gameObject);

        using (var holder = new RenderModelInterfaceHolder())
        {
            if (createComponents)
            {
                if (LoadComponents(holder, renderModelName))
                {
                    UpdateComponents(holder.instance);
                    return true;
                }

                Debug.Log("[" + gameObject.name + "] Render model does not support components, falling back to single mesh.");
            }

            if (!string.IsNullOrEmpty(renderModelName))
            {
                var model = models[renderModelName] as RenderModel;
                if (model == null || model.mesh == null)
                {
                    var renderModels = holder.instance;
                    if (renderModels == null)
                        return false;

                    if (verbose)
                        Debug.Log("Loading render model " + renderModelName);

                    model = LoadRenderModel(renderModels, renderModelName, renderModelName);
                    if (model == null)
                        return false;

                    models[renderModelName] = model;
                }

                gameObject.AddComponent<MeshFilter>().mesh = model.mesh;
                gameObject.AddComponent<MeshRenderer>().sharedMaterial = model.material;
                return true;
            }
        }

        return false;
    }

    RenderModel LoadRenderModel(CVRRenderModels renderModels, string renderModelName, string baseName)
    {
        var pRenderModel = System.IntPtr.Zero;

        EVRRenderModelError error;
        while ( true )
        {
            error = renderModels.LoadRenderModel_Async(renderModelName, ref pRenderModel);
            if (error != EVRRenderModelError.Loading)
                break;

            Sleep();
        }

        if (error != EVRRenderModelError.None)
        {
            Debug.LogError(string.Format("Failed to load render model {0} - {1}", renderModelName, error.ToString()));
            return null;
        }

        var renderModel = MarshalRenderModel(pRenderModel);

        var vertices = new Vector3[renderModel.unVertexCount];
        var normals = new Vector3[renderModel.unVertexCount];
        var uv = new Vector2[renderModel.unVertexCount];

        var type = typeof(RenderModel_Vertex_t);
        for (int iVert = 0; iVert < renderModel.unVertexCount; iVert++)
        {
            var ptr = new System.IntPtr(renderModel.rVertexData.ToInt64() + iVert * Marshal.SizeOf(type));
            var vert = (RenderModel_Vertex_t)Marshal.PtrToStructure(ptr, type);

            vertices[iVert] = new Vector3(vert.vPosition.v0, vert.vPosition.v1, -vert.vPosition.v2);
            normals[iVert] = new Vector3(vert.vNormal.v0, vert.vNormal.v1, -vert.vNormal.v2);
            uv[iVert] = new Vector2(vert.rfTextureCoord0, vert.rfTextureCoord1);
        }

        int indexCount = (int)renderModel.unTriangleCount * 3;
        var indices = new short[indexCount];
        Marshal.Copy(renderModel.rIndexData, indices, 0, indices.Length);

        var triangles = new int[indexCount];
        for (int iTri = 0; iTri < renderModel.unTriangleCount; iTri++)
        {
            triangles[iTri * 3 + 0] = (int)indices[iTri * 3 + 2];
            triangles[iTri * 3 + 1] = (int)indices[iTri * 3 + 1];
            triangles[iTri * 3 + 2] = (int)indices[iTri * 3 + 0];
        }

        var mesh = new Mesh();
        mesh.vertices = vertices;
        mesh.normals = normals;
        mesh.uv = uv;
        mesh.triangles = triangles;

#if (UNITY_5_4 || UNITY_5_3 || UNITY_5_2 || UNITY_5_1 || UNITY_5_0)
        mesh.Optimize();
#endif
        //mesh.hideFlags = HideFlags.DontUnloadUnusedAsset;

        // Check cache before loading texture.
        var material = materials[renderModel.diffuseTextureId] as Material;
        if (material == null || material.mainTexture == null)
        {
            var pDiffuseTexture = System.IntPtr.Zero;

            while (true)
            {
                error = renderModels.LoadTexture_Async(renderModel.diffuseTextureId, ref pDiffuseTexture);
                if (error != EVRRenderModelError.Loading)
                    break;

                Sleep();
            }

            if (error == EVRRenderModelError.None)
            {
                var diffuseTexture = MarshalRenderModel_TextureMap(pDiffuseTexture);
                var texture = new Texture2D(diffuseTexture.unWidth, diffuseTexture.unHeight, TextureFormat.RGBA32, false);
                if (SystemInfo.graphicsDeviceType == UnityEngine.Rendering.GraphicsDeviceType.Direct3D11)
                {
                    texture.Apply();

                    while (true)
                    {
                        error = renderModels.LoadIntoTextureD3D11_Async(renderModel.diffuseTextureId, texture.GetNativeTexturePtr());
                        if (error != EVRRenderModelError.Loading)
                            break;

                        Sleep();
                    }
                }
                else
                {
                    var textureMapData = new byte[diffuseTexture.unWidth * diffuseTexture.unHeight * 4]; // RGBA
                    Marshal.Copy(diffuseTexture.rubTextureMapData, textureMapData, 0, textureMapData.Length);

                    var colors = new Color32[diffuseTexture.unWidth * diffuseTexture.unHeight];
                    int iColor = 0;
                    for (int iHeight = 0; iHeight < diffuseTexture.unHeight; iHeight++)
                    {
                        for (int iWidth = 0; iWidth < diffuseTexture.unWidth; iWidth++)
                        {
                            var r = textureMapData[iColor++];
                            var g = textureMapData[iColor++];
                            var b = textureMapData[iColor++];
                            var a = textureMapData[iColor++];
                            colors[iHeight * diffuseTexture.unWidth + iWidth] = new Color32(r, g, b, a);
                        }
                    }

                    texture.SetPixels32(colors);
                    texture.Apply();
                }

                material = new Material(shader != null ? shader : Shader.Find("Standard"));
                material.mainTexture = texture;
                //material.hideFlags = HideFlags.DontUnloadUnusedAsset;

                materials[renderModel.diffuseTextureId] = material;

                renderModels.FreeTexture(pDiffuseTexture);
            }
            else
            {
                Debug.Log("Failed to load render model texture for render model " + renderModelName);
            }
        }

        // Delay freeing when we can since we'll often get multiple requests for the same model right
        // after another (e.g. two controllers or two basestations).
#if UNITY_EDITOR
        if (!Application.isPlaying)
            renderModels.FreeRenderModel(pRenderModel);
        else
#endif
            StartCoroutine(FreeRenderModel(pRenderModel));

        return new RenderModel(mesh, material);
    }

    IEnumerator FreeRenderModel(System.IntPtr pRenderModel)
    {
        yield return new WaitForSeconds(1.0f);

        using (var holder = new RenderModelInterfaceHolder())
        {
            var renderModels = holder.instance;
            renderModels.FreeRenderModel(pRenderModel);
        }
    }

    public Transform FindComponent(string componentName)
    {
        var t = transform;
        for (int i = 0; i < t.childCount; i++)
        {
            var child = t.GetChild(i);
            if (child.name == componentName)
                return child;
        }
        return null;
    }

    private void StripMesh(GameObject go)
    {
        var meshRenderer = go.GetComponent<MeshRenderer>();
        if (meshRenderer != null)
            DestroyImmediate(meshRenderer);

        var meshFilter = go.GetComponent<MeshFilter>();
        if (meshFilter != null)
            DestroyImmediate(meshFilter);
    }

    private bool LoadComponents(RenderModelInterfaceHolder holder, string renderModelName)
    {
        // Disable existing components (we will re-enable them if referenced by this new model).
        // Also strip mesh filter and renderer since these will get re-added if the new component needs them.
        var t = transform;
        for (int i = 0; i < t.childCount; i++)
        {
            var child = t.GetChild(i);
            child.gameObject.SetActive(false);
            StripMesh(child.gameObject);
        }

        // If no model specified, we're done; return success.
        if (string.IsNullOrEmpty(renderModelName))
            return true;

        var renderModels = holder.instance;
        if (renderModels == null)
            return false;

        var count = renderModels.GetComponentCount(renderModelName);
        if (count == 0)
            return false;

        for (int i = 0; i < count; i++)
        {
            var capacity = renderModels.GetComponentName(renderModelName, (uint)i, null, 0);
            if (capacity == 0)
                continue;

            var componentName = new System.Text.StringBuilder((int)capacity);
            if (renderModels.GetComponentName(renderModelName, (uint)i, componentName, capacity) == 0)
                continue;

            // Create (or reuse) a child object for this component (some components are dynamic and don't have meshes).
            t = FindComponent(componentName.ToString());
            if (t != null)
            {
                t.gameObject.SetActive(true);
            }
            else
            {
                t = new GameObject(componentName.ToString()).transform;
                t.parent = transform;
                t.gameObject.layer = gameObject.layer;

                // Also create a child 'attach' object for attaching things.
                var attach = new GameObject(k_localTransformName).transform;
                attach.parent = t;
                attach.localPosition = Vector3.zero;
                attach.localRotation = Quaternion.identity;
                attach.localScale = Vector3.one;
                attach.gameObject.layer = gameObject.layer;
            }

            // Reset transform.
            t.localPosition = Vector3.zero;
            t.localRotation = Quaternion.identity;
            t.localScale = Vector3.one;

            capacity = renderModels.GetComponentRenderModelName(renderModelName, componentName.ToString(), null, 0);
            if (capacity == 0)
                continue;

            var componentRenderModelName = new System.Text.StringBuilder((int)capacity);
            if (renderModels.GetComponentRenderModelName(renderModelName, componentName.ToString(), componentRenderModelName, capacity) == 0)
                continue;

            // Check the cache or load into memory.
            var model = models[componentRenderModelName] as RenderModel;
            if (model == null || model.mesh == null)
            {
                if (verbose)
                    Debug.Log("Loading render model " + componentRenderModelName);

                model = LoadRenderModel(renderModels, componentRenderModelName.ToString(), renderModelName);
                if (model == null)
                    continue;

                models[componentRenderModelName] = model;
            }

            t.gameObject.AddComponent<MeshFilter>().mesh = model.mesh;
            t.gameObject.AddComponent<MeshRenderer>().sharedMaterial = model.material;
        }

        return true;
    }

    SteamVR_Events.Action deviceConnectedAction, hideRenderModelsAction, modelSkinSettingsHaveChangedAction;

    SteamVR_RenderModel()
    {
        deviceConnectedAction = SteamVR_Events.DeviceConnectedAction(OnDeviceConnected);
        hideRenderModelsAction = SteamVR_Events.HideRenderModelsAction(OnHideRenderModels);
        modelSkinSettingsHaveChangedAction = SteamVR_Events.SystemAction(EVREventType.VREvent_ModelSkinSettingsHaveChanged, OnModelSkinSettingsHaveChanged);
    }

    void OnEnable()
    {
#if UNITY_EDITOR
        if (!Application.isPlaying)
            return;
#endif
        if (!string.IsNullOrEmpty(modelOverride))
        {
            Debug.Log(modelOverrideWarning);
            enabled = false;
            return;
        }

        var system = OpenVR.System;
        if (system != null && system.IsTrackedDeviceConnected((uint)index))
        {
            UpdateModel();
        }

        deviceConnectedAction.enabled = true;
        hideRenderModelsAction.enabled = true;
        modelSkinSettingsHaveChangedAction.enabled = true;
    }

    void OnDisable()
    {
#if UNITY_EDITOR
        if (!Application.isPlaying)
            return;
#endif
        deviceConnectedAction.enabled = false;
        hideRenderModelsAction.enabled = false;
        modelSkinSettingsHaveChangedAction.enabled = false;
    }

#if UNITY_EDITOR
    Hashtable values;
#endif
    void Update()
    {
#if UNITY_EDITOR
        if (!Application.isPlaying)
        {
            // See if anything has changed since this gets called whenever anything gets touched.
            var fields = GetType().GetFields(System.Reflection.BindingFlags.Instance | System.Reflection.BindingFlags.Public);

            bool modified = false;

            if (values == null)
            {
                modified = true;
            }
            else
            {
                foreach (var f in fields)
                {
                    if (!values.Contains(f))
                    {
                        modified = true;
                        break;
                    }

                    var v0 = values[f];
                    var v1 = f.GetValue(this);
                    if (v1 != null)
                    {
                        if (!v1.Equals(v0))
                        {
                            modified = true;
                            break;
                        }
                    }
                    else if (v0 != null)
                    {
                        modified = true;
                        break;
                    }
                }
            }

            if (modified)
            {
                if (renderModelName != modelOverride)
                {
                    renderModelName = modelOverride;
                    SetModel(modelOverride);
                }

                values = new Hashtable();
                foreach (var f in fields)
                    values[f] = f.GetValue(this);
            }

            return; // Do not update transforms (below) when not playing in Editor (to avoid keeping OpenVR running all the time).
        }
#endif
        // Update component transforms dynamically.
        if (updateDynamically)
            UpdateComponents(OpenVR.RenderModels);
    }

    Dictionary<int, string> nameCache;

    public void UpdateComponents(CVRRenderModels renderModels)
    {
        if (renderModels == null)
            return;

        var t = transform;
        if (t.childCount == 0)
            return;

        var controllerState = (index != SteamVR_TrackedObject.EIndex.None) ?
            SteamVR_Controller.Input((int)index).GetState() : new VRControllerState_t();

        if (nameCache == null)
            nameCache = new Dictionary<int, string>();

        for (int i = 0; i < t.childCount; i++)
        {
            var child = t.GetChild(i);

            // Cache names since accessing an object's name allocate memory.
            string name;
            if (!nameCache.TryGetValue(child.GetInstanceID(), out name))
            {
                name = child.name;
                nameCache.Add(child.GetInstanceID(), name);
            }

            var componentState = new RenderModel_ComponentState_t();
            if (!renderModels.GetComponentState(renderModelName, name, ref controllerState, ref controllerModeState, ref componentState))
                continue;

            var componentTransform = new SteamVR_Utils.RigidTransform(componentState.mTrackingToComponentRenderModel);
            child.localPosition = componentTransform.pos;
            child.localRotation = componentTransform.rot;

            var attach = child.Find(k_localTransformName);
            if (attach != null)
            {
                var attachTransform = new SteamVR_Utils.RigidTransform(componentState.mTrackingToComponentLocal);
                attach.position = t.TransformPoint(attachTransform.pos);
                attach.rotation = t.rotation * attachTransform.rot;
            }

            bool visible = (componentState.uProperties & (uint)EVRComponentProperty.IsVisible) != 0;
            if (visible != child.gameObject.activeSelf)
            {
                child.gameObject.SetActive(visible);
            }
        }
    }

    public void SetDeviceIndex(int index)
    {
        this.index = (SteamVR_TrackedObject.EIndex)index;
        modelOverride = "";

        if (enabled)
        {
            UpdateModel();
        }
    }

    private static void Sleep()
    {
#if !UNITY_METRO
        System.Threading.Thread.Sleep(1);
#endif
    }

    private RenderModel_t MarshalRenderModel(System.IntPtr pRenderModel)
    {
        if ((System.Environment.OSVersion.Platform == System.PlatformID.MacOSX) ||
            (System.Environment.OSVersion.Platform == System.PlatformID.Unix))
        {
            var packedModel = (RenderModel_t_Packed)Marshal.PtrToStructure(pRenderModel, typeof(RenderModel_t_Packed));
            RenderModel_t model = new RenderModel_t();
            packedModel.Unpack(ref model);
            return model;
        }
        else
        {
            return (RenderModel_t)Marshal.PtrToStructure(pRenderModel, typeof(RenderModel_t));
        }
    }

    private RenderModel_TextureMap_t MarshalRenderModel_TextureMap(System.IntPtr pRenderModel)
    {
        if ((System.Environment.OSVersion.Platform == System.PlatformID.MacOSX) ||
            (System.Environment.OSVersion.Platform == System.PlatformID.Unix))
        {
            var packedModel = (RenderModel_TextureMap_t_Packed)Marshal.PtrToStructure(pRenderModel, typeof(RenderModel_TextureMap_t_Packed));
            RenderModel_TextureMap_t model = new RenderModel_TextureMap_t();
            packedModel.Unpack(ref model);
            return model;
        }
        else
        {
            return (RenderModel_TextureMap_t)Marshal.PtrToStructure(pRenderModel, typeof(RenderModel_TextureMap_t));
        }
    }
}