From 9448e06ad8480292c5315e2f683c34e0d33b1f1c Mon Sep 17 00:00:00 2001
From: "Kristian S. Stangeland" <kr_stang@hotmail.com>
Date: Wed, 17 Jul 2013 20:59:09 +0200
Subject: [PATCH] Add the ability to store an NBT compound in a given
 ItemStack.

---
 .../protocol/wrappers/nbt/NbtFactory.java     | 46 +++++++++++++++----
 .../protocol/wrappers/nbt/NbtFactoryTest.java | 26 ++++++++++-
 2 files changed, 60 insertions(+), 12 deletions(-)

diff --git a/ProtocolLib/src/main/java/com/comphenix/protocol/wrappers/nbt/NbtFactory.java b/ProtocolLib/src/main/java/com/comphenix/protocol/wrappers/nbt/NbtFactory.java
index 0d72c602..97339f97 100644
--- a/ProtocolLib/src/main/java/com/comphenix/protocol/wrappers/nbt/NbtFactory.java
+++ b/ProtocolLib/src/main/java/com/comphenix/protocol/wrappers/nbt/NbtFactory.java
@@ -113,6 +113,22 @@ public class NbtFactory {
 		}
 	}
 	
+	/**
+	 * Set the NBT compound tag of a given item stack.
+	 * <p>
+	 * The item stack must be a wrapper for a CraftItemStack. Use 
+	 * {@link MinecraftReflection#getBukkitItemStack(ItemStack)} if not.
+	 * @param stack - the item stack.
+	 * @param compound - the new NBT compound.
+	 */
+	public static void setItemTag(ItemStack stack, NbtCompound compound) {
+		if (!MinecraftReflection.isCraftItemStack(stack))
+			throw new IllegalArgumentException("Stack must be a CraftItemStack.");
+		
+		StructureModifier<NbtBase<?>> modifier = getStackModifier(stack);
+		modifier.write(0, compound);
+	}
+	
 	/**
 	 * Construct a wrapper for an NBT tag stored (in memory) in an item stack. This is where
 	 * auxillary data such as enchanting, name and lore is stored. It doesn't include the items
@@ -127,16 +143,7 @@ public class NbtFactory {
 		if (!MinecraftReflection.isCraftItemStack(stack))
 			throw new IllegalArgumentException("Stack must be a CraftItemStack.");
 		
-		Object nmsStack = MinecraftReflection.getMinecraftItemStack(stack);
-		
-		if (itemStackModifier == null) {
-			itemStackModifier = new StructureModifier<Object>(nmsStack.getClass(), Object.class, false);
-		}
-		
-		// Use the first and best NBT tag
-		StructureModifier<NbtBase<?>> modifier = itemStackModifier.
-				withTarget(nmsStack).
-				withType(MinecraftReflection.getNBTBaseClass(), BukkitConverters.getNbtConverter());
+		StructureModifier<NbtBase<?>> modifier = getStackModifier(stack);
 		NbtBase<?> result = modifier.read(0);
 		
 		// Create the tag if it doesn't exist
@@ -147,6 +154,25 @@ public class NbtFactory {
 		return fromBase(result);
 	}
 	
+	/**
+	 * Retrieve a structure modifier that automatically marshalls between NBT wrappers and their NMS counterpart.
+	 * @param stack - the stack that will store the NBT compound.
+	 * @return The structure modifier.
+	 */
+	private static StructureModifier<NbtBase<?>> getStackModifier(ItemStack stack) {
+		Object nmsStack = MinecraftReflection.getMinecraftItemStack(stack);
+		
+		if (itemStackModifier == null) {
+			itemStackModifier = new StructureModifier<Object>(nmsStack.getClass(), Object.class, false);
+		}
+		
+		// Use the first and best NBT tag
+		return itemStackModifier.
+				withTarget(nmsStack).
+				withType(MinecraftReflection.getNBTBaseClass(), 
+						 BukkitConverters.getNbtConverter());
+	}
+	
 	/**
 	 * Initialize a NBT wrapper.
 	 * @param handle - the underlying net.minecraft.server object to wrap.
diff --git a/ProtocolLib/src/test/java/com/comphenix/protocol/wrappers/nbt/NbtFactoryTest.java b/ProtocolLib/src/test/java/com/comphenix/protocol/wrappers/nbt/NbtFactoryTest.java
index cd85bf7d..bb413974 100644
--- a/ProtocolLib/src/test/java/com/comphenix/protocol/wrappers/nbt/NbtFactoryTest.java
+++ b/ProtocolLib/src/test/java/com/comphenix/protocol/wrappers/nbt/NbtFactoryTest.java
@@ -26,21 +26,29 @@ import java.io.DataInputStream;
 import java.io.DataOutput;
 import java.io.DataOutputStream;
 
+import org.bukkit.Material;
+import org.bukkit.craftbukkit.v1_6_R2.inventory.CraftItemFactory;
+import org.bukkit.inventory.ItemStack;
 import org.junit.BeforeClass;
 import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+
 import com.comphenix.protocol.BukkitInitialization;
+import com.comphenix.protocol.utility.MinecraftReflection;
 import com.comphenix.protocol.wrappers.nbt.io.NbtBinarySerializer;
 
+@RunWith(org.powermock.modules.junit4.PowerMockRunner.class)
+@PrepareForTest(CraftItemFactory.class)
 public class NbtFactoryTest {
 	@BeforeClass
 	public static void initializeBukkit() throws IllegalAccessException {
-		BukkitInitialization.initializePackage();
+		BukkitInitialization.initializeItemMeta();
 	}
 	
 	@Test
 	public void testFromStream() {
 		WrappedCompound compound = WrappedCompound.fromName("tag");
-		
 		compound.put("name", "Test Testerson");
 		compound.put("age", 42);
 		
@@ -59,4 +67,18 @@ public class NbtFactoryTest {
 		assertEquals(compound.getInteger("age"), cloned.getInteger("age"));
 		assertEquals(compound.getList("nicknames"), cloned.getList("nicknames"));
 	}
+	
+	@Test
+	public void testItemTag() {
+		ItemStack test = new ItemStack(Material.GOLD_AXE);
+		ItemStack craftTest = MinecraftReflection.getBukkitItemStack(test);
+		
+		NbtCompound compound = NbtFactory.ofCompound("tag");
+		compound.put("name", "Test Testerson");
+		compound.put("age", 42);
+		
+		NbtFactory.setItemTag(craftTest, compound);
+		
+		assertEquals(compound, NbtFactory.fromItemTag(craftTest));
+	}
 }