chameleon/Assets/Scripts/Components/MeshPainter/MeshPainter.Gpu.cs

using UnityEngine;

public partial class MeshPainter
{
    // --- GPU backend ---

    /// <summary>GPU backend per map: a display RenderTexture + scratch for the dilation ping-pong.</summary>
    private class GpuMap
    {
        public RenderTexture display;   // painted into (work RT for the mask map)
        public RenderTexture temp;      // dilation ping-pong scratch
        public RenderTexture final;     // MSO only: packed mask map (R=metallic, G=occlusion, A=smoothness); null otherwise
        public bool drew;               // anything drawn this stroke?
        public bool sRGB;               // sRGB color map (albedo/emission) vs linear data map (MSO)
    }

    // MeshPaintBrush.shader pass indices: 0 = source-over brush, 1 = coverage seed.

    public Shader brushShader;               // optional; else Shader.Find("Hidden/MeshPaintBrush")
    public Shader dilateShader;              // optional; else Shader.Find("Hidden/MeshPaintDilate")
    public Shader packMsShader;              // optional; else Shader.Find("Hidden/MeshPaintPackMS")
    private GpuMap[] _gpu;
    private Material _brushMat;
    private Material _dilateMat;
    private Material _packMat;
    private bool _prevSRGB;

    // --- Shared backend resource pools ---------------------------------------------------------------
    // Every paint RT has the identical descriptor (textureSize² ARGBHalf linear), so a despawned avatar's
    // freed RTs can be reused by the next spawn instead of re-allocating ~70 MB. This closes the gap the
    // in-place reseed left: reseed reuses an EXISTING backend, but a Conversion hit (hider despawns ->
    // hunter respawns) tears one down and builds another - the observer stalled on that full free+alloc.
    // Pooling reclaims the just-freed set immediately. Bounded so a mass despawn (match end) does not park
    // every avatar's RTs; extras beyond the cap are released normally. Static, to survive the despawn->
    // respawn gap; flushed on quit. Unity's null-check safely skips any entry a fast-enter-playmode reload
    // left dangling.
    private static readonly System.Collections.Generic.Stack<RenderTexture> s_rtPool =
        new System.Collections.Generic.Stack<RenderTexture>();
    private const int RtPoolCap = 18; // ~2 avatars' worth (9 RTs each)
    private static bool s_poolQuitHooked;

    // Reused staging textures for seed uploads (one per colour space), so re-seeding on every spawn/apply
    // does not churn a 4 MB Texture2D per map. The linear flag must match the data (colour maps decode sRGB
    // on the blit), so the two spaces are kept apart; rebuilt if the paint size ever changes.
    private static Texture2D s_stageLinear;
    private static Texture2D s_stageSRGB;

