#include "class.h"
 #include "mux.h"
 
+struct typec_switch {
+       struct typec_switch_dev *sw_dev;
+};
+
 static int switch_fwnode_match(struct device *dev, const void *fwnode)
 {
-       if (!is_typec_switch(dev))
+       if (!is_typec_switch_dev(dev))
                return 0;
 
        return dev_fwnode(dev) == fwnode;
        dev = class_find_device(&typec_mux_class, NULL, fwnode,
                                switch_fwnode_match);
 
-       return dev ? to_typec_switch(dev) : ERR_PTR(-EPROBE_DEFER);
+       return dev ? to_typec_switch_dev(dev) : ERR_PTR(-EPROBE_DEFER);
 }
 
 /**
  */
 struct typec_switch *fwnode_typec_switch_get(struct fwnode_handle *fwnode)
 {
+       struct typec_switch_dev *sw_dev;
        struct typec_switch *sw;
 
-       sw = fwnode_connection_find_match(fwnode, "orientation-switch", NULL,
-                                         typec_switch_match);
-       if (!IS_ERR_OR_NULL(sw))
-               WARN_ON(!try_module_get(sw->dev.parent->driver->owner));
+       sw = kzalloc(sizeof(*sw), GFP_KERNEL);
+       if (!sw)
+               return ERR_PTR(-ENOMEM);
+
+       sw_dev = fwnode_connection_find_match(fwnode, "orientation-switch", NULL,
+                                             typec_switch_match);
+       if (IS_ERR_OR_NULL(sw_dev)) {
+               kfree(sw);
+               return ERR_CAST(sw_dev);
+       }
+
+       WARN_ON(!try_module_get(sw_dev->dev.parent->driver->owner));
+
+       sw->sw_dev = sw_dev;
 
        return sw;
 }
  */
 void typec_switch_put(struct typec_switch *sw)
 {
-       if (!IS_ERR_OR_NULL(sw)) {
-               module_put(sw->dev.parent->driver->owner);
-               put_device(&sw->dev);
-       }
+       struct typec_switch_dev *sw_dev;
+
+       if (IS_ERR_OR_NULL(sw))
+               return;
+
+       sw_dev = sw->sw_dev;
+
+       module_put(sw_dev->dev.parent->driver->owner);
+       put_device(&sw_dev->dev);
+       kfree(sw);
 }
 EXPORT_SYMBOL_GPL(typec_switch_put);
 
 static void typec_switch_release(struct device *dev)
 {
-       kfree(to_typec_switch(dev));
+       kfree(to_typec_switch_dev(dev));
 }
 
 const struct device_type typec_switch_dev_type = {
  * connector to the USB controllers. USB Type-C plugs can be inserted
  * right-side-up or upside-down.
  */
-struct typec_switch *
+struct typec_switch_dev *
 typec_switch_register(struct device *parent,
                      const struct typec_switch_desc *desc)
 {
-       struct typec_switch *sw;
+       struct typec_switch_dev *sw_dev;
        int ret;
 
        if (!desc || !desc->set)
                return ERR_PTR(-EINVAL);
 
-       sw = kzalloc(sizeof(*sw), GFP_KERNEL);
-       if (!sw)
+       sw_dev = kzalloc(sizeof(*sw_dev), GFP_KERNEL);
+       if (!sw_dev)
                return ERR_PTR(-ENOMEM);
 
-       sw->set = desc->set;
+       sw_dev->set = desc->set;
 
-       device_initialize(&sw->dev);
-       sw->dev.parent = parent;
-       sw->dev.fwnode = desc->fwnode;
-       sw->dev.class = &typec_mux_class;
-       sw->dev.type = &typec_switch_dev_type;
-       sw->dev.driver_data = desc->drvdata;
-       ret = dev_set_name(&sw->dev, "%s-switch", desc->name ? desc->name : dev_name(parent));
+       device_initialize(&sw_dev->dev);
+       sw_dev->dev.parent = parent;
+       sw_dev->dev.fwnode = desc->fwnode;
+       sw_dev->dev.class = &typec_mux_class;
+       sw_dev->dev.type = &typec_switch_dev_type;
+       sw_dev->dev.driver_data = desc->drvdata;
+       ret = dev_set_name(&sw_dev->dev, "%s-switch", desc->name ? desc->name : dev_name(parent));
        if (ret) {
-               put_device(&sw->dev);
+               put_device(&sw_dev->dev);
                return ERR_PTR(ret);
        }
 
-       ret = device_add(&sw->dev);
+       ret = device_add(&sw_dev->dev);
        if (ret) {
                dev_err(parent, "failed to register switch (%d)\n", ret);
-               put_device(&sw->dev);
+               put_device(&sw_dev->dev);
                return ERR_PTR(ret);
        }
 
-       return sw;
+       return sw_dev;
 }
 EXPORT_SYMBOL_GPL(typec_switch_register);
 
 int typec_switch_set(struct typec_switch *sw,
                     enum typec_orientation orientation)
 {
+       struct typec_switch_dev *sw_dev;
+
        if (IS_ERR_OR_NULL(sw))
                return 0;
 
-       return sw->set(sw, orientation);
+       sw_dev = sw->sw_dev;
+
+       return sw_dev->set(sw_dev, orientation);
 }
 EXPORT_SYMBOL_GPL(typec_switch_set);
 
 /**
  * typec_switch_unregister - Unregister USB Type-C orientation switch
- * @sw: USB Type-C orientation switch
+ * @sw_dev: USB Type-C orientation switch
  *
  * Unregister switch that was registered with typec_switch_register().
  */
-void typec_switch_unregister(struct typec_switch *sw)
+void typec_switch_unregister(struct typec_switch_dev *sw_dev)
 {
-       if (!IS_ERR_OR_NULL(sw))
-               device_unregister(&sw->dev);
+       if (!IS_ERR_OR_NULL(sw_dev))
+               device_unregister(&sw_dev->dev);
 }
 EXPORT_SYMBOL_GPL(typec_switch_unregister);
 
-void typec_switch_set_drvdata(struct typec_switch *sw, void *data)
+void typec_switch_set_drvdata(struct typec_switch_dev *sw_dev, void *data)
 {
-       dev_set_drvdata(&sw->dev, data);
+       dev_set_drvdata(&sw_dev->dev, data);
 }
 EXPORT_SYMBOL_GPL(typec_switch_set_drvdata);
 
-void *typec_switch_get_drvdata(struct typec_switch *sw)
+void *typec_switch_get_drvdata(struct typec_switch_dev *sw_dev)
 {
-       return dev_get_drvdata(&sw->dev);
+       return dev_get_drvdata(&sw_dev->dev);
 }
 EXPORT_SYMBOL_GPL(typec_switch_get_drvdata);
 
 /* ------------------------------------------------------------------------- */
 
+struct typec_mux {
+       struct typec_mux_dev *mux_dev;
+};
+
 static int mux_fwnode_match(struct device *dev, const void *fwnode)
 {
-       if (!is_typec_mux(dev))
+       if (!is_typec_mux_dev(dev))
                return 0;
 
        return dev_fwnode(dev) == fwnode;
        dev = class_find_device(&typec_mux_class, NULL, fwnode,
                                mux_fwnode_match);
 
-       return dev ? to_typec_mux(dev) : ERR_PTR(-EPROBE_DEFER);
+       return dev ? to_typec_mux_dev(dev) : ERR_PTR(-EPROBE_DEFER);
 }
 
 /**
 struct typec_mux *fwnode_typec_mux_get(struct fwnode_handle *fwnode,
                                       const struct typec_altmode_desc *desc)
 {
+       struct typec_mux_dev *mux_dev;
        struct typec_mux *mux;
 
-       mux = fwnode_connection_find_match(fwnode, "mode-switch", (void *)desc,
-                                          typec_mux_match);
-       if (!IS_ERR_OR_NULL(mux))
-               WARN_ON(!try_module_get(mux->dev.parent->driver->owner));
+       mux = kzalloc(sizeof(*mux), GFP_KERNEL);
+       if (!mux)
+               return ERR_PTR(-ENOMEM);
+
+       mux_dev = fwnode_connection_find_match(fwnode, "mode-switch", (void *)desc,
+                                              typec_mux_match);
+       if (IS_ERR_OR_NULL(mux_dev)) {
+               kfree(mux);
+               return ERR_CAST(mux_dev);
+       }
+
+       WARN_ON(!try_module_get(mux_dev->dev.parent->driver->owner));
+
+       mux->mux_dev = mux_dev;
 
        return mux;
 }
  */
 void typec_mux_put(struct typec_mux *mux)
 {
-       if (!IS_ERR_OR_NULL(mux)) {
-               module_put(mux->dev.parent->driver->owner);
-               put_device(&mux->dev);
-       }
+       struct typec_mux_dev *mux_dev;
+
+       if (IS_ERR_OR_NULL(mux))
+               return;
+
+       mux_dev = mux->mux_dev;
+       module_put(mux_dev->dev.parent->driver->owner);
+       put_device(&mux_dev->dev);
+       kfree(mux);
 }
 EXPORT_SYMBOL_GPL(typec_mux_put);
 
 int typec_mux_set(struct typec_mux *mux, struct typec_mux_state *state)
 {
+       struct typec_mux_dev *mux_dev;
+
        if (IS_ERR_OR_NULL(mux))
                return 0;
 
-       return mux->set(mux, state);
+       mux_dev = mux->mux_dev;
+
+       return mux_dev->set(mux_dev, state);
 }
 EXPORT_SYMBOL_GPL(typec_mux_set);
 
 static void typec_mux_release(struct device *dev)
 {
-       kfree(to_typec_mux(dev));
+       kfree(to_typec_mux_dev(dev));
 }
 
 const struct device_type typec_mux_dev_type = {
  * the pins on the connector need to be reconfigured. This function registers
  * multiplexer switches routing the pins on the connector.
  */
-struct typec_mux *
+struct typec_mux_dev *
 typec_mux_register(struct device *parent, const struct typec_mux_desc *desc)
 {
-       struct typec_mux *mux;
+       struct typec_mux_dev *mux_dev;
        int ret;
 
        if (!desc || !desc->set)
                return ERR_PTR(-EINVAL);
 
-       mux = kzalloc(sizeof(*mux), GFP_KERNEL);
-       if (!mux)
+       mux_dev = kzalloc(sizeof(*mux_dev), GFP_KERNEL);
+       if (!mux_dev)
                return ERR_PTR(-ENOMEM);
 
-       mux->set = desc->set;
+       mux_dev->set = desc->set;
 
-       device_initialize(&mux->dev);
-       mux->dev.parent = parent;
-       mux->dev.fwnode = desc->fwnode;
-       mux->dev.class = &typec_mux_class;
-       mux->dev.type = &typec_mux_dev_type;
-       mux->dev.driver_data = desc->drvdata;
-       ret = dev_set_name(&mux->dev, "%s-mux", desc->name ? desc->name : dev_name(parent));
+       device_initialize(&mux_dev->dev);
+       mux_dev->dev.parent = parent;
+       mux_dev->dev.fwnode = desc->fwnode;
+       mux_dev->dev.class = &typec_mux_class;
+       mux_dev->dev.type = &typec_mux_dev_type;
+       mux_dev->dev.driver_data = desc->drvdata;
+       ret = dev_set_name(&mux_dev->dev, "%s-mux", desc->name ? desc->name : dev_name(parent));
        if (ret) {
-               put_device(&mux->dev);
+               put_device(&mux_dev->dev);
                return ERR_PTR(ret);
        }
 
-       ret = device_add(&mux->dev);
+       ret = device_add(&mux_dev->dev);
        if (ret) {
                dev_err(parent, "failed to register mux (%d)\n", ret);
-               put_device(&mux->dev);
+               put_device(&mux_dev->dev);
                return ERR_PTR(ret);
        }
 
-       return mux;
+       return mux_dev;
 }
 EXPORT_SYMBOL_GPL(typec_mux_register);
 
 /**
  * typec_mux_unregister - Unregister Multiplexer Switch
- * @mux: USB Type-C Connector Multiplexer/DeMultiplexer
+ * @mux_dev: USB Type-C Connector Multiplexer/DeMultiplexer
  *
  * Unregister mux that was registered with typec_mux_register().
  */
-void typec_mux_unregister(struct typec_mux *mux)
+void typec_mux_unregister(struct typec_mux_dev *mux_dev)
 {
-       if (!IS_ERR_OR_NULL(mux))
-               device_unregister(&mux->dev);
+       if (!IS_ERR_OR_NULL(mux_dev))
+               device_unregister(&mux_dev->dev);
 }
 EXPORT_SYMBOL_GPL(typec_mux_unregister);
 
-void typec_mux_set_drvdata(struct typec_mux *mux, void *data)
+void typec_mux_set_drvdata(struct typec_mux_dev *mux_dev, void *data)
 {
-       dev_set_drvdata(&mux->dev, data);
+       dev_set_drvdata(&mux_dev->dev, data);
 }
 EXPORT_SYMBOL_GPL(typec_mux_set_drvdata);
 
-void *typec_mux_get_drvdata(struct typec_mux *mux)
+void *typec_mux_get_drvdata(struct typec_mux_dev *mux_dev)
 {
-       return dev_get_drvdata(&mux->dev);
+       return dev_get_drvdata(&mux_dev->dev);
 }
 EXPORT_SYMBOL_GPL(typec_mux_get_drvdata);