Improved Shapes ScenePointEditor.cs

ScenePointEditor.cs

using System;
using System.Collections.Generic;
using System.Linq;
using UnityEditor;
using UnityEditor.IMGUI.Controls;
using UnityEngine;
using Object = UnityEngine.Object;

// Shapes © Freya Holmér - https://twitter.com/FreyaHolmer/
// Website & Documentation - https://acegikmo.com/shapes/
//
// MODIFIED for Vectroid:
// - Removed the AddRemovePoints Tab-cycle mode entirely; add/remove now happens via
//   edge-hover ghost-button (Unity PolygonCollider2D-style insert) and the Delete key.
// - Added multiselect (shift-click to toggle selection), Ctrl/Cmd+D to duplicate
//   selected points, and a combined position handle for moving multi-selections.
// - hasAddRemoveMode now gates the new add/remove behavior so fixed-count shapes
//   (Line/Quad/Triangle) that explicitly set it to false still keep their counts.
namespace Shapes {

	public class ScenePointEditor : SceneEditGizmos {

		static bool isEditing;

		public bool hasAddRemoveMode = true;
		public bool hasEditThicknessMode = false;
		public bool hasEditColorMode = false;

		public bool useFlatThicknessHandles = false;

		public event Action<ShapeRenderer, int> onValuesChanged = delegate { };

		public bool[] colorEnabledArray = null;
		public bool[] positionEnabledArray = null;

		EditMode currentEditMode = EditMode.PositionHandles;
		readonly HashSet<int> selectedIndices = new HashSet<int>();

		const float EDGE_HOVER_PIXEL_THRESHOLD = 12f;
		const float SELECT_DOT_GUI_RADIUS = 7f;

		void GoToNextEditMode() {
			while( true ) {
				currentEditMode = (EditMode)( ( (int)currentEditMode + 1 ) % (int)EditMode.COUNT );
				if( CanEdit( currentEditMode ) == false ) continue;
				break;
			}
		}

		bool CanEdit( EditMode mode ) {
			switch( mode ) {
				case EditMode.EditThickness: return hasEditThicknessMode;
				case EditMode.EditColor:     return hasEditColorMode;
				default:                     return true;
			}
		}

		bool HasAnyExtraMode => hasEditThicknessMode || hasEditColorMode;

		enum EditMode {
			PositionHandles,
			EditThickness,
			EditColor,
			COUNT
		}

		SphereBoundsHandle discHandle = ShapesHandles.InitDiscHandle();

		public ScenePointEditor( Editor parentEditor ) => this.parentEditor = parentEditor;

		protected override bool IsEditing {
			get => isEditing;
			set => isEditing = value;
		}

		public bool DoSceneHandles( bool closed, Object component, List<PolylinePoint> points, Transform tf, float globalThicknessScale = 1f, Color globalColorTint = default ) {
			void SetPt( int i, Vector3 pt ) {
				PolylinePoint pp = points[i];
				pp.point = pt;
				points[i] = pp;
			}

			float GetThicknessWorld( int i ) {
				float localThickness = points[i].thickness * globalThicknessScale;
				return localThickness * tf.lossyScale.AvgComponentMagnitude();
			}

			void SetThicknessWorld( int i, float thicknessWorld ) {
				PolylinePoint pp = points[i];
				float localThickness = thicknessWorld / tf.lossyScale.AvgComponentMagnitude();
				pp.thickness = localThickness / globalThicknessScale;
				points[i] = pp;
			}


			Color GetNetColor( int i ) => points[i].color * globalColorTint;
			Color GetColor( int i ) => points[i].color;

			void SetColor( int i, Color color ) {
				PolylinePoint p = points[i];
				p.color = color;
				points[i] = p;
			}

			return DoSceneHandles( closed, component, points, tf, i => points[i].point, p => p.point, SetPt, PolylinePoint.Lerp, GetThicknessWorld, SetThicknessWorld, GetNetColor, GetColor, SetColor, () => points.Count );
		}

		public bool DoSceneHandles( bool closed, Object component, List<Vector2> points, Transform tf ) {
			return DoSceneHandles( closed, component, points, tf, i => points[i], p => p, ( i, p ) => points[i] = p, Vector2.LerpUnclamped );
		}

		public bool DoSceneHandles( bool closed, Object component, List<Vector3> points, Transform tf ) {
			return DoSceneHandles( closed, component, points, tf, i => points[i], p => p, ( i, p ) => points[i] = p, Vector3.LerpUnclamped );
		}

