using System;
using System.Collections.Generic;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Rendering;
namespace UnityEngine.PathTracing.Core
{
// Container type that stores a set of arbitrarily sized textures in a 2D texture array,
// where each slice is a quadtree. Supports add/update/remove operations.
internal class TextureSlotAllocator : IDisposable
{
public readonly struct TextureLocation
{
public readonly int AtlasIndex; // Which atlas is the texture in
public readonly TextureQuadTree.TextureNode TextureNode; // Which position in the atlas is the texture
public TextureLocation(int atlasIndex, TextureQuadTree.TextureNode textureNode)
{
AtlasIndex = atlasIndex;
TextureNode = textureNode;
}
public static readonly TextureLocation Invalid = new(-1, null);
public bool IsValid => AtlasIndex >= 0 && TextureNode != null;
}
private readonly int _size;
private readonly GraphicsFormat _format;
private readonly FilterMode _filterMode;
private readonly CommandBuffer _cmd;
private RenderTexture _atlases;
private readonly HashSet<int> _freeAtlases = new();
private TextureQuadTree[] _textureQuadTrees = Array.Empty<TextureQuadTree>();
// Texture sizes stored in each quadtree node don't necessarily match the size of the quadtree node.
// We keep track of the real texture size here, so we can write to the atlas without stretching.
private readonly Dictionary<TextureLocation, Vector2Int> _textureSizes = new();
private readonly int _tempTextureID = Shader.PropertyToID("_TempSlotRT");
public RenderTexture Texture => _atlases;
public TextureSlotAllocator(int size, GraphicsFormat format, FilterMode filterMode)
{
_size = size;
_format = format;
_filterMode = filterMode;
_cmd = new CommandBuffer();
_cmd.name = $"Texture Slot Allocator ({size}, {format})";
// Initial allocation, in order to have a texture atlas in place even if we don't add anything to it.
// Note: This is required when binding the textures to ray tracing shaders (compute shaders seem to be OK with null textures)
ResizeAtlas(1);
}
public TextureLocation AddTexture(Texture texture, Vector2 scale, Vector2 offset)
{
// Find the max dimension
int size = Mathf.Max(texture.width, texture.height);
// Try to find an appropriately sized atlas in the free list
int atlasIndex = -1;
TextureQuadTree atlas = null;
foreach (int freeIndex in _freeAtlases)
{
if (_textureQuadTrees[freeIndex].HasSpaceForTexture(size))
{
atlasIndex = freeIndex;
atlas = _textureQuadTrees[atlasIndex];
break;
}
}
// If we couldn't find one, create a new atlas
if (atlasIndex == -1 || atlas == null)
{
// Create a new atlas
atlasIndex = _textureQuadTrees.Length;
Array.Resize(ref _textureQuadTrees, atlasIndex + 1);
atlas = new TextureQuadTree(_size);
_textureQuadTrees[atlasIndex] = atlas;
ResizeAtlas(atlasIndex + 1);
// Add it to the free list
_freeAtlases.Add(atlasIndex);
}
// Place in atlas
if (!atlas.AddTexture(size, out var textureNode))
{
return TextureLocation.Invalid;
}
// We might have filled up the atlas. If so, we remove it from the free list, if it is there.
if (atlas.IsFull)
{
_freeAtlases.Remove(atlasIndex);
}
// Blit the texture into the atlas
TextureLocation location = new TextureLocation(atlasIndex, textureNode);
UpdateTexture(in location, texture, scale, offset);
return location;
}
public void UpdateTexture(in TextureLocation location, Texture texture, Vector2 scale, Vector2 offset)
{
TextureQuadTree.TextureNode node = location.TextureNode;
int atlasIndex = location.AtlasIndex;
Debug.Assert(atlasIndex < _atlases.volumeDepth, "The texture atlas is too small.");
Vector2 textureScale = new Vector2(
node.Size / (float)Mathf.Min(texture.width, node.Size),
node.Size / (float)Mathf.Min(texture.height, node.Size));
_cmd.GetTemporaryRT(_tempTextureID, new RenderTextureDescriptor(node.Size, node.Size, _format, 0));
_cmd.Blit(texture, _tempTextureID, textureScale * scale, offset);
_cmd.CopyTexture(_tempTextureID, 0, 0, 0, 0, node.Size, node.Size, _atlases, atlasIndex, 0, node.PosX, node.PosY);
_cmd.ReleaseTemporaryRT(_tempTextureID);
Graphics.ExecuteCommandBuffer(_cmd);
_cmd.Clear();
// Keep track of texture size
_textureSizes[location] = new Vector2Int(texture.width, texture.height);
}
public void RemoveTexture(in TextureLocation location)
{
Debug.Assert(location.AtlasIndex < _textureQuadTrees.Length);
var atlas = _textureQuadTrees[location.AtlasIndex];
atlas.RemoveTexture(location.TextureNode);
// After removal, the atlas can't be full, so make sure it's in the free list
_freeAtlases.Add(location.AtlasIndex);
// No need to hold on to texture size anymore
_textureSizes.Remove(location);
// We could choose to shrink the atlas here to save memory,
// but don't do so currently due to the overhead of resizing the atlas.
}
public Vector2Int GetTextureSize(in TextureLocation location)
{
return new Vector2Int(
Mathf.Min(_textureSizes[location].x, location.TextureNode.Size),
Mathf.Min(_textureSizes[location].y, location.TextureNode.Size));
}
public void GetScaleAndOffset(in TextureLocation location, out Vector2 scale, out Vector2 offset)
{
Debug.Assert(location.AtlasIndex < _textureQuadTrees.Length);
Vector2Int textureSize = GetTextureSize(location);
scale = new Vector2(textureSize.x / (float)_size, textureSize.y / (float)_size);
offset = new Vector2(location.TextureNode.PosX / (float)_size, location.TextureNode.PosY / (float)_size);
}
private void ResizeAtlas(int sliceCount)
{
// Create new atlas array
var texture = new RenderTexture(new RenderTextureDescriptor(_size, _size)
{
dimension = TextureDimension.Tex2DArray,
depthBufferBits = 0,
volumeDepth = sliceCount,
msaaSamples = 1,
vrUsage = VRTextureUsage.OneEye,
graphicsFormat = _format,
enableRandomWrite = true,
})
{
name = $"Texture Atlas (Format = {_format})",
hideFlags = HideFlags.DontSaveInEditor,
wrapMode = TextureWrapMode.Clamp,
wrapModeU = TextureWrapMode.Clamp,
wrapModeV = TextureWrapMode.Clamp,
filterMode = _filterMode,
};
texture.Create();
if (_atlases != null)
{
var previousRT = RenderTexture.active;
// Blit old atlases into new one
for (int i = 0; i < Mathf.Min(sliceCount, _atlases.volumeDepth); i++)
{
Graphics.Blit(_atlases, texture, i, i);
}
RenderTexture.active = previousRT;
// Release old atlases
_atlases.Release();
CoreUtils.Destroy(_atlases);
}
_atlases = texture;
}
public void Dispose()
{
_cmd?.Dispose();
if (_atlases != null)
{
_atlases.Release();
CoreUtils.Destroy(_atlases);
}
}
}
}