Code/Network/WaterNetwork.cs

Component that represents the scene water network: stores nodes and rivers, provides graph queries and edit operations (add/remove/split/merge/reverse rivers and nodes), maintains indexes and cached curve tables, and drives per-frame ticking and rendering updates.

File Access
using System;
using System.Collections.Generic;
using System.Linq;
using Sandbox;

namespace RedSnail.WaterTool2;

/// <summary>
/// The scene's water: a graph of nodes joined by rivers, and the lakes those rivers run into and out of. It is the
/// single source of truth for every river, confluence and lake in it; every surface is generated from it, piece by
/// piece, and rebuilt only where the graph changed. Connected water is built as one surface: a river's last row of
/// vertices is the very row its confluence or lake starts from, and every vertex carries the state of the water there
/// (its flow, its wave strength, its foam), so the water runs from one piece into the next without a seam.
/// </summary>
/// <remarks>
/// Runs everywhere, a dedicated server included: the pieces are what swimming, drowning and buoyancy ask about. Only
/// the GPU side is skipped where nothing renders.
/// </remarks>
[Title("Water Network")]
[Icon("water")]
public sealed partial class WaterNetwork : Component, Component.ExecuteInEditor
{
	/// <summary>
	/// The network gameplay asks about. Set by the network that was enabled last (so the play scene's network takes
	/// over from the editor scene's as play starts), and picked up again by any live network when it goes.
	/// </summary>
	[SkipHotload] public static WaterNetwork Current { get; private set; }

	/// <summary>
	/// Undo/redo and scene loads assign whole new lists here instead of mutating the existing ones, so the setters are
	/// the one place that reliably sees "the graph changed" and asks for a reconcile pass.
	/// </summary>
	[Property, Hide]
	public List<WaterNode> Nodes
	{
		get;
		set
		{
			field = value ?? new List<WaterNode>();
			MarkDirty();
		}
	} = [];

	[Property, Hide]
	public List<WaterRiver> Rivers
	{
		get;
		set
		{
			field = value ?? new List<WaterRiver>();
			MarkDirty();
		}
	} = [];

	/// <summary>
	/// Shared counter for node, river and lake ids. Serialized so ids stay unique across undo and reloads.
	/// </summary>
	[Property, Hide] public int NextId { get; set; } = 1;



	/// <summary>
	/// Every network asks the scene for "the" network, so a second one would be found by nobody in particular. The
	/// one that arrives second takes itself out again.
	/// </summary>
	protected override void OnAwake()
	{
		var existing = Scene.Get<WaterNetwork>();

		if (!existing.IsValid() || existing == this)
			return;

		Log.Warning($"Only one {nameof(WaterNetwork)} can exist in a scene: the one added to \"{GameObject.Name}\" has been removed. Use the one on \"{existing.GameObject.Name}\".");

		Destroy();
	}



	protected override void OnEnabled()
	{
		Current = this;

		MarkDirty();
	}



	protected override void OnDisabled()
	{
		ReleaseRendering();
		ReleasePieces();
		ReleaseRipples();

		if (Current == this)
			Current = null;
	}



	/// <summary>
	/// Undo/redo restore the graph by deserializing into this component; the list setters catch most of it, this catches
	/// the rest (and inspector edits).
	/// </summary>
	protected override void OnValidate()
	{
		MarkDirty();
	}



	protected override void OnUpdate()
	{
		// Leaving play mode takes the play scene's network away; the editor's carries on.
		if (!Current.IsValid())
			Current = this;

		Tick();
		UpdateRipples();

		if (Application.IsHeadless)
			return;

		UpdateRendering();
	}



	private readonly Dictionary<int, WaterNode> m_NodeIndex = new();
	private readonly Dictionary<int, WaterRiver> m_RiverIndex = new();
	private readonly Dictionary<int, List<WaterRiver>> m_NodeRivers = new();
	private static readonly List<WaterRiver> s_NoRivers = new();

	private int m_GraphRevision;
	private int m_IndexRevision = -1;

	/// <summary>
	/// Says the graph's shape changed, so the lookups below are rebuilt before they are next used.
	/// </summary>
	private void InvalidateIndex()
	{
		m_GraphRevision++;
	}