		public bool DoSceneHandles( bool closed, Object component, List<Vector3> points, List<Color> colors, Transform tf ) {
			void SetColor( int i, Color c ) {
				colors[i] = c;
				onValuesChanged( component as ShapeRenderer, i );
			}

			return DoSceneHandles( closed, component, points, tf, i => points[i], p => p, ( i, p ) => points[i] = p, Vector3.LerpUnclamped, null, null, i => colors[i], i => colors[i], SetColor, () => colors.Count );
		}

		bool TextureButton( Vector3 worldPos, Texture2D tex, float scale, bool fade = true ) {
			Rect r = new Rect( 0, 0, tex.width * scale, tex.height * scale );
			r.center = HandleUtility.WorldToGUIPoint( worldPos );

			Vector2 mousePos = Event.current.mousePosition;

			if( fade ) {
				float t = Mathf.InverseLerp( 200, 64, Vector2.Distance( mousePos, r.center ) );
				float a = Mathf.Lerp( 0.3f, 1f, t );
				GUI.color = new Color( 1, 1, 1, a );
			}

			bool pressed = GUI.Button( r, tex, GUIStyle.none );
			if( fade )
				GUI.color = Color.white;
			return pressed;
		}

		bool DoSceneHandles<T>( bool closed,
								Object component,
								List<T> points,
								Transform tf,
								Func<int, Vector3> GetPtAt,
								Func<T, Vector3> GetPt,
								Action<int, Vector3> SetPt,
								Func<T, T, float, T> Lerp,
								Func<int, float> GetThicknessWorld = null,
								Action<int, float> SetThicknessWorld = null,
								Func<int, Color> GetNetColor = null,
								Func<int, Color> GetColor = null,
								Action<int, Color> SetColor = null,
								Func<int> ColorCount = null ) {
			CheckForCancelEditAction();
			if( IsHoldingAlt )
				return false;

			bool changed = false;

			Vector3 GetWorldPt( int i ) => tf.TransformPoint( GetPtAt( i ) );

			if( !isEditing )
				return false;

			// Drop selection indices that no longer exist (e.g. after external edits)
			selectedIndices.RemoveWhere( idx => idx < 0 || idx >= points.Count );

			if( Event.current.isKey && Event.current.keyCode == KeyCode.Tab && HasAnyExtraMode ) {
				if( Event.current.type == EventType.KeyDown )
					GoToNextEditMode();
				Event.current.Use();
			}

			// Mode label overlay
			if( Selection.gameObjects.Length > 0 && Selection.gameObjects[0] == ( (Component)component ).gameObject ) {
				if( Event.current.type == EventType.MouseMove )
					SceneView.lastActiveSceneView.Repaint();
				if( HasAnyExtraMode ) {
					Handles.BeginGUI();
					Vector2 mousePos = Event.current.mousePosition;
					Rect r = new Rect( mousePos.x + 32, mousePos.y, Screen.width, 128 );
					string label = "Press Tab to cycle modes:";
					void SelectLabel( string str, EditMode mode, bool exists = true ) {
						if( !exists ) return;
						label += mode == currentEditMode ? "\n> " + str : "\n  " + str;
					}
					SelectLabel( "position", EditMode.PositionHandles );
					SelectLabel( "thickness", EditMode.EditThickness, hasEditThicknessMode );
					SelectLabel( "color", EditMode.EditColor, hasEditColorMode );
					GUI.Label( r, label );
					Handles.EndGUI();
				}
			}

			Quaternion handleRotation = Tools.pivotRotation == PivotRotation.Global ? Quaternion.identity : tf.rotation;

			if( currentEditMode == EditMode.EditThickness ) {
				Transform camTf = SceneView.lastActiveSceneView.camera.transform;
				Vector3 camPos = camTf.position;
				Vector3 camUp = camTf.up;
				for( int i = 0; i < points.Count; i++ ) {
					discHandle.radius = GetThicknessWorld( i ) / 2f;
					discHandle.center = Vector3.zero;
					Vector3 discPos = GetWorldPt( i );
					Vector3 dirToCamera = discPos - camPos;
					Quaternion discRot = useFlatThicknessHandles ? tf.rotation : Quaternion.LookRotation( dirToCamera, camUp );
					Matrix4x4 mtx = Matrix4x4.TRS( discPos, discRot, Vector3.one );
					using( var chChk = new EditorGUI.ChangeCheckScope() ) {
						using( new Handles.DrawingScope( ShapesHandles.GetHandleColor( GetNetColor( i ) ), mtx ) )
							discHandle.DrawHandle();
						if( chChk.changed ) {
							changed = true;
							Undo.RecordObject( component, "edit thickness" );
							SetThicknessWorld( i, discHandle.radius * 2 );
							break;
						}
					}
				}
			} else if( currentEditMode == EditMode.EditColor ) {
				Handles.BeginGUI();
				for( int i = 0; i < ColorCount(); i++ ) {
					if( colorEnabledArray != null && colorEnabledArray[i] == false ) continue;
					Vector3 ptWorld = GetWorldPt( i );
					Color col = GetColor( i );
					col.a = 1f;
					GUI.color = col;
					if( TextureButton( ptWorld, UIAssets.Instance.pointEditColor, 0.5f, fade: false ) ) {
						int captured = i;
						ShapesUI.ShowColorPicker( OnColorChanged, GetColor( i ) );
						void OnColorChanged( Color c ) {
							Undo.RecordObject( component, "modify color" );
							SetColor( captured, c );
							( component as ShapeRenderer )?.UpdateAllMaterialProperties();
							( component as ShapeRenderer )?.UpdateMesh( force: true );
							ShapesUI.RepaintAllSceneViews();
						}
					}
				}
				GUI.color = Color.white;
				Handles.EndGUI();
			} else if( currentEditMode == EditMode.PositionHandles ) {
				changed |= DoPositionMode( closed, component, points, tf, GetWorldPt, GetPt, SetPt, Lerp );
			}

			return changed;
		}

