blk: Support iteration

It is useful to be able to iterate over block devices. Typically there
are fixed and removable devices. For security reasons it is sometimes
useful to ignore removable devices since they are under user control.

Add iterators which support selecting the block-device type.

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

3 files changed, 160 insertions, 0 deletions

diff --git a/drivers/block/blk-uclass.c b/drivers/block/blk-uclass.c
index dfc0d46970..83682dcc18 100644
--- a/drivers/block/blk-uclass.c
+++ b/drivers/block/blk-uclass.c
@@ -540,6 +540,55 @@ int blk_next_free_devnum(enum if_type if_type)
 	return ret + 1;
 }
 
+static int blk_flags_check(struct udevice *dev, enum blk_flag_t req_flags)
+{
+	const struct blk_desc *desc = dev_get_uclass_plat(dev);
+	enum blk_flag_t flags;
+
+	flags = desc->removable ? BLKF_REMOVABLE : BLKF_FIXED;
+
+	return flags & req_flags ? 0 : 1;
+}
+
+int blk_first_device_err(enum blk_flag_t flags, struct udevice **devp)
+{
+	int ret;
+
+	for (ret = uclass_first_device_err(UCLASS_BLK, devp);
+	     !ret;
+	     ret = uclass_next_device_err(devp)) {
+		if (!blk_flags_check(*devp, flags))
+			return 0;
+	}
+
+	return -ENODEV;
+}
+
+int blk_next_device_err(enum blk_flag_t flags, struct udevice **devp)
+{
+	int ret;
+
+	for (ret = uclass_next_device_err(devp);
+	     !ret;
+	     ret = uclass_next_device_err(devp)) {
+		if (!blk_flags_check(*devp, flags))
+			return 0;
+	}
+
+	return -ENODEV;
+}
+
+int blk_count_devices(enum blk_flag_t flag)
+{
+	struct udevice *dev;
+	int count = 0;
+
+	blk_foreach_probe(flag, dev)
+		count++;
+
+	return count;
+}
+
 static int blk_claim_devnum(enum if_type if_type, int devnum)
 {
 	struct udevice *dev;
diff --git a/include/blk.h b/include/blk.h
index c4401b0025..19bab081c2 100644
--- a/include/blk.h
+++ b/include/blk.h
@@ -19,6 +19,8 @@ typedef ulong lbaint_t;
 #define LBAF "%" LBAFlength "x"
 #define LBAFU "%" LBAFlength "u"
 
+struct udevice;
+
 /* Interface types: */
 enum if_type {
 	IF_TYPE_UNKNOWN = 0,
@@ -683,4 +685,58 @@ const char *blk_get_if_type_name(enum if_type if_type);
 int blk_common_cmd(int argc, char *const argv[], enum if_type if_type,
 		   int *cur_devnump);
 
+enum blk_flag_t {
+	BLKF_FIXED	= 1 << 0,
+	BLKF_REMOVABLE	= 1 << 1,
+	BLKF_BOTH	= BLKF_FIXED | BLKF_REMOVABLE,
+};
+
+/**
+ * blk_first_device_err() - Get the first block device
+ *
+ * The device returned is probed if necessary, and ready for use
+ *
+ * @flags: Indicates type of device to return
+ * @devp: Returns pointer to the first device in that uclass, or NULL if none
+ * @return 0 if found, -ENODEV if not found, other -ve on error
+ */
+int blk_first_device_err(enum blk_flag_t flags, struct udevice **devp);
+
+/**
+ * blk_next_device_err() - Get the next block device
+ *
+ * The device returned is probed if necessary, and ready for use
+ *
+ * @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 uclass if no error occurred, or -ENODEV if
+ * there is no next device.
+ * @return 0 if found, -ENODEV if not found, other -ve on error
+ */
+int blk_next_device_err(enum blk_flag_t flags, struct udevice **devp);
+
+/**
+ * blk_foreach_probe() - Helper function to iteration through block devices
+ *
+ * This creates a for() loop which works through the available devices in
+ * a uclass in order from start to end. Devices are probed if necessary,
+ * and ready for use.
+ *
+ * @flags: Indicates type of device to return
+ * @dev: struct udevice * to hold the current device. Set to NULL when there
+ * are no more devices.
+ */
+#define blk_foreach_probe(flags, pos)	\
+	for (int _ret = blk_first_device_err(flags, &(pos)); \
+	     !_ret && pos; \
+	     _ret = blk_next_device_err(flags, &(pos)))
+
+/**
+ * blk_count_devices() - count the number of devices of a particular type
+ *
+ * @flags: Indicates type of device to find
+ * @return number of devices matching those flags
+ */
+int blk_count_devices(enum blk_flag_t flag);
+
 #endif
diff --git a/test/dm/blk.c b/test/dm/blk.c
index b7f4304e9e..deccf05289 100644
--- a/test/dm/blk.c
+++ b/test/dm/blk.c
@@ -162,3 +162,58 @@ static int dm_test_blk_get_from_parent(struct unit_test_state *uts)
 	return 0;
 }
 DM_TEST(dm_test_blk_get_from_parent, UT_TESTF_SCAN_PDATA | UT_TESTF_SCAN_FDT);
+
+/* Test iteration through block devices */
+static int dm_test_blk_iter(struct unit_test_state *uts)
+{
+	struct udevice *dev;
+	int i;
+
+	/*
+	 * See sandbox test.dts - it has:
+	 *
+	 *   mmc0 - removable
+	 *   mmc1 - removable
+	 *   mmc2 - fixed
+	 */
+	ut_assertok(blk_first_device_err(BLKF_FIXED, &dev));
+	ut_asserteq_str("mmc2.blk", dev->name);
+	ut_asserteq(-ENODEV, blk_next_device_err(BLKF_FIXED, &dev));
+
+	ut_assertok(blk_first_device_err(BLKF_REMOVABLE, &dev));
+	ut_asserteq_str("mmc1.blk", dev->name);
+	ut_assertok(blk_next_device_err(BLKF_REMOVABLE, &dev));
+	ut_asserteq_str("mmc0.blk", dev->name);
+	ut_asserteq(-ENODEV, blk_next_device_err(BLKF_REMOVABLE, &dev));
+
+	ut_assertok(blk_first_device_err(BLKF_BOTH, &dev));
+	ut_asserteq_str("mmc2.blk", dev->name);
+	ut_assertok(blk_next_device_err(BLKF_BOTH, &dev));
+	ut_asserteq_str("mmc1.blk", dev->name);
+	ut_assertok(blk_next_device_err(BLKF_BOTH, &dev));
+	ut_asserteq_str("mmc0.blk", dev->name);
+	ut_asserteq(-ENODEV, blk_next_device_err(BLKF_FIXED, &dev));
+
+	ut_asserteq(1, blk_count_devices(BLKF_FIXED));
+	ut_asserteq(2, blk_count_devices(BLKF_REMOVABLE));
+	ut_asserteq(3, blk_count_devices(BLKF_BOTH));
+
+	i = 0;
+	blk_foreach_probe(BLKF_FIXED, dev)
+		ut_asserteq_str((i++, "mmc2.blk"), dev->name);
+	ut_asserteq(1, i);
+
+	i = 0;
+	blk_foreach_probe(BLKF_REMOVABLE, dev)
+		ut_asserteq_str(i++ ? "mmc0.blk" : "mmc1.blk", dev->name);
+	ut_asserteq(2, i);
+
+	i = 0;
+	blk_foreach_probe(BLKF_BOTH, dev)
+		ut_asserteq_str((++i == 1 ? "mmc2.blk" : i == 2 ?
+			"mmc1.blk" : "mmc0.blk"), dev->name);
+	ut_asserteq(3, i);
+
+	return 0;
+}
+DM_TEST(dm_test_blk_iter, UT_TESTF_SCAN_PDATA | UT_TESTF_SCAN_FDT);