	/// <summary>
	/// Nodes and rivers by id, and the rivers at each node: everything asks the graph these questions constantly.
	/// </summary>
	private void EnsureIndex()
	{
		if (m_IndexRevision == m_GraphRevision)
			return;

		m_IndexRevision = m_GraphRevision;

		m_NodeIndex.Clear();
		m_RiverIndex.Clear();
		m_NodeRivers.Clear();

		foreach (var node in Nodes)
		{
			if (node is not null)
				m_NodeIndex[node.Id] = node;
		}

		foreach (var river in Rivers)
		{
			if (river is null)
				continue;

			m_RiverIndex[river.Id] = river;

			Attach(river.StartNodeId, river);

			if (river.EndNodeId != river.StartNodeId)
				Attach(river.EndNodeId, river);
		}
	}



	private void Attach(int _NodeId, WaterRiver _River)
	{
		if (!m_NodeRivers.TryGetValue(_NodeId, out var list))
		{
			list = new List<WaterRiver>();
			m_NodeRivers[_NodeId] = list;
		}

		list.Add(_River);
	}



	public WaterNode FindNode(int _Id)
	{
		EnsureIndex();

		return m_NodeIndex.GetValueOrDefault(_Id);
	}



	public WaterRiver FindRiver(int _Id)
	{
		EnsureIndex();

		return m_RiverIndex.GetValueOrDefault(_Id);
	}



	public WaterRiver FindRiverBetween(int _NodeA, int _NodeB)
	{
		if (_NodeA == _NodeB)
			return null;

		foreach (var river in GetRiversAt(_NodeA))
		{
			if (river.Touches(_NodeB))
				return river;
		}

		return null;
	}



	/// <summary>
	/// The rivers at a node. The list belongs to the network: treat it as read only, and copy it before changing the
	/// graph while walking it.
	/// </summary>
	public IReadOnlyList<WaterRiver> GetRiversAt(int _NodeId)
	{
		EnsureIndex();

		return m_NodeRivers.TryGetValue(_NodeId, out var rivers) ? rivers : s_NoRivers;
	}



	public int GetDegree(int _NodeId)
	{
		return GetRiversAt(_NodeId).Count;
	}



	public WaterNode FindNearestNode(Vector3 _LocalPosition, float _MaxDistance)
	{
		WaterNode nearest = null;
		float nearestDistance = _MaxDistance;

		foreach (var node in Nodes)
		{
			float distance = Vector3.DistanceBetween(node.Position, _LocalPosition);

			if (distance > nearestDistance)
				continue;

			nearest = node;
			nearestDistance = distance;
		}

		return nearest;
	}



	public WaterNode AddNode(Vector3 _LocalPosition)
	{
		var node = new WaterNode { Id = NextId++, Position = _LocalPosition };

		Nodes.Add(node);
		MarkDirty();

		return node;
	}



	/// <summary>
	/// Joins two nodes with a river flowing from the first to the second. Returns null when the join makes no sense
	/// (same node, unknown node, a river already runs between them, or either node opens into a lake that already has
	/// its river).
	/// </summary>
	public WaterRiver AddRiver(int _UpstreamNodeId, int _DownstreamNodeId, float _Width, float _Speed, float _Waves, float _Depth)
	{
		if (_UpstreamNodeId == _DownstreamNodeId)
			return null;

		var start = FindNode(_UpstreamNodeId);
		var end = FindNode(_DownstreamNodeId);

		if (start is null || end is null)
			return null;

		if (FindRiverBetween(_UpstreamNodeId, _DownstreamNodeId) is not null)
			return null;

		// A lake opening or a river mouth carries exactly one river: the lake, or the sea, is the rest of the junction.
		if (IsEndTaken(_UpstreamNodeId) || IsEndTaken(_DownstreamNodeId))
			return null;

		var river = new WaterRiver
		{
			Id = NextId++,
			StartNodeId = _UpstreamNodeId,
			EndNodeId = _DownstreamNodeId,
			StartWidth = _Width,
			EndWidth = _Width,
			Speed = _Speed,
			Waves = _Waves,
			Depth = _Depth,
			StartTangent = StraightHandle(start.Position, end.Position),
			EndTangent = StraightHandle(end.Position, start.Position)
		};

		Rivers.Add(river);
		InvalidateIndex();

		RefreshTangentsAround(_UpstreamNodeId);
		RefreshTangentsAround(_DownstreamNodeId);
		MarkDirty();

		return river;
	}



