mmc: report capacity for the selected partition

Enhance the MMC core to calculate the size of each MMC partition, and
update mmc->capacity whenever a partition is selected. This causes:

mmc dev 0 1 ; mmcinfo

... to report the size of the currently selected partition, rather than
always reporting the size of the user partition.

Signed-off-by: Stephen Warren <swarren@nvidia.com>
Signed-off-by: Andy Fleming <afleming@freescale.com>
diff --git a/drivers/mmc/mmc.c b/drivers/mmc/mmc.c
index fe83934..e6a296a 100644
--- a/drivers/mmc/mmc.c
+++ b/drivers/mmc/mmc.c
@@ -702,16 +702,49 @@
 	return 0;
 }
 
+static int mmc_set_capacity(struct mmc *mmc, int part_num)
+{
+	switch (part_num) {
+	case 0:
+		mmc->capacity = mmc->capacity_user;
+		break;
+	case 1:
+	case 2:
+		mmc->capacity = mmc->capacity_boot;
+		break;
+	case 3:
+		mmc->capacity = mmc->capacity_rpmb;
+		break;
+	case 4:
+	case 5:
+	case 6:
+	case 7:
+		mmc->capacity = mmc->capacity_gp[part_num - 4];
+		break;
+	default:
+		return -1;
+	}
+
+	mmc->block_dev.lba = lldiv(mmc->capacity, mmc->read_bl_len);
+
+	return 0;
+}
+
 int mmc_switch_part(int dev_num, unsigned int part_num)
 {
 	struct mmc *mmc = find_mmc_device(dev_num);
+	int ret;
 
 	if (!mmc)
 		return -1;
 
-	return mmc_switch(mmc, EXT_CSD_CMD_SET_NORMAL, EXT_CSD_PART_CONF,
-			  (mmc->part_config & ~PART_ACCESS_MASK)
-			  | (part_num & PART_ACCESS_MASK));
+	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);
 }
 
 int mmc_getcd(struct mmc *mmc)
@@ -919,7 +952,7 @@
 
 static int mmc_startup(struct mmc *mmc)
 {
-	int err;
+	int err, i;
 	uint mult, freq;
 	u64 cmult, csize, capacity;
 	struct mmc_cmd cmd;
@@ -1037,8 +1070,12 @@
 		cmult = (mmc->csd[2] & 0x00038000) >> 15;
 	}
 
-	mmc->capacity = (csize + 1) << (cmult + 2);
-	mmc->capacity *= mmc->read_bl_len;
+	mmc->capacity_user = (csize + 1) << (cmult + 2);
+	mmc->capacity_user *= mmc->read_bl_len;
+	mmc->capacity_boot = 0;
+	mmc->capacity_rpmb = 0;
+	for (i = 0; i < 4; i++)
+		mmc->capacity_gp[i] = 0;
 
 	if (mmc->read_bl_len > MMC_MAX_BLOCK_LEN)
 		mmc->read_bl_len = MMC_MAX_BLOCK_LEN;
@@ -1077,7 +1114,7 @@
 					| ext_csd[EXT_CSD_SEC_CNT + 3] << 24;
 			capacity *= MMC_MAX_BLOCK_LEN;
 			if ((capacity >> 20) > 2 * 1024)
-				mmc->capacity = capacity;
+				mmc->capacity_user = capacity;
 		}
 
 		switch (ext_csd[EXT_CSD_REV]) {
@@ -1119,8 +1156,25 @@
 		if ((ext_csd[EXT_CSD_PARTITIONING_SUPPORT] & PART_SUPPORT) ||
 		    ext_csd[EXT_CSD_BOOT_MULT])
 			mmc->part_config = ext_csd[EXT_CSD_PART_CONF];
+
+		mmc->capacity_boot = ext_csd[EXT_CSD_BOOT_MULT] << 17;
+
+		mmc->capacity_rpmb = ext_csd[EXT_CSD_RPMB_MULT] << 17;
+
+		for (i = 0; i < 4; i++) {
+			int idx = EXT_CSD_GP_SIZE_MULT + i * 3;
+			mmc->capacity_gp[i] = (ext_csd[idx + 2] << 16) +
+				(ext_csd[idx + 1] << 8) + ext_csd[idx];
+			mmc->capacity_gp[i] *=
+				ext_csd[EXT_CSD_HC_ERASE_GRP_SIZE];
+			mmc->capacity_gp[i] *= ext_csd[EXT_CSD_HC_WP_GRP_SIZE];
+		}
 	}
 
+	err = mmc_set_capacity(mmc, mmc->part_num);
+	if (err)
+		return err;
+
 	if (IS_SD(mmc))
 		err = sd_change_freq(mmc);
 	else
diff --git a/include/mmc.h b/include/mmc.h
index 566db59..ea198d8 100644
--- a/include/mmc.h
+++ b/include/mmc.h
@@ -158,7 +158,9 @@
 /*
  * EXT_CSD fields
  */
+#define EXT_CSD_GP_SIZE_MULT		143	/* R/W */
 #define EXT_CSD_PARTITIONING_SUPPORT	160	/* RO */
+#define EXT_CSD_RPMB_MULT		168	/* RO */
 #define EXT_CSD_ERASE_GROUP_DEF		175	/* R/W */
 #define EXT_CSD_PART_CONF		179	/* R/W */
 #define EXT_CSD_BUS_WIDTH		183	/* R/W */
@@ -166,6 +168,7 @@
 #define EXT_CSD_REV			192	/* RO */
 #define EXT_CSD_CARD_TYPE		196	/* RO */
 #define EXT_CSD_SEC_CNT			212	/* RO, 4 bytes */
+#define EXT_CSD_HC_WP_GRP_SIZE		221	/* RO */
 #define EXT_CSD_HC_ERASE_GRP_SIZE	224	/* RO */
 #define EXT_CSD_BOOT_MULT		226	/* RO */
 
@@ -263,6 +266,10 @@
 	uint write_bl_len;
 	uint erase_grp_size;
 	u64 capacity;
+	u64 capacity_user;
+	u64 capacity_boot;
+	u64 capacity_rpmb;
+	u64 capacity_gp[4];
 	block_dev_desc_t block_dev;
 	int (*send_cmd)(struct mmc *mmc,
 			struct mmc_cmd *cmd, struct mmc_data *data);