forked from BilalY/Rasagar
273 lines
11 KiB
C#
273 lines
11 KiB
C#
|
using System;
|
||
|
|
using System.Collections.Generic;
|
||
|
|
using UnityEngine.Assertions;
|
||
|
|
using UnityEngine.Rendering;
|
||
|
|
using Unity.Mathematics;
|
||
|
|
|
||
|
|
#if UNITY_EDITOR
|
||
|
|
using UnityEditor;
|
||
|
|
#endif
|
||
|
|
|
||
|
|
namespace UnityEngine.Rendering.UnifiedRayTracing
|
||
|
|
{
|
||
|
|
internal class AccelStructAdapter : IDisposable
|
||
|
|
{
|
||
|
|
private IRayTracingAccelStruct _accelStruct;
|
||
|
|
AccelStructInstances _instances;
|
||
|
|
|
||
|
|
internal AccelStructInstances Instances { get => _instances; }
|
||
|
|
|
||
|
|
struct InstanceIDs
|
||
|
|
{
|
||
|
|
public int InstanceID;
|
||
|
|
public int AccelStructID;
|
||
|
|
}
|
||
|
|
|
||
|
|
private readonly Dictionary<int, InstanceIDs[]> _objectHandleToInstances = new();
|
||
|
|
|
||
|
|
public AccelStructAdapter(IRayTracingAccelStruct accelStruct, GeometryPool geometryPool)
|
||
|
|
{
|
||
|
|
_accelStruct = accelStruct;
|
||
|
|
_instances = new AccelStructInstances(geometryPool);
|
||
|
|
}
|
||
|
|
|
||
|
|
public AccelStructAdapter(IRayTracingAccelStruct accelStruct, RayTracingResources resources)
|
||
|
|
: this(accelStruct, new GeometryPool(GeometryPoolDesc.NewDefault(), resources.geometryPoolKernels, resources.copyBuffer))
|
||
|
|
{ }
|
||
|
|
|
||
|
|
public IRayTracingAccelStruct GetAccelerationStructure()
|
||
|
|
{
|
||
|
|
return _accelStruct;
|
||
|
|
}
|
||
|
|
|
||
|
|
public GeometryPool GeometryPool => _instances.geometryPool;
|
||
|
|
|
||
|
|
public void Bind(CommandBuffer cmd, string propertyName, IRayTracingShader shader)
|
||
|
|
{
|
||
|
|
shader.SetAccelerationStructure(cmd, propertyName, _accelStruct);
|
||
|
|
_instances.Bind(cmd, shader);
|
||
|
|
}
|
||
|
|
|
||
|
|
public void Dispose()
|
||
|
|
{
|
||
|
|
_instances?.Dispose();
|
||
|
|
_instances = null;
|
||
|
|
_accelStruct?.Dispose();
|
||
|
|
_accelStruct = null;
|
||
|
|
_objectHandleToInstances.Clear();
|
||
|
|
}
|
||
|
|
|
||
|
|
public void AddInstance(int objectHandle, Component meshRendererOrTerrain, uint[] perSubMeshMask, uint[] perSubMeshMaterialIDs)
|
||
|
|
{
|
||
|
|
if (meshRendererOrTerrain is Terrain terrain)
|
||
|
|
{
|
||
|
|
Debug.Assert(terrain.enabled, "Terrains are expected to be enabled.");
|
||
|
|
var terrainDesc = new TerrainDesc();
|
||
|
|
terrainDesc.terrain = terrain;
|
||
|
|
terrainDesc.mask = perSubMeshMask[0];
|
||
|
|
terrainDesc.materialID = perSubMeshMaterialIDs[0];
|
||
|
|
AddInstance(objectHandle, terrainDesc);
|
||
|
|
}
|
||
|
|
else
|
||
|
|
{
|
||
|
|
var meshRenderer = (MeshRenderer)meshRendererOrTerrain;
|
||
|
|
Debug.Assert(meshRenderer.enabled, "Mesh renderers are expected to be enabled.");
|
||
|
|
Debug.Assert(!meshRenderer.isPartOfStaticBatch, "Mesh renderers are expected to not be part of static batch.");
|
||
|
|
var mesh = meshRenderer.GetComponent<MeshFilter>().sharedMesh;
|
||
|
|
AddInstance(objectHandle, mesh, meshRenderer.transform.localToWorldMatrix, perSubMeshMask, perSubMeshMaterialIDs);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
public void AddInstance(int objectHandle, Mesh mesh, Matrix4x4 localToWorldMatrix, uint[] perSubMeshMask, uint[] perSubMeshMaterialIDs)
|
||
|
|
{
|
||
|
|
int subMeshCount = mesh.subMeshCount;
|
||
|
|
|
||
|
|
var instances = new InstanceIDs[subMeshCount];
|
||
|
|
for (int i = 0; i < subMeshCount; ++i)
|
||
|
|
{
|
||
|
|
var instanceDesc = new MeshInstanceDesc(mesh, i)
|
||
|
|
{
|
||
|
|
localToWorldMatrix = localToWorldMatrix,
|
||
|
|
mask = perSubMeshMask[i],
|
||
|
|
};
|
||
|
|
|
||
|
|
instances[i].InstanceID = _instances.AddInstance(instanceDesc, perSubMeshMaterialIDs[i]);
|
||
|
|
instanceDesc.instanceID = (uint)instances[i].InstanceID;
|
||
|
|
instances[i].AccelStructID = _accelStruct.AddInstance(instanceDesc);
|
||
|
|
}
|
||
|
|
|
||
|
|
_objectHandleToInstances.Add(objectHandle, instances);
|
||
|
|
}
|
||
|
|
|
||
|
|
private void AddInstance(int objectHandle, TerrainDesc terrainDesc)
|
||
|
|
{
|
||
|
|
List<InstanceIDs> instanceHandles = new List<InstanceIDs>();
|
||
|
|
|
||
|
|
AddHeightmap(terrainDesc, ref instanceHandles);
|
||
|
|
AddTrees(terrainDesc, ref instanceHandles);
|
||
|
|
|
||
|
|
_objectHandleToInstances.Add(objectHandle, instanceHandles.ToArray());
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
void AddHeightmap(TerrainDesc terrainDesc, ref List<InstanceIDs> instanceHandles)
|
||
|
|
{
|
||
|
|
var terrainMesh = TerrainToMesh.Convert(terrainDesc.terrain);
|
||
|
|
var instanceDesc = new MeshInstanceDesc(terrainMesh);
|
||
|
|
instanceDesc.localToWorldMatrix = terrainDesc.localToWorldMatrix;
|
||
|
|
instanceDesc.mask = terrainDesc.mask;
|
||
|
|
instanceDesc.enableTriangleCulling = terrainDesc.enableTriangleCulling;
|
||
|
|
instanceDesc.frontTriangleCounterClockwise = terrainDesc.frontTriangleCounterClockwise;
|
||
|
|
|
||
|
|
instanceHandles.Add(AddInstance(instanceDesc, terrainDesc.materialID));
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
void AddTrees(TerrainDesc terrainDesc, ref List<InstanceIDs> instanceHandles)
|
||
|
|
{
|
||
|
|
TerrainData terrainData = terrainDesc.terrain.terrainData;
|
||
|
|
float4x4 terrainLocalToWorld = terrainDesc.localToWorldMatrix;
|
||
|
|
float3 positionScale = new float3((float)terrainData.heightmapResolution, 1.0f, (float)terrainData.heightmapResolution) * terrainData.heightmapScale;
|
||
|
|
float3 positionOffset = new float3(terrainLocalToWorld[3].x, terrainLocalToWorld[3].y, terrainLocalToWorld[3].z);
|
||
|
|
|
||
|
|
foreach (var treeInstance in terrainData.treeInstances)
|
||
|
|
{
|
||
|
|
var localToWorld = Matrix4x4.TRS(
|
||
|
|
positionOffset + new float3(treeInstance.position) * positionScale,
|
||
|
|
Quaternion.AngleAxis(treeInstance.rotation, Vector3.up),
|
||
|
|
new Vector3(treeInstance.widthScale, treeInstance.heightScale, treeInstance.widthScale));
|
||
|
|
|
||
|
|
var prefab = terrainData.treePrototypes[treeInstance.prototypeIndex].prefab;
|
||
|
|
|
||
|
|
GameObject go = prefab.gameObject;
|
||
|
|
if (prefab.TryGetComponent<LODGroup>(out var lodGroup))
|
||
|
|
{
|
||
|
|
var groups = lodGroup.GetLODs();
|
||
|
|
if (groups.Length != 0 && groups[0].renderers.Length != 0)
|
||
|
|
go = (groups[0].renderers[0] as MeshRenderer).gameObject;
|
||
|
|
}
|
||
|
|
if (!go.TryGetComponent<MeshFilter>(out var filter))
|
||
|
|
continue;
|
||
|
|
|
||
|
|
var mesh = filter.sharedMesh;
|
||
|
|
for (int i = 0; i < mesh.subMeshCount; ++i)
|
||
|
|
{
|
||
|
|
var instanceDesc = new MeshInstanceDesc(mesh, i);
|
||
|
|
instanceDesc.localToWorldMatrix = localToWorld;
|
||
|
|
instanceDesc.mask = terrainDesc.mask;
|
||
|
|
instanceDesc.enableTriangleCulling = terrainDesc.enableTriangleCulling;
|
||
|
|
instanceDesc.frontTriangleCounterClockwise = terrainDesc.frontTriangleCounterClockwise;
|
||
|
|
instanceHandles.Add(AddInstance(instanceDesc, terrainDesc.materialID));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
InstanceIDs AddInstance(MeshInstanceDesc instanceDesc, uint materialID)
|
||
|
|
{
|
||
|
|
InstanceIDs res = new InstanceIDs();
|
||
|
|
res.InstanceID = _instances.AddInstance(instanceDesc, materialID);
|
||
|
|
instanceDesc.instanceID = (uint)res.InstanceID;
|
||
|
|
res.AccelStructID = _accelStruct.AddInstance(instanceDesc);
|
||
|
|
|
||
|
|
return res;
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
public void RemoveInstance(int objectHandle)
|
||
|
|
{
|
||
|
|
bool success = _objectHandleToInstances.TryGetValue(objectHandle, out var instances);
|
||
|
|
Assert.IsTrue(success);
|
||
|
|
|
||
|
|
foreach (var instance in instances)
|
||
|
|
{
|
||
|
|
_instances.RemoveInstance(instance.InstanceID);
|
||
|
|
_accelStruct.RemoveInstance(instance.AccelStructID);
|
||
|
|
}
|
||
|
|
|
||
|
|
_objectHandleToInstances.Remove(objectHandle);
|
||
|
|
}
|
||
|
|
|
||
|
|
public void UpdateInstanceTransform(int objectHandle, Matrix4x4 localToWorldMatrix)
|
||
|
|
{
|
||
|
|
bool success = _objectHandleToInstances.TryGetValue(objectHandle, out var instances);
|
||
|
|
Assert.IsTrue(success);
|
||
|
|
|
||
|
|
foreach(var instance in instances)
|
||
|
|
{
|
||
|
|
_instances.UpdateInstanceTransform(instance.InstanceID, localToWorldMatrix);
|
||
|
|
_accelStruct.UpdateInstanceTransform(instance.AccelStructID, localToWorldMatrix);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
public void UpdateInstanceMaterialIDs(int objectHandle, uint[] perSubMeshMaterialIDs)
|
||
|
|
{
|
||
|
|
bool success = _objectHandleToInstances.TryGetValue(objectHandle, out var instances);
|
||
|
|
Assert.IsTrue(success);
|
||
|
|
Assert.IsTrue(perSubMeshMaterialIDs.Length >= instances.Length);
|
||
|
|
int i = 0;
|
||
|
|
foreach (var instance in instances)
|
||
|
|
{
|
||
|
|
_instances.UpdateInstanceMaterialID(instance.InstanceID, perSubMeshMaterialIDs[i++]);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
public void UpdateInstanceMask(int objectHandle, uint[] perSubMeshMask)
|
||
|
|
{
|
||
|
|
bool success = _objectHandleToInstances.TryGetValue(objectHandle, out var instances);
|
||
|
|
Assert.IsTrue(success);
|
||
|
|
Assert.IsTrue(perSubMeshMask.Length >= instances.Length);
|
||
|
|
int i = 0;
|
||
|
|
foreach (var instance in instances)
|
||
|
|
{
|
||
|
|
_instances.UpdateInstanceMask(instance.InstanceID, perSubMeshMask[i]);
|
||
|
|
_accelStruct.UpdateInstanceMask(instance.AccelStructID, perSubMeshMask[i]);
|
||
|
|
i++;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
public void Build(CommandBuffer cmd, ref GraphicsBuffer scratchBuffer)
|
||
|
|
{
|
||
|
|
RayTracingHelper.ResizeScratchBufferForBuild(_accelStruct, ref scratchBuffer);
|
||
|
|
_accelStruct.Build(cmd, scratchBuffer);
|
||
|
|
}
|
||
|
|
|
||
|
|
public void NextFrame()
|
||
|
|
{
|
||
|
|
_instances.NextFrame();
|
||
|
|
}
|
||
|
|
|
||
|
|
public bool GetInstanceIDs(int rendererID, out int[] instanceIDs)
|
||
|
|
{
|
||
|
|
if (!_objectHandleToInstances.TryGetValue(rendererID, out InstanceIDs[] instIDs))
|
||
|
|
{
|
||
|
|
// This should never happen as long as the renderer was already added to the acceleration structure
|
||
|
|
instanceIDs = null;
|
||
|
|
return false;
|
||
|
|
}
|
||
|
|
instanceIDs = Array.ConvertAll(instIDs, item => item.InstanceID);
|
||
|
|
return true;
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|
||
|
|
|
||
|
|
internal struct TerrainDesc
|
||
|
|
{
|
||
|
|
public Terrain terrain;
|
||
|
|
public Matrix4x4 localToWorldMatrix;
|
||
|
|
public uint mask;
|
||
|
|
public uint materialID;
|
||
|
|
public bool enableTriangleCulling;
|
||
|
|
public bool frontTriangleCounterClockwise;
|
||
|
|
|
||
|
|
public TerrainDesc(Terrain terrain)
|
||
|
|
{
|
||
|
|
this.terrain = terrain;
|
||
|
|
localToWorldMatrix = Matrix4x4.identity;
|
||
|
|
mask = 0xFFFFFFFF;
|
||
|
|
materialID = 0;
|
||
|
|
enableTriangleCulling = true;
|
||
|
|
frontTriangleCounterClockwise = false;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|