		// Position-handle mode with multiselect, edge-insert, delete, duplicate.
		bool DoPositionMode<T>( bool closed,
								Object component,
								List<T> points,
								Transform tf,
								Func<int, Vector3> GetWorldPt,
								Func<T, Vector3> GetPt,
								Action<int, Vector3> SetPt,
								Func<T, T, float, T> Lerp ) {
			bool changed = false;
			Event e = Event.current;
			Quaternion handleRotation = Tools.pivotRotation == PivotRotation.Global ? Quaternion.identity : tf.rotation;
			int minPoints = closed ? 3 : 2;

			// --- Edge-hover ghost (insert) ---
			int hoverEdgeStart = -1;
			Vector3 hoverInsertWorld = default;
			float hoverInsertT = 0f;
			if( hasAddRemoveMode && points.Count >= 2 ) {
				Vector2 mouseGui = e.mousePosition;
				float bestDist = EDGE_HOVER_PIXEL_THRESHOLD;
				int edgeCount = closed ? points.Count : points.Count - 1;
				for( int i = 0; i < edgeCount; i++ ) {
					int j = ( i + 1 ) % points.Count;
					Vector3 wa = GetWorldPt( i );
					Vector3 wb = GetWorldPt( j );
					Vector2 ga = HandleUtility.WorldToGUIPoint( wa );
					Vector2 gb = HandleUtility.WorldToGUIPoint( wb );
					Vector2 ab = gb - ga;
					float lenSq = ab.sqrMagnitude;
					if( lenSq < 0.0001f ) continue;
					float t = Mathf.Clamp01( Vector2.Dot( mouseGui - ga, ab ) / lenSq );
					Vector2 closestGui = ga + t * ab;
					// Skip if closer to either endpoint (those are point selectors)
					if( Vector2.Distance( closestGui, ga ) < SELECT_DOT_GUI_RADIUS * 1.5f ) continue;
					if( Vector2.Distance( closestGui, gb ) < SELECT_DOT_GUI_RADIUS * 1.5f ) continue;
					float dist = Vector2.Distance( mouseGui, closestGui );
					if( dist < bestDist ) {
						bestDist = dist;
						hoverEdgeStart = i;
						hoverInsertT = t;
						hoverInsertWorld = Vector3.LerpUnclamped( wa, wb, t );
					}
				}
			}

			// --- Selection dots (per point, click to select / shift-click to toggle) ---
			for( int i = 0; i < points.Count; i++ ) {
				if( positionEnabledArray != null && positionEnabledArray[i] == false ) continue;
				Vector3 ptWorld = GetWorldPt( i );
				bool sel = selectedIndices.Contains( i );
				float handleSize = HandleUtility.GetHandleSize( ptWorld ) * 0.06f;
				Color dotColor = sel ? new Color( 1f, 0.7f, 0.1f ) : new Color( 0.2f, 0.6f, 1f );
				using( new Handles.DrawingScope( dotColor ) ) {
					if( Handles.Button( ptWorld, Quaternion.identity, handleSize, handleSize * 1.6f, Handles.DotHandleCap ) ) {
						bool shift = ( e.modifiers & EventModifiers.Shift ) != 0;
						if( shift ) {
							if( !selectedIndices.Add( i ) )
								selectedIndices.Remove( i );
						} else {
							selectedIndices.Clear();
							selectedIndices.Add( i );
						}
						SceneView.RepaintAll();
					}
				}
			}

			// --- Edge insert ghost button (drawn after dots so it overlays correctly) ---
			if( hoverEdgeStart >= 0 ) {
				float ghostSize = HandleUtility.GetHandleSize( hoverInsertWorld ) * 0.06f;
				using( new Handles.DrawingScope( new Color( 0.4f, 1f, 0.4f, 0.8f ) ) ) {
					if( Handles.Button( hoverInsertWorld, Quaternion.identity, ghostSize, ghostSize * 1.6f, Handles.SphereHandleCap ) ) {
						Undo.RecordObject( component, "insert point" );
						T newPt = Lerp( points[hoverEdgeStart], points[( hoverEdgeStart + 1 ) % points.Count], hoverInsertT );
						int insertAt = hoverEdgeStart + 1;
						points.Insert( insertAt, newPt );
						selectedIndices.Clear();
						selectedIndices.Add( insertAt );
						changed = true;
						SceneView.RepaintAll();
					}
				}
				// Repaint while moving so the ghost tracks the mouse smoothly
				if( e.type == EventType.MouseMove )
					SceneView.lastActiveSceneView.Repaint();
			}

			// --- Movement handles ---
			if( selectedIndices.Count > 1 ) {
				// Combined handle at centroid of selected
				Vector3 centroid = Vector3.zero;
				foreach( int idx in selectedIndices ) centroid += GetWorldPt( idx );
				centroid /= selectedIndices.Count;
				EditorGUI.BeginChangeCheck();
				Vector3 newCentroid = Handles.PositionHandle( centroid, handleRotation );
				if( EditorGUI.EndChangeCheck() ) {
					Vector3 deltaWorld = newCentroid - centroid;
					if( deltaWorld.sqrMagnitude > 0f ) {
						Undo.RecordObject( component, "move points" );
						foreach( int idx in selectedIndices ) {
							Vector3 nw = GetWorldPt( idx ) + deltaWorld;
							SetPt( idx, tf.InverseTransformPoint( nw ) );
						}
						changed = true;
					}
				}
			} else {
				// Single position handle on each point that's selected; if none selected,
				// show a position handle on every point so click-and-drag still works
				// without explicit selection (matches the original behavior).
				bool anySelected = selectedIndices.Count == 1;
				for( int i = 0; i < points.Count; i++ ) {
					if( positionEnabledArray != null && positionEnabledArray[i] == false ) continue;
					if( anySelected && !selectedIndices.Contains( i ) ) continue;
					Vector3 ptWorld = GetWorldPt( i );
					EditorGUI.BeginChangeCheck();
					Vector3 newPosWorld = Handles.PositionHandle( ptWorld, handleRotation );
					if( EditorGUI.EndChangeCheck() ) {
						changed = true;
						Undo.RecordObject( component, "modify points" );
						SetPt( i, tf.InverseTransformPoint( newPosWorld ) );
					}
				}
			}

			// --- Keyboard shortcuts ---
			if( e.type == EventType.KeyDown && hasAddRemoveMode && selectedIndices.Count > 0 ) {
				bool ctrl = ( e.modifiers & ( EventModifiers.Control | EventModifiers.Command ) ) != 0;
				if( e.keyCode == KeyCode.Delete || e.keyCode == KeyCode.Backspace ) {
					int allowedRemovals = Mathf.Max( 0, points.Count - minPoints );
					if( allowedRemovals > 0 ) {
						Undo.RecordObject( component, "delete points" );
						var toRemove = selectedIndices.OrderByDescending( x => x ).Take( allowedRemovals ).ToList();
						foreach( int idx in toRemove )
							points.RemoveAt( idx );
						selectedIndices.Clear();
						changed = true;
					}
					e.Use();
					SceneView.RepaintAll();
				} else if( ctrl && e.keyCode == KeyCode.D ) {
					Undo.RecordObject( component, "duplicate points" );
					var sortedAsc = selectedIndices.OrderBy( x => x ).ToList();
					// Insert highest-first so earlier indices don't shift
					for( int i = sortedAsc.Count - 1; i >= 0; i-- ) {
						int idx = sortedAsc[i];
						points.Insert( idx + 1, points[idx] );
					}
					selectedIndices.Clear();
					for( int i = 0; i < sortedAsc.Count; i++ )
						selectedIndices.Add( sortedAsc[i] + 1 + i );
					changed = true;
					e.Use();
					SceneView.RepaintAll();
				} else if( e.keyCode == KeyCode.A && ctrl ) {
					selectedIndices.Clear();
					for( int i = 0; i < points.Count; i++ )
						selectedIndices.Add( i );
					e.Use();
					SceneView.RepaintAll();
				}
			}

			return changed;
		}

	}

}