Bypassing Unity Undo System's Memory Limitation

By Myonmyon

It is said that in the crumbled city of Uniograad, there is a catacomb buried deep under the debris. Whenever wind travels through the ruins, the empty chambers of the catacomb would produce a wailing sound. And if you dare to dig into the catacomb, you would find… the unspeakable abomination lurking in the walls…

Which is Unity’s Undo system.

Usually this has not much of impact if you don’t do a very specific thing, but if you dare to crank Terrain resolution to 4096, the undo system simply can’t handle it. When you paint texture, chances are you are unable to undo.

The reason is simple, undo system has a built-in buffer size limitation. If your object is too large and surpasses this size, the undo system will not record it. A terrain with 4096 resolution and 16 layers will have 4 Alpha Map each has 4096x4096 in size, 4 channels, and that would give you something like 1 GB per undo step. It is still reasonable that Unity limited the size because this would quickly blow up your RAM or VRAM, but being unable to do even one undo is kinda… extreme.

So how to solve this problem?

Well, this complaint has been on Unity Forum for… more than 8 years! And Unity still hasn’t exposed this buffer size limitation. Also, the Undo system is written in C++, so you cannot just use reflection to simply adjust a value.

A dead end, and another day to curse Unity for being unable to solve a simple problem.

But really, we can do nothing about it?

The Detour

Well, there is a way to hack our way out of this… It is still more than necessary, but if you are so desperate, this might be the only way out.

First, when you paint, the Terrain Tools system will invoke a method in Undo system, requesting the Undo system to save the current state of the objects. And these objects, for painting terrain texture, are simply 4 texture 2D. The instance id of these textures does not change, which means no matter how many strokes you make on a terrain, the 4 textures all keep their instance id. This information is useful for restoring the state of the alpha maps.

Now… if only we can hijack this invocation…

Oh, we do have something. There is a plugin called Harmony that can intercept a method call and run some code before actually calling the method. And you can also use the injected code to decide whether to run the original method. This “some code” is called a “prefix” in Harmony terminology.

But! There is a caveat, Harmony can’t take extern methods. And the bad news is that most of the methods in the Undo system are extern. The good news is, the call made by Terrain Tools system, lands on a non extern method.

So, heh, after intercepting the objects, we need to create copies for them. It is essentially copying data from one texture to another. And according to this blogpost, the best way to do this is to use Graphics.CopyTexture, which takes no noticeable time. However, this means the data lives in GPU, so you probably need to be careful about it, as the textures are fairly large.

To restore the textures, first we will need to find alpha maps from TerrainData that have matching instance id as your recorded textures (you did record the instance id or a reference to the original texture as well, right? Because creating a new texture will have a different instance id). Then… you will need to use SetAlphaMaps to update the original alpha maps. This method requires a float[,,] to work, but we don’t have it yet…

The best approach I found is to use GetPixels on the copied texture to download all pixel data from GPU to your RAM, and use a Parallel.For to create the 3 dimensional float array. This should reduce the undo operation time to about 1~2 seconds.

Oh, and finally… It is sad that we can’t intercept the PerformUndo call since it is extern, so I enforced Alt+Z to undo when painting texture. This can be done with a simple [EditorMenu] attribute.

Anyway, this is not perfect, but still this is the best I can do if Unity keeps hiding that precious buffer size limitation.

Example Code

using System;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Threading.Tasks;
using HarmonyLib;
using UnityEditor;
using UnityEditor.SceneManagement;
using UnityEngine;
using Object = UnityEngine.Object;

namespace UndoHijack
{
    public static class UndoHijackSys
    {
        private static Harmony _harmony;