    /// <summary>
    /// Sets up the GPU backend (RenderTexture set + materials per map). Returns true on
    /// success; on missing shaders it falls back to the CPU backend and returns false so
    /// the caller leaves _gpuActive=false (and useGpu in sync) instead of crashing later.
    /// </summary>
    private bool SetupGpu(Color32[][] canonical)
    {
        if (!brushShader) brushShader = Shader.Find("Hidden/MeshPaintBrush");
        if (!dilateShader) dilateShader = Shader.Find("Hidden/MeshPaintDilate");
        if (!packMsShader) packMsShader = Shader.Find("Hidden/MeshPaintPackMS");

        if (!brushShader || !dilateShader || !packMsShader)
        {
            Debug.LogError("[MeshPainter] GPU shaders not found (Hidden/MeshPaintBrush, " +
                           "Hidden/MeshPaintDilate or Hidden/MeshPaintPackMS). Falling back to CPU painting.", this);
            useGpu = false;
            SetupCpu(canonical);
            return false;
        }

        _brushMat = new Material(brushShader) { hideFlags = HideFlags.HideAndDontSave };
        _dilateMat = new Material(dilateShader) { hideFlags = HideFlags.HideAndDontSave };
        _packMat = new Material(packMsShader) { hideFlags = HideFlags.HideAndDontSave };

        _gpu = new GpuMap[MAP_COUNT];
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            bool linear = MapIsLinear(mi);
            var gm = new GpuMap
            {
                sRGB = !linear,
                display = CreateRT("MeshPainter_DisplayRT_" + mi),
                temp = CreateRT("MeshPainter_TempRT_" + mi),
                final = (mi == MAP_MSO) ? CreateRT("MeshPainter_FinalRT_" + mi) : null
            };
            _gpu[mi] = gm;

            if (!DeferInitialSeed) SeedGpuMap(mi, canonical); // deferred: RTs stay unseeded until the snapshot re-seed

            // Color maps render to the display RT; the MS map exposes its packed final RT.
            AssignMapTexture(mi, gm.final ? gm.final : gm.display);
        }
        return true;
    }

    /// <summary>Seeds a GPU map's display RT from canonical data (coverage alpha = UV islands), then packs MS to final.</summary>
    private void SeedGpuMap(int mi, Color32[][] canonical)
    {
        var src = ResolveSeed(mi, canonical);
        var gm = _gpu[mi];

        if (mi == MAP_MSO)
        {
            // Canonical (R=metallic, G=occlusion, A=smoothness) -> work layout
            // (R=metallic, G=smoothness, B=occlusion). Alpha = coverage, stamped below.
            var work = new Color32[src.Length];
            for (int i = 0; i < src.Length; i++) work[i] = new Color32(src[i].r, src[i].a, src[i].g, 0);
            UploadColorsToRT(work, gm.display, linear: true);
            StampCoverage(gm.display);
            PackMs(gm); // build the final mask RT once so the material shows the seeded values
        }
        else if (mi == MAP_FUR)
        {
            // Fur R IS data (strand length), not coverage. Upload the seeded/loaded R verbatim (linear
            // data map) and skip the coverage stamp, so unpainted texels stay bald (R 0) instead of being
            // marked full length. (Dilation is also skipped for fur in EndGpuStroke.)
            UploadColorsToRT((Color32[])src.Clone(), gm.display, linear: true);
        }
        else
        {
            UploadColorsToRT((Color32[])src.Clone(), gm.display, linear: gm.sRGB == false);
            StampCoverage(gm.display);
        }
    }

    /// <summary>
    /// Re-seeds the ALREADY-built GPU maps from canonical data without recreating the RenderTextures or
    /// materials: each map's display RT is re-uploaded + re-stamped (MSO re-packed) exactly as SetupGpu
    /// does, but the existing RTs (same size/format) are reused, so the fur/albedo RTs a sibling
    /// FurShellRenderer samples keep their references. Avoids the per-apply ~70 MB RT alloc/free churn a
    /// full RebuildBackend incurs at 1024². Callers ensure _gpuActive && _gpu != null.
    /// </summary>
    private void ReseedGpuInPlace(Color32[][] canonical)
    {
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            SeedGpuMap(mi, canonical);
            var gm = _gpu[mi];
            AssignMapTexture(mi, gm.final ? gm.final : gm.display);
        }
    }

    /// <summary>
    /// Stamps the dilation coverage guard into a map's alpha: rasterizes the mesh with the brush
    /// shader's coverage pass (pass 1), writing alpha = 1 on exactly the texels the GPU brush can
    /// paint. This replaces the CPU island-mask approximation, whose generous edge tolerance marked
    /// sub-pixel island borders as protected even though the brush never reaches them - so they
    /// could never be overpainted (stayed frozen at their seeded value) while the dilation guarded
    /// them. With a GPU-accurate guard those border texels fall in the gutter (alpha 0) and dilation
    /// overpaints them normally. Texels stay alpha 0 until reached, identical to the brush.
    /// </summary>
    private void StampCoverage(RenderTexture rt)
    {
        if (!_brushMat || !_bakedMesh) return;
        var prev = RenderTexture.active;
        Graphics.SetRenderTarget(rt);
        if (_brushMat.SetPass(1))
        {
            DrawBakedMeshNow();
        }
        else
        {
            Debug.LogError("[MeshPainter] Brush coverage pass (pass 1) missing - dilation guard not seeded.", this);
        }
        RenderTexture.active = prev;
    }

    private void TeardownGpu()
    {
        if (_gpu != null)
        {
            foreach (var gm in _gpu)
            {
                if (gm == null) continue;
                RecycleRT(gm.display);
                RecycleRT(gm.temp);
                RecycleRT(gm.final);
            }
            _gpu = null;
        }
        if (_brushMat) { Destroy(_brushMat); _brushMat = null; }
        if (_dilateMat) { Destroy(_dilateMat); _dilateMat = null; }
        if (_packMat) { Destroy(_packMat); _packMat = null; }
    }

    // Returns an RT to the shared pool for the next spawn to reuse, or releases it if the pool is full.
    // Deliberately NOT Released before pooling - Release frees the GPU memory the reuse is meant to keep.
    private static void RecycleRT(RenderTexture rt)
    {
        if (!rt) return;
        if (s_rtPool.Count < RtPoolCap) s_rtPool.Push(rt);
        else { rt.Release(); Object.Destroy(rt); }
    }

    // Flushes the pooled RTs + staging textures. Hooked to Application.quitting (does not fire on editor
    // play-mode stop, so the capped set lingers until the next domain reload there - a bounded ~150 MB).
    private static void ClearBackendPool()
    {
        while (s_rtPool.Count > 0)
        {
            var rt = s_rtPool.Pop();
            if (rt) { rt.Release(); Object.Destroy(rt); }
        }
        if (s_stageLinear) { Object.Destroy(s_stageLinear); s_stageLinear = null; }
        if (s_stageSRGB) { Object.Destroy(s_stageSRGB); s_stageSRGB = null; }
    }

    private RenderTexture CreateRT(string rtName)
    {
        // Reuse a pooled RT of the matching size before allocating (see s_rtPool). A pooled RT is still
        // Created (never Released while pooled); its stale content is fully overwritten by SeedGpuMap.
        while (s_rtPool.Count > 0)
        {
            var pooled = s_rtPool.Pop();
            if (pooled && pooled.width == textureSize && pooled.height == textureSize)
            {
                pooled.name = rtName;
                return pooled;
            }
            if (pooled) { pooled.Release(); Destroy(pooled); } // size mismatch / stale -> drop it
        }

        if (!s_poolQuitHooked) { Application.quitting += ClearBackendPool; s_poolQuitHooked = true; }

        // 16-bit half-float, NOT ARGB32: repeated low-opacity source-over only converges to the
        // target colour if the accumulation buffer has sub-8-bit precision. At 8 bpp the per-stroke
        // decrement dst*alpha rounds back to dst once it drops below half an LSB, so painting e.g.
        // black at low opacity freezes at a grey floor (~0.5/alpha in 255ths) and never reaches 0.
        // Half-float removes that floor. Storage is linear; the sRGB round-trip for colour maps
        // happens at the byte boundary (seed decodes on upload, ReadRtPixels re-encodes on readback),
        // so canonical bytes and CPU/GPU parity are unchanged - only the accumulation gained bits.
        var rt = new RenderTexture(textureSize, textureSize, 0, RenderTextureFormat.ARGBHalf,
            RenderTextureReadWrite.Linear)
        {
            name = rtName,
            wrapMode = TextureWrapMode.Clamp,
            filterMode = FilterMode.Bilinear,
            antiAliasing = 1,
            useMipMap = false,
            autoGenerateMips = false
        };
        rt.Create();
        return rt;
    }

    /// <summary>Writes Color32 RGB into a RenderTexture via a temporary texture. Alpha (the dilation
    /// coverage marker) is cleared to 0 here and stamped afterwards by StampCoverage, so the guard matches
    /// the GPU rasterizer exactly (see SeedGpuMap). Fur carries no coverage (its length lives in R), so a
    /// zero alpha is correct there too.</summary>
    private void UploadColorsToRT(Color32[] colors, RenderTexture rt, bool linear)
    {
        var tmp = AcquireStage(linear); // reused per colour space, not allocated per call
        for (int i = 0; i < colors.Length; i++) colors[i].a = 0;
        tmp.SetPixels32(colors);
        tmp.Apply(false);

        bool prev = GL.sRGBWrite;
        // Paint RTs store LINEAR half-float, so never sRGB-encode on write. The source tex's own
        // colour space (sRGB for colour maps) still decodes to linear when the blit samples it.
        GL.sRGBWrite = false;
        Graphics.Blit(tmp, rt);
        GL.sRGBWrite = prev;
    }

    /// <summary>A reusable staging Texture2D per colour space (see s_stageLinear/s_stageSRGB): seeding uploads
    /// through it instead of churning a 4 MB texture per map. The linear flag must match the data (colour maps
    /// decode sRGB when the blit samples the staging tex), so the two spaces are kept separate; rebuilt if the
    /// paint size ever changes. Main-thread only, so the shared instances never overlap.</summary>
    private Texture2D AcquireStage(bool linear)
    {
        ref Texture2D slot = ref linear ? ref s_stageLinear : ref s_stageSRGB;
        if (!slot || slot.width != textureSize || slot.height != textureSize)
        {
            if (slot) Destroy(slot);
            slot = new Texture2D(textureSize, textureSize, TextureFormat.RGBA32, false, linear)
            {
                name = linear ? "MeshPainter_StageLinear" : "MeshPainter_StageSRGB",
                hideFlags = HideFlags.HideAndDontSave
            };
        }
        return slot;
    }

    /// <summary>
    /// Reads a paint RT back into 8-bit canonical bytes. The paint RTs are linear 16-bit half (for
    /// blend precision), so this first blits down into a throwaway 8-bit RT using the SAME encoding
    /// the old ARGB32 RTs had - sRGB for colour maps (GL.sRGBWrite re-encodes), linear for data maps -
    /// then reads that. The resulting bytes are bit-for-bit what the old 8-bit path (and the CPU
    /// backend) produced for the same colour; only the accumulation gained precision.
    /// </summary>
    private Color32[] ReadRtPixels(RenderTexture rt, bool linear)
    {
        var prevActive = RenderTexture.active;

        var tmp = RenderTexture.GetTemporary(textureSize, textureSize, 0, RenderTextureFormat.ARGB32,
            linear ? RenderTextureReadWrite.Linear : RenderTextureReadWrite.sRGB);
        bool prevSRGB = GL.sRGBWrite;
        GL.sRGBWrite = !linear; // colour maps: re-encode linear->sRGB on write; data maps: linear
        Graphics.Blit(rt, tmp);
        GL.sRGBWrite = prevSRGB;

        RenderTexture.active = tmp;
        var tex = new Texture2D(textureSize, textureSize, TextureFormat.RGBA32, false, linear);
        tex.ReadPixels(new Rect(0, 0, textureSize, textureSize), 0, 0, false);
        tex.Apply(false);

        RenderTexture.active = prevActive;
        RenderTexture.ReleaseTemporary(tmp);
        var px = tex.GetPixels32();
        Destroy(tex);
        return px;
    }

    private void PackMs(GpuMap gm)
    {
        bool prev = GL.sRGBWrite;
        GL.sRGBWrite = false; // linear data, no conversion
        Graphics.Blit(gm.display, gm.final, _packMat, 0);
        GL.sRGBWrite = prev;
    }

    // --- GPU painting ---

    private void BeginGpuStroke()
    {
        _prevSRGB = GL.sRGBWrite;
        for (int mi = 0; mi < MAP_COUNT; mi++)
            if (IsChannelActive(mi)) _gpu[mi].drew = false;
    }

    private void DrawBrushGpu(Vector3 pLocal, float radiusLocal, Vector3 u, Vector3 v, Vector3 n)
    {
        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            if (!IsChannelActive(mi)) continue;
            var gm = _gpu[mi];

            Graphics.SetRenderTarget(gm.display);
            GL.sRGBWrite = false; // linear half RTs: blend natively in linear, no sRGB encode

            // Stamp placement is shared by every draw on this map.
            _brushMat.SetTexture(MaskTexId, _stampMask);
            _brushMat.SetVector(BrushCenterId, pLocal);
            _brushMat.SetFloat(BrushRadiusId, radiusLocal);
            _brushMat.SetVector(BrushUId, u);
            _brushMat.SetVector(BrushVId, v);
            _brushMat.SetVector(BrushNormalId, n);

            if (mi == MAP_MSO)
            {
                // Metallic/smoothness and occlusion share one texture but paint with their OWN
                // opacity, so each is a separate stamp (own _BrushColor.a + per-channel write
                // mask). Untouched channels keep their value while coverage (A) accumulates.
                if (paintMetallicSmoothness)
                    StampMso(Mathf.Clamp01(metallicSmoothnessOpacity), 13f); // R|G|A (skip B = occlusion)
                if (paintOcclusion)
                    StampMso(Mathf.Clamp01(occlusionOpacity), 3f);           // B|A (occlusion only)
            }
            else if (mi == MAP_FUR)
            {
                // Fur stores ONLY strand length in R (the colour comes from the albedo map at render). Gated by
                // its own toggle: the length eases toward furLength with the brush's alpha as the deposit
                // weight. Reuse the source-over pass (0) on R only, so R = lerp(dst.R, furLength, op) with
                // op = luma*mask.a*brushColor.a. SetVector + EncodeData (not SetColor) keeps furLength as raw
                // linear data - the fur RT is a linear data map the shell samples directly for length.
                float fl = EncodeData(furLength);
                _brushMat.SetVector(BrushColorId, new Vector4(fl, fl, fl, Mathf.Clamp01(brushColor.a)));
                _brushMat.SetFloat(ColorWriteMaskId, 8f); // R only (bits: R=8, G=4, B=2, A=1)
                DrawStamp();
            }
            else
            {
                // Albedo / emission: full RGBA, opacity carried by the colour's own alpha (source-over).
                Color paintCol = mi == MAP_ALBEDO ? brushColor : emissionColor;
                _brushMat.SetColor(BrushColorId, paintCol);
                _brushMat.SetFloat(ColorWriteMaskId, 15f);
                DrawStamp();
            }

            gm.drew = true;
        }
    }

    /// <summary>
    /// Issues one mask-map stamp: raw (metallic, smoothness, occlusion) data with the given paint
    /// opacity in alpha and a per-channel colour write mask, so metallic/smoothness and occlusion
    /// stay independently paintable with their own opacity. EncodeData pre-compensates the Color
    /// property's sRGB->linear upload so the values round-trip exactly (CPU parity); alpha is not converted.
    /// </summary>
    private void StampMso(float opacity, float colorMask)
    {
        _brushMat.SetVector(BrushColorId, new Vector4(
            EncodeData(metallic), EncodeData(smoothness), EncodeData(occlusion), opacity));
        _brushMat.SetFloat(ColorWriteMaskId, colorMask);
        DrawStamp();
    }

    /// <summary>Draws the baked mesh once per sub-mesh with the current brush pass (UV-space rasterization).</summary>
    private void DrawStamp()
    {
        if (!_brushMat.SetPass(0)) return;
        DrawBakedMeshNow();
    }

    /// <summary>Draws the baked mesh once per sub-mesh with the pass the caller just set.</summary>
    private void DrawBakedMeshNow()
    {
        int sub = _bakedMesh.subMeshCount;
        for (int i = 0; i < sub; i++)
            Graphics.DrawMeshNow(_bakedMesh, Matrix4x4.identity, i);
    }

    private void EndGpuStroke()
    {
        Graphics.SetRenderTarget(null);

        for (int mi = 0; mi < MAP_COUNT; mi++)
        {
            if (!IsChannelActive(mi)) continue;
            var gm = _gpu[mi];

            // Edge padding against UV seams: dilate the display RT into temp, then swap. Skipped for fur:
            // its alpha is length (not a coverage guard), so dilation has no valid source mask and the
            // display RT reference must stay stable for FurMapTexture (no ping-pong swap).
            if (gm.drew && paddingTexels > 0 && mi != MAP_FUR)
            {
                GL.sRGBWrite = false;
                _dilateMat.SetInt(PadId, paddingTexels);
                Graphics.Blit(gm.display, gm.temp, _dilateMat, 0);

                var swap = gm.display; gm.display = gm.temp; gm.temp = swap;
                if (!gm.final) AssignMapTexture(mi, gm.display);
            }

            // MS: pack the work layout (R=metallic, G=smoothness) into the URP final layout.
            if (gm.final) PackMs(gm);
        }

        GL.sRGBWrite = _prevSRGB;
    }

    /// <summary>
    /// Pre-compensates a raw data value (0..1) for the brush's Color property, which Unity
    /// converts sRGB->linear on upload in a linear project. Applying the inverse here makes
    /// the value round-trip exactly (so metallic/smoothness/occlusion are stored faithfully
    /// and match the CPU byte writes). In gamma color space no conversion happens, so pass through.
    /// </summary>
    private static float EncodeData(float v)
    {
        v = Mathf.Clamp01(v);
        return QualitySettings.activeColorSpace == ColorSpace.Linear ? Mathf.LinearToGammaSpace(v) : v;
    }
}