#include "tb.h"
 
+#define CAP_OFFSET_MAX         0xff
+#define VSE_CAP_OFFSET_MAX     0xffff
 
 struct tb_cap_any {
        union {
        };
 } __packed;
 
-static bool tb_cap_is_basic(struct tb_cap_any *cap)
-{
-       /* basic.cap is u8. This checks only the lower 8 bit of cap. */
-       return cap->basic.cap != 5;
-}
-
-static bool tb_cap_is_long(struct tb_cap_any *cap)
+/**
+ * tb_port_find_cap() - Find port capability
+ * @port: Port to find the capability for
+ * @cap: Capability to look
+ *
+ * Returns offset to start of capability or %-ENOENT if no such
+ * capability was found. Negative errno is returned if there was an
+ * error.
+ */
+int tb_port_find_cap(struct tb_port *port, enum tb_port_cap cap)
 {
-       return !tb_cap_is_basic(cap)
-              && cap->extended_short.next == 0
-              && cap->extended_short.length == 0;
-}
+       u32 offset;
 
-static enum tb_cap tb_cap(struct tb_cap_any *cap)
-{
-       if (tb_cap_is_basic(cap))
-               return cap->basic.cap;
+       /*
+        * DP out adapters claim to implement TMU capability but in
+        * reality they do not so we hard code the adapter specific
+        * capability offset here.
+        */
+       if (port->config.type == TB_TYPE_DP_HDMI_OUT)
+               offset = 0x39;
        else
-               /* extended_short/long have cap at the same offset. */
-               return cap->extended_short.cap;
+               offset = 0x1;
+
+       do {
+               struct tb_cap_any header;
+               int ret;
+
+               ret = tb_port_read(port, &header, TB_CFG_PORT, offset, 1);
+               if (ret)
+                       return ret;
+
+               if (header.basic.cap == cap)
+                       return offset;
+
+               offset = header.basic.next;
+       } while (offset);
+
+       return -ENOENT;
 }
 
