Editor/Tools/PaintTool.cs

An editor tool for painting subterrain textures. It handles brush preview, mouse dragging, dispatches a compute shader to modify the terrain control map, tracks a dirty region and creates undo actions by copying control map pixels before and after painting.

File AccessNative Interop
using System;
using Sandbox;
using Editor.TerrainEditor;

namespace Editor.SubTerrain;

public enum SubTerrainPaintLayer
{
	Base = 0,
	Overlay = 1
}

[Title( "Paint Texture" )]
[Icon( "brush" )]
[Alias( "subterrain_paint" )]
[Group( "1" )]
[Order( 3 )]
public class SubTerrainPaintTool : EditorTool
{
	readonly SubTerrainTool _parent;
	bool _dragging;
	RectInt _dirtyRegion;

	public SubTerrainPaintTool( SubTerrainTool parent )
	{
		_parent = parent;
	}

	public static SubTerrainPaintLayer ActiveLayer { get; set; } = SubTerrainPaintLayer.Base;

	public override void OnEnabled()
	{
		AllowGameObjectSelection = false;
	}

	public override void OnUpdate()
	{
		var terrain = GetSelectedComponent<Terrain>();
		if ( !terrain.IsValid() )
		{
			_parent.ClearStatus();
			return;
		}

		if ( !terrain.RayIntersects( Gizmo.CurrentRay, Gizmo.RayDepth, out var hitPosition ) )
		{
			_parent.ClearStatus();
			return;
		}

		var tx = terrain.WorldTransform;
		_parent.UpdateStatus( terrain, hitPosition );

		var previewTransform = new Transform( tx.PointToWorld( hitPosition ), tx.Rotation );
		_parent.DrawBrushPreview( previewTransform );

		if ( Gizmo.IsLeftMouseDown )
		{
			bool shouldPaint = !_dragging || !Application.CursorDelta.IsNearZeroLength;

			if ( !_dragging )
			{
				_dragging = true;

				var uv = new Vector2( hitPosition.x, hitPosition.y ) / terrain.Storage.TerrainSize;
				var x = (int)Math.Floor( terrain.Storage.Resolution * uv.x );
				var y = (int)Math.Floor( terrain.Storage.Resolution * uv.y );
				_dirtyRegion = new( new Vector2Int( x, y ) );
			}

			if ( shouldPaint )
			{
				OnPaint( terrain, hitPosition );
			}
		}
		else if ( _dragging )
		{
			_dragging = false;
			OnPaintEnded( terrain );
		}
	}

	void OnPaint( Terrain terrain, Vector3 hitPosition )
	{
		var brushTexture = _parent.SelectedBrushTexture;
		if ( brushTexture is null )
			return;

		var hitUV = new Vector2( hitPosition.x, hitPosition.y ) / terrain.Storage.TerrainSize;
		int size = (int)Math.Floor( _parent.BrushSettings.Size * 2.0f / terrain.Storage.TerrainSize * terrain.Storage.Resolution );
		size = Math.Max( size, 1 );

		var cs = new ComputeShader( "terrain/cs_subterrain_splat" );
		cs.Attributes.Set( "ControlMap", terrain.ControlMap );
		cs.Attributes.Set( "ControlUV", hitUV );
		cs.Attributes.Set( "BrushStrength", _parent.BrushSettings.Opacity * (Gizmo.IsCtrlPressed ? -1.0f : 1.0f) );
		cs.Attributes.Set( "BrushSize", size );
		cs.Attributes.Set( "Brush", brushTexture );
		cs.Attributes.Set( "SplatChannel", PaintTextureTool.SplatChannel );
		cs.Attributes.Set( "PaintLayer", (int)ActiveLayer );

		cs.Dispatch( size, size, 1 );

		var x = (int)Math.Floor( terrain.Storage.Resolution * hitUV.x ) - size / 2;
		var y = (int)Math.Floor( terrain.Storage.Resolution * hitUV.y ) - size / 2;
		_dirtyRegion.Add( new RectInt( x, y, size + 1, size + 1 ) );
	}

	void OnPaintEnded( Terrain terrain )
	{
		_dirtyRegion.Left = Math.Clamp( _dirtyRegion.Left, 0, terrain.Storage.Resolution - 1 );
		_dirtyRegion.Right = Math.Clamp( _dirtyRegion.Right, 0, terrain.Storage.Resolution - 1 );
		_dirtyRegion.Top = Math.Clamp( _dirtyRegion.Top, 0, terrain.Storage.Resolution - 1 );
		_dirtyRegion.Bottom = Math.Clamp( _dirtyRegion.Bottom, 0, terrain.Storage.Resolution - 1 );

		var dirtyRegion = _dirtyRegion;

		static uint[] CopyRegion( uint[] data, int stride, RectInt rect )
		{
			uint[] region = new uint[rect.Width * rect.Height];

			for ( int y = 0; y < rect.Height; y++ )
			{
				for ( int x = 0; x < rect.Width; x++ )
				{
					region[x + y * rect.Width] = data[rect.Left + x + (rect.Top + y) * stride];
				}
			}

			return region;
		}

		var regionBefore = CopyRegion( terrain.Storage.ControlMap, terrain.Storage.Resolution, dirtyRegion );
		terrain.SyncCPUTexture( Terrain.SyncFlags.Control, dirtyRegion );
		var regionAfter = CopyRegion( terrain.Storage.ControlMap, terrain.Storage.Resolution, dirtyRegion );

		Action CreateUndoAction( uint[] region ) => () =>
		{
			if ( !terrain.IsValid() )
				return;

			for ( int y = 0; y < dirtyRegion.Height; y++ )
			{
				for ( int x = 0; x < dirtyRegion.Width; x++ )
				{
					terrain.Storage.ControlMap[dirtyRegion.Left + x + (dirtyRegion.Top + y) * terrain.Storage.Resolution] = region[x + y * dirtyRegion.Width];
				}
			}

			terrain.SyncGPUTexture();
		};

		SceneEditorSession.Active.UndoSystem.Insert( $"Terrain {DisplayInfo.For( this ).Name}",
			CreateUndoAction( regionBefore ),
			CreateUndoAction( regionAfter ) );
	}
}