mmc: restore capacity when switching to partition 0

The capacity and lba for an MMC device with part_num 0 reflects the
whole device.  When mmc_switch_part() successfully switches to a
partition, the capacity is changed to that partition.  As partition 0
does not physically exist, attempts to switch back to the whole device
will indicate an error, but the capacity setting for the whole device
must still be restored to match the partition.

Signed-off-by: Peter A. Bigot <pab@pabigot.com>
Tested-by: Tom Rini <trini@ti.com>
Acked-by: Pantelis Antoniou <panto@antoniou-consulting.com>
diff --git a/drivers/mmc/mmc.c b/drivers/mmc/mmc.c
index 52a8e36..e4a28d4 100644
--- a/drivers/mmc/mmc.c
+++ b/drivers/mmc/mmc.c
@@ -594,10 +594,15 @@
 	ret = mmc_switch(mmc, EXT_CSD_CMD_SET_NORMAL, EXT_CSD_PART_CONF,
 			 (mmc->part_config & ~PART_ACCESS_MASK)
 			 | (part_num & PART_ACCESS_MASK));
-	if (ret)
-		return ret;
 
-	return mmc_set_capacity(mmc, part_num);
+	/*
+	 * Set the capacity if the switch succeeded or was intended
+	 * to return to representing the raw device.
+	 */
+	if ((ret == 0) || ((ret == -ENODEV) && (part_num == 0)))
+		ret = mmc_set_capacity(mmc, part_num);
+
+	return ret;
 }
 
 int mmc_getcd(struct mmc *mmc)