Fix weapon set allocation issues + add colour to nodes

spec/System/TestPassiveSpec_spec.lua

@@ -43,4 +43,86 @@ describe("TestPassiveSpec", function()
 
 		assert.is_true(ok, err)
 	end)
+
+	local function allocNode(spec, nodeId, allocMode)
+		local node = spec.nodes[nodeId]
+		spec.allocMode = allocMode
+		spec:AllocNode(node)
+		assert.are.equals(allocMode, node.allocMode)
+		return node
+	end
+
+	it("normal passive allocation promotes the shortest path instead of using a longer detour", function()
+		local spec = build.spec
+		allocNode(spec, 56651, 0)
+		local weaponSetNode = allocNode(spec, 38143, 1)
+		assert.are.equals("Strength", weaponSetNode.dn)
+
+		local reachableNode = spec.nodes[43923]
+		assert.are.equals("Accuracy", reachableNode.dn)
+
+		spec.allocMode = 0
+		spec:AllocNode(reachableNode)
+
+		assert.True(reachableNode.alloc)
+		assert.are.equals(0, reachableNode.allocMode)
+		assert.are.equals(0, weaponSetNode.allocMode)
+	end)
+
+	it("normal passive allocation promotes the weapon-set chain behind the path root", function()
+		local spec = build.spec
+		allocNode(spec, 56651, 0)
+		allocNode(spec, 35324, 0)
+		allocNode(spec, 35660, 1)
+		allocNode(spec, 18548, 1)
+
+		local promotedNode = spec.nodes[28992]
+		assert.are.equals("Honed Instincts", promotedNode.dn)
+
+		spec.allocMode = 0
+		spec:AllocNode(promotedNode)
+
+		assert.True(promotedNode.alloc)
+		assert.are.equals(0, promotedNode.allocMode)
+		assert.are.equals(0, spec.nodes[18548].allocMode)
+		assert.are.equals(0, spec.nodes[35660].allocMode)
+	end)
+
+	it("weapon-set allocation cannot originate from another weapon set path", function()
+		local spec = build.spec
+		allocNode(spec, 56651, 0)
+		allocNode(spec, 35324, 0)
+		allocNode(spec, 35660, 1)
+		allocNode(spec, 18548, 1)
+
+		local weaponSet2Node = spec.nodes[28992]
+		assert.are.equals("Honed Instincts", weaponSet2Node.dn)
+
+		spec.allocMode = 2
+		spec:AllocNode(weaponSet2Node)
+
+		assert.True(weaponSet2Node.alloc)
+		assert.are.equals(2, weaponSet2Node.allocMode)
+		assert.are.equals(1, spec.nodes[18548].allocMode)
+		assert.are.equals(1, spec.nodes[35660].allocMode)
+	end)
+
+	it("normal passives cannot stay connected through weapon-set-only paths", function()
+		local spec = build.spec
+		allocNode(spec, 56651, 0)
+		allocNode(spec, 35234, 0)
+		allocNode(spec, 6789, 0)
+		allocNode(spec, 4313, 0)
+		allocNode(spec, 28992, 0)
+		allocNode(spec, 35660, 1)
+		allocNode(spec, 18548, 1)
+
+		spec:DeallocNode(spec.nodes[4313])
+
+		assert.are_not.equals(true, spec.nodes[28992].alloc)
+		assert.True(spec.nodes[35660].alloc)
+		assert.True(spec.nodes[18548].alloc)
+		assert.are.equals(1, spec.nodes[35660].allocMode)
+		assert.are.equals(1, spec.nodes[18548].allocMode)
+	end)
 end)

src/Classes/PassiveSpec.lua

@@ -72,7 +72,7 @@ function PassiveSpecClass:Init(treeVersion, convert)
 	-- Keys are node IDs, values are nodes
 	self.allocNodes = { }
 
