Preserve main skill selection across character reimports

spec/System/TestImportReimport_spec.lua

@@ -55,6 +55,22 @@ describe("TestImportReimport", function()
 		}
 	end
 
+	local function pasteSocketGroups(mainSocketGroup, ...)
+		for _, lines in ipairs({ ... }) do
+			build.skillsTab:PasteSocketGroup(table.concat(lines, "\n") .. "\n")
+		end
+		build.mainSocketGroup = mainSocketGroup
+		runCallback("OnFrame")
+	end
+
+	local function makeImportItemWithGems(itemTypeLine, inventoryId, ...)
+		local socketedItems = { }
+		for index, gemName in ipairs({ ... }) do
+			socketedItems[index] = makeSocketedGemEntry(index - 1, gemName:match(" Support$") ~= nil, gemName, 20)
+		end
+		return makeImportItem(itemTypeLine, inventoryId, socketedItems, "test-import-item-" .. inventoryId)
+	end
+
 	local function reimportSocketedItemsWithOptions(itemTypeLine, inventoryId, socketedItems, clearItems)
 		build.importTab:ImportItemsAndSkills(buildImportPayload({
 			makeImportItem(itemTypeLine, inventoryId, socketedItems),
@@ -228,6 +244,52 @@ Blight 20/0  1
 		assert.is_false(groupsBySlot.Gloves.enabled)
 	end)
 
+	it("preserves the main socket group by skill when supports change", function()
+		pasteSocketGroups(1,
+			{ "Slot: Body Armour", "Cleave 20/0  1", "Added Fire Damage 20/0 DISABLED 1" }
+		)
+
+		build.importTab:ImportItemsAndSkills(buildImportPayload({
+			makeImportItemWithGems("Rawhide Gloves", "Gloves", "Cleave"),
+			makeImportItemWithGems("Plate Vest", "BodyArmour", "Cleave", "Faster Attacks Support"),
+		}), false, true, true)
+		runCallback("OnFrame")
+
+		local mainSocketGroup = build.skillsTab.socketGroupList[build.mainSocketGroup]
+		assert.are.equal("Body Armour", mainSocketGroup.slot)
+		assert.is_true(mainSocketGroup.gemList[2].enabled)
+	end)
+
+	it("preserves a support-granted main skill when supports change", function()
+		pasteSocketGroups(1,
+			{ "Slot: Body Armour", "Cleave 20/0  1", "Prismatic Burst 20/0  1", "Added Fire Damage 20/0  1" }
+		)
+		build.skillsTab.socketGroupList[1].mainActiveSkill = 2
+
+		build.importTab:ImportItemsAndSkills(buildImportPayload({
+			makeImportItemWithGems("Plate Vest", "BodyArmour", "Cleave", "Prismatic Burst Support", "Faster Attacks Support"),
+		}), false, true, true)
+		runCallback("OnFrame")
+
+		local mainSocketGroup = build.skillsTab.socketGroupList[build.mainSocketGroup]
+		assert.are.equal("PrismaticBurst", mainSocketGroup.displaySkillList[mainSocketGroup.mainActiveSkill].activeEffect.grantedEffect.id)
+	end)
+
+	it("guesses a new main socket group when the previous main skill is gone", function()
+		pasteSocketGroups(1,
+			{ "Slot: Body Armour", "Cleave 20/0  1" },
+			{ "Slot: Boots", "Frostblink 20/0  1" }
+		)
+
+		build.importTab:ImportItemsAndSkills(buildImportPayload({
+			makeImportItemWithGems("Iron Greaves", "Boots", "Frostblink"),
+			makeImportItemWithGems("Rawhide Gloves", "Gloves", "Fireball", "Added Fire Damage Support"),
+		}), false, true, true)
+		runCallback("OnFrame")
+
+		assert.are.equal("Gloves", build.skillsTab.socketGroupList[build.mainSocketGroup].slot)
+	end)
+
 	it("preserves skill part selection when reimporting items and skills", function()
 		assertReimportPreservesSkillSubstate("Helmet", "Iron Hat", "Helm", "Blight", "skillPart", 2)
 	end)

src/Classes/ImportTab.lua

@@ -1267,6 +1267,22 @@ local function getSocketGroupReimportKey(socketGroup)
 	}, SOCKET_GROUP_REIMPORT_KEY_SEPARATOR)
 end
 