	/// <summary>
	/// Inserts a node on a river without changing its shape: the river keeps its id and becomes the upstream half, a
	/// new river carries on from the node. A split landing within <paramref name="_MinEndDistance"/> of an end returns
	/// that end's node instead, so no sliver rivers get made.
	/// </summary>
	public WaterNode SplitRiver(int _RiverId, float _T, float _MinEndDistance = 25.0f)
	{
		var river = FindRiver(_RiverId);

		if (river is null || !TryGetCurve(river, out var curve))
			return null;

		var table = GetCurveTable(river);
		float distance = _T * table.Length;

		if (distance <= _MinEndDistance)
			return FindNode(river.StartNodeId);

		if (table.Length - distance <= _MinEndDistance)
			return FindNode(river.EndNodeId);

		curve.Split(_T, out var first, out var second);

		// The handles at the split come from the subdivision, which is the only way the halves keep tracing the old
		// river, so they are hand-placed from the node's point of view.
		var node = new WaterNode { Id = NextId++, Position = first.P3, AutoTangents = false, CornerRadius = FindNode(river.StartNodeId)?.CornerRadius ?? 150.0f };
		Nodes.Add(node);

		int endNodeId = river.EndNodeId;
		float widthAtSplit = new RiverWidth(river.StartWidth, river.EndWidth, table.Length).At(distance);
		float endWidth = river.EndWidth;

		river.EndNodeId = node.Id;
		river.EndWidth = widthAtSplit;
		river.StartTangent = first.P1 - first.P0;
		river.EndTangent = first.P2 - first.P3;

		Rivers.Add(new WaterRiver
		{
			Id = NextId++,
			StartNodeId = node.Id,
			EndNodeId = endNodeId,
			StartWidth = widthAtSplit,
			EndWidth = endWidth,
			Speed = river.Speed,
			Waves = river.Waves,
			Depth = river.Depth,
			StartTangent = second.P1 - second.P0,
			EndTangent = second.P2 - second.P3
		});

		MarkDirty();

		return node;
	}



	/// <summary>
	/// Switches a node between automatic and hand-placed tangents. Going manual straightens every handle at the node
	/// (a sharp bend) unless <paramref name="_Straighten"/> is off, which just freezes the handles as they are.
	/// </summary>
	public void SetAutoTangents(int _NodeId, bool _Auto, bool _Straighten = true)
	{
		var node = FindNode(_NodeId);

		if (node is null)
			return;

		node.AutoTangents = _Auto;

		if (!_Auto && _Straighten)
		{
			foreach (var river in GetRiversAt(_NodeId))
			{
				var other = FindNode(river.OtherNodeId(_NodeId));

				if (other is not null)
					SetHandleInternal(river, _NodeId, StraightHandle(node.Position, other.Position));
			}
		}

		RefreshTangentsAround(_NodeId);
		MarkDirty();
	}



	/// <summary>
	/// The Bezier handle a river has at one of its nodes, pointing from the node into the river.
	/// </summary>
	public static Vector3 GetHandle(WaterRiver _River, int _NodeId)
	{
		return _River.StartNodeId == _NodeId ? _River.StartTangent : _River.EndTangent;
	}



	/// <summary>
	/// Places one handle by hand. With <paramref name="_Mirror"/>, a bend keeps its other handle opposite (its own
	/// length kept) so the river stays smooth through it.
	/// </summary>
	public void SetHandle(int _RiverId, int _NodeId, Vector3 _Handle, bool _Mirror)
	{
		var river = FindRiver(_RiverId);
		var node = FindNode(_NodeId);

		if (river is null || node is null || !river.Touches(_NodeId))
			return;

		node.AutoTangents = false;
		SetHandleInternal(river, _NodeId, _Handle);

		if (_Mirror && !_Handle.IsNearZeroLength)
		{
			var attached = GetRiversAt(_NodeId).ToList();

			if (attached.Count == 2)
			{
				var other = attached[0] == river ? attached[1] : attached[0];
				float otherLength = GetHandle(other, _NodeId).Length;

				SetHandleInternal(other, _NodeId, -_Handle.Normal * otherLength);
			}
		}

		MarkDirty();
	}



	private static void SetHandleInternal(WaterRiver _River, int _NodeId, Vector3 _Handle)
	{
		if (_River.StartNodeId == _NodeId)
			_River.StartTangent = _Handle;
		else
			_River.EndTangent = _Handle;
	}