-	-- List of nodes allocated in subgraphs; used to maintain allocation when loading, and when rebuilding subgraphs
+	-- List of nodes allocated in subgraphs, used to maintain allocation when loading, and when rebuilding subgraphs
 	self.allocSubgraphNodes = { }
 
 	-- List of cluster nodes to allocate
@@ -813,19 +813,114 @@ end
 -- Allocate the given node, if possible, and all nodes along the path to the node
 -- An alternate path to the node may be provided, otherwise the default path will be used
 -- The path must always contain the given node, as will be the case for the default path
+function PassiveSpecClass:CanPathThroughAllocMode(allocMode, node)
+	-- Normal allocation can only use normal nodes, weapon set allocation can also use its own set.
+	local nodeMode = node.allocMode or 0
+	return nodeMode == 0 or allocMode > 0 and nodeMode == allocMode
+end
+
+function PassiveSpecClass:GetAllocationPath(target, allocMode, allocatedOnly, skipRoot)
+	-- Build a shortest path using only roots and allocated nodes that are valid for the requested alloc mode.
+	local visited = { }
+	local prev = { }
+	local queue = { }
+	local o = 1
+	for _, node in pairs(self.allocNodes) do
+		-- Start from every compatible allocated node so the first hit is the shortest valid route.
+		if node ~= skipRoot and self:CanPathThroughAllocMode(allocatedOnly and 0 or allocMode, node) then
+			visited[node] = true
+			queue[#queue + 1] = node
+		end
+	end
+	while queue[o] do
+		local node = queue[o]
+		o = o + 1
+		if node == target then
+			local path = { }
+			local current = target
+			-- Rebuild the path backwards from the target to the allocated root.
+			while current and prev[current] do
+				t_insert(path, current)
+				current = prev[current]
+			end
+			-- Keep the root outside the path list so hover rendering can preview the first connector.
+			path.root = current
+			return path
+		end
+		if node.unlockConstraint then
+			for _, nodeId in ipairs(node.unlockConstraint.nodes) do
+				if not self.nodes[nodeId].alloc then
+					goto continue
+				end
+			end
+		end
+		for _, other in ipairs(node.linked) do
+			local canPath = true
+			if other.unlockConstraint then
+				for _, nodeId in ipairs(other.unlockConstraint.nodes) do
+					if not self.nodes[nodeId].alloc then
+						canPath = false
+						break
+					end
+				end
+			end
+			local otherMode = other.allocMode or 0
+			local canVisit = allocatedOnly and other.alloc and self:CanPathThroughAllocMode(allocMode, other) and (otherMode > 0 or other.type ~= "ClassStart" and other.type ~= "AscendClassStart") or not allocatedOnly and (not other.alloc or self:CanPathThroughAllocMode(allocMode, other))
+			-- Keep the normal tree, weapon set 1, and weapon set 2 as separate pathing graphs.
+			if canPath and canVisit and not visited[other] and node.type ~= "Mastery" and other.type ~= "ClassStart" and other.type ~= "AscendClassStart" and (node.ascendancyName == other.ascendancyName or (not prev[node] and not other.ascendancyName)) then
+				visited[other] = true
+				prev[other] = node
+				queue[#queue + 1] = other
+			end
+		end
+		::continue::
+	end
+end
+
+function PassiveSpecClass:GetEffectiveAllocationPath(node, altPath)
+	local path = altPath or node.path
+	if altPath or #node.intuitiveLeapLikesAffecting > 0 then
+		return path
+	end
+	local pathRoot = node.pathRoot
+	local pathRootMode = pathRoot and pathRoot.allocMode or 0
+	if pathRootMode == 0 then
+		return path
+	end
+	if self.allocMode == 0 and pathRoot.alloc then
+		-- Normal allocation through a weapon-set branch promotes that branch back to normal pathing.
+		path = { unpack(path) }
+		local rootPath = self:GetAllocationPath(pathRoot, pathRootMode, true, pathRoot)
+		path.root = rootPath and rootPath.root
+		for _, pathNode in ipairs(rootPath or { pathRoot }) do
+			if not isValueInArray(path, pathNode) then
+				t_insert(path, pathNode)
+			end
+		end
+	elseif self.allocMode > 0 and pathRootMode ~= self.allocMode then
+		-- Weapon-set paths cannot start from the other weapon set, so recalculate from compatible roots.
+		path = self:GetAllocationPath(node, self.allocMode)
+	end
+	return path
+end
+
 function PassiveSpecClass:AllocNode(node, altPath)
 	if not node.path then
 		-- Node cannot be connected to the tree as there is no possible path
 		return
 	end
+	local path = self:GetEffectiveAllocationPath(node, altPath)
+	if not path then
+		return
+	end
 
 	-- Allocate all nodes along the path
 	if #node.intuitiveLeapLikesAffecting > 0 then
 		node.alloc = true
 		node.allocMode = (node.ascendancyName or node.type == "Keystone" or node.type == "Socket" or node.containJewelSocket) and 0 or self.allocMode
 		self.allocNodes[node.id] = node
 	else
-		for _, pathNode in ipairs(altPath or node.path) do
+		for _, pathNode in ipairs(path) do
 			pathNode.alloc = true
 			pathNode.allocMode = (node.ascendancyName or pathNode.type == "Keystone" or pathNode.type == "Socket" or pathNode.containJewelSocket) and 0 or self.allocMode
 			-- set path attribute nodes to latest chosen attribute or default to Strength if allocating before choosing an attribute
@@ -913,7 +1008,8 @@ end
 
 -- Attempt to find a class start node starting from the given node
 -- Unless noAscent == true it will also look for an ascendancy class start node
-function PassiveSpecClass:FindStartFromNode(node, visited, noAscend)
+function PassiveSpecClass:FindStartFromNode(node, visited, noAscend, allocMode)
+	allocMode = allocMode or node.allocMode or 0
 	-- Mark the current node as visited so we don't go around in circles
 	node.visited = true
 	t_insert(visited, node)
@@ -923,9 +1019,9 @@ function PassiveSpecClass:FindStartFromNode(node, visited, noAscend)
 		--  - the other node is a start node, or
 		--  - there is a path to a start node through the other node which didn't pass through any nodes which have already been visited
 		local startIndex = #visited + 1
-		if other.alloc and
+		if other.alloc and self:CanPathThroughAllocMode(allocMode, other) and
 		  (other.type == "ClassStart" or other.type == "AscendClassStart" or
-		    (not other.visited and node.type ~= "Mastery" and self:FindStartFromNode(other, visited, noAscend))
+		    (not other.visited and node.type ~= "Mastery" and self:FindStartFromNode(other, visited, noAscend, allocMode))
 		  ) then
 			if node.ascendancyName and not other.ascendancyName then
 				-- Pathing out of Ascendant, un-visit the outside nodes
@@ -955,6 +1051,7 @@ end
 function PassiveSpecClass:BuildPathFromNode(root)
 	root.pathDist = 0
 	root.path = { }
+	root.pathRoot = root
 	local queue = { root }
 	local o, i = 1, 2 -- Out, in
 	while o < i do
@@ -995,13 +1092,14 @@ function PassiveSpecClass:BuildPathFromNode(root)
 			if not other.pathDist then
 				ConPrintTable(other, true)
 			end
-			if node.type ~= "Mastery" and other.type ~= "ClassStart" and other.type ~= "AscendClassStart" and other.pathDist > curDist and (node.ascendancyName == other.ascendancyName or (curDist == 0 and not other.ascendancyName)) and canPath then
+			if node.type ~= "Mastery" and other.type ~= "ClassStart" and other.type ~= "AscendClassStart" and (not other.alloc or self:CanPathThroughAllocMode(root.allocMode or 0, other)) and other.pathDist > curDist and (node.ascendancyName == other.ascendancyName or (curDist == 0 and not other.ascendancyName)) and canPath then
 				-- The shortest path to the other node is through the current node
 				other.pathDist = curDist
 				if not other.alloc then
 					other.pathDist = other.pathDist + 1
 				end
 				other.path = wipeTable(other.path)
+				other.pathRoot = root
 				other.path[1] = other
 				for i, n in ipairs(node.path) do
 					other.path[i+1] = n
@@ -1469,13 +1567,13 @@ function PassiveSpecClass:BuildAllDependsAndPaths()
 		node.connectedToStart = false
 		local anyStartFound = (node.type == "ClassStart" or node.type == "AscendClassStart")
 		for _, other in ipairs(node.linked) do
-			if other.alloc and not isValueInArray(node.depends, other) then
+			if other.alloc and self:CanPathThroughAllocMode(node.allocMode or 0, other) and not isValueInArray(node.depends, other) then
 				-- The other node is allocated and isn't already dependent on this node, so try and find a path to a start node through it
 				if other.type == "ClassStart" or other.type == "AscendClassStart" then
 					-- Well that was easy!
 					anyStartFound = true
 					node.connectedToStart = true
-				elseif self:FindStartFromNode(other, visited) then
+				elseif self:FindStartFromNode(other, visited, nil, node.allocMode or 0) then
 					-- We found a path through the other node, therefore the other node cannot be dependent on this node
 					anyStartFound = true
 					node.connectedToStart = true
@@ -1497,13 +1595,13 @@ function PassiveSpecClass:BuildAllDependsAndPaths()
 							local otherPath = false
 							local allocatedLinkCount = 0
 							for _, linkedNode in ipairs(n.linked) do
-								if linkedNode.alloc then
+								if linkedNode.alloc and self:CanPathThroughAllocMode(n.allocMode or 0, linkedNode) then
 									allocatedLinkCount = allocatedLinkCount + 1
 								end
 							end
 							if allocatedLinkCount > 1 then
 								for _, linkedNode in ipairs(n.linked) do
-									if linkedNode.alloc and not depIds[linkedNode.id] then
+									if linkedNode.alloc and self:CanPathThroughAllocMode(n.allocMode or 0, linkedNode) and not depIds[linkedNode.id] then
 										otherPath = true
 									end
 								end
@@ -1674,6 +1772,7 @@ function PassiveSpecClass:BuildAllDependsAndPaths()
 	for id, node in pairs(self.nodes) do
 		node.pathDist = (node.alloc and #node.intuitiveLeapLikesAffecting == 0) and 0 or 1000
 		node.path = nil
+		node.pathRoot = nil
 		if node.isJewelSocket or node.expansionJewel then
 			node.distanceToClassStart = 0
 		end

src/Classes/PassiveTreeView.lua

@@ -328,9 +328,14 @@ function PassiveTreeViewClass:Draw(build, viewPort, inputEvents)
 		-- Use the node's own path and dependence list
 		hoverPath = { }
 		if #hoverNode.intuitiveLeapLikesAffecting == 0 then
-			for _, pathNode in pairs(hoverNode.path) do
+			-- Use the same effective path as allocation so weapon-set promotion previews match clicks.
+			local path = (hoverNode.alloc and hoverNode.path or spec:GetEffectiveAllocationPath(hoverNode)) or { }
+			for _, pathNode in ipairs(path) do
 				hoverPath[pathNode] = true
 			end
+			if path.root then
+				hoverPath[path.root] = true
+			end
 		end
 		hoverDep = { }
 		for _, depNode in pairs(hoverNode.depends) do
@@ -640,13 +645,33 @@ function PassiveTreeViewClass:Draw(build, viewPort, inputEvents)
 	local function setConnectorColor(r, g, b)
 		connectorColor[1], connectorColor[2], connectorColor[3] = r, g, b
 	end
+	local function nodeIsHoverPathEndpoint(node)
+		if node == hoverNode or hoverPath[node] then
+			return true
+		end
+		if node.alloc then
+			local mode = node.allocMode or 0
+			return mode == 0 or mode == spec.allocMode
+		end
+	end
+	local function nodeWillChangeAllocMode(node)
+		-- Hovering a normal allocation can preview weapon-set nodes being promoted back to normal.
+		return hoverNode and hoverPath and hoverPath[node] and not hoverNode.alloc and spec.allocMode == 0 and (node.allocMode or 0) > 0
+	end
+	local function nodeWillAllocateWithAllocMode(node)
+		-- Unallocated hover/trace path nodes should preview with the selected weapon-set tint.
+		return hoverPath and hoverPath[node] and (self.traceMode or hoverNode and not hoverNode.alloc) and not node.alloc and spec.allocMode > 0 and not node.ascendancyName and node.type ~= "Keystone" and node.type ~= "Socket" and not node.containJewelSocket
+	end
 	local function getState(n1, n2)
 		-- Determine the connector state
 		local state = "Normal"
-		if n1.alloc and n2.alloc then
+		local mode1, mode2 = n1.allocMode or 0, n2.allocMode or 0
+		if hoverPath and nodeIsHoverPathEndpoint(n1) and nodeIsHoverPathEndpoint(n2) and hoverPath[n1] and hoverPath[n2] and (nodeWillChangeAllocMode(n1) or nodeWillChangeAllocMode(n2)) then
+			state = "Intermediate"
+		elseif n1.alloc and n2.alloc and (mode1 == 0 or mode2 == 0 or mode1 == mode2) then
 			state = "Active"
 		elseif hoverPath then
-			if (n1.alloc or n1 == hoverNode or hoverPath[n1]) and (n2.alloc or n2 == hoverNode or hoverPath[n2]) then
+			if nodeIsHoverPathEndpoint(n1) and nodeIsHoverPathEndpoint(n2) and (not n1.alloc or not n2.alloc or hoverPath[n1] and hoverPath[n2]) then
 				state = "Intermediate"
 			end
 		end
@@ -831,6 +856,8 @@ function PassiveTreeViewClass:Draw(build, viewPort, inputEvents)
 
 		local base, overlay, effect
 		local isAlloc = node.alloc or build.calcsTab.mainEnv.grantedPassives[nodeId] or (compareNode and compareNode.alloc)
+		local allocMode = nodeWillAllocateWithAllocMode(node) and spec.allocMode or node.allocMode
+		local allocModeColor = not self.showHeatMap and not launch.devModeAlt and not compareNode and allocMode and allocMode > 0 and not nodeWillChangeAllocMode(node)
 		SetDrawLayer(nil, 25)
 		if node.type == "ClassStart" then
 			overlay = nil
@@ -1000,6 +1027,9 @@ function PassiveTreeViewClass:Draw(build, viewPort, inputEvents)
 		end
 
 		if overlay then
+			if allocModeColor then
+				SetDrawColor(unpack(hexToRGB(colorCodes[allocMode == 1 and "NEGATIVE" or "POSITIVE"]:sub(3))))
+			end
 			-- Draw overlay
 			if node.type ~= "ClassStart" and node.type ~= "AscendClassStart" then
 				if hoverNode and hoverNode ~= node then
@@ -1052,6 +1082,8 @@ function PassiveTreeViewClass:Draw(build, viewPort, inputEvents)
 			self:DrawAsset(overlayImage, scrX, scrY, scale)
 			if not self.showHeatMap and not launch.devModeAlt and not node.alloc and (node.type == "AscendClassStart" or node.type == "ClassStart") then
 				SetDrawColor(1, 1, 1)
+			elseif allocModeColor then
+				SetDrawColor(1, 1, 1)
 			end
 		end
 		if self.searchStrResults[nodeId] then