[PATCH V2] dm: core: Add late driver remove option

[Marek Vasut <marek.vasut@gmail.com>]

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>
Cc: Simon Glass <sjg@chromium.org>
Cc: Tom Rini <trini@konsulko.com>
---
V2: Fix DM tests
---
 arch/arm/lib/bootm.c         |  1 +
 drivers/core/device-remove.c | 11 ++++++++---
 drivers/core/root.c          |  2 ++
 drivers/core/uclass.c        | 32 +++++++++++++++++++++++++-------
 include/dm/device.h          |  4 ++++
 include/dm/uclass-internal.h |  3 ++-
 test/dm/core.c               | 21 ++++++++++++---------
 test/dm/test-main.c          | 30 +++++++++++++++++-------------
 8 files changed, 71 insertions(+), 33 deletions(-)

diff --git a/arch/arm/lib/bootm.c b/arch/arm/lib/bootm.c
index 1206e306db..f9091a3d41 100644
--- a/arch/arm/lib/bootm.c
+++ b/arch/arm/lib/bootm.c
@@ -120,6 +120,7 @@ static void announce_and_cleanup(int fake)
 	 * of DMA operation or releasing device internal buffers.
 	 */
 	dm_remove_devices_flags(DM_REMOVE_ACTIVE_ALL);
+	dm_remove_devices_flags(DM_REMOVE_ACTIVE_ALL | DM_REMOVE_LATE);
 
 	cleanup_before_linux();
 }
diff --git a/drivers/core/device-remove.c b/drivers/core/device-remove.c
index efdb0f2905..07b241b6bb 100644
--- a/drivers/core/device-remove.c
+++ b/drivers/core/device-remove.c
@@ -172,14 +172,19 @@ int device_remove(struct udevice *dev, uint flags)
 	drv = dev->driver;
 	assert(drv);
 
-	ret = uclass_pre_remove_device(dev);
-	if (ret)
-		return ret;
+	if (!(flags & DM_REMOVE_LATE)) {
+		ret = uclass_pre_remove_device(dev);
+		if (ret)
+			return ret;
+	}
 
 	ret = device_chld_remove(dev, NULL, flags);
 	if (ret)
 		goto err;
 
+	if (!(flags & DM_REMOVE_LATE) && (drv->flags & DM_FLAG_REMOVE_LATE))
+		return 0;
+
 	/*
 	 * Remove the device if called with the "normal" remove flag set,
 	 * or if the remove flag matches any of the drivers remove flags
diff --git a/drivers/core/root.c b/drivers/core/root.c
index 0726be6b79..21f3054125 100644
--- a/drivers/core/root.c
+++ b/drivers/core/root.c
@@ -168,6 +168,7 @@ int dm_init(bool of_live)
 int dm_uninit(void)
 {
 	device_remove(dm_root(), DM_REMOVE_NORMAL);
+	device_remove(dm_root(), DM_REMOVE_LATE);
 	device_unbind(dm_root());
 	gd->dm_root = NULL;
 
@@ -393,6 +394,7 @@ struct acpi_ops root_acpi_ops = {
 U_BOOT_DRIVER(root_driver) = {
 	.name	= "root_driver",
 	.id	= UCLASS_ROOT,
+	.flags	= DM_FLAG_REMOVE_LATE,
 	ACPI_OPS_PTR(&root_acpi_ops)
 };
 
diff --git a/drivers/core/uclass.c b/drivers/core/uclass.c
index c3f1b73cd6..ac474d3ff8 100644
--- a/drivers/core/uclass.c
+++ b/drivers/core/uclass.c
@@ -104,10 +104,28 @@ fail_mem:
 	return ret;
 }
 
-int uclass_destroy(struct uclass *uc)
+int uclass_find_device_by_drv_flag(struct uclass *uc, const unsigned int flag,
+				   const bool neg, struct udevice **devp)
+{
+	struct udevice *dev;
+
+	*devp = NULL;
+	uclass_foreach_dev(dev, uc) {
+		if ((neg && (dev->driver->flags & flag)) ||
+		    (!neg && !(dev->driver->flags & flag))) {
+			*devp = dev;
+			return 0;
+		}
+	}
+
+	return -ENODEV;
+}
+
+int uclass_destroy(struct uclass *uc, unsigned int flag)
 {
 	struct uclass_driver *uc_drv;
 	struct udevice *dev;
+	bool late = flag & DM_REMOVE_LATE;
 	int ret;
 
 	/*
@@ -116,15 +134,15 @@ int uclass_destroy(struct uclass *uc)
 	 * unbound (by the recursion in the call to device_unbind() below).
 	 * We can loop until the list is empty.
 	 */
