diff --git a/Android.mk b/Android.mk
index 14873fb..3e664cc 100644
--- a/Android.mk
+++ b/Android.mk
@@ -5,7 +5,7 @@
 LOCAL_C_INCLUDES += $(TARGET_OUT_HEADERS)/gpt-utils/inc
 LOCAL_CFLAGS += -Wall -Werror
 LOCAL_SHARED_LIBRARIES += liblog librecovery_updater_msm libcutils
-LOCAL_SRC_FILES := boot_control.c
+LOCAL_SRC_FILES := boot_control.cpp
 LOCAL_MODULE_RELATIVE_PATH := hw
 LOCAL_MODULE := bootctrl.$(TARGET_BOARD_PLATFORM)
 include $(BUILD_SHARED_LIBRARY)
diff --git a/boot_control.c b/boot_control.cpp
similarity index 72%
rename from boot_control.c
rename to boot_control.cpp
index d5e2e4b..539c52c 100644
--- a/boot_control.c
+++ b/boot_control.cpp
@@ -26,6 +26,13 @@
  * OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
  * IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  */
+#include <map>
+#include <list>
+#include <string>
+#include <vector>
+#ifdef __cplusplus
+extern "C" {
+#endif
 #include <errno.h>
 #define LOG_TAG "bootcontrolhal"
 #include <cutils/log.h>
@@ -46,6 +53,18 @@
 #define LUN_NAME_END_LOC 14
 #define BOOT_SLOT_PROP "ro.boot.slot_suffix"
 
+#define SLOT_ACTIVE 1
+#define SLOT_INACTIVE 2
+#define UPDATE_SLOT(pentry, guid, slot_state) ({ \
+		memcpy(pentry, guid, TYPE_GUID_SIZE); \
+		if (slot_state == SLOT_ACTIVE)\
+			*(pentry + AB_FLAG_OFFSET) = AB_SLOT_ACTIVE_VAL; \
+		else if (slot_state == SLOT_INACTIVE) \
+		*(pentry + AB_FLAG_OFFSET)  = (*(pentry + AB_FLAG_OFFSET)& \
+			~AB_PARTITION_ATTR_SLOT_ACTIVE); \
+		})
+
+using namespace std;
 const char *slot_suffix_arr[] = {
 	AB_SLOT_A_SUFFIX,
 	AB_SLOT_B_SUFFIX,
@@ -169,7 +188,7 @@
 				"%s%s",
 				ptn_list[i],
 				slot);
-		disk = gpt_disk_alloc(disk);
+		disk = gpt_disk_alloc();
 		if (!disk) {
 			ALOGE("%s: Failed to alloc disk struct",
 					__func__);
@@ -294,36 +313,18 @@
 	return 0;
 }
 
