Protecting the render tree

src/bunit.core/Extensions/WaitForHelpers/WaitForHelperLoggerExtensions.cs

@@ -5,22 +5,22 @@ namespace Bunit.Extensions.WaitForHelpers;
 internal static class WaitForHelperLoggerExtensions
 {
 	private static readonly Action<ILogger, int, Exception?> CheckingWaitCondition
-		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(1, "OnAfterRender"), "Checking the wait condition for component {Id}.");
+		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(1, "CheckingWaitCondition"), "Checking the wait condition for component {Id}.");
 
 	private static readonly Action<ILogger, int, Exception?> CheckCompleted
-		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(2, "OnAfterRender"), "The check completed successfully for component {Id}.");
+		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(2, "CheckCompleted"), "The check completed successfully for component {Id}.");
 
 	private static readonly Action<ILogger, int, Exception?> CheckFailed
-		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(3, "OnAfterRender"), "The check failed for component {Id}.");
+		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(3, "CheckFailed"), "The check failed for component {Id}.");
 
 	private static readonly Action<ILogger, int, Exception> CheckThrow
-		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(4, "OnAfterRender"), "The checker for component {Id} throw an exception.");
+		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(4, "CheckThrow"), "The checker for component {Id} throw an exception.");
 
 	private static readonly Action<ILogger, int, Exception?> WaiterTimedOut
-		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(10, "OnTimeout"), "The waiter for component {Id} timed out.");
+		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(10, "WaiterTimedOut"), "The waiter for component {Id} timed out.");
 
 	private static readonly Action<ILogger, int, Exception?> WaiterDisposed
-		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(20, "OnTimeout"), "The waiter for component {Id} disposed.");
+		= LoggerMessage.Define<int>(LogLevel.Debug, new EventId(20, "WaiterDisposed"), "The waiter for component {Id} disposed.");
 
 	internal static void LogCheckingWaitCondition<T>(this ILogger<WaitForHelper<T>> logger, int componentId)
 	{

src/bunit.core/Rendering/RenderEvent.cs

@@ -5,24 +5,15 @@ namespace Bunit.Rendering;
 /// </summary>
 public sealed class RenderEvent
 {
-	private readonly RenderBatch renderBatch;
+	private readonly Dictionary<int, Status> statuses = new();
+
+	internal IReadOnlyDictionary<int, Status> Statuses => statuses;
 
 	/// <summary>
 	/// Gets a collection of <see cref="ArrayRange{RenderTreeFrame}"/>, accessible via the ID
 	/// of the component they are created by.
 	/// </summary>
-	public RenderTreeFrameDictionary Frames { get; }
-
-	/// <summary>
-	/// Initializes a new instance of the <see cref="RenderEvent"/> class.
-	/// </summary>
-	/// <param name="renderBatch">The <see cref="RenderBatch"/> update from the render event.</param>
-	/// <param name="frames">The <see cref="RenderTreeFrameDictionary"/> from the current render.</param>
-	internal RenderEvent(RenderBatch renderBatch, RenderTreeFrameDictionary frames)
-	{
-		this.renderBatch = renderBatch;
-		Frames = frames;
-	}
+	public RenderTreeFrameDictionary Frames { get; } = new();
 
 	/// <summary>
 	/// Gets the render status for a <paramref name="renderedComponent"/>.
@@ -34,86 +25,57 @@ internal RenderEvent(RenderBatch renderBatch, RenderTreeFrameDictionary frames)
 		if (renderedComponent is null)
 			throw new ArgumentNullException(nameof(renderedComponent));
 
-		var result = (Rendered: false, Changed: false, Disposed: false);
+		return statuses.TryGetValue(renderedComponent.ComponentId, out var status)
+			? (status.Rendered, status.Changed, status.Disposed)
+			: (Rendered: false, Changed: false, Disposed: false);
+	}
 
-		if (DidComponentDispose(renderedComponent))
-		{
-			result.Disposed = true;
-		}
-		else
+	internal Status GetOrCreateStatus(int componentId)
+	{
+		if (!statuses.TryGetValue(componentId, out var status))
 		{
-			(result.Rendered, result.Changed) = GetRenderAndChangeStatus(renderedComponent);
+			status = new();
+			statuses[componentId] = status;
 		}
+		return status;
+	}
 
-		return result;
+	internal void SetDisposed(int componentId)
+	{
+		GetOrCreateStatus(componentId).Disposed = true;
 	}
 
-	private bool DidComponentDispose(IRenderedFragmentBase renderedComponent)
+	internal void SetUpdated(int componentId, bool hasChanges)
 	{
-		for (var i = 0; i < renderBatch.DisposedComponentIDs.Count; i++)
-		{
-			if (renderBatch.DisposedComponentIDs.Array[i].Equals(renderedComponent.ComponentId))
-			{
-				return true;
-			}
-		}
+		var status = GetOrCreateStatus(componentId);
+		status.Rendered = true;
+		status.Changed = hasChanges;
+	}
 
-		return false;
+	internal void SetUpdatedApplied(int componentId)
+	{
+		GetOrCreateStatus(componentId).UpdatesApplied = true;
 	}
 
-	/// <summary>
-	/// This method determines if the <paramref name="renderedComponent"/> or any of the
-	/// components underneath it in the render tree rendered and whether they they changed
-	/// their render tree during render.
-	///
-	/// It does this by getting the status from the <paramref name="renderedComponent"/>,
-	/// then from all its children, using a recursive pattern, where the internal methods
-	/// GetStatus and GetStatusFromChildren call each other until there are no more children,
-	/// or both a render and a change is found.
-	/// </summary>
-	private (bool Rendered, bool HasChanges) GetRenderAndChangeStatus(IRenderedFragmentBase renderedComponent)
+	internal void AddFrames(int componentId, ArrayRange<RenderTreeFrame> frames)
 	{
-		var result = (Rendered: false, HasChanges: false);
+		Frames.Add(componentId, frames);
+		GetOrCreateStatus(componentId).FramesLoaded = true;
+	}
 
-		GetStatus(renderedComponent.ComponentId);
+	internal record class Status
+	{
+		public bool Rendered { get; set; }
 
-		return result;
+		public bool Changed { get; set; }
 
-		void GetStatus(int componentId)
-		{
-			for (var i = 0; i < renderBatch.UpdatedComponents.Count; i++)
-			{
-				ref var update = ref renderBatch.UpdatedComponents.Array[i];
-				if (update.ComponentId == componentId)
-				{
-					result.Rendered = true;
-					result.HasChanges = update.Edits.Count > 0;
-					break;
-				}
-			}
-
-			if (!result.HasChanges)
-			{
-				GetStatusFromChildren(componentId);
-			}
-		}
+		public bool Disposed { get; set; }
 
-		void GetStatusFromChildren(int componentId)
-		{
-			var frames = Frames[componentId];
-			for (var i = 0; i < frames.Count; i++)
-			{
-				ref var frame = ref frames.Array[i];
-				if (frame.FrameType == RenderTreeFrameType.Component)
-				{
-					GetStatus(frame.ComponentId);
-
-					if (result.HasChanges)
-					{
-						break;
-					}
-				}
-			}
-		}
+		public bool UpdatesApplied { get; set; }
+
+		public bool FramesLoaded { get; set; }
+
+		public bool UpdateNeeded => Rendered || Changed;
 	}
 }
+