Fix legacy DPMS changes with MST

index 7780567aa669..58b9554711c7 100644
@@ -3129,6 +3129,82 @@ static const struct drm_private_state_funcs mst_state_funcs = {
 	.atomic_destroy_state = drm_dp_mst_destroy_state,
 };
 
+static bool
+drm_dp_mst_connector_still_exists(struct drm_connector *connector,
+				  struct drm_dp_mst_topology_mgr *mgr,
+				  struct drm_dp_mst_branch *mstb)
+{
+	struct drm_dp_mst_port *port, *tmp;
+	bool exists = false;
+
+	mstb = drm_dp_get_validated_mstb_ref(mgr, mstb);
+	if (!mstb)
+		return false;
+
+	list_for_each_entry_safe(port, tmp, &mstb->ports, next) {
+		port = drm_dp_get_validated_port_ref(mgr, port);
+		if (!port)
+			continue;
+
+		exists = (port->connector == connector ||
+			  (port->mstb &&
+			   drm_dp_mst_connector_still_exists(connector, mgr,
+							     port->mstb)));
+
+		drm_dp_put_port(port);
+		if (exists)
+			break;
+	}
+
+	drm_dp_put_mst_branch_device(mstb);
+	return exists;
+}
+
+/**
+ * drm_dp_mst_connector_atomic_check - Helper for validating a new atomic
+ *                                     state on an MST connector
+ * @connector: drm connector
+ * @connector_state: the new atomic state of @connector
+ * @mgr: the MST topology mgr for @connector
+ *
+ * This function performs various atomic checks that apply to all drivers
+ * using the DRM DP MST helpers. This should be called by all drivers at the
+ * start of the atomic_check function for their MST connectors.
+ *
+ * Return 0 for success, or negative error code on failure.
+ */
+int
+drm_dp_mst_connector_atomic_check(struct drm_connector *connector,
+				  struct drm_connector_state *connector_state,
+				  struct drm_dp_mst_topology_mgr *mgr)
+{
+	struct drm_atomic_state *state = connector_state->state;
+	struct drm_crtc *crtc = connector_state->crtc;
+	struct drm_crtc_state *new_crtc_state;
+
+	if (!crtc)
+		return 0;
+
+	new_crtc_state = drm_atomic_get_new_crtc_state(state, crtc);
+	if (!new_crtc_state)
+		return 0;
+
+	if (!drm_atomic_crtc_needs_modeset(new_crtc_state) ||
+	    !new_crtc_state->active)
+		return 0;
+
+	/* Make sure that the port for this MST connector still exists */
+	if (!drm_dp_mst_connector_still_exists(connector, mgr,
+					       mgr->mst_primary)) {
+		DRM_DEBUG_ATOMIC("[CONNECTOR:%d:%s] has disappeared from the MST topology\n",
+				 connector->base.id, connector->name);
+		return -EINVAL;
+	}
+
+	return 0;
+}
+EXPORT_SYMBOL(drm_dp_mst_connector_atomic_check);
+
 /**
  * drm_atomic_get_mst_topology_state: get MST topology state
  *

index 7f78d26a0766..8e33c2c85d1e 100644
@@ -625,6 +625,9 @@ void drm_dp_mst_topology_mgr_suspend(struct drm_dp_mst_topology_mgr *mgr);
 int drm_dp_mst_topology_mgr_resume(struct drm_dp_mst_topology_mgr *mgr);
 struct drm_dp_mst_topology_state *drm_atomic_get_mst_topology_state(struct drm_atomic_state *state,
 								    struct drm_dp_mst_topology_mgr *mgr);
+int drm_dp_mst_connector_atomic_check(struct drm_connector *connector,
+				      struct drm_connector_state *connector_state,
+				      struct drm_dp_mst_topology_mgr *mgr);
 int drm_dp_atomic_find_vcpi_slots(struct drm_atomic_state *state,
 				  struct drm_dp_mst_topology_mgr *mgr,
 				  struct drm_dp_mst_port *port, int pbn);

index 9da0bdfe1e1c..8d6f6ee9bc75 100644
@@ -881,22 +881,16 @@ nv50_mstc_atomic_best_encoder(struct drm_connector *connector,
 {
 	struct nv50_head *head = nv50_head(connector_state->crtc);
 	struct nv50_mstc *mstc = nv50_mstc(connector);
-	if (mstc->port) {
-		struct nv50_mstm *mstm = mstc->mstm;
-		return &mstm->msto[head->base.index]->encoder;
-	}
-	return NULL;
+
+	return &mstc->mstm->msto[head->base.index]->encoder;
 }
 
 static struct drm_encoder *
 nv50_mstc_best_encoder(struct drm_connector *connector)
 {
 	struct nv50_mstc *mstc = nv50_mstc(connector);
-	if (mstc->port) {
-		struct nv50_mstm *mstm = mstc->mstm;
-		return &mstm->msto[0]->encoder;
-	}
-	return NULL;
+
+	return &mstc->mstm->msto[0]->encoder;
 }
 
 static enum drm_mode_status

@@ -926,10 +920,21 @@ nv50_mstc_get_modes(struct drm_connector *connector)
 	return ret;
 }
 
+static int
+nv50_mstc_atomic_check(struct drm_connector *connector,
+		       struct drm_connector_state *state)
+{
+	struct nv50_mstc *mstc = nv50_mstc(connector);
+
+	return drm_dp_mst_connector_atomic_check(connector, state,
+						 &mstc->mstm->mgr);
+}
+
 static const struct drm_connector_helper_funcs
 nv50_mstc_help = {
 	.get_modes = nv50_mstc_get_modes,
 	.mode_valid = nv50_mstc_mode_valid,
+	.atomic_check = nv50_mstc_atomic_check,
 	.best_encoder = nv50_mstc_best_encoder,
 	.atomic_best_encoder = nv50_mstc_atomic_best_encoder,
 };

index 4ecd65375603..fcb9b87b9339 100644
@@ -311,9 +311,8 @@ static int intel_dp_mst_get_ddc_modes(struct drm_connector *connector)
 	struct edid *edid;
 	int ret;
 
-	if (!intel_dp) {
+	if (intel_connector->mst_port_gone)
 		return intel_connector_update_modes(connector, NULL);
-	}
 
 	edid = drm_dp_mst_get_edid(connector, &intel_dp->mst_mgr, intel_connector->port);
 	ret = intel_connector_update_modes(connector, edid);

@@ -328,9 +327,10 @@ intel_dp_mst_detect(struct drm_connector *connector, bool force)
 	struct intel_connector *intel_connector = to_intel_connector(connector);
 	struct intel_dp *intel_dp = intel_connector->mst_port;
 
-	if (!intel_dp)
+	if (intel_connector->mst_port_gone)
 		return connector_status_disconnected;
-	return drm_dp_mst_detect_port(connector, &intel_dp->mst_mgr, intel_connector->port);
+	return drm_dp_mst_detect_port(connector, &intel_dp->mst_mgr,
+				      intel_connector->port);
 }
 
 static void

@@ -370,7 +370,7 @@ intel_dp_mst_mode_valid(struct drm_connector *connector,
 	int bpp = 24; /* MST uses fixed bpp */
 	int max_rate, mode_rate, max_lanes, max_link_clock;
 
