usb_mux: retimer: mux as chained mux and retimer

This makes retimers appear as generic muxes.  By allowing a
chain of muxes they can be stacked up to the new configurations
that zork requires and will continue to work as they did before
on configurations that only have a single mux.

The code used to have two different arrays, 1) muxes and 2)
retimers.  On one of the zork configurations the processor
MUX stopped being the primary mux and the retimer took its
place.  In a different configuration of that same platform
it left the primary and secondary alone but the mux_set
FLIP operation had to be ignored.  Since the same
interfaces needed to be available for both it stopped making
sense to have two different structures and two different
methods of handling them.  This consolodates the two into
one.

The platforms that do not have retimers, this change will
not make any difference. For platforms like zork, it will
remove the retimers and make them chained muxes.  So
testing on trembyle makes sense to verify,

BUG=b:147593660
BRANCH=none
TEST=verify USB still works on trembyle

Change-Id: I286cf1e302f9bd3dd7e81098ec08514a2a009fe3
Signed-off-by: Denis Brockus <dbrockus@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/ec/+/2066794
Commit-Queue: Jett Rink <jettrink@chromium.org>
Reviewed-by: Edward Hill <ecgh@chromium.org>
Reviewed-by: Jett Rink <jettrink@chromium.org>

10 files changed, 182 insertions, 195 deletions

diff --git a/driver/retimer/bb_retimer.c b/driver/retimer/bb_retimer.c
index 61a84c4116..08afbc7308 100644
--- a/driver/retimer/bb_retimer.c
+++ b/driver/retimer/bb_retimer.c
@@ -27,7 +27,8 @@
 /**
  * Utility functions
  */