	/// <summary>
	/// Moves a node. A node a river opens into a lake (or the sea) at stays on that lake's shore (or that coast): it
	/// slides round it instead.
	/// </summary>
	public void MoveNode(int _NodeId, Vector3 _LocalPosition)
	{
		var node = FindNode(_NodeId);

		if (node is null)
			return;

		Vector3 target = _LocalPosition;

		if (LakeAtOpening(_NodeId) is { } lake && GetShore(lake) is { } shore)
		{
			shore.At(shore.NearestDistance(_LocalPosition), out Vector3 onShore, out _, out _);
			target = onShore;
		}

		// A drag that is holding still would otherwise rebuild the pieces every frame over float noise.
		if (Vector3.DistanceBetween(node.Position, target) < 0.01f)
			return;

		node.Position = target;

		RefreshTangentsAround(_NodeId);
		MarkDirty();
	}



	/// <summary>
	/// Removes a river. A node left with nothing attached goes with it (and stops being a lake opening), the others
	/// get their smoothing refreshed.
	/// </summary>
	public bool RemoveRiver(int _RiverId)
	{
		var river = FindRiver(_RiverId);

		if (river is null)
			return false;

		Rivers.Remove(river);
		InvalidateIndex();

		foreach (int nodeId in new[] { river.StartNodeId, river.EndNodeId })
		{
			if (GetDegree(nodeId) == 0)
				RemoveNodeInternal(nodeId);
			else
				RefreshTangentsAround(nodeId);
		}

		MarkDirty();

		return true;
	}



	public bool RemoveNode(int _NodeId)
	{
		if (FindNode(_NodeId) is null)
			return false;

		foreach (var river in GetRiversAt(_NodeId).ToList())
			RemoveRiver(river.Id);

		// RemoveRiver already dropped it once it had no rivers left; this covers a node that never had any.
		RemoveNodeInternal(_NodeId);
		MarkDirty();

		return true;
	}



	private void RemoveNodeInternal(int _NodeId)
	{
		Nodes.RemoveAll(node => node.Id == _NodeId);

		foreach (var lake in Lakes)
			lake.Openings.Remove(_NodeId);

		InvalidateIndex();
	}



	/// <summary>
	/// A node that opens into a lake, or into the sea through a coast, and already has its one river.
	/// </summary>
	private bool IsEndTaken(int _NodeId)
	{
		return LakeAtOpening(_NodeId) is not null && GetDegree(_NodeId) > 0;
	}



	/// <summary>
	/// Takes a node a river simply passes through out, and joins its two rivers into one. The merged river keeps the
	/// handles the two had at their far ends, stretched to the full length (which undoes a split exactly on a straight
	/// river and closely on a curved one). Null when the node is not a plain pass-through (a confluence, a river end, a
	/// lake opening, or two rivers flowing against each other), in which case nothing changes.
	/// </summary>
	public WaterRiver DissolveNode(int _NodeId)
	{
		var node = FindNode(_NodeId);

		if (node is null || LakeAtOpening(_NodeId) is not null)
			return null;

		var attached = GetRiversAt(_NodeId).ToList();

		if (attached.Count != 2)
			return null;

		// The water has to run through the node: one river arriving, the other leaving.
		var upstream = attached.FirstOrDefault(river => river.EndNodeId == _NodeId);
		var downstream = attached.FirstOrDefault(river => river.StartNodeId == _NodeId);

		if (upstream is null || downstream is null || upstream == downstream)
			return null;

		int startId = upstream.StartNodeId;
		int endId = downstream.EndNodeId;

		if (startId == endId || FindRiverBetween(startId, endId) is not null)
			return null;

		float firstLength = GetCurveTable(upstream)?.Length ?? 0.0f;
		float secondLength = GetCurveTable(downstream)?.Length ?? 0.0f;
		float total = firstLength + secondLength;
		var main = secondLength > firstLength ? downstream : upstream;

		var merged = new WaterRiver
		{
			Id = NextId++,
			StartNodeId = startId,
			EndNodeId = endId,
			StartWidth = upstream.StartWidth,
			EndWidth = downstream.EndWidth,
			Speed = main.Speed,
			Waves = main.Waves,
			Depth = main.Depth,
			StartTangent = upstream.StartTangent * (firstLength > 0.0f ? total / firstLength : 1.0f),
			EndTangent = downstream.EndTangent * (secondLength > 0.0f ? total / secondLength : 1.0f)
		};

		Rivers.Remove(upstream);
		Rivers.Remove(downstream);
		Nodes.Remove(node);
		Rivers.Add(merged);
		InvalidateIndex();

		RefreshTangentsAround(startId);
		RefreshTangentsAround(endId);
		MarkDirty();

		return merged;
	}