-static unsigned get_current_active_slot(struct boot_control_module *module)
+static int boot_control_check_slot_sanity(struct boot_control_module *module,
+		unsigned slot)
 {
-	uint32_t num_slots = 0;
-	char bootPartition[MAX_GPT_NAME_SIZE + 1];
-	unsigned i = 0;
-	if (!module) {
-		ALOGE("%s: Invalid argument", __func__);
-		goto error;
+	if (!module)
+		return -1;
+	uint32_t num_slots = get_number_slots(module);
+	if ((num_slots < 1) || (slot > num_slots - 1)) {
+		ALOGE("Invalid slot number");
+		return -1;
 	}
-	num_slots = get_number_slots(module);
-	if (num_slots <= 1) {
-		//Slot 0 is the only slot around.
-		return 0;
-	}
-	//Iterate through a list of partitons named as boot+suffix
-	//and see which one is currently active.
-	for (i = 0; slot_suffix_arr[i] != NULL ; i++) {
-		memset(bootPartition, '\0', sizeof(bootPartition));
-		snprintf(bootPartition, sizeof(bootPartition) - 1,
-				"boot%s",
-				slot_suffix_arr[i]);
-		if (get_partition_attribute(bootPartition,
-					ATTR_SLOT_ACTIVE) == 1)
-			return i;
-	}
-error:
-	//The HAL spec requires that we return a number between
-	//0 to num_slots - 1. Since something went wrong here we
-	//are just going to return the default slot.
 	return 0;
+
 }
 
 int mark_boot_successful(struct boot_control_module *module)
@@ -346,133 +347,121 @@
 
 const char *get_suffix(struct boot_control_module *module, unsigned slot)
 {
-	unsigned num_slots = 0;
-	if (!module) {
-		ALOGE("%s: Invalid arg", __func__);
-	}
-	num_slots = get_number_slots(module);
-	if (num_slots < 1 || slot > num_slots - 1)
+	if (boot_control_check_slot_sanity(module, slot) != 0)
 		return NULL;
 	else
 		return slot_suffix_arr[slot];
 }
 
-int set_active_boot_slot(struct boot_control_module *module, unsigned slot)
+
+//Return a gpt disk structure representing the disk that holds
+//partition.
+static struct gpt_disk* boot_ctl_get_disk_info(char *partition)
 {
-	const char ptn_list[][MAX_GPT_NAME_SIZE] = { AB_PTN_LIST };
+	struct gpt_disk *disk = NULL;
+	if (!partition)
+		return NULL;
+	disk = gpt_disk_alloc();
+	if (!disk) {
+		ALOGE("%s: Failed to alloc disk",
+				__func__);
+		goto error;
+	}
+	if (gpt_disk_get_disk_info(partition, disk)) {
+		ALOGE("failed to get disk info for %s",
+				partition);
+		goto error;
+	}
+	return disk;
+error:
+	if (disk)
+		gpt_disk_free(disk);
+	return NULL;
+}
+
+//The argument here is a vector of partition names(including the slot suffix)
+//that lie on a single disk
+static int boot_ctl_set_active_slot_for_partitions(vector<string> part_list,
+		unsigned slot)
+{
+	char buf[PATH_MAX] = {0};
+	struct gpt_disk *disk = NULL;
 	char slotA[MAX_GPT_NAME_SIZE + 1] = {0};
 	char slotB[MAX_GPT_NAME_SIZE + 1] = {0};
 	char active_guid[TYPE_GUID_SIZE + 1] = {0};
 	char inactive_guid[TYPE_GUID_SIZE + 1] = {0};
-	struct gpt_disk *disk = NULL;
-	//Pointer to partition entry of current 'A' partition
+	//Pointer to the partition entry of current 'A' partition
 	uint8_t *pentryA = NULL;
 	uint8_t *pentryA_bak = NULL;
 	//Pointer to partition entry of current 'B' partition
 	uint8_t *pentryB = NULL;
 	uint8_t *pentryB_bak = NULL;
-	uint8_t *slot_info = NULL;
-	uint32_t i;
-	int rc = -1;
-	char buf[PATH_MAX] = {0};
 	struct stat st;
-	unsigned num_slots = 0;
-	unsigned current_slot = 0;
-	int is_ufs = gpt_utils_is_ufs_device();
+	vector<string>::iterator partition_iterator;
 
-	if (!module) {
-		ALOGE("%s: Invalid arg", __func__);
-		goto error;
-	}
-	num_slots = get_number_slots(module);
-	if ((num_slots < 1) || (slot > num_slots - 1)) {
-		ALOGE("%s: Unable to get num slots/Invalid slot value",
-				__func__);
-		goto error;
-	}
-	current_slot = get_current_active_slot(module);
-	if (current_slot == slot) {
-		//Nothing to do here. Just return
-		return 0;
-	}
-	for (i=0; i < ARRAY_SIZE(ptn_list); i++) {
-		//XBL is handled differrently for ufs devices
-		if (is_ufs && !strncmp(ptn_list[i], PTN_XBL, strlen(PTN_XBL)))
-				continue;
-		memset(buf, '\0', sizeof(buf));
+	for (partition_iterator = part_list.begin();
+			partition_iterator != part_list.end();
+			partition_iterator++) {
+		//Chop off the slot suffix from the partition name to
+		//make the string easier to work with.
+		string prefix = *partition_iterator;
+		if (prefix.size() < (strlen(AB_SLOT_A_SUFFIX) + 1)) {
+			ALOGE("Invalid partition name: %s", prefix.c_str());
+			goto error;
+		}
+		prefix.resize(prefix.size() - strlen(AB_SLOT_A_SUFFIX));
 		//Check if A/B versions of this ptn exist
-		snprintf(buf, sizeof(buf) - 1,
-                                        "%s/%s%s",
-                                        BOOT_DEV_DIR,
-                                        ptn_list[i],
-					AB_SLOT_A_SUFFIX
-					);
-		if (stat(buf, &st)) {
-			//partition does not have _a version
-			continue;
-		}
-		memset(buf, '\0', sizeof(buf));
-		snprintf(buf, sizeof(buf) - 1,
-                                        "%s/%s%s",
-                                        BOOT_DEV_DIR,
-                                        ptn_list[i],
-					AB_SLOT_B_SUFFIX
-					);
-		if (stat(buf, &st)) {
-			//partition does not have _a version
-			continue;
-		}
-		disk = gpt_disk_alloc();
-		if (!disk)
-			goto error;
-		memset(slotA, 0, sizeof(slotA));
-		memset(slotB, 0, sizeof(slotB));
-		snprintf(slotA, sizeof(slotA) - 1, "%s%s",
-				ptn_list[i],
+		snprintf(buf, sizeof(buf) - 1, "%s/%s%s", BOOT_DEV_DIR,
+				prefix.c_str(),
 				AB_SLOT_A_SUFFIX);
-		snprintf(slotB, sizeof(slotB) - 1,"%s%s",
-				ptn_list[i],
+		if (stat(buf, &st))
+			continue;
+		memset(buf, '\0', sizeof(buf));
+		snprintf(buf, sizeof(buf) - 1, "%s/%s%s", BOOT_DEV_DIR,
+				prefix.c_str(),
 				AB_SLOT_B_SUFFIX);
-		//It is assumed that both the A and B slots reside on the
-		//same physical disk
-		if (gpt_disk_get_disk_info(slotA, disk))
-			goto error;
-		//Get partition entry for slot A from primary table
+		if (stat(buf, &st))
+			continue;
+		memset(slotA, 0, sizeof(slotA));
+		memset(slotB, 0, sizeof(slotA));
+		snprintf(slotA, sizeof(slotA) - 1, "%s%s", prefix.c_str(),
+				AB_SLOT_A_SUFFIX);
+		snprintf(slotB, sizeof(slotB) - 1,"%s%s", prefix.c_str(),
+				AB_SLOT_B_SUFFIX);
+		//Get the disk containing the partitions that were passed in.
+		//All partitions passed in must lie on the same disk.
+		if (!disk) {
+			disk = boot_ctl_get_disk_info(slotA);
+			if (!disk)
+				goto error;
+		}
+		//Get partition entry for slot A & B from the primary
+		//and backup tables.
 		pentryA = gpt_disk_get_pentry(disk, slotA, PRIMARY_GPT);
-		//Get partition entry for slot A from backup table
 		pentryA_bak = gpt_disk_get_pentry(disk, slotA, SECONDARY_GPT);
-		//Get partition entry for slot B from primary table
 		pentryB = gpt_disk_get_pentry(disk, slotB, PRIMARY_GPT);
-		//Get partition entry for slot B from backup table
 		pentryB_bak = gpt_disk_get_pentry(disk, slotB, SECONDARY_GPT);
 		if ( !pentryA || !pentryA_bak || !pentryB || !pentryB_bak) {
-			//Something has gone wrong here.We know that we have
-			//_a and _b versions of this partition due to the
-			//check at the start of the loop so none of these
-			//should be NULL.
+			//None of these should be NULL since we have already
+			//checked for A & B versions earlier.
 			ALOGE("Slot pentries for %s not found.",
-					ptn_list[i]);
+					prefix.c_str());
 			goto error;
 		}
 		memset(active_guid, '\0', sizeof(active_guid));
 		memset(inactive_guid, '\0', sizeof(inactive_guid));
 		if (get_partition_attribute(slotA, ATTR_SLOT_ACTIVE) == 1) {
 			//A is the current active slot
-			memcpy((void*)active_guid,
-					(const void*)pentryA,
+			memcpy((void*)active_guid, (const void*)pentryA,
 					TYPE_GUID_SIZE);
-			memcpy((void*)inactive_guid,
-					(const void*)pentryB,
+			memcpy((void*)inactive_guid,(const void*)pentryB,
 					TYPE_GUID_SIZE);
-
 		} else if (get_partition_attribute(slotB,
 					ATTR_SLOT_ACTIVE) == 1) {
 			//B is the current active slot
-			memcpy((void*)active_guid,
-					(const void*)pentryB,
+			memcpy((void*)active_guid, (const void*)pentryB,
 					TYPE_GUID_SIZE);
-			memcpy((void*)inactive_guid,
-					(const void*)pentryA,
+			memcpy((void*)inactive_guid, (const void*)pentryA,
 					TYPE_GUID_SIZE);
 		} else {
 			ALOGE("Both A & B are inactive..Aborting");
@@ -481,64 +470,97 @@
 		if (!strncmp(slot_suffix_arr[slot], AB_SLOT_A_SUFFIX,
 					strlen(AB_SLOT_A_SUFFIX))){
 			//Mark A as active in primary table
-			memcpy(pentryA, active_guid, TYPE_GUID_SIZE);
-			slot_info = pentryA + AB_FLAG_OFFSET;
-			*slot_info = AB_SLOT_ACTIVE_VAL;
-
+			UPDATE_SLOT(pentryA, active_guid, SLOT_ACTIVE);
 			//Mark A as active in backup table
-			memcpy(pentryA_bak, active_guid, TYPE_GUID_SIZE);
-			slot_info = pentryA_bak + AB_FLAG_OFFSET;
-			*slot_info = AB_SLOT_ACTIVE_VAL;
-
+			UPDATE_SLOT(pentryA_bak, active_guid, SLOT_ACTIVE);
 			//Mark B as inactive in primary table
-			memcpy(pentryB, inactive_guid, TYPE_GUID_SIZE);
-			slot_info = pentryB + AB_FLAG_OFFSET;
-			*slot_info = *(slot_info) &
-				~AB_PARTITION_ATTR_SLOT_ACTIVE;
-
+			UPDATE_SLOT(pentryB, inactive_guid, SLOT_INACTIVE);
 			//Mark B as inactive in backup table
-			memcpy(pentryB_bak, inactive_guid, TYPE_GUID_SIZE);
-			slot_info = pentryB_bak + AB_FLAG_OFFSET;
-			*slot_info = *(slot_info) &
-				~AB_PARTITION_ATTR_SLOT_ACTIVE;
+			UPDATE_SLOT(pentryB_bak, inactive_guid, SLOT_INACTIVE);
 		} else if (!strncmp(slot_suffix_arr[slot], AB_SLOT_B_SUFFIX,
 					strlen(AB_SLOT_B_SUFFIX))){
 			//Mark B as active in primary table
-			memcpy(pentryB, active_guid, TYPE_GUID_SIZE);
-			slot_info = pentryB + AB_FLAG_OFFSET;
-			*slot_info = AB_SLOT_ACTIVE_VAL;
-
+			UPDATE_SLOT(pentryB, active_guid, SLOT_ACTIVE);
 			//Mark B as active in backup table
-			memcpy(pentryB_bak, active_guid, TYPE_GUID_SIZE);
-			slot_info = pentryB_bak + AB_FLAG_OFFSET;
-			*slot_info = AB_SLOT_ACTIVE_VAL;
-
+			UPDATE_SLOT(pentryB_bak, active_guid, SLOT_ACTIVE);
 			//Mark A as inavtive in primary table
-			memcpy(pentryA, inactive_guid, TYPE_GUID_SIZE);
-			slot_info = pentryA + AB_FLAG_OFFSET;
-			*slot_info = *(slot_info) &
-				~AB_PARTITION_ATTR_SLOT_ACTIVE;
-
+			UPDATE_SLOT(pentryA, inactive_guid, SLOT_INACTIVE);
 			//Mark A as inactive in backup table
-			memcpy(pentryA_bak, inactive_guid, TYPE_GUID_SIZE);
-			slot_info = pentryA_bak + AB_FLAG_OFFSET;
-			*slot_info = *(slot_info) &
-				~AB_PARTITION_ATTR_SLOT_ACTIVE;
+			UPDATE_SLOT(pentryA_bak, inactive_guid, SLOT_INACTIVE);
 		} else {
 			//Something has gone terribly terribly wrong
 			ALOGE("%s: Unknown slot suffix!", __func__);
 			goto error;
 		}
-		if (gpt_disk_update_crc(disk) != 0) {
-			ALOGE("%s: Failed to update gpt_disk crc", __func__);
-			goto error;
+		if (disk) {
+			if (gpt_disk_update_crc(disk) != 0) {
+				ALOGE("%s: Failed to update gpt_disk crc",
+						__func__);
+				goto error;
+			}
 		}
-		if (gpt_disk_commit(disk) != 0) {
-			ALOGE("%s: Failed to commit disk info", __func__);
+	}
+	//write updated content to disk
+	if (disk) {
+		if (gpt_disk_commit(disk)) {
+			ALOGE("Failed to commit disk entry");
 			goto error;
 		}
 		gpt_disk_free(disk);
-		disk = NULL;
+	}
+	return 0;
+
+error:
+	if (disk)
+		gpt_disk_free(disk);
+	return -1;
+}
+
+int set_active_boot_slot(struct boot_control_module *module, unsigned slot)
+{
+	map<string, vector<string>> ptn_map;
+	vector<string> ptn_vec;
+	const char ptn_list[][MAX_GPT_NAME_SIZE] = { AB_PTN_LIST };
+	uint32_t i;
+	int rc = -1;
+	int is_ufs = gpt_utils_is_ufs_device();
+	map<string, vector<string>>::iterator map_iter;
+	vector<string>::iterator string_iter;
+
+	if (boot_control_check_slot_sanity(module, slot)) {
+		ALOGE("%s: Bad arguments", __func__);
+		goto error;
+	}
+	//The partition list just contains prefixes(without the _a/_b) of the
+	//partitions that support A/B. In order to get the layout we need the
+	//actual names. To do this we append the slot suffix to every member
+	//in the list.
+	for (i = 0; i < ARRAY_SIZE(ptn_list); i++) {
+		//XBL is handled differrently for ufs devices so ignore it
+		if (is_ufs && !strncmp(ptn_list[i], PTN_XBL, strlen(PTN_XBL)))
+				continue;
+		//The partition list will be the list of _a partitions
+		string cur_ptn = ptn_list[i];
+		cur_ptn.append(AB_SLOT_A_SUFFIX);
+		ptn_vec.push_back(cur_ptn);
+
+	}
+	//The partition map gives us info in the following format:
+	// [path_to_block_device_1]--><partitions on device 1>
+	// [path_to_block_device_2]--><partitions on device 2>
+	// ...
+	// ...
+	// eg:
+	// [/dev/block/sdb]---><system, boot, rpm, tz,....>
+	if (gpt_utils_get_partition_map(ptn_vec, ptn_map)) {
+		ALOGE("%s: Failed to get partition map",
+				__func__);
+		goto error;
+	}
+	for (map_iter = ptn_map.begin(); map_iter != ptn_map.end(); map_iter++){
+		if (map_iter->second.size() < 1)
+			continue;
+		boot_ctl_set_active_slot_for_partitions(map_iter->second, slot);
 	}
 	if (is_ufs) {
 		if (!strncmp(slot_suffix_arr[slot], AB_SLOT_A_SUFFIX,
@@ -562,22 +584,13 @@
 	}
 	return 0;
 error:
-	if (disk)
-		gpt_disk_free(disk);
 	return -1;
 }
 
 int set_slot_as_unbootable(struct boot_control_module *module, unsigned slot)
 {
-	unsigned num_slots = 0;
-	if (!module) {
-		ALOGE("%s: Invalid argument", __func__);
-		goto error;
-	}
-	num_slots = get_number_slots(module);
-	if (num_slots < 1 || slot > num_slots - 1) {
-		ALOGE("%s: Unable to get num_slots/Invalid slot value",
-				__func__);
+	if (boot_control_check_slot_sanity(module, slot) != 0) {
+		ALOGE("%s: Argument check failed", __func__);
 		goto error;
 	}
 	if (update_slot_attribute(slot_suffix_arr[slot],
@@ -592,17 +605,11 @@
 
 int is_slot_bootable(struct boot_control_module *module, unsigned slot)
 {
-	unsigned num_slots = 0;
 	int attr = 0;
 	char bootPartition[MAX_GPT_NAME_SIZE + 1] = {0};
-	if (!module) {
-		ALOGE("%s: Invalid argument", __func__);
-		goto error;
-	}
-	num_slots = get_number_slots(module);
-	if (num_slots < 1 || slot > num_slots - 1) {
-		ALOGE("%s: Unable to get num_slots/Invalid slot value",
-				__func__);
+
+	if (boot_control_check_slot_sanity(module, slot) != 0) {
+		ALOGE("%s: Argument check failed", __func__);
 		goto error;
 	}
 	snprintf(bootPartition,
@@ -617,17 +624,11 @@
 
 int is_slot_marked_successful(struct boot_control_module *module, unsigned slot)
 {
-	unsigned num_slots = 0;
 	int attr = 0;
 	char bootPartition[MAX_GPT_NAME_SIZE + 1] = {0};
-	if (!module) {
-		ALOGE("%s: Invalid argument", __func__);
-		goto error;
-	}
-	num_slots = get_number_slots(module);
-	if (num_slots < 1 || slot > num_slots - 1) {
-		ALOGE("%s: Unable to get num_slots/Invalid slot value",
-				__func__);
+
+	if (boot_control_check_slot_sanity(module, slot) != 0) {
+		ALOGE("%s: Argument check failed", __func__);
 		goto error;
 	}
 	snprintf(bootPartition,
@@ -664,3 +665,6 @@
 	.getSuffix = get_suffix,
 	.isSlotMarkedSuccessful = is_slot_marked_successful,
 };
+#ifdef __cplusplus
+}
+#endif