Restore cluster allocations correctly during undo and redo

spec/System/TestClusterJewelGraphs_spec.lua

@@ -112,4 +112,137 @@ Added Small Passive Skills grant: 12% increased Chaos Damage]])
 		spec:DeallocNode(nestedSocket)
 		assert.are.equal(1, countSubgraphs(spec))
 	end)
+
+	it("restores all cluster allocations on the first undo after loading a build", function()
+		local spec = build.spec
+		local outerSocket = getOuterSocket(spec)
+		spec:AllocNode(outerSocket)
+		spec.jewels[outerSocket.id] = addItem([[Rarity: RARE
+New Item
+Large Cluster Jewel
+Cluster Jewel Skill: affliction_chaos_damage
+Cluster Jewel Node Count: 8
+Implicits: 3
+Adds 8 Passive Skills
+2 Added Passive Skills are Jewel Sockets
+Added Small Passive Skills grant: 12% increased Chaos Damage]]).id
+		spec:BuildClusterJewelGraphs()
+		local _, graph = next(spec.subGraphs)
+		local nestedSocketId
+		for _, node in ipairs(graph.nodes) do
+			if node.id then
+				spec:AllocNode(spec.nodes[node.id])
+				if node.type == "Socket" then
+					nestedSocketId = node.id
+				end
+			end
+		end
+		spec.jewels[assert(nestedSocketId)] = addKeystoneCluster().id
+		spec:BuildClusterJewelGraphs()
+		for _, subgraph in pairs(spec.subGraphs) do
+			for _, node in ipairs(subgraph.nodes) do
+				if node.id then
+					spec:AllocNode(spec.nodes[node.id])
+				end
+			end
+		end
+		runCallback("OnFrame")
+		local function allocatedIds()
+			local ids = { }
+			for id in pairs(build.spec.allocNodes) do
+				table.insert(ids, id)
+			end
+			table.sort(ids)
+			return ids
+		end
+		local before = allocatedIds()
+		loadBuildFromXML(build:SaveDB("code"))
+		spec = build.spec
+		assert.are.same(before, allocatedIds())
+		assert.are.equal(2, countSubgraphs(spec))
+		local leaf
+		for _, node in pairs(spec.allocNodes) do
+			if node.expansionSkill and node.type == "Normal" and #node.depends == 1 then
+				leaf = node
+				break
+			end
+		end
+		spec:DeallocNode(assert(leaf))
+		spec:AddUndoState()
+		runCallback("OnFrame")
+		local after = allocatedIds()
+		assert.are.equal(#before - 1, #after)
+		spec:Undo()
+		runCallback("OnFrame")
+		assert.are.same(before, allocatedIds())
+		spec:Redo()
+		runCallback("OnFrame")
+		assert.are.same(after, allocatedIds())
+
+		-- Removing a parent socket temporarily removes its nested cluster graph.
+		spec:DeallocNode(spec.nodes[nestedSocketId])
+		spec:AddUndoState()
+		runCallback("OnFrame")
+		local withoutNested = allocatedIds()
+		assert.are.equal(1, countSubgraphs(spec))
+		spec:Undo()
+		runCallback("OnFrame")
+		assert.are.same(after, allocatedIds())
+		assert.are.equal(2, countSubgraphs(spec))
+		spec:Redo()
+		runCallback("OnFrame")
+		assert.are.same(withoutNested, allocatedIds())
+		spec:Undo()
+		runCallback("OnFrame")
+		assert.are.same(after, allocatedIds())
+	end)
+
+	it("does not revive deallocated cluster nodes when undoing later tree edits", function()
+		local spec = build.spec
+		local socket = getOuterSocket(spec)
+		spec:AllocNode(socket)
+		spec.jewels[socket.id] = addKeystoneCluster().id
+		spec.extended_hashes = { 123 }
+		spec.jewel_data = {
+			[socket.id] = {
+				subgraph = {
+					groups = { keystone = { proxy = socket.expansionJewel.proxy, nodes = { "123" } } },
+					nodes = { ["123"] = { group = "keystone", isKeystone = true, orbitIndex = 0 } },
+				},
+			},
+		}
+		spec:BuildClusterJewelGraphs()
+		local _, graph = next(spec.subGraphs)
+		local clusterId = graph.nodes[1].id
+		spec:AllocNode(spec.nodes[clusterId])
+		-- Rebuilding also occurs when loading equipped cluster jewels.
+		spec:BuildClusterJewelGraphs()
+		spec:ResetUndo()
+		spec:DeallocNode(spec.nodes[clusterId])
+		spec:AddUndoState()
+
+		for step = 1, 20 do
+			local candidates = { }
+			for id, node in pairs(spec.nodes) do
+				if node.type == "Normal" and not node.expansionSkill and not node.alloc and node.path and #node.path > 0 then
+					table.insert(candidates, id)
+				end
+			end
+			table.sort(candidates)
+			spec:AllocNode(spec.nodes[assert(candidates[1])])
+			spec:AddUndoState()
+		end
+		assert.is_nil(spec.allocNodes[clusterId])
+		spec:Undo()
+		assert.is_nil(spec.allocNodes[clusterId])
+		spec:Redo()
+		assert.is_nil(spec.allocNodes[clusterId])
+		for step = 1, 21 do
+			spec:Undo()
+		end
+		assert.is_truthy(spec.allocNodes[clusterId])
+		spec:Redo()
+		assert.is_nil(spec.allocNodes[clusterId])
+	end)
+
 end)