+local function findSocketGroupBySkillId(socketGroupList, skillId, preferredSlot)
+	local fallbackIndex
+	for index, socketGroup in ipairs(socketGroupList) do
+		for _, gem in ipairs(socketGroup.gemList) do
+			if gem.skillId == skillId then
+				if socketGroup.slot == preferredSlot then
+					return index
+				end
+				fallbackIndex = fallbackIndex or index
+				break
+			end
+		end
+	end
+	return fallbackIndex
+end
+
 local function snapshotSocketGroupReimportState(socketGroup, isMainGroup)
 	local gemStates = { }
 	for gemIndex, gem in ipairs(socketGroup.gemList) do
@@ -1390,7 +1406,18 @@ function ImportTabClass:ImportItemsAndSkills(charData, clearItems, clearSkills,
 	local mainSkillEmpty = #self.build.skillsTab.socketGroupList == 0
 	local skillOrder
 	local preservedSocketGroupStateByKey
+	local preservedMainSkillId
+	local preservedMainSkillSlot
+	local preservedMainActiveSkillId
 	if clearSkills then
+		local mainSocketGroup = self.build.skillsTab.socketGroupList[self.build.mainSocketGroup]
+		local mainActiveSkill = mainSocketGroup and mainSocketGroup.displaySkillList
+			and mainSocketGroup.displaySkillList[mainSocketGroup.mainActiveSkill or 1]
+		local mainActiveEffect = mainActiveSkill and mainActiveSkill.activeEffect
+		local mainSkillSource = mainActiveEffect and mainActiveEffect.srcInstance
+		preservedMainSkillId = mainSkillSource and mainSkillSource.skillId
+		preservedMainSkillSlot = mainSocketGroup and mainSocketGroup.slot
+		preservedMainActiveSkillId = mainActiveEffect and mainActiveEffect.grantedEffect and mainActiveEffect.grantedEffect.id
 		skillOrder = { }
 		preservedSocketGroupStateByKey = { }
 		for _, socketGroup in ipairs(self.build.skillsTab.socketGroupList) do
@@ -1472,6 +1499,15 @@ function ImportTabClass:ImportItemsAndSkills(charData, clearItems, clearSkills,
 		end
 		if restoredMainSocketGroup then
 			self.build.mainSocketGroup = restoredMainSocketGroup
+		elseif not mainSkillEmpty then
+			local matchingMainSocketGroup = preservedMainSkillId
+				and findSocketGroupBySkillId(self.build.skillsTab.socketGroupList, preservedMainSkillId, preservedMainSkillSlot)
+			if matchingMainSocketGroup then
+				self.build.mainSocketGroup = matchingMainSocketGroup
+				self.build.skillsTab.socketGroupList[matchingMainSocketGroup].pendingMainActiveSkillId = preservedMainActiveSkillId
+			else
+				self.build.mainSocketGroup = self:GuessMainSocketGroup()
+			end
 		end
 	end
 	if mainSkillEmpty then

src/Modules/CalcSetup.lua

@@ -1772,6 +1772,15 @@ function calcs.initEnv(build, mode, override, specEnv)
 					env.player.modDB:AddMod(modLib.setSource(value.value, groupCfg.slotName or ""))
 				end
 
+				if index == env.mainSocketGroup and env.mode == "MAIN" and group.pendingMainActiveSkillId then
+					for activeSkillIndex, activeSkill in ipairs(socketGroupSkillList) do
+						if activeSkill.activeEffect.grantedEffect.id == group.pendingMainActiveSkillId then
+							group.mainActiveSkill = activeSkillIndex
+							break
+						end
+					end
+				end
+
 				if index == env.mainSocketGroup and #socketGroupSkillList > 0 then
 					-- Select the main skill from this socket group
 					local activeSkillIndex
@@ -1789,6 +1798,7 @@ function calcs.initEnv(build, mode, override, specEnv)
 			end
 
 			if env.mode == "MAIN" then
+				group.pendingMainActiveSkillId = nil
 				-- Create display label for the socket group if the user didn't specify one
 				if group.label and group.label:match("%S") then
 					group.displayLabel = group.label