drm/amd/display: Fix up plane_states add/remove logic

Our plane_states array trimming logic was faulty, we should be starting
to shuffle from the plane that was just released, not from the very
beginning of the array.

Also get rid of a leftover line that was setting the plane state at the
stream index to null, leading to issues. Also move the
dc_plane_state_retain call to where we reference plane_state, in case we
do hit the error case where we can't get a free pipe.

Signed-off-by: Andrew Jiang <Andrew.Jiang@amd.com>
Reviewed-by: Tony Cheng <Tony.Cheng@amd.com>
Acked-by: Harry Wentland <Harry.Wentland@amd.com>
Signed-off-by: Alex Deucher <alexander.deucher@amd.com>

drivers/gpu/drm/amd/display/dc/core/dc_resource.c

@@ -1073,9 +1073,6 @@ bool dc_add_plane_to_context(
 		return false;
 	}
 
-	/* retain new surfaces */
-	dc_plane_state_retain(plane_state);
-
 	free_pipe = acquire_free_pipe_for_stream(context, pool, stream);
 
 #if defined(CONFIG_DRM_AMD_DC_DCN1_0)
@@ -1085,11 +1082,11 @@ bool dc_add_plane_to_context(
 			free_pipe = &context->res_ctx.pipe_ctx[pipe_idx];
 	}
 #endif
-	if (!free_pipe) {
-		stream_status->plane_states[i] = NULL;
+	if (!free_pipe)
 		return false;
-	}
 
+	/* retain new surfaces */
+	dc_plane_state_retain(plane_state);
 	free_pipe->plane_state = plane_state;
 
 	if (head_pipe != free_pipe) {
@@ -1181,8 +1178,8 @@ bool dc_remove_plane_from_context(
 
 	stream_status->plane_count--;
 
-	/* Trim back arrays */
-	for (i = 0; i < stream_status->plane_count; i++)
+	/* Start at the plane we've just released, and move all the planes one index forward to "trim" the array */
+	for (; i < stream_status->plane_count; i++)
 		stream_status->plane_states[i] = stream_status->plane_states[i + 1];
 
 	stream_status->plane_states[stream_status->plane_count] = NULL;