Fix torn notes not combining when moved into inventory

Source/inv.cpp

@@ -140,8 +140,85 @@ namespace {
 
 OptionalOwnedClxSpriteList pInvCels;
 
-/**
- * @brief Adds an item to a player's InvGrid array
+bool IsTornNaKrulNote(_item_indexes id)
+{
+	return IsAnyOf(id, IDI_NOTE1, IDI_NOTE2, IDI_NOTE3);
+}
+
+bool HasAllTornNaKrulNotes(const Player &player)
+{
+	return HasInventoryItemWithId(player, IDI_NOTE1)
+	    && HasInventoryItemWithId(player, IDI_NOTE2)
+	    && HasInventoryItemWithId(player, IDI_NOTE3);
+}
+
+void ConvertToFullNaKrulNote(Item &item)
+{
+	item = {};
+	GetItemAttrs(item, IDI_FULLNOTE, 16);
+	SetupItem(item);
+}
+
+std::array<_item_indexes, 2> GetOtherTornNaKrulNotes(_item_indexes preservedNoteId)
+{
+	assert(IsTornNaKrulNote(preservedNoteId));
+
+	switch (preservedNoteId) {
+	case IDI_NOTE1:
+		return { IDI_NOTE2, IDI_NOTE3 };
+	case IDI_NOTE2:
+		return { IDI_NOTE1, IDI_NOTE3 };
+	case IDI_NOTE3:
+		return { IDI_NOTE1, IDI_NOTE2 };
+	default:
+		app_fatal("Unexpected Na-Krul note id");
+	}
+}
+
+int TryCombineNaKrulNoteAfterInventoryInsert(Player &player, int insertedInvIndex)
+{
+	if (insertedInvIndex < 0 || insertedInvIndex >= player._pNumInv) {
+		return insertedInvIndex;
+	}
+
+	const _item_indexes insertedId = player.InvList[insertedInvIndex].IDidx;
+	if (!IsTornNaKrulNote(insertedId) || !HasAllTornNaKrulNotes(player)) {
+		return insertedInvIndex;
+	}
+
+	player.Say(HeroSpeech::JustWhatIWasLookingFor, 10);
+
+	std::array<int, 2> removedNoteIndices {};
+	size_t removeCount = 0;
+	for (const _item_indexes note : GetOtherTornNaKrulNotes(insertedId)) {
+		for (int i = 0; i < player._pNumInv; i++) {
+			if (player.InvList[i].IDidx == note) {
+				removedNoteIndices[removeCount++] = i;
+				break;
+			}
+		}
+	}
+
+	if (removeCount != removedNoteIndices.size()) {
+		return insertedInvIndex;
+	}
+
+	std::sort(removedNoteIndices.begin(), removedNoteIndices.end(), std::greater<int>());
+	for (const int removedNoteIndex : removedNoteIndices) {
+		player.RemoveInvItem(removedNoteIndex, false);
+		if (removedNoteIndex < insertedInvIndex) {
+			insertedInvIndex--;
+		}
+	}
+
+	Item &combinedNote = player.InvList[insertedInvIndex];
+	ConvertToFullNaKrulNote(combinedNote);
+	combinedNote.updateRequiredStatsCacheForPlayer(player);
+	return insertedInvIndex;
+}
+
+/**
+ * @brief Adds an item to a player's InvGrid array
  * @param player The player reference
  * @param invGridIndex Item's position in InvGrid (this should be the item's topleft grid tile)
  * @param invListIndex The item's InvList index (it's expected this already has +1 added to it since InvGrid can't store a 0 index)
@@ -498,7 +575,7 @@ bool ChangeInvItem(Player &player, int slot, Size itemSize)
 		if (prevItemId == 0) {
 			player.InvList[player._pNumInv] = player.HoldItem.pop();
 			player._pNumInv++;
-			prevItemId = player._pNumInv;
+			prevItemId = TryCombineNaKrulNoteAfterInventoryInsert(player, player._pNumInv - 1) + 1;
 		} else {
 			const int invIndex = prevItemId - 1;
 			if (player.HoldItem._itype == ItemType::Gold)
@@ -512,8 +589,11 @@ bool ChangeInvItem(Player &player, int slot, Size itemSize)
 				if (itemIndex == -prevItemId)
 					itemIndex = 0;
 			}
+			prevItemId = TryCombineNaKrulNoteAfterInventoryInsert(player, invIndex) + 1;
 		}
 
+		itemSize = GetInventorySize(player.InvList[prevItemId - 1]);
+
 		AddItemToInvGrid(player, slot - SLOTXY_INV_FIRST, prevItemId, itemSize, &player == MyPlayer);
 	}
 
@@ -957,31 +1037,19 @@ void CheckInvCut(Player &player, Point cursorPosition, bool automaticMove, bool
 
 void TryCombineNaKrulNotes(Player &player, Item &noteItem)
 {
-	const int idx = noteItem.IDidx;
-	const _item_indexes notes[] = { IDI_NOTE1, IDI_NOTE2, IDI_NOTE3 };
-
-	if (IsNoneOf(idx, IDI_NOTE1, IDI_NOTE2, IDI_NOTE3)) {
-		return;
-	}
-
-	for (const _item_indexes note : notes) {
-		if (idx != note && !HasInventoryItemWithId(player, note)) {
-			return; // the player doesn't have all notes
-		}
+	const _item_indexes noteId = noteItem.IDidx;
+	if (!IsTornNaKrulNote(noteId) || !HasAllTornNaKrulNotes(player)) {
+		return; // the player doesn't have all notes
 	}
 
 	MyPlayer->Say(HeroSpeech::JustWhatIWasLookingFor, 10);
 
-	for (const _item_indexes note : notes) {
-		if (idx != note) {
-			RemoveInventoryItemById(player, note);
-		}
+	for (const _item_indexes note : GetOtherTornNaKrulNotes(noteId)) {
+		RemoveInventoryItemById(player, note);
 	}
 
 	const Point position = noteItem.position; // copy the position to restore it after re-initialising the item
-	noteItem = {};
-	GetItemAttrs(noteItem, IDI_FULLNOTE, 16);
-	SetupItem(noteItem);
+	ConvertToFullNaKrulNote(noteItem);
 	noteItem.position = position; // this ensures CleanupItem removes the entry in the dropped items lookup table
 }
 
@@ -1392,16 +1460,20 @@ bool CanFitItemInInventory(const Player &player, const Item &item)
 	return static_cast<bool>(FindSlotForItem(player, GetInventorySize(item)));
 }
 
-bool AutoPlaceItemInInventory(Player &player, const Item &item, bool sendNetworkMessage)
+bool AutoPlaceItemInInventory(Player &player, const Item &item, bool sendNetworkMessage, InventoryInsertSemantics semantics)
 {
 	const Size itemSize = GetInventorySize(item);
 	std::optional<int> targetSlot = FindSlotForItem(player, itemSize);
 
 	if (targetSlot) {
 		player.InvList[player._pNumInv] = item;
 		player._pNumInv++;
+		int invIndex = player._pNumInv - 1;
+		if (semantics == InventoryInsertSemantics::PlayerAction) {
+			invIndex = TryCombineNaKrulNoteAfterInventoryInsert(player, invIndex);
+		}
 
-		AddItemToInvGrid(player, *targetSlot, player._pNumInv, itemSize, sendNetworkMessage);
+		AddItemToInvGrid(player, *targetSlot, invIndex + 1, GetInventorySize(player.InvList[invIndex]), sendNetworkMessage);
 		player.CalcScrolls();
 
 		return true;
@@ -1459,7 +1531,7 @@ void ReorganizeInventory(Player &player)
 	bool reorganizationFailed = false;
 	for (const int index : sortedIndices) {
 		const Item &item = tempStorage[index];
-		if (!AutoPlaceItemInInventory(player, item, false)) {
+		if (!AutoPlaceItemInInventory(player, item, false, InventoryInsertSemantics::InternalRebuild)) {
 			reorganizationFailed = true;
 			break;
 		}

Source/inv.h

@@ -74,14 +74,19 @@ enum inv_xy_slot : uint8_t {
 };
 
 enum item_color : uint8_t {
-	// clang-format off
-	ICOL_YELLOW = PAL16_YELLOW + 5,
-	ICOL_WHITE  = PAL16_GRAY   + 5,
-	ICOL_BLUE   = PAL16_BLUE   + 5,
+	// clang-format off
+	ICOL_YELLOW = PAL16_YELLOW + 5,
+	ICOL_WHITE  = PAL16_GRAY   + 5,
+	ICOL_BLUE   = PAL16_BLUE   + 5,
 	ICOL_RED    = PAL16_RED    + 5,
 	// clang-format on
 };
 
+enum class InventoryInsertSemantics {
+	PlayerAction,
+	InternalRebuild,
+};
+
 extern bool invflag;
 extern const Rectangle InvRect[NUM_XY_SLOTS];
 
@@ -148,15 +153,16 @@ bool AutoEquip(Player &player, const Item &item, bool persistItem = true, bool s
  */
 bool CanFitItemInInventory(const Player &player, const Item &item);
 
-/**
- * @brief Attempts to place the given item in the specified player's inventory.
- * @param player The player whose inventory will be used.
- * @param item The item to be placed.
- * @param sendNetworkMessage Set to true if you want a network message to be generated if the item is persisted.
- * Should only be set if a local player is placing an item in a play session (not when creating a new game)
- * @return 'True' if the item was placed on the player's inventory and 'False' otherwise.
+/**
+ * @brief Attempts to place the given item in the specified player's inventory.
+ * @param player The player whose inventory will be used.
+ * @param item The item to be placed.
+ * @param sendNetworkMessage Set to true if you want a network message to be generated if the item is persisted.
+ * Should only be set if a local player is placing an item in a play session (not when creating a new game)
+ * @param semantics Distinguishes player-facing item insertion from internal inventory rebuilds.
+ * @return 'True' if the item was placed on the player's inventory and 'False' otherwise.
  */
-bool AutoPlaceItemInInventory(Player &player, const Item &item, bool sendNetworkMessage = false);
+bool AutoPlaceItemInInventory(Player &player, const Item &item, bool sendNetworkMessage = false, InventoryInsertSemantics semantics = InventoryInsertSemantics::PlayerAction);
 
 /**
  * @brief Checks whether the given item can be placed on the specified player's belt. Returns 'True' when the item can be placed

test/inv_test.cpp

@@ -1,3 +1,5 @@
+#include <algorithm>
+
 #include <gtest/gtest.h>
 
 #include "cursor.h"
@@ -64,6 +66,39 @@ void clear_inventory()
 	MyPlayer->_pNumInv = 0;
 }
 
+void place_inventory_item(int invIndex, int gridIndex, _item_indexes itemId)
+{
+	Item &item = MyPlayer->InvList[invIndex];
+	InitializeItem(item, itemId);
+	item.updateRequiredStatsCacheForPlayer(*MyPlayer);
+	MyPlayer->InvGrid[gridIndex] = invIndex + 1;
+	MyPlayer->_pNumInv = std::max(MyPlayer->_pNumInv, invIndex + 1);
+}
+
+int count_inventory_items_with_id(_item_indexes itemId)
+{
+	int count = 0;
+	for (int i = 0; i < MyPlayer->_pNumInv; i++) {
+		if (MyPlayer->InvList[i].IDidx == itemId) {
+			count++;
+		}
+	}
+
+	return count;
+}
+
+int count_positive_inv_grid_slots()
+{
+	int count = 0;
+	for (const int8_t cell : MyPlayer->InvGrid) {
+		if (cell > 0) {
+			count++;
+		}
+	}
+
+	return count;
+}
+
 // Test that the scroll is used in the inventory in correct conditions
 TEST_F(InvTest, UseScroll_from_inventory)
 {
@@ -385,5 +420,76 @@ TEST_F(InvTest, ItemSizeLastDiabloItem)
 	EXPECT_EQ(GetInventorySize(testItem), Size(2, 3));
 }
 
+TEST_F(InvTest, AutoPlaceItemInInventoryCombinesInsertedNaKrulNote)
+{
+	if (!gbIsHellfire) return;
+	SNetInitializeProvider(SELCONN_LOOPBACK, nullptr);
+
+	clear_inventory();
+	place_inventory_item(0, 0, IDI_NOTE2);
+	place_inventory_item(1, 1, IDI_NOTE3);
+
+	Item insertedNote {};
+	InitializeItem(insertedNote, IDI_NOTE1);
+	insertedNote.updateRequiredStatsCacheForPlayer(*MyPlayer);
+
+	ASSERT_TRUE(AutoPlaceItemInInventory(*MyPlayer, insertedNote));
+	EXPECT_EQ(MyPlayer->_pNumInv, 1);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_FULLNOTE), 1);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE1), 0);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE2), 0);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE3), 0);
+	EXPECT_EQ(MyPlayer->InvGrid[0], 0);
+	EXPECT_EQ(MyPlayer->InvGrid[1], 0);
+	EXPECT_EQ(MyPlayer->InvGrid[2], 1);
+}
+
+TEST_F(InvTest, AutoPlaceItemInInventoryCombinesInsertedDuplicateNaKrulNote)
+{
+	if (!gbIsHellfire) return;
+	SNetInitializeProvider(SELCONN_LOOPBACK, nullptr);
+
+	clear_inventory();
+	place_inventory_item(0, 0, IDI_NOTE1);
+	place_inventory_item(1, 1, IDI_NOTE2);
+	place_inventory_item(2, 2, IDI_NOTE3);
+
+	Item insertedNote {};
+	InitializeItem(insertedNote, IDI_NOTE1);
+	insertedNote.updateRequiredStatsCacheForPlayer(*MyPlayer);
+
+	ASSERT_TRUE(AutoPlaceItemInInventory(*MyPlayer, insertedNote));
+	EXPECT_EQ(MyPlayer->_pNumInv, 2);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_FULLNOTE), 1);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE1), 1);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE2), 0);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE3), 0);
+	EXPECT_EQ(count_positive_inv_grid_slots(), 2);
+	EXPECT_EQ(MyPlayer->InvGrid[0], 1);
+	EXPECT_EQ(MyPlayer->InvGrid[3], 2);
+	EXPECT_EQ(MyPlayer->InvList[0].IDidx, IDI_NOTE1);
+	EXPECT_EQ(MyPlayer->InvList[1].IDidx, IDI_FULLNOTE);
+}
+
+TEST_F(InvTest, ReorganizeInventoryDoesNotCombineNaKrulNotes)
+{
+	if (!gbIsHellfire) return;
+	SNetInitializeProvider(SELCONN_LOOPBACK, nullptr);
+
+	clear_inventory();
+	place_inventory_item(0, 0, IDI_NOTE1);
+	place_inventory_item(1, 1, IDI_NOTE2);
+	place_inventory_item(2, 2, IDI_NOTE3);
+
+	ReorganizeInventory(*MyPlayer);
+
+	EXPECT_EQ(MyPlayer->_pNumInv, 3);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_FULLNOTE), 0);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE1), 1);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE2), 1);
+	EXPECT_EQ(count_inventory_items_with_id(IDI_NOTE3), 1);
+	EXPECT_EQ(count_positive_inv_grid_slots(), 3);
+}
+
 } // namespace
 } // namespace devilution