-	while (!list_empty(&uc->dev_head)) {
-		dev = list_first_entry(&uc->dev_head, struct udevice,
-				       uclass_node);
-		ret = device_remove(dev, DM_REMOVE_NORMAL | DM_REMOVE_NO_PD);
-		if (ret)
+	while (!uclass_find_device_by_drv_flag(uc, DM_FLAG_REMOVE_LATE, late, &dev)) {
+		ret = device_remove(dev, flag | DM_REMOVE_NO_PD);
+		if (ret) {
 			return log_msg_ret("remove", ret);
+		}
 		ret = device_unbind(dev);
-		if (ret)
+		if (ret) {
 			return log_msg_ret("unbind", ret);
+		}
 	}
 
 	uc_drv = uc->uc_drv;
diff --git a/include/dm/device.h b/include/dm/device.h
index 953706cf52..7b1db252bf 100644
--- a/include/dm/device.h
+++ b/include/dm/device.h
@@ -73,6 +73,8 @@ struct driver_info;
  */
 #define DM_FLAG_REMOVE_WITH_PD_ON	(1 << 13)
 
+#define DM_FLAG_REMOVE_LATE		(1 << 14)
+
 /*
  * One or multiple of these flags are passed to device_remove() so that
  * a selective device removal as specified by the remove-stage and the
@@ -95,6 +97,8 @@ enum {
 
 	/* Don't power down any attached power domains */
 	DM_REMOVE_NO_PD		= 1 << 1,
+
+	DM_REMOVE_LATE		= 1 << 2,
 };
 
 /**
diff --git a/include/dm/uclass-internal.h b/include/dm/uclass-internal.h
index 6e3f15c2b0..b5926b0f5c 100644
--- a/include/dm/uclass-internal.h
+++ b/include/dm/uclass-internal.h
@@ -247,8 +247,9 @@ struct uclass *uclass_find(enum uclass_id key);
  * Destroy a uclass and all its devices
  *
  * @uc: uclass to destroy
+ * @flag: driver flags (DM_REMOVE_NORMAL or DM_REMOVE_LATE)
  * @return 0 on success, -ve on error
  */
-int uclass_destroy(struct uclass *uc);
+int uclass_destroy(struct uclass *uc, unsigned int flag);
 
 #endif
diff --git a/test/dm/core.c b/test/dm/core.c
index d20c48443f..0000b10c5e 100644
--- a/test/dm/core.c
+++ b/test/dm/core.c
@@ -81,16 +81,19 @@ void dm_leak_check_start(struct unit_test_state *uts)
 int dm_leak_check_end(struct unit_test_state *uts)
 {
 	struct mallinfo end;
-	int id, diff;
+	int i, id, diff;
 
 	/* Don't delete the root class, since we started with that */
-	for (id = UCLASS_ROOT + 1; id < UCLASS_COUNT; id++) {
-		struct uclass *uc;
-
-		uc = uclass_find(id);
-		if (!uc)
-			continue;
-		ut_assertok(uclass_destroy(uc));
+	for (i = 0; i < 2; i++) {
+		for (id = UCLASS_ROOT + 1; id < UCLASS_COUNT; id++) {
+			struct uclass *uc;
+
+			uc = uclass_find(id);
+			if (!uc)
+				continue;
+			ut_assertok(uclass_destroy(uc,
+				    i ? DM_REMOVE_LATE : DM_REMOVE_NORMAL));
+		}
 	}
 
 	end = mallinfo();
@@ -513,7 +516,7 @@ static int dm_test_uclass(struct unit_test_state *uts)
 	ut_asserteq(0, dm_testdrv_op_count[DM_TEST_OP_DESTROY]);
 	ut_assert(uc->priv);
 
-	ut_assertok(uclass_destroy(uc));
+	ut_assertok(uclass_destroy(uc, DM_REMOVE_LATE));
 	ut_asserteq(1, dm_testdrv_op_count[DM_TEST_OP_INIT]);
 	ut_asserteq(1, dm_testdrv_op_count[DM_TEST_OP_DESTROY]);
 
diff --git a/test/dm/test-main.c b/test/dm/test-main.c
index 53e5ca321f..2002480bf5 100644
--- a/test/dm/test-main.c
+++ b/test/dm/test-main.c
@@ -58,19 +58,23 @@ static int do_autoprobe(struct unit_test_state *uts)
 
 static int dm_test_destroy(struct unit_test_state *uts)
 {
-	int id;
-
-	for (id = 0; id < UCLASS_COUNT; id++) {
-		struct uclass *uc;
-
-		/*
-		 * If the uclass doesn't exist we don't want to create it. So
-		 * check that here before we call uclass_find_device().
-		 */
-		uc = uclass_find(id);
-		if (!uc)
-			continue;
-		ut_assertok(uclass_destroy(uc));
+	int i, id;
+
+	for (i = 0; i < 2; i++) {
+		for (id = 0; id < UCLASS_COUNT; id++) {
+			struct uclass *uc;
+
+			/*
+			 * If the uclass doesn't exist we don't want to
+			 * create it. So check that here before we call
+			 * uclass_find_device().
+			 */
+			uc = uclass_find(id);
+			if (!uc)
+				continue;
+			ut_assertok(uclass_destroy(uc,
+				    i ? DM_REMOVE_LATE : DM_REMOVE_NORMAL));
+		}
 	}
 
 	return 0;
-- 
2.27.0