dm: blk: Expand iteration and add tests

Add some functions which support iteration before probing. Also add tests
for the functions.

Signed-off-by: Simon Glass <sjg@chromium.org>

drivers/block/blk-uclass.c

@@ -552,6 +552,30 @@ static int blk_flags_check(struct udevice *dev, enum blk_flag_t req_flags)
 	return flags & req_flags ? 0 : 1;
 }
 
+int blk_find_first(enum blk_flag_t flags, struct udevice **devp)
+{
+	int ret;
+
+	for (ret = uclass_find_first_device(UCLASS_BLK, devp);
+	     *devp && !blk_flags_check(*devp, flags);
+	     ret = uclass_find_next_device(devp))
+		return 0;
+
+	return -ENODEV;
+}
+
+int blk_find_next(enum blk_flag_t flags, struct udevice **devp)
+{
+	int ret;
+
+	for (ret = uclass_find_next_device(devp);
+	     *devp && !blk_flags_check(*devp, flags);
+	     ret = uclass_find_next_device(devp))
+		return 0;
+
+	return -ENODEV;
+}
+
 int blk_first_device_err(enum blk_flag_t flags, struct udevice **devp)
 {
 	int ret;

include/blk.h

@@ -731,6 +731,51 @@ int blk_first_device_err(enum blk_flag_t flags, struct udevice **devp);
  */
 int blk_next_device_err(enum blk_flag_t flags, struct udevice **devp);
 
+/**
+ * blk_find_first() - Return the first matching block device
+ * @flags: Indicates type of device to return
+ * @devp:	Returns pointer to device, or NULL on error
+ *
+ * The device is not prepared for use - this is an internal function.
+ * The function uclass_get_device_tail() can be used to probe the device.
+ *
+ * Note that some devices are considered removable until they have been probed
+ *
+ * @return 0 if found, -ENODEV if not found
+ */
+int blk_find_first(enum blk_flag_t flags, struct udevice **devp);
+
+/**
+ * blk_find_next() - Return the next matching block device
+ * @flags: Indicates type of device to return
+ * @devp: On entry, pointer to device to lookup. On exit, returns pointer
+ * to the next device in the same uclass, or NULL if none
+ *
+ * The device is not prepared for use - this is an internal function.
+ * The function uclass_get_device_tail() can be used to probe the device.
+ *
+ * Note that some devices are considered removable until they have been probed
+ *
+ * @return 0 if found, -ENODEV if not found
+ */
+int blk_find_next(enum blk_flag_t flags, struct udevice **devp);
+
+/**
+ * blk_foreach() - iterate through block devices
+ *
+ * This creates a for() loop which works through the available block devices in
+ * order from start to end.
+ *
+ * If for some reason the uclass cannot be found, this does nothing.
+ *
+ * @_flags: Indicates type of device to return
+ * @_pos: struct udevice * to hold the current device. Set to NULL when there
+ * are no more devices.
+ */
+#define blk_foreach(_flags, _pos) \
+	for (int _ret = blk_find_first(_flags, &_pos); !_ret && _pos; \
+	     _ret = blk_find_next(_flags, &_pos))
+
 /**
  * blk_foreach_probe() - Helper function to iteration through block devices
  *

test/dm/blk.c

@@ -217,3 +217,114 @@ static int dm_test_blk_iter(struct unit_test_state *uts)
 	return 0;
 }
 DM_TEST(dm_test_blk_iter, UT_TESTF_SCAN_PDATA | UT_TESTF_SCAN_FDT);
+
+/* Test finding fixed/removable block devices */
+static int dm_test_blk_flags(struct unit_test_state *uts)
+{
+	struct udevice *dev;
+
+	/* Iterate through devices without probing them */
+	ut_assertok(blk_find_first(BLKF_BOTH, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc2.blk", dev->name);
+
+	ut_assertok(blk_find_next(BLKF_BOTH, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc1.blk", dev->name);
+
+	ut_assertok(blk_find_next(BLKF_BOTH, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc0.blk", dev->name);
+
+	ut_asserteq(-ENODEV, blk_find_next(BLKF_BOTH, &dev));
+	ut_assertnull(dev);
+
+	/* All devices are removable until probed */
+	ut_asserteq(-ENODEV, blk_find_first(BLKF_FIXED, &dev));
+
+	ut_assertok(blk_find_first(BLKF_REMOVABLE, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc2.blk", dev->name);
+
+	/* Now probe them and iterate again */
+	ut_assertok(blk_first_device_err(BLKF_BOTH, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc2.blk", dev->name);
+
+	ut_assertok(blk_next_device_err(BLKF_BOTH, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc1.blk", dev->name);
+
+	ut_assertok(blk_next_device_err(BLKF_BOTH, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc0.blk", dev->name);
+
+	ut_asserteq(-ENODEV, blk_next_device_err(BLKF_BOTH, &dev));
+
+	/* Look only for fixed devices */
+	ut_assertok(blk_first_device_err(BLKF_FIXED, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc2.blk", dev->name);
+
+	ut_asserteq(-ENODEV, blk_next_device_err(BLKF_FIXED, &dev));
+
+	/* Look only for removable devices */
+	ut_assertok(blk_first_device_err(BLKF_REMOVABLE, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc1.blk", dev->name);
+
+	ut_assertok(blk_next_device_err(BLKF_REMOVABLE, &dev));
+	ut_assertnonnull(dev);
+	ut_asserteq_str("mmc0.blk", dev->name);
+
+	ut_asserteq(-ENODEV, blk_next_device_err(BLKF_REMOVABLE, &dev));
+
+	return 0;
+}
+DM_TEST(dm_test_blk_flags, UT_TESTF_SCAN_PDATA | UT_TESTF_SCAN_FDT);
+
+/* Test blk_foreach() and friend */
+static int dm_test_blk_foreach(struct unit_test_state *uts)
+{
+	struct udevice *dev;
+	int found;
+
+	/* Test blk_foreach() - use the 3rd bytes of the name (0/1/2) */
+	found = 0;
+	blk_foreach(BLKF_BOTH, dev)
+		found |= 1 << dectoul(&dev->name[3], NULL);
+	ut_asserteq(7, found);
+
+	/* All devices are removable until probed */
+	found = 0;
+	blk_foreach(BLKF_FIXED, dev)
+		found |= 1 << dectoul(&dev->name[3], NULL);
+	ut_asserteq(0, found);
+
+	found = 0;
+	blk_foreach(BLKF_REMOVABLE, dev)
+		found |= 1 << dectoul(&dev->name[3], NULL);
+	ut_asserteq(7, found);
+
+	/* Now try again with the probing functions */
+	found = 0;
+	blk_foreach_probe(BLKF_BOTH, dev)
+		found |= 1 << dectoul(&dev->name[3], NULL);
+	ut_asserteq(7, found);
+	ut_asserteq(3, blk_count_devices(BLKF_BOTH));
+
+	found = 0;
+	blk_foreach_probe(BLKF_FIXED, dev)
+		found |= 1 << dectoul(&dev->name[3], NULL);
+	ut_asserteq(4, found);
+	ut_asserteq(1, blk_count_devices(BLKF_FIXED));
+
+	found = 0;
+	blk_foreach_probe(BLKF_REMOVABLE, dev)
+		found |= 1 << dectoul(&dev->name[3], NULL);
+	ut_asserteq(3, found);
+	ut_asserteq(2, blk_count_devices(BLKF_REMOVABLE));
+
+	return 0;
+}
+DM_TEST(dm_test_blk_foreach, UT_TESTF_SCAN_PDATA | UT_TESTF_SCAN_FDT);