	/// <summary>
	/// Turns a river round so its water flows the other way: its ends, handles and widths swap.
	/// </summary>
	public void ReverseRiver(int _RiverId)
	{
		var river = FindRiver(_RiverId);

		if (river is null)
			return;

		(river.StartNodeId, river.EndNodeId) = (river.EndNodeId, river.StartNodeId);
		(river.StartTangent, river.EndTangent) = (river.EndTangent, river.StartTangent);
		(river.StartWidth, river.EndWidth) = (river.EndWidth, river.StartWidth);

		InvalidateIndex();
		MarkDirty();
	}



	/// <summary>
	/// Turns round the whole run of rivers a river belongs to: every river joined to it through plain bends, up to the
	/// confluences, lakes and ends either side. That is what reversing "a river" means to whoever drew it as one.
	/// </summary>
	public void ReverseRun(int _RiverId)
	{
		foreach (var river in GetRun(_RiverId))
			ReverseRiver(river.Id);
	}



	/// <summary>
	/// The rivers joined to one through plain bends (nodes with exactly two rivers that are not lake openings).
	/// </summary>
	public List<WaterRiver> GetRun(int _RiverId)
	{
		var run = new List<WaterRiver>();
		var first = FindRiver(_RiverId);

		if (first is null)
			return run;

		var seen = new HashSet<int> { first.Id };
		var queue = new Queue<WaterRiver>();

		queue.Enqueue(first);

		while (queue.Count > 0)
		{
			var river = queue.Dequeue();
			run.Add(river);

			foreach (int nodeId in new[] { river.StartNodeId, river.EndNodeId })
			{
				if (GetDegree(nodeId) != 2 || LakeAtOpening(nodeId) is not null)
					continue;

				foreach (var other in GetRiversAt(nodeId))
				{
					if (seen.Add(other.Id))
						queue.Enqueue(other);
				}
			}
		}

		return run;
	}



	/// <summary>
	/// The curve of a river in network space. A river opening into a lake leaves its shore square on, so the two
	/// surfaces meet cleanly.
	/// </summary>
	public bool TryGetCurve(WaterRiver _River, out WaterCurve _Curve)
	{
		var start = FindNode(_River.StartNodeId);
		var end = FindNode(_River.EndNodeId);

		if (start is null || end is null)
		{
			_Curve = default;
			return false;
		}

		Vector3 startHandle = _River.StartTangent;
		Vector3 endHandle = _River.EndTangent;

		if (TryLakeOutward(_River.StartNodeId, out Vector3 startOutward))
			startHandle = startOutward * MathF.Max(1.0f, startHandle.Length);

		if (TryLakeOutward(_River.EndNodeId, out Vector3 endOutward))
			endHandle = endOutward * MathF.Max(1.0f, endHandle.Length);

		_Curve = new WaterCurve(start.Position, start.Position + startHandle, end.Position + endHandle, end.Position);

		return true;
	}



	private readonly Dictionary<int, WaterCurveTable> m_CurveTables = new();

	/// <summary>
	/// The arc-length table of a river, rebuilt only when its curve actually changed.
	/// </summary>
	public WaterCurveTable GetCurveTable(WaterRiver _River)
	{
		if (_River is null || !TryGetCurve(_River, out var curve))
			return null;

		if (m_CurveTables.TryGetValue(_River.Id, out var table) && table.Curve.Equals(curve))
			return table;

		table = WaterCurveTable.Create(curve);
		m_CurveTables[_River.Id] = table;

		return table;
	}



	private void PruneCurveTables()
	{
		var live = new HashSet<int>(Rivers.Select(river => river.Id));
		var stale = m_CurveTables.Keys.Where(id => !live.Contains(id)).ToList();

		foreach (int id in stale)
			m_CurveTables.Remove(id);
	}



	/// <summary>
	/// The width profile a river is built with, as the last reconcile worked it out (a river through a bend shares its
	/// width there with the next). Its own widths when nothing has been worked out yet.
	/// </summary>
	public RiverWidth WidthOf(WaterRiver _River)
	{
		if (m_Widths.TryGetValue(_River.Id, out var width))
			return width;

		return new RiverWidth(_River.StartWidth, _River.EndWidth, GetCurveTable(_River)?.Length ?? 1.0f);
	}
}