-static u32 tb_cap_next(struct tb_cap_any *cap, u32 offset)
+static int tb_switch_find_cap(struct tb_switch *sw, enum tb_switch_cap cap)
 {
-       int next;
-       if (offset == 1) {
-               /*
-                * The first pointer is part of the switch header and always
-                * a simple pointer.
-                */
-               next = cap->basic.next;
-       } else {
-               /*
-                * Somehow Intel decided to use 3 different types of capability
-                * headers. It is not like anyone could have predicted that
-                * single byte offsets are not enough...
-                */
-               if (tb_cap_is_basic(cap))
-                       next = cap->basic.next;
-               else if (!tb_cap_is_long(cap))
-                       next = cap->extended_short.next;
-               else
-                       next = cap->extended_long.next;
+       int offset = sw->config.first_cap_offset;
+
+       while (offset > 0 && offset < CAP_OFFSET_MAX) {
+               struct tb_cap_any header;
+               int ret;
+
+               ret = tb_sw_read(sw, &header, TB_CFG_SWITCH, offset, 1);
+               if (ret)
+                       return ret;
+
+               if (header.basic.cap == cap)
+                       return offset;
+
+               offset = header.basic.next;
        }
-       /*
-        * "Hey, we could terminate some capability lists with a null offset
-        *  and others with a pointer to the last element." - "Great idea!"
-        */
-       if (next == offset)
-               return 0;
-       return next;
+
+       return -ENOENT;
 }
 
 /**
- * tb_find_cap() - find a capability
+ * tb_switch_find_vse_cap() - Find switch vendor specific capability
+ * @sw: Switch to find the capability for
+ * @vsec: Vendor specific capability to look
  *
- * Return: Returns a positive offset if the capability was found and 0 if not.
- * Returns an error code on failure.
+ * Functions enumerates vendor specific capabilities (VSEC) of a switch
+ * and returns offset when capability matching @vsec is found. If no
+ * such capability is found returns %-ENOENT. In case of error returns
+ * negative errno.
  */
-int tb_find_cap(struct tb_port *port, enum tb_cfg_space space, enum tb_cap cap)
+int tb_switch_find_vse_cap(struct tb_switch *sw, enum tb_switch_vse_cap vsec)
 {
-       u32 offset = 1;
        struct tb_cap_any header;
-       int res;
-       int retries = 10;
-       while (retries--) {
-               res = tb_port_read(port, &header, space, offset, 1);
-               if (res) {
-                       /* Intel needs some help with linked lists. */
-                       if (space == TB_CFG_PORT && offset == 0xa
-                           && port->config.type == TB_TYPE_DP_HDMI_OUT) {
-                               offset = 0x39;
-                               continue;
-                       }
-                       return res;
-               }
-               if (offset != 1) {
-                       if (tb_cap(&header) == cap)
+       int offset;
+
+       offset = tb_switch_find_cap(sw, TB_SWITCH_CAP_VSE);
+       if (offset < 0)
+               return offset;
+
+       while (offset > 0 && offset < VSE_CAP_OFFSET_MAX) {
+               int ret;
+
+               ret = tb_sw_read(sw, &header, TB_CFG_SWITCH, offset, 2);
+               if (ret)
+                       return ret;
+
+               /*
+                * Extended vendor specific capabilities come in two
+                * flavors: short and long. The latter is used when
+                * offset is over 0xff.
+                */
+               if (offset >= CAP_OFFSET_MAX) {
+                       if (header.extended_long.vsec_id == vsec)
                                return offset;
-                       if (tb_cap_is_long(&header)) {
-                               /* tb_cap_extended_long is 2 dwords */
-                               res = tb_port_read(port, &header, space,
-                                                  offset, 2);
-                               if (res)
-                                       return res;
-                       }
+                       offset = header.extended_long.next;
+               } else {
+                       if (header.extended_short.vsec_id == vsec)
+                               return offset;
+                       if (!header.extended_short.length)
+                               return -ENOENT;
+                       offset = header.extended_short.next;
                }
-               offset = tb_cap_next(&header, offset);
-               if (!offset)
-                       return 0;
        }
-       tb_port_WARN(port,
-                    "run out of retries while looking for cap %#x in config space %d, last offset: %#x\n",
-                    cap, space, offset);
-       return -EIO;
+
+       return -ENOENT;
 }
 
  */
 #define TB_MAX_CONFIG_RW_LENGTH 60
 
-enum tb_cap {
-       TB_CAP_PHY              = 0x0001,
-       TB_CAP_TIME1            = 0x0003,
-       TB_CAP_PCIE             = 0x0004,
-       TB_CAP_I2C              = 0x0005,
-       TB_CAP_PLUG_EVENTS      = 0x0105, /* also EEPROM */
-       TB_CAP_TIME2            = 0x0305,
-       TB_CAP_IECS             = 0x0405,
-       TB_CAP_LINK_CONTROLLER  = 0x0605, /* also IECS */
+enum tb_switch_cap {
+       TB_SWITCH_CAP_VSE               = 0x05,
+};
+
+enum tb_switch_vse_cap {
+       TB_VSE_CAP_PLUG_EVENTS          = 0x01, /* also EEPROM */
+       TB_VSE_CAP_TIME2                = 0x03,
+       TB_VSE_CAP_IECS                 = 0x04,
+       TB_VSE_CAP_LINK_CONTROLLER      = 0x06, /* also IECS */
+};
+
+enum tb_port_cap {
+       TB_PORT_CAP_PHY                 = 0x01,
+       TB_PORT_CAP_TIME1               = 0x03,
+       TB_PORT_CAP_ADAP                = 0x04,
+       TB_PORT_CAP_VSE                 = 0x05,
 };
 
 enum tb_port_state {
        u8 cap; /* if cap == 0x05 then we have a extended capability */
 } __packed;
 
+/**
+ * struct tb_cap_extended_short - Switch extended short capability
+ * @next: Pointer to the next capability. If @next and @length are zero
+ *       then we have a long cap.
+ * @cap: Base capability ID (see &enum tb_switch_cap)
+ * @vsec_id: Vendor specific capability ID (see &enum switch_vse_cap)
+ * @length: Length of this capability
+ */
 struct tb_cap_extended_short {
-       u8 next; /* if next and length are zero then we have a long cap */
-       enum tb_cap cap:16;
+       u8 next;
+       u8 cap;
+       u8 vsec_id;
        u8 length;
 } __packed;
 
+/**
+ * struct tb_cap_extended_long - Switch extended long capability
+ * @zero1: This field should be zero
+ * @cap: Base capability ID (see &enum tb_switch_cap)
+ * @vsec_id: Vendor specific capability ID (see &enum switch_vse_cap)
+ * @zero2: This field should be zero
+ * @next: Pointer to the next capability
+ * @length: Length of this capability
+ */
 struct tb_cap_extended_long {
        u8 zero1;
-       enum tb_cap cap:16;
+       u8 cap;
+       u8 vsec_id;
        u8 zero2;
        u16 next;
        u16 length;