dm: core: Add late driver remove option

Add another flag to the DM core which could be assigned to drivers and
which makes those drivers call their remove callbacks last, just before
booting OS and after all the other drivers finished with their remove
callbacks. This is necessary for things like clock drivers, where the
other drivers might depend on the clock driver in their remove callbacks.
Prime example is the mmc subsystem, which can reconfigure a card from HS
mode to slower modes in the remove callback and for that it needs to
reconfigure the controller clock.

Signed-off-by: Marek Vasut <marek.vasut+renesas@gmail.com>
Signed-off-by: Simon Glass <sjg@chromium.org>
diff --git a/test/dm/core.c b/test/dm/core.c
index 1f5ca57..bfd6565 100644
--- a/test/dm/core.c
+++ b/test/dm/core.c
@@ -72,6 +72,14 @@
 	.name = "test_act_dma_drv",
 };
 
+static struct driver_info driver_info_vital_clk = {
+	.name = "test_vital_clk_drv",
+};
+
+static struct driver_info driver_info_act_dma_vital_clk = {
+	.name = "test_act_dma_vital_clk_drv",
+};
+
 void dm_leak_check_start(struct unit_test_state *uts)
 {
 	uts->start = mallinfo();
@@ -883,6 +891,92 @@
 }
 DM_TEST(dm_test_remove_active_dma, 0);
 
+/* Test removal of 'vital' devices */
+static int dm_test_remove_vital(struct unit_test_state *uts)
+{
+	struct dm_test_state *dms = uts->priv;
+	struct udevice *normal, *dma, *vital, *dma_vital;
+
+	/* Skip the behaviour in test_post_probe() */
+	dms->skip_post_probe = 1;
+
+	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_manual,
+					&normal));
+	ut_assertnonnull(normal);
+
+	ut_assertok(device_bind_by_name(dms->root, false, &driver_info_act_dma,
+					&dma));
+	ut_assertnonnull(dma);
+
+	ut_assertok(device_bind_by_name(dms->root, false,
+					&driver_info_vital_clk, &vital));
+	ut_assertnonnull(vital);
+
+	ut_assertok(device_bind_by_name(dms->root, false,
+					&driver_info_act_dma_vital_clk,
+					&dma_vital));
+	ut_assertnonnull(dma_vital);
+
+	/* Probe the devices */
+	ut_assertok(device_probe(normal));
+	ut_assertok(device_probe(dma));
+	ut_assertok(device_probe(vital));
+	ut_assertok(device_probe(dma_vital));
+
+	/* Check that devices are active right now */
+	ut_asserteq(true, device_active(normal));
+	ut_asserteq(true, device_active(dma));
+	ut_asserteq(true, device_active(vital));
+	ut_asserteq(true, device_active(dma_vital));
+
+	/* Remove active devices via selective remove flag */
+	dm_remove_devices_flags(DM_REMOVE_NON_VITAL | DM_REMOVE_ACTIVE_ALL);
+
+	/*
+	 * Check that this only has an effect on the dma device, since two
+	 * devices are vital and the third does not have active DMA
+	 */
+	ut_asserteq(true, device_active(normal));
+	ut_asserteq(false, device_active(dma));
+	ut_asserteq(true, device_active(vital));
+	ut_asserteq(true, device_active(dma_vital));
+
+	/* Remove active devices via selective remove flag */
+	ut_assertok(device_probe(dma));
+	dm_remove_devices_flags(DM_REMOVE_ACTIVE_ALL);
+
+	/* This should have affected both active-dma devices */
+	ut_asserteq(true, device_active(normal));
+	ut_asserteq(false, device_active(dma));
+	ut_asserteq(true, device_active(vital));
+	ut_asserteq(false, device_active(dma_vital));
+
+	/* Remove non-vital devices */
+	ut_assertok(device_probe(dma));
+	ut_assertok(device_probe(dma_vital));
+	dm_remove_devices_flags(DM_REMOVE_NON_VITAL);
+
+	/* This should have affected only non-vital devices */
+	ut_asserteq(false, device_active(normal));
+	ut_asserteq(false, device_active(dma));
+	ut_asserteq(true, device_active(vital));
+	ut_asserteq(true, device_active(dma_vital));
+
+	/* Remove vital devices via normal remove flag */
+	ut_assertok(device_probe(normal));
+	ut_assertok(device_probe(dma));
+	dm_remove_devices_flags(DM_REMOVE_NORMAL);
+
+	/* Check that all devices are inactive right now */
+	ut_asserteq(false, device_active(normal));
+	ut_asserteq(false, device_active(dma));
+	ut_asserteq(false, device_active(vital));
+	ut_asserteq(false, device_active(dma_vital));
+
+	return 0;
+}
+DM_TEST(dm_test_remove_vital, 0);
+
 static int dm_test_uclass_before_ready(struct unit_test_state *uts)
 {
 	struct uclass *uc;
diff --git a/test/dm/test-driver.c b/test/dm/test-driver.c
index a67f5d3..ca7626a 100644
--- a/test/dm/test-driver.c
+++ b/test/dm/test-driver.c
@@ -170,3 +170,25 @@
 	.unbind	= test_manual_unbind,
 	.flags	= DM_FLAG_ACTIVE_DMA,
 };
+
+U_BOOT_DRIVER(test_vital_clk_drv) = {
+	.name	= "test_vital_clk_drv",
+	.id	= UCLASS_TEST,
+	.ops	= &test_manual_ops,
+	.bind	= test_manual_bind,
+	.probe	= test_manual_probe,
+	.remove	= test_manual_remove,
+	.unbind	= test_manual_unbind,
+	.flags	= DM_FLAG_VITAL,
+};
+
+U_BOOT_DRIVER(test_act_dma_vital_clk_drv) = {
+	.name	= "test_act_dma_vital_clk_drv",
+	.id	= UCLASS_TEST,
+	.ops	= &test_manual_ops,
+	.bind	= test_manual_bind,
+	.probe	= test_manual_probe,
+	.remove	= test_manual_remove,
+	.unbind	= test_manual_unbind,
+	.flags	= DM_FLAG_VITAL | DM_FLAG_ACTIVE_DMA,
+};