dm: pci: Add APIs to find next capability and extended capability

This introduces two new APIs dm_pci_find_next_capability() and
dm_pci_find_next_ext_capability() to get PCI capability address
and PCI express extended capability address for a given PCI device
starting from a given offset.

Signed-off-by: Bin Meng <bmeng.cn@gmail.com>
Reviewed-by: Simon Glass <sjg@chromium.org>
diff --git a/drivers/pci/pci-uclass.c b/drivers/pci/pci-uclass.c
index da49c96e..0c52337 100644
--- a/drivers/pci/pci-uclass.c
+++ b/drivers/pci/pci-uclass.c
@@ -1344,26 +1344,14 @@
 	return dm_pci_bus_to_virt(dev, pci_bus_addr, flags, 0, MAP_NOCACHE);
 }
 
-int dm_pci_find_capability(struct udevice *dev, int cap)
+static int _dm_pci_find_next_capability(struct udevice *dev, u8 pos, int cap)
 {
-	u16 status;
-	u8 header_type;
 	int ttl = PCI_FIND_CAP_TTL;
 	u8 id;
 	u16 ent;
-	u8 pos;
-
-	dm_pci_read_config16(dev, PCI_STATUS, &status);
-	if (!(status & PCI_STATUS_CAP_LIST))
-		return 0;
-
-	dm_pci_read_config8(dev, PCI_HEADER_TYPE, &header_type);
-	if ((header_type & 0x7f) == PCI_HEADER_TYPE_CARDBUS)
-		pos = PCI_CB_CAPABILITY_LIST;
-	else
-		pos = PCI_CAPABILITY_LIST;
 
 	dm_pci_read_config8(dev, pos, &pos);
+
 	while (ttl--) {
 		if (pos < PCI_STD_HEADER_SIZEOF)
 			break;
@@ -1381,7 +1369,32 @@
 	return 0;
 }
 
-int dm_pci_find_ext_capability(struct udevice *dev, int cap)
+int dm_pci_find_next_capability(struct udevice *dev, u8 start, int cap)
+{
+	return _dm_pci_find_next_capability(dev, start + PCI_CAP_LIST_NEXT,
+					    cap);
+}
+
+int dm_pci_find_capability(struct udevice *dev, int cap)
+{
+	u16 status;
+	u8 header_type;
+	u8 pos;
+
+	dm_pci_read_config16(dev, PCI_STATUS, &status);
+	if (!(status & PCI_STATUS_CAP_LIST))
+		return 0;
+
+	dm_pci_read_config8(dev, PCI_HEADER_TYPE, &header_type);
+	if ((header_type & 0x7f) == PCI_HEADER_TYPE_CARDBUS)
+		pos = PCI_CB_CAPABILITY_LIST;
+	else
+		pos = PCI_CAPABILITY_LIST;
+
+	return _dm_pci_find_next_capability(dev, pos, cap);
+}
+
+int dm_pci_find_next_ext_capability(struct udevice *dev, int start, int cap)
 {
 	u32 header;
 	int ttl;
@@ -1390,6 +1403,9 @@
 	/* minimum 8 bytes per capability */
 	ttl = (PCI_CFG_SPACE_EXP_SIZE - PCI_CFG_SPACE_SIZE) / 8;
 
+	if (start)
+		pos = start;
+
 	dm_pci_read_config32(dev, pos, &header);
 	/*
 	 * If we have no capabilities, this is indicated by cap ID,
@@ -1412,6 +1428,11 @@
 	return 0;
 }
 
+int dm_pci_find_ext_capability(struct udevice *dev, int cap)
+{
+	return dm_pci_find_next_ext_capability(dev, 0, cap);
+}
+
 UCLASS_DRIVER(pci) = {
 	.id		= UCLASS_PCI,
 	.name		= "pci",