-	if (!intel_dp)
+	if (intel_connector->mst_port_gone)
 		return MODE_ERROR;
 
 	if (mode->flags & DRM_MODE_FLAG_DBLSCAN)

@@ -402,7 +402,7 @@ static struct drm_encoder *intel_mst_atomic_best_encoder(struct drm_connector *c
 	struct intel_dp *intel_dp = intel_connector->mst_port;
 	struct intel_crtc *crtc = to_intel_crtc(state->crtc);
 
-	if (!intel_dp)
+	if (intel_connector->mst_port_gone)
 		return NULL;
 	return &intel_dp->mst_encoders[crtc->pipe]->base.base;
 }

@@ -514,7 +514,7 @@ static void intel_dp_destroy_mst_connector(struct drm_dp_mst_topology_mgr *mgr,
 						   connector);
 	/* prevent race with the check in ->detect */
 	drm_modeset_lock(&connector->dev->mode_config.connection_mutex, NULL);
-	intel_connector->mst_port = NULL;
+	intel_connector->mst_port_gone = true;
 	drm_modeset_unlock(&connector->dev->mode_config.connection_mutex);
 
 	drm_connector_put(connector);

index 8fc61e96754f..87ce772ae7f8 100644
@@ -409,6 +409,7 @@ struct intel_connector {
 	void *port; /* store this opaque as its illegal to dereference it */
 
 	struct intel_dp *mst_port;
+	bool mst_port_gone;
 
 	/* Work struct to schedule a uevent on link train failure */
 	struct work_struct modeset_retry_work;

index fcb9b87b9339..a366f32b048a 100644
@@ -42,7 +42,7 @@ static bool intel_dp_mst_compute_config(struct intel_encoder *encoder,
 		to_intel_connector(conn_state->connector);
 	struct drm_atomic_state *state = pipe_config->base.state;
 	int bpp;
-	int lane_count, slots;
+	int lane_count, slots = 0;
 	const struct drm_display_mode *adjusted_mode = &pipe_config->base.adjusted_mode;
 	int mst_pbn;
 	bool reduce_m_n = drm_dp_has_quirk(&intel_dp->desc,

@@ -76,11 +76,16 @@ static bool intel_dp_mst_compute_config(struct intel_encoder *encoder,
 	mst_pbn = drm_dp_calc_pbn_mode(adjusted_mode->crtc_clock, bpp);
 	pipe_config->pbn = mst_pbn;
 
-	slots = drm_dp_atomic_find_vcpi_slots(state, &intel_dp->mst_mgr,
-					      connector->port, mst_pbn);
-	if (slots < 0) {
-		DRM_DEBUG_KMS("failed finding vcpi slots:%d\n", slots);
-		return false;
+	if (!connector->mst_port_gone) {
+		slots = drm_dp_atomic_find_vcpi_slots(state,
+						      &intel_dp->mst_mgr,
+						      connector->port,
+						      mst_pbn);
+		if (slots < 0) {
+			DRM_DEBUG_KMS("failed finding vcpi slots:%d\n",
+				      slots);
+			return false;
+		}
 	}
 
 	intel_link_compute_m_n(bpp, lane_count,

index a366f32b048a..2b798d4592f0 100644
@@ -106,14 +106,21 @@ static bool intel_dp_mst_compute_config(struct intel_encoder *encoder,
 }
 
 static int intel_dp_mst_atomic_check(struct drm_connector *connector,
-		struct drm_connector_state *new_conn_state)
+				     struct drm_connector_state *new_conn_state)
 {
 	struct drm_atomic_state *state = new_conn_state->state;
 	struct drm_connector_state *old_conn_state;
 	struct drm_crtc *old_crtc;
 	struct drm_crtc_state *crtc_state;
+	struct drm_dp_mst_topology_mgr *mgr =
+		&to_intel_connector(connector)->mst_port->mst_mgr;
 	int slots, ret = 0;
 
+	ret = drm_dp_mst_connector_atomic_check(connector, new_conn_state,
+						mgr);
+	if (ret)
+		return ret;
+
 	old_conn_state = drm_atomic_get_old_connector_state(state, connector);
 	old_crtc = old_conn_state->crtc;
 	if (!old_crtc)

@@ -122,12 +129,6 @@ static int intel_dp_mst_atomic_check(struct drm_connector *connector,
 	crtc_state = drm_atomic_get_new_crtc_state(state, old_crtc);
 	slots = to_intel_crtc_state(crtc_state)->dp_m_n.tu;
 	if (drm_atomic_crtc_needs_modeset(crtc_state) && slots > 0) {
-		struct drm_dp_mst_topology_mgr *mgr;
-		struct drm_encoder *old_encoder;
-
-		old_encoder = old_conn_state->best_encoder;
-		mgr = &enc_to_mst(old_encoder)->primary->dp.mst_mgr;
-
 		ret = drm_dp_atomic_release_vcpi_slots(state, mgr, slots);
 		if (ret)
 			DRM_DEBUG_KMS("failed releasing %d vcpi slots:%d\n", slots, ret);

@@ -407,8 +408,6 @@ static struct drm_encoder *intel_mst_atomic_best_encoder(struct drm_connector *c
 	struct intel_dp *intel_dp = intel_connector->mst_port;
 	struct intel_crtc *crtc = to_intel_crtc(state->crtc);
 
-	if (intel_connector->mst_port_gone)
-		return NULL;
 	return &intel_dp->mst_encoders[crtc->pipe]->base.base;
 }
 

index 9a300732ba37..d011a39f17b2 100644
@@ -294,10 +294,22 @@ static struct drm_encoder *dm_mst_best_encoder(struct drm_connector *connector)
 	return &amdgpu_dm_connector->mst_encoder->base;
 }
 
+static int
+amdgpu_dm_mst_connector_atomic_check(struct drm_connector *connector,
+				     struct drm_connector_state *new_cstate)
+{
+	struct amdgpu_dm_connector *aconnector =
+		to_amdgpu_dm_connector(connector);
+
+	return drm_dp_mst_connector_atomic_check(connector, new_cstate,
+						 &aconnector->mst_mgr);
+}
+
 static const struct drm_connector_helper_funcs dm_dp_mst_connector_helper_funcs = {
 	.get_modes = dm_dp_mst_get_modes,
 	.mode_valid = amdgpu_dm_connector_mode_valid,
 	.best_encoder = dm_mst_best_encoder,
+	.atomic_check = amdgpu_dm_mst_connector_atomic_check,
 };
 
 static void amdgpu_dm_encoder_destroy(struct drm_encoder *encoder)