SkinnedMeshCollider.cs

using System;
using System.Threading;
using Unity.Collections;
using Unity.Jobs;
using UnityEngine;
using UnityEngine.Rendering;

/// <summary>
/// SkinnedMeshRendererの変形済み頂点をGPUから読み戻してMeshColliderを更新する。
/// Unity 6 (6000.x) 以降専用。Awaitable / EntityId / Physics.BakeMeshのジョブ実行を使う
/// </summary>
[RequireComponent(typeof(SkinnedMeshRenderer))]
public class SkinnedMeshCollider : MonoBehaviour
	{
		//座標のみ (float3) なので12バイト
		private const int PositionStride = 12;

		[SerializeField] private SkinnedMeshRenderer _skinnedMeshRenderer;
		[SerializeField] private MeshCollider _meshCollider;

		private static readonly int VertsId = Shader.PropertyToID("Verts");
		private static readonly int PositionsId = Shader.PropertyToID("Positions");
		private static readonly int LocalToWorldId = Shader.PropertyToID("LocalToWorld");
		private static readonly int SrcStrideId = Shader.PropertyToID("SrcStride");
		private static readonly int SrcPositionOffsetId = Shader.PropertyToID("SrcPositionOffset");
		private static readonly int VertexCountId = Shader.PropertyToID("VertexCount");

		private ComputeBuffer _positionBuffer;
		private int _kernel;
		private int _dispatchCount;
		private int _vertexCount;
		private bool _layoutValidated;
		private bool _loggedMissingBuffer;

		private ComputeShader _bakePointComputeShader;

		private Mesh _mesh;

		//読み戻し先。毎フレーム確保せず使い回す
		private NativeArray<Vector3> _readbackBuffer;
		private bool _readbackInFlight;

		//cooking用
		private EntityId _meshEntityId;
		private bool _convex;
		private MeshColliderCookingOptions _cookingOptions;
		private JobHandle _bakeHandle;
		private bool _bakeScheduled;

		/// <summary>
		/// PhysXのtriangle mesh cookingをワーカースレッドで走らせる。
		/// Physics.BakeMeshは内部でP/Invokeするのでburstコンパイルはできないが、通常のジョブからは呼べる
		/// </summary>
		private struct BakeMeshJob : IJob
		{
			public EntityId MeshId;
			public bool Convex;
			public MeshColliderCookingOptions CookingOptions;

			public void Execute()
			{
				Physics.BakeMesh(MeshId, Convex, CookingOptions);
			}
		}

		private void Start()
		{
			if (!Initialize())
			{
				enabled = false;
				return;
			}

			//破棄時はdestroyCancellationTokenで畳まれるので、戻り値は待たない
			_ = RunAsync();
		}

		private bool Initialize()
		{
			if (_skinnedMeshRenderer == null)
			{
				_skinnedMeshRenderer = GetComponent<SkinnedMeshRenderer>();
			}

			if (_meshCollider == null)
			{
				_meshCollider = GetComponent<MeshCollider>();
			}

			if (_skinnedMeshRenderer == null || _meshCollider == null)
			{
				return false;
			}

			var sharedMesh = _skinnedMeshRenderer.sharedMesh;
			if (sharedMesh == null || sharedMesh.vertexCount == 0)
			{
				return false;
			}

			//rootBoneはLocalToWorldの算出に使うので、無ければ動かせない
			if (_skinnedMeshRenderer.rootBone == null)
			{
				Debug.LogWarning($"{name}: rootBoneが未設定のため無効化する", this);
				return false;
			}

			//頂点バッファをRawとして読む要求は、GetVertexBufferより先に出しておく。下記マニュアル参照
			//https://docs.unity3d.com/6000.6/Documentation/ScriptReference/SkinnedMeshRenderer-vertexBufferTarget.html
			_skinnedMeshRenderer.vertexBufferTarget |= GraphicsBuffer.Target.Raw;

			//インデックスバッファと初期姿勢を得るために一度だけBakeする
			_mesh = new Mesh();
			_mesh.name = "skinmesh";
			_skinnedMeshRenderer.BakeMesh(_mesh);

			//meshes.verticesはアクセスごとに配列を新規確保して返すので、ループ内で引かない
			var bakedPositions = _mesh.vertices;
			_vertexCount = bakedPositions.Length;

			//衝突判定に法線とUVは不要。座標のみのレイアウトにして読み戻し量を32B/頂点から12B/頂点へ落とす
			_mesh.SetVertexBufferParams(_vertexCount,
				new VertexAttributeDescriptor(VertexAttribute.Position, VertexAttributeFormat.Float32, 3));
			_mesh.SetVertexBufferData(bakedPositions, 0, 0, _vertexCount);
			_mesh.MarkDynamic();

			_positionBuffer = new ComputeBuffer(_vertexCount, PositionStride);
			_positionBuffer.SetData(bakedPositions, 0, 0, _vertexCount);

			_readbackBuffer = new NativeArray<Vector3>(_vertexCount, Allocator.Persistent,
				NativeArrayOptions.UninitializedMemory);

			//cookingの引数はcolliderの設定に合わせる。
			//既定値でbakeすると、colliderが非既定のcookingOptionsを持つ場合にcacheが一致せず再cookになる
			_meshEntityId = _mesh.GetEntityId();
			_convex = _meshCollider.convex;
			_cookingOptions = _meshCollider.cookingOptions;

			//computeshaderをロード。
			var computeShaderAsset = Resources.Load<ComputeShader>("SkinMeshtoMesh");
			if (computeShaderAsset == null)
			{
				Debug.LogError($"{name}: Resources/SkinMeshtoMesh.computeが見つからないため無効化する", this);
				return false;
			}

			//SetBufferの状態はComputeShaderごとに持つので、インスタンスごとに複製する
			_bakePointComputeShader = Instantiate(computeShaderAsset);
			_kernel = _bakePointComputeShader.FindKernel("CSMain");
			_bakePointComputeShader.GetKernelThreadGroupSizes(_kernel, out var threadX, out _, out _);
			_dispatchCount = (int)(((uint)_vertexCount + threadX - 1) / threadX);

			_bakePointComputeShader.SetBuffer(_kernel, PositionsId, _positionBuffer);

			//Vertsのレイアウトは初回にバッファが取れた時点で実測する (ValidateAndBindLayout)
			return true;
		}

		/// <summary>
		/// 1フレーム1周の更新ループ。
		/// AsyncGPUReadbackRequest.doneを毎フレーム見る状態機械をawaitで置き換えている
		/// </summary>
		private async Awaitable RunAsync()
		{
			var token = destroyCancellationToken;

			try
			{
				while (true)
				{
					await Awaitable.NextFrameAsync(token);

					if (_skinnedMeshRenderer == null || _bakePointComputeShader == null)
					{
						return;
					}

					var dispatched = TryDispatch();

					//前の周回で投げたcookingをここで回収する。
					//Schedule直後にawaitして待つとその分1フレーム余計にかかるので、
					//_meshを触らない区間 (NextFrame + Dispatch) を挟んでから回収する
					if (_bakeScheduled)
					{
						_bakeHandle.Complete();
						_bakeScheduled = false;

						//cook済みのデータがMeshにcacheされているので、代入時に再cookは走らない
						_meshCollider.sharedMesh = _mesh;
					}

					if (!dispatched)
					{
						//レイアウトが未対応でValidateAndBindLayoutが無効化した場合は、ループごと畳む
						if (!enabled)
						{
							return;
						}

						continue;
					}

					//読み戻し先を使い回すので、GetData<T>()が返す一時バッファとそのコピーが要らない
					_readbackInFlight = true;
					var request = await AsyncGPUReadback.RequestIntoNativeArrayAsync(ref _readbackBuffer,
						_positionBuffer);
					_readbackInFlight = false;

					//このawaitはtokenを取らないので、読み戻し中に破棄されてもここまで進む。
					//Releaseが済んだ後の継続で破棄済みのMeshを触らないよう、ここで畳む
					token.ThrowIfCancellationRequested();

					if (request.hasError)
					{
						continue;
					}

					//SetVertexBufferDataだけでMeshがdirtyになり、collision meshの再生成が必要になる。
					//法線はcolliderに不要なのでRecalculateNormalsは呼ばない
					_mesh.SetVertexBufferData(_readbackBuffer, 0, 0, _vertexCount);

					//PhysXに増分cookingは無く、頂点を触ればmidphaseの再構築が要る。
					//消せないのでメインスレッドから追い出す
					_bakeHandle = new BakeMeshJob
					{
						MeshId = _meshEntityId,
						Convex = _convex,
						CookingOptions = _cookingOptions
					}.Schedule();
					_bakeScheduled = true;

					//ここでCompleteするとメインスレッドで待つことになり、同期cookingと変わらない。
					//回収は次の周回の頭で行う
				}
			}
			catch (OperationCanceledException)
			{
				//破棄された。Releaseで後始末する
			}
		}

		/// <summary>
		/// 変形済み頂点バッファからcompute shaderで座標だけを抜き出す。
		/// バッファがまだ無ければfalseを返す
		/// </summary>
		private bool TryDispatch()
		{
			var buffer = _skinnedMeshRenderer.GetVertexBuffer();
			if (buffer == null)
			{
				//初回skinning前・画面外(updateWhenOffscreen == false)・cached mesh未準備のいずれか。
				//常に更新したい場合はupdateWhenOffscreenをtrueにする
				if (!_loggedMissingBuffer)
				{
					_loggedMissingBuffer = true;
					Debug.Log($"{name}: 変形済み頂点バッファが未生成。skinningされるまで待つ", this);
				}

				return false;
			}

			try
			{
				if (!_layoutValidated && !ValidateAndBindLayout(buffer))
				{
					enabled = false;
					return false;
				}

				_bakePointComputeShader.SetBuffer(_kernel, VertsId, buffer);
				_bakePointComputeShader.SetMatrix(LocalToWorldId,
					_skinnedMeshRenderer.worldToLocalMatrix * _skinnedMeshRenderer.rootBone.localToWorldMatrix);
				_bakePointComputeShader.Dispatch(_kernel, _dispatchCount, 1, 1);
			}
			finally
			{
				//GetVertexBufferは呼び出しごとにwrapperを新規生成する。
				//wrapperはm_ReferencesBuffer=trueで作られるので、破棄してもrenderer側の実バッファは消えない
				buffer.Dispose();
			}

			return true;
		}

		/// <summary>
		/// 変形済み頂点バッファのstrideを実測して、compute shaderへ渡す。
		/// stride 40の決め打ちはしない (normal / tangentの有無で変わるため)
		/// </summary>
		private bool ValidateAndBindLayout(GraphicsBuffer buffer)
		{
			var reportedStride = buffer.stride;

			//Rawビューではstrideが4で報告されることがあるので、総バイト数からも算出して妥当な方を採る
			var totalBytes = (long)buffer.count * reportedStride;
			var derivedStride = _vertexCount > 0 ? (int)(totalBytes / _vertexCount) : 0;
			var srcStride = (reportedStride >= PositionStride && reportedStride % 4 == 0)
				? reportedStride
				: derivedStride;

			Debug.Log($"{name}: deformed vertex buffer stride={reportedStride} count={buffer.count} " +
			          $"derived={derivedStride} -> 使用stride={srcStride}", this);

			//ByteAddressBuffer.Load3は4バイト境界を要求する
			if (srcStride < PositionStride || srcStride % 4 != 0)
			{
				Debug.LogError(
					$"{name}: 変形済み頂点バッファのstrideを決定できない (stride={reportedStride}, count={buffer.count})。無効化する",
					this);
				return false;
			}

			var sharedMesh = _skinnedMeshRenderer.sharedMesh;
			if (sharedMesh.GetVertexAttributeFormat(VertexAttribute.Position) != VertexAttributeFormat.Float32 ||
			    sharedMesh.GetVertexAttributeDimension(VertexAttribute.Position) != 3)
			{
				Debug.LogError($"{name}: PositionがFloat32x3ではないため未対応。無効化する", this);
				return false;
			}

			_bakePointComputeShader.SetInt(SrcStrideId, srcStride);
			//skinning出力のhot streamはposition / normal / tangentの順なので、positionは先頭
			_bakePointComputeShader.SetInt(SrcPositionOffsetId, 0);
			_bakePointComputeShader.SetInt(VertexCountId, _vertexCount);

			_layoutValidated = true;
			return true;
		}

		private void Release()
		{
			//ジョブがMeshを掴んでいる間に破棄しない
			if (_bakeScheduled)
			{
				_bakeHandle.Complete();
				_bakeScheduled = false;
			}

			//destroyCancellationTokenはawaitを畳むだけで、GPU側の読み戻し要求までは取り消さない。
			//書き込み先の_readbackBufferを解放する前に、要求の完了を待つ必要がある
			if (_readbackInFlight)
			{
				AsyncGPUReadback.WaitAllRequests();
				_readbackInFlight = false;
			}

			if (_readbackBuffer.IsCreated)
			{
				_readbackBuffer.Dispose();
			}

			if (_positionBuffer != null)
			{
				_positionBuffer.Release();
				_positionBuffer = null;
			}

			if (_bakePointComputeShader != null)
			{
				Destroy(_bakePointComputeShader);
				_bakePointComputeShader = null;
			}

			if (_mesh != null)
			{
				Destroy(_mesh);
				_mesh = null;
			}

			_skinnedMeshRenderer = null;
			_meshCollider = null;
		}

		private void OnDestroy()
		{
			Release();
		}
	}