        [MenuItem("Help/Hacks/Hijack Undo")]
        [InitializeOnLoadMethod]
        public static void Init()
        {
            _harmony = new Harmony("UndoHijack");

            var original = typeof(Undo).GetMethod("RegisterCompleteObjectUndo",
                new[] { typeof(UnityEngine.Object[]), typeof(string) });
            var prefix = typeof(UndoHijackSys).GetMethod("HijackTexturePaintingUndo",
                BindingFlags.NonPublic | BindingFlags.Static);
            _harmony.Patch(original, prefix: new HarmonyMethod(prefix));
        }

        internal class UndoElement: IDisposable
        {
            public int instanceId;
            public Texture2D values;

            public void Dispose()
            {
                Object.DestroyImmediate(values);
            }
        }

        public const int maxUndo = 6;
        private static List<UndoElement[]> _undo = new();

        private static bool HijackTexturePaintingUndo(UnityEngine.Object[] objectsToUndo, string name)
        {
            if (name != "PaintTextureTool - Texture") return true;
            if (_undo.Count >= maxUndo)
            {
                var excessiveStep = _undo[0];
                _undo.RemoveAt(0);
                foreach (var element in excessiveStep)
                {
                    element.Dispose();
                }
            }

            _undo.Add(objectsToUndo.Select(x => CreateUndoStep(x as Texture2D)).ToArray());
            Debug.Log($"Adding undo ({objectsToUndo.Length})");
            return false;
        }

        private static UndoElement CreateUndoStep(Texture2D t2d)
        {
            var result = new UndoElement();
            result.instanceId = t2d.GetInstanceID();
            var width = t2d.width;
            var height = t2d.height;
            result.values = new Texture2D(width,height);
            Graphics.CopyTexture(t2d, result.values);
            return result;
        }

        [MenuItem("Help/Hacks/Alt Undo &z")]
        private static void Undo()
        {
            if (_undo.Count < 1)
            {
                Debug.LogWarning("Texture paint undo threshold reached.");
                return;
            }
            var stepToUndo = _undo[^1];
            var step = stepToUndo.ToDictionary(x => x.instanceId, x => x.values);
            var scene = EditorSceneManager.GetActiveScene();
            var terrains = scene.GetRootGameObjects()
                .SelectMany(x => x.GetComponentsInChildren<Terrain>()).ToList();
            foreach (var t in terrains)
            {
                /*
                 *  Conclusion after conducting experiments:
                 *  A terrain-data's Alphamap Texture's instance id does not change.
                 *  BUT! These Alphamaps are created on the fly in the get method.
                 */
                // Only compare the first. If the first doesn't match, none will match.
                var allAlphamapTexures = t.terrainData.alphamapTextures;
                var alphamapTexture = t.terrainData.alphamapTextures[0];
                var id = alphamapTexture.GetInstanceID();
                if (!step.ContainsKey(id)) continue;
                // write back
                var width = t.terrainData.alphamapWidth;
                var height = t.terrainData.alphamapHeight;
                var map = new float[width, height, t.terrainData.terrainLayers.Length];
                var counter = 0;
                foreach (var texture in allAlphamapTexures)
                {
                    var instanceID = texture.GetInstanceID();
                    var undoTexture = step[instanceID];
                    WriteToArray(undoTexture, map, width, height, counter);
                    counter++;
                }

                t.terrainData.SetAlphamaps(0, 0, map);
            }

            foreach (var element in stepToUndo)
            {
                element.Dispose();
            }
            _undo.RemoveAt(_undo.Count - 1);
        }

        private static void WriteToArray(Texture2D t2d, float[,,] map, int width, int height, int layer)
        {
            var res = t2d.GetPixels();
            var result = Parallel.For(0, width, i =>
            {
                for (int j = 0; j < height; j++)
                {
                    map[i, j, layer * 4 + 0] = res[i * width + j].r;
                    map[i, j, layer * 4 + 1] = res[i * width + j].g;
                    map[i, j, layer * 4 + 2] = res[i * width + j].b;
                    map[i, j, layer * 4 + 3] = res[i * width + j].a;
                }
            });
        }

        
    }
}
Share: X (Twitter) Facebook LinkedIn