-static int bb_retimer_read(int port, const uint8_t offset, uint32_t *data)
+static int bb_retimer_read(const struct usb_mux *me,
+			   const uint8_t offset, uint32_t *data)
 {
 	int rv;
 	uint8_t buf[BB_RETIMER_READ_SIZE];
@@ -39,8 +40,7 @@ static int bb_retimer_read(int port, const uint8_t offset, uint32_t *data)
 	 * byte[1:4] : Data [LSB -> MSB]
 	 * Stop
 	 */
-	rv = i2c_xfer(usb_retimers[port].i2c_port,
-		      usb_retimers[port].i2c_addr_flags,
+	rv = i2c_xfer(me->i2c_port, me->i2c_addr_flags,
 		      &offset, 1, buf, BB_RETIMER_READ_SIZE);
 	if (rv)
 		return rv;
@@ -52,7 +52,8 @@ static int bb_retimer_read(int port, const uint8_t offset, uint32_t *data)
 	return EC_SUCCESS;
 }
 
-static int bb_retimer_write(int port, const uint8_t offset, uint32_t data)
+static int bb_retimer_write(const struct usb_mux *me,
+			    const uint8_t offset, uint32_t data)
 {
 	uint8_t buf[BB_RETIMER_WRITE_SIZE];
 
@@ -71,14 +72,14 @@ static int bb_retimer_write(int port, const uint8_t offset, uint32_t data)
 	buf[4] = (data >> 16) & 0xFF;
 	buf[5] = (data >> 24) & 0xFF;
 
-	return i2c_xfer(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags,
+	return i2c_xfer(me->i2c_port,
+			me->i2c_addr_flags,
 			buf, BB_RETIMER_WRITE_SIZE, NULL, 0);
 }
 
-static void bb_retimer_power_handle(int port, int on_off)
+static void bb_retimer_power_handle(const struct usb_mux *me, int on_off)
 {
-	const struct bb_usb_control *control = &bb_controls[port];
+	const struct bb_usb_control *control = &bb_controls[me->usb_port];
 
 	/* handle retimer's power domain */
 
@@ -110,10 +111,11 @@ static void bb_retimer_power_handle(int port, int on_off)
 /**
  * Driver interface functions
  */
-static int retimer_set_state(int port, mux_state_t mux_state)
+static int retimer_set_state(const struct usb_mux *me, mux_state_t mux_state)
 {
 	uint32_t set_retimer_con = 0;
 	uint8_t dp_pin_mode;
+	int port = me->usb_port;
 	union tbt_mode_resp_cable cable_resp;
 	union tbt_mode_resp_device dev_resp;
 
@@ -280,30 +282,30 @@ static int retimer_set_state(int port, mux_state_t mux_state)
 						cable_resp.tbt_cable_speed);
 	}
 	/* Writing the register4 */
-	return bb_retimer_write(port, BB_RETIMER_REG_CONNECTION_STATE,
+	return bb_retimer_write(me, BB_RETIMER_REG_CONNECTION_STATE,
 			set_retimer_con);
 }
 
-static int retimer_low_power_mode(int port)
+static int retimer_low_power_mode(const struct usb_mux *me)
 {
-	bb_retimer_power_handle(port, 0);
+	bb_retimer_power_handle(me, 0);
 	return EC_SUCCESS;
 }
 
-static int retimer_init(int port)
+static int retimer_init(const struct usb_mux *me)
 {
 	int rv;
 	uint32_t data;
 
-	bb_retimer_power_handle(port, 1);
+	bb_retimer_power_handle(me, 1);
 
-	rv = bb_retimer_read(port, BB_RETIMER_REG_VENDOR_ID, &data);
+	rv = bb_retimer_read(me, BB_RETIMER_REG_VENDOR_ID, &data);
 	if (rv)
 		return rv;
 	if (data != BB_RETIMER_VENDOR_ID)
 		return EC_ERROR_UNKNOWN;
 
-	rv = bb_retimer_read(port, BB_RETIMER_REG_DEVICE_ID, &data);
+	rv = bb_retimer_read(me, BB_RETIMER_REG_DEVICE_ID, &data);
 	if (rv)
 		return rv;
 
@@ -313,7 +315,7 @@ static int retimer_init(int port)
 	return EC_SUCCESS;
 }
 
-const struct usb_retimer_driver bb_usb_retimer = {
+const struct usb_mux_driver bb_usb_retimer = {
 	.init = retimer_init,
 	.set = retimer_set_state,
 	.enter_low_power_mode = retimer_low_power_mode,
@@ -324,14 +326,24 @@ static int console_command_bb_retimer(int argc, char **argv)
 {
 	char rw, *e;
 	int rv, port, reg, data, val;
+	const struct usb_mux *mux;
 
 	if (argc < 4)
 		return EC_ERROR_PARAM_COUNT;
 
 	/* Get port number */
 	port = strtoi(argv[1], &e, 0);
-	if (*e || port < 0 || port > board_get_usb_pd_port_count() ||
-		!usb_retimers[port].driver)
+	if (*e || port < 0 || port > board_get_usb_pd_port_count())
+		return EC_ERROR_PARAM1;
+
+	mux = &usb_muxes[port];
+	while (mux) {
+		if (mux->driver == &bb_usb_retimer)
+			break;
+		mux = mux->next_mux;
+	}
+
+	if (!mux)
 		return EC_ERROR_PARAM1;
 
 	/* Validate r/w selection */
@@ -345,16 +357,16 @@ static int console_command_bb_retimer(int argc, char **argv)
 		return EC_ERROR_PARAM3;
 
 	if (rw == 'r')
-		rv = bb_retimer_read(port, reg, &data);
+		rv = bb_retimer_read(mux, reg, &data);
 	else {
 		/* Get value to be written */
 		val = strtoi(argv[4], &e, 0);
 		if (*e || val < 0)
 			return EC_ERROR_PARAM4;
 
-		rv = bb_retimer_write(port, reg, val);
+		rv = bb_retimer_write(mux, reg, val);
 		if (rv == EC_SUCCESS) {
-			rv = bb_retimer_read(port, reg, &data);
+			rv = bb_retimer_read(mux, reg, &data);
 			if (rv == EC_SUCCESS && data != val)
 				rv = EC_ERROR_UNKNOWN;
 		}
diff --git a/driver/retimer/bb_retimer.h b/driver/retimer/bb_retimer.h
index 3b34b72e3d..868cb5ee7c 100644
--- a/driver/retimer/bb_retimer.h
+++ b/driver/retimer/bb_retimer.h
@@ -40,7 +40,7 @@
 #define BB_RETIMER_TBT_CABLE_GENERATION(x)	(((x) & 0x3) << 28)
 
 /* Supported USB retimer drivers */
-extern const struct usb_retimer_driver bb_usb_retimer;
+extern const struct usb_mux_driver bb_usb_retimer;
 
 /* Retimer driver hardware specific controls */
 struct bb_usb_control {
diff --git a/driver/retimer/nb7v904m.c b/driver/retimer/nb7v904m.c
index 81fd5abc93..e1cedede95 100644
--- a/driver/retimer/nb7v904m.c
+++ b/driver/retimer/nb7v904m.c
@@ -15,28 +15,28 @@
 #define CPRINTS(format, args...) cprints(CC_USB, format, ## args)
 #define CPRINTF(format, args...) cprintf(CC_USB, format, ## args)
 
-static int nb7v904m_write(int port, int offset, int data)
+static int nb7v904m_write(const struct usb_mux *me, int offset, int data)
 {
-	return i2c_write8(usb_retimers[port].i2c_port,
-			  usb_retimers[port].i2c_addr_flags,
+	return i2c_write8(me->i2c_port,
+			  me->i2c_addr_flags,
 			  offset, data);
 
 }
 
-static int nb7v904m_read(int port, int offset, int *regval)
+static int nb7v904m_read(const struct usb_mux *me, int offset, int *regval)
 {
-	return i2c_read8(usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags,
+	return i2c_read8(me->i2c_port,
+			 me->i2c_addr_flags,
 			 offset, regval);
 
 }
 
-static int set_low_power_mode(int port, bool enable)
+static int set_low_power_mode(const struct usb_mux *me, bool enable)
 {
 	int regval;
 	int rv;
 
-	rv = nb7v904m_read(port, NB7V904M_REG_GEN_DEV_SETTINGS, &regval);
+	rv = nb7v904m_read(me, NB7V904M_REG_GEN_DEV_SETTINGS, &regval);
 	if (rv)
 		return rv;
 
@@ -45,28 +45,29 @@ static int set_low_power_mode(int port, bool enable)
 	else
 		regval &= ~NB7V904M_CHIP_EN;
 
-	return nb7v904m_write(port, NB7V904M_REG_GEN_DEV_SETTINGS, regval);
+	return nb7v904m_write(me, NB7V904M_REG_GEN_DEV_SETTINGS, regval);
 }
 
-static int nb7v904m_enter_low_power_mode(int port)
+static int nb7v904m_enter_low_power_mode(const struct usb_mux *me)
 {
-	int rv = set_low_power_mode(port, 1);
+	int rv = set_low_power_mode(me, 1);
 
 	if (rv)
-		CPRINTS("C%d: NB7V904M: Failed to enter low power mode!", port);
+		CPRINTS("C%d: NB7V904M: Failed to enter low power mode!",
+			me->usb_port);
 	return rv;
 }
 
-static int nb7v904m_init(int port)
+static int nb7v904m_init(const struct usb_mux *me)
 {
-	int rv = set_low_power_mode(port, 0);
+	int rv = set_low_power_mode(me, 0);
 
 	if (rv)
-		CPRINTS("C%d: NB7V904M: init failed!", port);
+		CPRINTS("C%d: NB7V904M: init failed!", me->usb_port);
 	return rv;
 }
 
-static int nb7v904m_set_mux(int port, mux_state_t mux_state)
+static int nb7v904m_set_mux(const struct usb_mux *me, mux_state_t mux_state)
 {
 	int rv = EC_SUCCESS;
 	int regval;
@@ -74,17 +75,17 @@ static int nb7v904m_set_mux(int port, mux_state_t mux_state)
 
 	/* Turn off redriver if it's not needed at all. */
 	if (mux_state == USB_PD_MUX_NONE)
-		return nb7v904m_enter_low_power_mode(port);
+		return nb7v904m_enter_low_power_mode(me);
 
-	rv = nb7v904m_init(port);
+	rv = nb7v904m_init(me);
 	if (rv)
 		return rv;
 
 	/* Clear operation mode field */
-	rv = nb7v904m_read(port, NB7V904M_REG_GEN_DEV_SETTINGS, &regval);
+	rv = nb7v904m_read(me, NB7V904M_REG_GEN_DEV_SETTINGS, &regval);
 	if (rv) {
-		CPRINTS("C%d %s: Failed to obtain dev settings!", port,
-			__func__);
+		CPRINTS("C%d %s: Failed to obtain dev settings!",
+			me->usb_port, __func__);
 		return rv;
 	}
 	regval &= ~NB7V904M_OP_MODE_MASK;
@@ -108,14 +109,14 @@ static int nb7v904m_set_mux(int port, mux_state_t mux_state)
 
 	if (mux_state & USB_PD_MUX_DP_ENABLED) {
 		/* Connect AUX */
-		rv = nb7v904m_write(port, NB7V904M_REG_AUX_CH_CTRL, flipped ?
+		rv = nb7v904m_write(me, NB7V904M_REG_AUX_CH_CTRL, flipped ?
 				    NB7V904M_AUX_CH_FLIPPED :
 				    NB7V904M_AUX_CH_NORMAL);
 		/* Enable all channels for DP */
 		regval |= NB7V904M_CH_EN_MASK;
 	} else {
 		/* Disconnect AUX since it's not being used. */
-		rv = nb7v904m_write(port, NB7V904M_REG_AUX_CH_CTRL,
+		rv = nb7v904m_write(me, NB7V904M_REG_AUX_CH_CTRL,
 				    NB7V904M_AUX_CH_HI_Z);
 
 		/* Disable the unused channels to save power */
@@ -129,14 +130,14 @@ static int nb7v904m_set_mux(int port, mux_state_t mux_state)
 		}
 	}
 
-	rv |= nb7v904m_write(port, NB7V904M_REG_GEN_DEV_SETTINGS, regval);
+	rv |= nb7v904m_write(me, NB7V904M_REG_GEN_DEV_SETTINGS, regval);
 	if (rv)
-		CPRINTS("C%d: %s failed!", port, __func__);
+		CPRINTS("C%d: %s failed!", me->usb_port, __func__);
 
 	return rv;
 }
 
-const struct usb_retimer_driver nb7v904m_usb_redriver_drv = {
+const struct usb_mux_driver nb7v904m_usb_redriver_drv = {
 	.enter_low_power_mode = &nb7v904m_enter_low_power_mode,
 	.init = &nb7v904m_init,
 	.set = &nb7v904m_set_mux,
diff --git a/driver/retimer/nb7v904m.h b/driver/retimer/nb7v904m.h
index 4306615825..8bb43f79eb 100644
--- a/driver/retimer/nb7v904m.h
+++ b/driver/retimer/nb7v904m.h
@@ -37,5 +37,5 @@
 #define NB7V904M_AUX_CH_FLIPPED  BIT(0)
 #define NB7V904M_AUX_CH_HI_Z     BIT(1)
 
-extern const struct usb_retimer_driver nb7v904m_usb_redriver_drv;
+extern const struct usb_mux_driver nb7v904m_usb_redriver_drv;
 #endif /* __CROS_EC_USB_REDRIVER_NB7V904M_H */
diff --git a/driver/retimer/pi3dpx1207.c b/driver/retimer/pi3dpx1207.c
index 74eb80aac1..9254c3d4bb 100644
--- a/driver/retimer/pi3dpx1207.c
+++ b/driver/retimer/pi3dpx1207.c
@@ -21,8 +21,7 @@ static uint8_t buf[PI3DPX1207_NUM_REGISTERS];
 /**
  * Local utility functions
  */
-static int pi3dpx1207_i2c_write(int i2c_port,
-				uint16_t addr_flags,
+static int pi3dpx1207_i2c_write(const struct usb_mux *me,
 				uint8_t offset,
 				uint8_t val)
 {
@@ -45,7 +44,7 @@ static int pi3dpx1207_i2c_write(int i2c_port,
 		attempt = 0;
 		do {
 			attempt++;
-			rv = i2c_xfer(i2c_port, addr_flags,
+			rv = i2c_xfer(me->i2c_port, me->i2c_addr_flags,
 				      NULL, 0, buf, offset);
 		} while ((rv != EC_SUCCESS) && (attempt < I2C_MAX_RETRIES));
 	}
@@ -56,15 +55,16 @@ static int pi3dpx1207_i2c_write(int i2c_port,
 		attempt = 0;
 		do {
 			attempt++;
-			rv = i2c_xfer(i2c_port, addr_flags,
+			rv = i2c_xfer(me->i2c_port, me->i2c_addr_flags,
 				      buf, offset + 1, NULL, 0);
 		} while ((rv != EC_SUCCESS) && (attempt < I2C_MAX_RETRIES));
 	}
 	return rv;
 }
 
-static void pi3dpx1207_shutoff_power(int port)
+static void pi3dpx1207_shutoff_power(const struct usb_mux *me)
 {
+	const int port = me->usb_port;
 	const int gpio_enable = pi3dpx1207_controls[port].enable_gpio;
 	const int gpio_dp_enable = pi3dpx1207_controls[port].dp_enable_gpio;
 
@@ -75,27 +75,26 @@ static void pi3dpx1207_shutoff_power(int port)
 /**
  * Driver interface code
  */
-static int pi3dpx1207_init(int port)
+static int pi3dpx1207_init(const struct usb_mux *me)
 {
+	const int port = me->usb_port;
 	const int gpio_enable = pi3dpx1207_controls[port].enable_gpio;
 
 	gpio_or_ioex_set_level(gpio_enable, 1);
 	return EC_SUCCESS;
 }
 
-static int pi3dpx1207_enter_low_power_mode(int port)
+static int pi3dpx1207_enter_low_power_mode(const struct usb_mux *me)
 {
-	pi3dpx1207_shutoff_power(port);
+	pi3dpx1207_shutoff_power(me);
 	return EC_SUCCESS;
 }
 
-static int pi3dpx1207_set_mux(int port, mux_state_t mux_state)
+static int pi3dpx1207_set_mux(const struct usb_mux *me, mux_state_t mux_state)
 {
 	int rv = EC_SUCCESS;
 	uint8_t mode_val = PI3DPX1207_MODE_WATCHDOG_EN;
-
-	const int i2c_port = usb_retimers[port].i2c_port;
-	const uint16_t i2c_addr_flags = usb_retimers[port].i2c_addr_flags;
+	const int port = me->usb_port;
 	const int gpio_enable = pi3dpx1207_controls[port].enable_gpio;
 	const int gpio_dp_enable = pi3dpx1207_controls[port].dp_enable_gpio;
 
@@ -127,18 +126,16 @@ static int pi3dpx1207_set_mux(int port, mux_state_t mux_state)
 	}
 	/* Nothing enabled, power down the retimer */
 	else {
-		pi3dpx1207_shutoff_power(port);
+		pi3dpx1207_shutoff_power(me);
 		return EC_SUCCESS;
 	}
 
 	/* Write the retimer config byte */
-	rv = pi3dpx1207_i2c_write(i2c_port, i2c_addr_flags,
-				  PI3DPX1207_MODE_OFFSET,
-				  mode_val);
+	rv = pi3dpx1207_i2c_write(me, PI3DPX1207_MODE_OFFSET, mode_val);
 	return rv;
 }
 
-const struct usb_retimer_driver pi3dpx1207_usb_retimer = {
+const struct usb_mux_driver pi3dpx1207_usb_retimer = {
 	.init = pi3dpx1207_init,
 	.set = pi3dpx1207_set_mux,
 	.enter_low_power_mode = pi3dpx1207_enter_low_power_mode,
diff --git a/driver/retimer/pi3dpx1207.h b/driver/retimer/pi3dpx1207.h
index 53b849b081..2e3405d1aa 100644
--- a/driver/retimer/pi3dpx1207.h
+++ b/driver/retimer/pi3dpx1207.h
@@ -40,7 +40,7 @@
 #define PI3DPX1207_MODE_CONF_USB_SUPER		0xC0
 
 /* Supported USB retimer drivers */
-extern const struct usb_retimer_driver pi3dpx1207_usb_retimer;
+extern const struct usb_mux_driver pi3dpx1207_usb_retimer;
 
 /* Retimer driver hardware specific controls */
 struct pi3dpx1207_usb_control {
diff --git a/driver/retimer/ps8802.c b/driver/retimer/ps8802.c
index 94ea2e5797..c4a9a5478b 100644
--- a/driver/retimer/ps8802.c
+++ b/driver/retimer/ps8802.c
@@ -16,47 +16,47 @@
 #define PS8802_DEBUG 0
 #define PS8802_I2C_WAKE_DELAY 500
 
-int ps8802_i2c_read(int port, int page, int offset, int *data)
+int ps8802_i2c_read(const struct usb_mux *me, int page, int offset, int *data)
 {
 	int rv;
 
-	rv = i2c_read8(usb_retimers[port].i2c_port,
-		       usb_retimers[port].i2c_addr_flags + page,
+	rv = i2c_read8(me->i2c_port,
+		       me->i2c_addr_flags + page,
 		       offset, data);
 
 	if (PS8802_DEBUG)
 		ccprintf("%s(%d:0x%02X, 0x%02X) =>0x%02X\n", __func__,
-			 usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+			 me->i2c_port,
+			 me->i2c_addr_flags + page,
 			 offset, *data);
 
 	return rv;
 }
 
-int ps8802_i2c_write(int port, int page, int offset, int data)
+int ps8802_i2c_write(const struct usb_mux *me, int page, int offset, int data)
 {
 	int rv;
 	int pre_val, post_val;
 
 	if (PS8802_DEBUG)
-		i2c_read8(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags + page,
+		i2c_read8(me->i2c_port,
+			me->i2c_addr_flags + page,
 			offset, &pre_val);
 
-	rv = i2c_write8(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags + page,
+	rv = i2c_write8(me->i2c_port,
+			me->i2c_addr_flags + page,
 			offset, data);
 
 	if (PS8802_DEBUG) {
-		i2c_read8(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags + page,
+		i2c_read8(me->i2c_port,
+			me->i2c_addr_flags + page,
 			offset, &post_val);
 
 		ccprintf("%s(%d:0x%02X, 0x%02X, 0x%02X) "
 			"0x%02X=>0x%02X\n",
 			 __func__,
-			 usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+			 me->i2c_port,
+			 me->i2c_addr_flags + page,
 			 offset, data,
 			 pre_val, post_val);
 	}
@@ -64,30 +64,31 @@ int ps8802_i2c_write(int port, int page, int offset, int data)
 	return rv;
 }
 
-int ps8802_i2c_write16(int port, int page, int offset, int data)
+int ps8802_i2c_write16(const struct usb_mux *me, int page, int offset,
+		       int data)
 {
 	int rv;
 	int pre_val, post_val;
 
 	if (PS8802_DEBUG)
-		i2c_read16(usb_retimers[port].i2c_port,
-			   usb_retimers[port].i2c_addr_flags + page,
+		i2c_read16(me->i2c_port,
+			   me->i2c_addr_flags + page,
 			   offset, &pre_val);
 
-	rv = i2c_write16(usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+	rv = i2c_write16(me->i2c_port,
+			 me->i2c_addr_flags + page,
 			 offset, data);
 
 	if (PS8802_DEBUG) {
-		i2c_read16(usb_retimers[port].i2c_port,
-			   usb_retimers[port].i2c_addr_flags + page,
+		i2c_read16(me->i2c_port,
+			   me->i2c_addr_flags + page,
 			   offset, &post_val);
 
 		ccprintf("%s(%d:0x%02X, 0x%02X, 0x%04X) "
 			 "0x%04X=>0x%04X\n",
 			 __func__,
-			 usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+			 me->i2c_port,
+			 me->i2c_addr_flags + page,
 			 offset, data,
 			 pre_val, post_val);
 	}
@@ -95,33 +96,33 @@ int ps8802_i2c_write16(int port, int page, int offset, int data)
 	return rv;
 }
 
-int ps8802_i2c_field_update8(int port, int page, int offset,
+int ps8802_i2c_field_update8(const struct usb_mux *me, int page, int offset,
 			     uint8_t field_mask, uint8_t set_value)
 {
 	int rv;
 	int pre_val, post_val;
 
 	if (PS8802_DEBUG)
-		i2c_read8(usb_retimers[port].i2c_port,
-			  usb_retimers[port].i2c_addr_flags + page,
+		i2c_read8(me->i2c_port,
+			  me->i2c_addr_flags + page,
 			  offset, &pre_val);
 
-	rv = i2c_field_update8(usb_retimers[port].i2c_port,
-			       usb_retimers[port].i2c_addr_flags + page,
+	rv = i2c_field_update8(me->i2c_port,
+			       me->i2c_addr_flags + page,
 			       offset,
 			       field_mask,
 			       set_value);
 
 	if (PS8802_DEBUG) {
-		i2c_read8(usb_retimers[port].i2c_port,
-			  usb_retimers[port].i2c_addr_flags + page,
+		i2c_read8(me->i2c_port,
+			  me->i2c_addr_flags + page,
 			  offset, &post_val);
 
 		ccprintf("%s(%d:0x%02X, 0x%02X, 0x%02X, 0x%02X) "
 			 "0x%02X=>0x%02X\n",
 			 __func__,
-			 usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+			 me->i2c_port,
+			 me->i2c_addr_flags + page,
 			 offset, field_mask, set_value,
 			 pre_val, post_val);
 	}
@@ -129,33 +130,33 @@ int ps8802_i2c_field_update8(int port, int page, int offset,
 	return rv;
 }
 
-int ps8802_i2c_field_update16(int port, int page, int offset,
+int ps8802_i2c_field_update16(const struct usb_mux *me, int page, int offset,
 			     uint16_t field_mask, uint16_t set_value)
 {
 	int rv;
 	int pre_val, post_val;
 
 	if (PS8802_DEBUG)
-		i2c_read16(usb_retimers[port].i2c_port,
-			   usb_retimers[port].i2c_addr_flags + page,
+		i2c_read16(me->i2c_port,
+			   me->i2c_addr_flags + page,
 			   offset, &pre_val);
 
-	rv = i2c_field_update16(usb_retimers[port].i2c_port,
-				usb_retimers[port].i2c_addr_flags + page,
+	rv = i2c_field_update16(me->i2c_port,
+				me->i2c_addr_flags + page,
 				offset,
 				field_mask,
 				set_value);
 
 	if (PS8802_DEBUG) {
-		i2c_read16(usb_retimers[port].i2c_port,
-			   usb_retimers[port].i2c_addr_flags + page,
+		i2c_read16(me->i2c_port,
+			   me->i2c_addr_flags + page,
 			   offset, &post_val);
 
 		ccprintf("%s(%d:0x%02X, 0x%02X, 0x%02X, 0x%04X) "
 			 "0x%04X=>0x%04X\n",
 			 __func__,
-			 usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+			 me->i2c_port,
+			 me->i2c_addr_flags + page,
 			 offset, field_mask, set_value,
 			 pre_val, post_val);
 	}
@@ -168,14 +169,14 @@ int ps8802_i2c_field_update16(int port, int page, int offset,
  * From Application Note: 1) Activate by reading any Page 2 register. 2) Wait
  * 500 microseconds. 3) After 5 seconds idle, PS8802 will return to standby.
  */
-int ps8802_i2c_wake(int port)
+int ps8802_i2c_wake(const struct usb_mux *me)
 {
 	int data;
 	int rv = EC_ERROR_UNKNOWN;
 
 	/* If in standby, first read will fail, second should succeed. */
 	for (int i = 0; i < 2; i++) {
-		rv = ps8802_i2c_read(port,
+		rv = ps8802_i2c_read(me,
 				     PS8802_REG_PAGE2,
 				     PS8802_REG2_MODE,
 				     &data);
@@ -188,39 +189,39 @@ int ps8802_i2c_wake(int port)
 	return rv;
 }
 
-int ps8802_detect(int port)
+int ps8802_detect(const struct usb_mux *me)
 {
 	int rv = EC_ERROR_NOT_POWERED;
 
 	/* Detected if we are powered and can read the device */
 	if (!chipset_in_state(CHIPSET_STATE_HARD_OFF))
-		rv = ps8802_i2c_wake(port);
+		rv = ps8802_i2c_wake(me);
 
 	return rv;
 }
 
-static int ps8802_init(int port)
+static int ps8802_init(const struct usb_mux *me)
 {
 	return EC_SUCCESS;
 }
 
-static int ps8802_set_mux(int port, mux_state_t mux_state)
+static int ps8802_set_mux(const struct usb_mux *me, mux_state_t mux_state)
 {
 	int val;
 	int rv;
 
 	if (chipset_in_state(CHIPSET_STATE_HARD_OFF))
 		return (mux_state == USB_PD_MUX_NONE) ? EC_SUCCESS
-						     : EC_ERROR_NOT_POWERED;
+						      : EC_ERROR_NOT_POWERED;
 
 	/* Make sure the PS8802 is awake */
-	rv = ps8802_i2c_wake(port);
+	rv = ps8802_i2c_wake(me);
 	if (rv)
 		return rv;
 
 	if (PS8802_DEBUG)
 		ccprintf("%s(%d, 0x%02X) %s %s %s\n",
-			 __func__, port, mux_state,
+			 __func__, me->usb_port, mux_state,
 			 (mux_state & USB_PD_MUX_USB_ENABLED)	? "USB" : "",
 			 (mux_state & USB_PD_MUX_DP_ENABLED)	? "DP" : "",
 			 (mux_state & USB_PD_MUX_POLARITY_INVERTED)
@@ -239,24 +240,15 @@ static int ps8802_set_mux(int port, mux_state_t mux_state)
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		val |= PS8802_MODE_FLIP_ENABLE;
 
-	rv = ps8802_i2c_write(port,
+	rv = ps8802_i2c_write(me,
 			      PS8802_REG_PAGE2,
 			      PS8802_REG2_MODE,
 			      val);
-	if (rv)
-		return rv;
-
-	/* Board specific retimer mux tuning */
-	if (usb_retimers[port].tune) {
-		rv = usb_retimers[port].tune(port, mux_state);
-		if (rv)
-			return rv;
-	}
 
 	return rv;
 }
 
-static int ps8802_get_mux(int port, mux_state_t *mux_state)
+static int ps8802_get_mux(const struct usb_mux *me, mux_state_t *mux_state)
 {
 	int rv;
 	int val;
@@ -266,11 +258,11 @@ static int ps8802_get_mux(int port, mux_state_t *mux_state)
 	if (chipset_in_state(CHIPSET_STATE_HARD_OFF))
 		return EC_ERROR_NOT_POWERED;
 
-	rv = ps8802_i2c_wake(port);
+	rv = ps8802_i2c_wake(me);
 	if (rv)
 		return rv;
 
-	rv = ps8802_i2c_read(port,
+	rv = ps8802_i2c_read(me,
 			     PS8802_REG_PAGE2,
 			     PS8802_REG2_MODE,
 			     &val);
@@ -287,17 +279,6 @@ static int ps8802_get_mux(int port, mux_state_t *mux_state)
 	return rv;
 }
 
-/*
- * PS8802 can look like a retimer or a MUX. So create both tables
- * and use them as needed, until retimers become a type of MUX and
- * then we will only need one of these tables.
- *
- * TODO(b:147593660) Cleanup of retimers as muxes in a more
- * generalized mechanism
- */
-const struct usb_retimer_driver ps8802_usb_retimer = {
-	.set = ps8802_set_mux,
-};
 const struct usb_mux_driver ps8802_usb_mux_driver = {
 	.init = ps8802_init,
 	.set = ps8802_set_mux,
diff --git a/driver/retimer/ps8802.h b/driver/retimer/ps8802.h
index 7c1b5c0cc8..2e4a301275 100644
--- a/driver/retimer/ps8802.h
+++ b/driver/retimer/ps8802.h
@@ -4,6 +4,7 @@
  *
  * PS8802 retimer.
  */
+#include "usb_mux.h"
 
 #ifndef __CROS_EC_USB_RETIMER_PS8802_H
 #define __CROS_EC_USB_RETIMER_PS8802_H
@@ -65,17 +66,17 @@
 
 
 extern const struct usb_mux_driver ps8802_usb_mux_driver;
-extern const struct usb_retimer_driver ps8802_usb_retimer;
 
-int ps8802_i2c_wake(int port);
-int ps8802_detect(int port);
+int ps8802_i2c_wake(const struct usb_mux *me);
+int ps8802_detect(const struct usb_mux *me);
 
-int ps8802_i2c_read(int port, int page, int offset, int *data);
-int ps8802_i2c_write(int port, int page, int offset, int data);
-int ps8802_i2c_write16(int port, int page, int offset, int data);
-int ps8802_i2c_field_update8(int port, int page, int offset,
+int ps8802_i2c_read(const struct usb_mux *me, int page, int offset, int *data);
+int ps8802_i2c_write(const struct usb_mux *me, int page, int offset, int data);
+int ps8802_i2c_write16(const struct usb_mux *me, int page, int offset,
+			int data);
+int ps8802_i2c_field_update8(const struct usb_mux *me, int page, int offset,
 			     uint8_t field_mask, uint8_t set_value);
-int ps8802_i2c_field_update16(int port, int page, int offset,
+int ps8802_i2c_field_update16(const struct usb_mux *me, int page, int offset,
 			     uint16_t field_mask, uint16_t set_value);
 
 #endif /* __CROS_EC_USB_RETIMER_PS8802_H */
diff --git a/driver/retimer/ps8818.c b/driver/retimer/ps8818.c
index 10dc70dca0..e70b8bf986 100644
--- a/driver/retimer/ps8818.c
+++ b/driver/retimer/ps8818.c
@@ -16,47 +16,47 @@
 
 #define PS8818_DEBUG 0
 
-int ps8818_i2c_read(int port, int page, int offset, int *data)
+int ps8818_i2c_read(const struct usb_mux *me, int page, int offset, int *data)
 {
 	int rv;
 
-	rv = i2c_read8(usb_retimers[port].i2c_port,
-		       usb_retimers[port].i2c_addr_flags + page,
+	rv = i2c_read8(me->i2c_port,
+		       me->i2c_addr_flags + page,
 		       offset, data);
 
 	if (PS8818_DEBUG)
 		ccprintf("%s(%d:0x%02X, 0x%02X) =>0x%02X\n", __func__,
-			 usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+			 me->usb_port,
+			 me->i2c_addr_flags + page,
 			 offset, *data);
 
 	return rv;
 }
 
-int ps8818_i2c_write(int port, int page, int offset, int data)
+int ps8818_i2c_write(const struct usb_mux *me, int page, int offset, int data)
 {
 	int rv;
 	int pre_val, post_val;
 
 	if (PS8818_DEBUG)
-		i2c_read8(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags + page,
+		i2c_read8(me->i2c_port,
+			me->i2c_addr_flags + page,
 			offset, &pre_val);
 
-	rv = i2c_write8(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags + page,
+	rv = i2c_write8(me->i2c_port,
+			me->i2c_addr_flags + page,
 			offset, data);
 
 	if (PS8818_DEBUG) {
-		i2c_read8(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags + page,
+		i2c_read8(me->i2c_port,
+			me->i2c_addr_flags + page,
 			offset, &post_val);
 
 		ccprintf("%s(%d:0x%02X, 0x%02X, 0x%02X) "
 			"0x%02X=>0x%02X\n",
 			 __func__,
-			 usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+			 me->usb_port,
+			 me->i2c_addr_flags + page,
 			 offset, data,
 			 pre_val, post_val);
 	}
@@ -64,33 +64,33 @@ int ps8818_i2c_write(int port, int page, int offset, int data)
 	return rv;
 }
 
-int ps8818_i2c_field_update8(int port, int page, int offset,
+int ps8818_i2c_field_update8(const struct usb_mux *me, int page, int offset,
 			     uint8_t field_mask, uint8_t set_value)
 {
 	int rv;
 	int pre_val, post_val;
 
 	if (PS8818_DEBUG)
-		i2c_read8(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags + page,
+		i2c_read8(me->i2c_port,
+			me->i2c_addr_flags + page,
 			offset, &pre_val);
 
-	rv = i2c_field_update8(usb_retimers[port].i2c_port,
-			       usb_retimers[port].i2c_addr_flags + page,
+	rv = i2c_field_update8(me->i2c_port,
+			       me->i2c_addr_flags + page,
 			       offset,
 			       field_mask,
 			       set_value);
 
 	if (PS8818_DEBUG) {
-		i2c_read8(usb_retimers[port].i2c_port,
-			usb_retimers[port].i2c_addr_flags + page,
+		i2c_read8(me->i2c_port,
+			me->i2c_addr_flags + page,
 			offset, &post_val);
 
 		ccprintf("%s(%d:0x%02X, 0x%02X, 0x%02X, 0x%02X) "
 			 "0x%02X=>0x%02X\n",
 			 __func__,
-			 usb_retimers[port].i2c_port,
-			 usb_retimers[port].i2c_addr_flags + page,
+			 me->usb_port,
+			 me->i2c_addr_flags + page,
 			 offset, field_mask, set_value,
 			 pre_val, post_val);
 	}
@@ -98,14 +98,14 @@ int ps8818_i2c_field_update8(int port, int page, int offset,
 	return rv;
 }
 
-int ps8818_detect(int port)
+int ps8818_detect(const struct usb_mux *me)
 {
 	int rv = EC_ERROR_NOT_POWERED;
 	int val;
 
 	/* Detected if we are powered and can read the device */
 	if (!chipset_in_state(CHIPSET_STATE_HARD_OFF))
-		rv = ps8818_i2c_read(port,
+		rv = ps8818_i2c_read(me,
 				     PS8818_REG_PAGE0,
 				     PS8818_REG0_FLIP,
 				     &val);
@@ -113,7 +113,7 @@ int ps8818_detect(int port)
 	return rv;
 }
 
-static int ps8818_set_mux(int port, mux_state_t mux_state)
+static int ps8818_set_mux(const struct usb_mux *me, mux_state_t mux_state)
 {
 	int rv;
 	int val = 0;
@@ -124,7 +124,7 @@ static int ps8818_set_mux(int port, mux_state_t mux_state)
 
 	if (PS8818_DEBUG)
 		ccprintf("%s(%d, 0x%02X) %s %s %s\n",
-			 __func__, port, mux_state,
+			 __func__, me->usb_port, mux_state,
 			 (mux_state & USB_PD_MUX_USB_ENABLED)	? "USB" : "",
 			 (mux_state & USB_PD_MUX_DP_ENABLED)	? "DP" : "",
 			 (mux_state & USB_PD_MUX_POLARITY_INVERTED)
@@ -136,7 +136,7 @@ static int ps8818_set_mux(int port, mux_state_t mux_state)
 	if (mux_state & USB_PD_MUX_DP_ENABLED)
 		val |= PS8818_MODE_DP_ENABLE;
 
-	rv = ps8818_i2c_field_update8(port,
+	rv = ps8818_i2c_field_update8(me,
 				PS8818_REG_PAGE0,
 				PS8818_REG0_MODE,
 				PS8818_MODE_NON_RESERVED_MASK,
@@ -149,7 +149,7 @@ static int ps8818_set_mux(int port, mux_state_t mux_state)
 	if (mux_state & USB_PD_MUX_POLARITY_INVERTED)
 		val |= PS8818_FLIP_CONFIG;
 
-	rv = ps8818_i2c_field_update8(port,
+	rv = ps8818_i2c_field_update8(me,
 				PS8818_REG_PAGE0,
 				PS8818_REG0_FLIP,
 				PS8818_FLIP_NON_RESERVED_MASK,
@@ -162,24 +162,15 @@ static int ps8818_set_mux(int port, mux_state_t mux_state)
 	if (mux_state & USB_PD_MUX_DP_ENABLED)
 		val |= PS8818_DPHPD_PLUGGED;
 
-	rv = ps8818_i2c_field_update8(port,
+	rv = ps8818_i2c_field_update8(me,
 				PS8818_REG_PAGE0,
 				PS8818_REG0_DPHPD_CONFIG,
 				PS8818_DPHPD_NON_RESERVED_MASK,
 				val);
-	if (rv)
-		return rv;
-
-	/* Board specific retimer mux tuning */
-	if (usb_retimers[port].tune) {
-		rv = usb_retimers[port].tune(port, mux_state);
-		if (rv)
-			return rv;
-	}
 
 	return rv;
 }
 
-const struct usb_retimer_driver ps8818_usb_retimer = {
+const struct usb_mux_driver ps8818_usb_retimer_driver = {
 	.set = ps8818_set_mux,
 };
diff --git a/driver/retimer/ps8818.h b/driver/retimer/ps8818.h
index a28d04a38b..c761284cf8 100644
--- a/driver/retimer/ps8818.h
+++ b/driver/retimer/ps8818.h
@@ -4,6 +4,7 @@
  *
  * PS8818 retimer.
  */
+#include "usb_mux.h"
 
 #ifndef __CROS_EC_USB_RETIMER_PS8818_H
 #define __CROS_EC_USB_RETIMER_PS8818_H
@@ -77,13 +78,16 @@
 #define PS8818_STATUS_NORMAL_OPERATION		BIT(7)
 #define PS8818_STATUS_10_GBPS			BIT(5)
 
-extern const struct usb_retimer_driver ps8818_usb_retimer;
+extern const struct usb_mux_driver ps8818_usb_retimer_driver;
 
-int ps8818_detect(int port);
+int ps8818_detect(const struct usb_mux *me);
 
-int ps8818_i2c_read(int port, int page, int offset, int *data);
-int ps8818_i2c_write(int port, int page, int offset, int data);
-int ps8818_i2c_field_update8(int port, int page, int offset,
+int ps8818_i2c_read(const struct usb_mux *me,
+		    int page, int offset, int *data);
+int ps8818_i2c_write(const struct usb_mux *me,
+		     int page, int offset, int data);
+int ps8818_i2c_field_update8(const struct usb_mux *me,
+			     int page, int offset,
 			     uint8_t field_mask, uint8_t set_value);
 
 #endif /* __CROS_EC_USB_RETIMER_PS8818_H */