src/Classes/PassiveSpec.lua

@@ -191,7 +191,6 @@ function PassiveSpecClass:Load(xml, dbFileName)
 	elseif url then
 		self:DecodeURL(url)
 	end
-	self:ResetUndo()
 end
 
 function PassiveSpecClass:Save(xml)
@@ -250,6 +249,8 @@ end
 
 function PassiveSpecClass:PostLoad()
 	self:BuildClusterJewelGraphs()
+	-- Capture the initial history only after saved cluster allocations have been restored.
+	self:ResetUndo()
 end
 
 -- Import passive spec from the provided class IDs and node hash list
@@ -1803,7 +1804,8 @@ function PassiveSpecClass:EndLegacyClusterHashConversion()
 	self.legacyClusterNodeMapReverse = nil
 end
 
-function PassiveSpecClass:BuildClusterJewelGraphs()
+---@param skipImportedAllocations boolean?
+function PassiveSpecClass:BuildClusterJewelGraphs(skipImportedAllocations)
 	local needsLegacyClusterHashConversion = self:BeginLegacyClusterHashConversion()
 
 	-- Mark that path building should clear out stale references to cluster nodes
@@ -1831,7 +1833,7 @@ function PassiveSpecClass:BuildClusterJewelGraphs()
 
 	local importedGroups = { }
 	local importedNodes = { }
-	if self.jewel_data then
+	if self.jewel_data and not skipImportedAllocations then
 		for _, value in pairs(self.jewel_data) do
 			if value.subgraph then
 				for groupId, groupData in pairs(value.subgraph.groups) do
@@ -2383,7 +2385,12 @@ function PassiveSpecClass:CreateUndoState()
 end
 
 function PassiveSpecClass:RestoreUndoState(state, treeVersion)
+	-- The snapshot owns allocations; do not reapply nodes remembered from earlier cluster rebuilds.
+	wipeTable(self.allocExtendedNodes)
 	self:ImportFromNodeList(nil, state.classId, state.ascendClassId, state.secondaryAscendClassId, state.hashList, state.hashOverrides, state.masteryEffects, treeVersion or state.treeVersion)
+	if next(self.subGraphs) or next(self.allocSubgraphNodes) then
+		self:BuildClusterJewelGraphs(true)
+	end
 	self:SetWindowTitleWithBuildClass()
 end