Ensure custom ROM is copied in crafting

Fixes:
 - Turtle upgrades
 - Pocket upgrades (modems and custom)
 - Normal to advanced upgrade

Also fixes upgraded pocket computers being craftable with a modem

src/main/java/org/squiddev/cctweaks/core/asm/ASMTransformer.java

@@ -100,6 +100,7 @@ public boolean matches(String className) {
 			),
 
 			// Pocket upgrades
+			new PreventModemUpgrade(),
 			new ClassMerger(
 				"dan200.computercraft.shared.pocket.items.ItemPocketComputer",
 				"org.squiddev.cctweaks.core.patch.ItemPocketComputer_Patch"
@@ -175,6 +176,7 @@ public boolean matches(String className) {
 
 			// Custom ROM booting
 			new SetCustomRom(),
+			new CopyRom(),
 			new ClassMerger(
 				"dan200.computercraft.shared.computer.items.ComputerItemFactory",
 				"org.squiddev.cctweaks.core.patch.ComputerItemFactory_Patch"

src/main/java/org/squiddev/cctweaks/core/asm/CopyRom.java

@@ -0,0 +1,57 @@
+package org.squiddev.cctweaks.core.asm;
+
+import dan200.computercraft.shared.pocket.recipes.PocketComputerUpgradeRecipe;
+import dan200.computercraft.shared.turtle.recipes.TurtleUpgradeRecipe;
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.tree.InsnList;
+import org.objectweb.asm.tree.InsnNode;
+import org.objectweb.asm.tree.MethodInsnNode;
+import org.squiddev.patcher.transformer.IPatcher;
+import org.squiddev.patcher.visitors.FindingVisitor;
+
+import static org.objectweb.asm.Opcodes.*;
+
+/**
+ * Modifies {@link TurtleUpgradeRecipe} and {@link PocketComputerUpgradeRecipe} to set the ROM id.
+ */
+public class CopyRom implements IPatcher {
+	@Override
+	public boolean matches(String className) {
+		return className.equals("dan200.computercraft.shared.turtle.recipes.TurtleUpgradeRecipe")
+			|| className.equals("dan200.computercraft.shared.pocket.recipes.PocketComputerUpgradeRecipe");
+	}
+
+	@Override
+	public ClassVisitor patch(String className, ClassVisitor delegate) throws Exception {
+		if (className.endsWith("TurtleUpgradeRecipe")) {
+			return new FindingVisitor(
+				delegate,
+				new MethodInsnNode(INVOKESTATIC, "dan200/computercraft/shared/turtle/items/TurtleItemFactory", "create", "(ILjava/lang/String;Ldan200/computercraft/shared/util/Colour;Ldan200/computercraft/shared/computer/core/ComputerFamily;Ldan200/computercraft/api/turtle/ITurtleUpgrade;Ldan200/computercraft/api/turtle/ITurtleUpgrade;ILnet/minecraft/util/ResourceLocation;)Lnet/minecraft/item/ItemStack;", false),
+				new InsnNode(ARETURN)
+			) {
+				@Override
+				public void handle(InsnList nodes, MethodVisitor visitor) {
+					nodes.getFirst().accept(visitor);
+					visitor.visitVarInsn(ALOAD, 3);
+					visitor.visitMethodInsn(INVOKESTATIC, "org/squiddev/cctweaks/core/rom/CraftingSetRom", "copyRom", "(Lnet/minecraft/item/ItemStack;Lnet/minecraft/item/ItemStack;)Lnet/minecraft/item/ItemStack;", false);
+					visitor.visitInsn(ARETURN);
+				}
+			}.onMethod("func_77572_b").onMethod("getCraftingResult").mustFind().once();
+		} else {
+			return new FindingVisitor(
+				delegate,
+				new MethodInsnNode(INVOKESTATIC, "dan200/computercraft/shared/pocket/items/PocketComputerItemFactory", "create", "(ILjava/lang/String;Ldan200/computercraft/shared/computer/core/ComputerFamily;Z)Lnet/minecraft/item/ItemStack;", false),
+				new InsnNode(ARETURN)
+			) {
+				@Override
+				public void handle(InsnList nodes, MethodVisitor visitor) {
+					nodes.getFirst().accept(visitor);
+					visitor.visitVarInsn(ALOAD, 2);
+					visitor.visitMethodInsn(INVOKESTATIC, "org/squiddev/cctweaks/core/rom/CraftingSetRom", "copyRom", "(Lnet/minecraft/item/ItemStack;Lnet/minecraft/item/ItemStack;)Lnet/minecraft/item/ItemStack;", false);
+					visitor.visitInsn(ARETURN);
+				}
+			}.onMethod("func_77572_b").onMethod("getCraftingResult").mustFind().once();
+		}
+	}
+}

src/main/java/org/squiddev/cctweaks/core/asm/PreventModemUpgrade.java

@@ -0,0 +1,46 @@
+package org.squiddev.cctweaks.core.asm;
+
+import org.objectweb.asm.ClassVisitor;
+import org.objectweb.asm.Label;
+import org.objectweb.asm.MethodVisitor;
+import org.objectweb.asm.tree.InsnList;
+import org.objectweb.asm.tree.JumpInsnNode;
+import org.objectweb.asm.tree.MethodInsnNode;
+import org.squiddev.patcher.transformer.IPatcher;
+import org.squiddev.patcher.visitors.FindingVisitor;
+
+import static org.objectweb.asm.Opcodes.*;
+
+/**
+ * Prevents {@link dan200.computercraft.shared.pocket.recipes.PocketComputerUpgradeRecipe} from working when there is a
+ * custom upgrade.
+ */
+public class PreventModemUpgrade implements IPatcher {
+	@Override
+	public boolean matches(String className) {
+		return className.equals("dan200.computercraft.shared.pocket.recipes.PocketComputerUpgradeRecipe");
+	}
+
+	@Override
+	public ClassVisitor patch(String className, ClassVisitor delegate) throws Exception {
+		return new FindingVisitor(delegate,
+			new MethodInsnNode(INVOKEVIRTUAL, "dan200/computercraft/shared/pocket/items/ItemPocketComputer", "getHasModem", "(Lnet/minecraft/item/ItemStack;)Z", false),
+			new JumpInsnNode(IFEQ, null)
+		) {
+			@Override
+			public void handle(InsnList nodes, MethodVisitor visitor) {
+				Label exit = new Label();
+				Label cont = ((JumpInsnNode) nodes.get(1)).label.getLabel();
+
+				visitor.visitMethodInsn(INVOKEVIRTUAL, "dan200/computercraft/shared/pocket/items/ItemPocketComputer", "getHasModem", "(Lnet/minecraft/item/ItemStack;)Z", false);
+				visitor.visitJumpInsn(IFNE, exit);
+
+				visitor.visitVarInsn(ALOAD, 2);
+				visitor.visitMethodInsn(INVOKESTATIC, "org/squiddev/cctweaks/core/pocket/PocketHooks", "hasPocketUpgrade", "(Lnet/minecraft/item/ItemStack;)Z", false);
+				visitor.visitJumpInsn(IFEQ, cont);
+
+				visitor.visitLabel(exit);
+			}
+		};
+	}
+}

src/main/java/org/squiddev/cctweaks/core/asm/SetCustomRom.java

@@ -35,8 +35,7 @@ public void handle(InsnList nodes, MethodVisitor visitor) {
 
 				visitor.visitVarInsn(ALOAD, 0);
 				visitor.visitFieldInsn(GETFIELD, "dan200/computercraft/shared/computer/blocks/TileComputerBase", "hasDisk", "Z");
-				visitor.visitInsn(ICONST_0);
-				visitor.visitJumpInsn(IF_ICMPEQ, finish);
+				visitor.visitJumpInsn(IFEQ, finish);
 
 				visitor.visitVarInsn(ALOAD, 2);
 				visitor.visitVarInsn(ALOAD, 0);

src/main/java/org/squiddev/cctweaks/core/pocket/CraftingPocketUpgrade.java

@@ -36,7 +36,7 @@ public ItemStack getCraftingResult(InventoryCrafting inventory) {
 				if (pocket != null) return null;
 
 				// If we have an upgrade then abort
-				if (stack.hasTagCompound() && stack.getTagCompound().getShort("upgrade") != 0) return null;
+				if (PocketHooks.hasPocketUpgrade(stack)) return null;
 
 				pocket = stack;
 			} else {

src/main/java/org/squiddev/cctweaks/core/pocket/PocketHooks.java

@@ -49,4 +49,8 @@ public static void create(IInventory inventory, ItemStack stack, ServerComputer
 	public static void destroy(ServerComputer computer) {
 		upgradeInstance.remove(computer);
 	}
+
+	public static boolean hasPocketUpgrade(ItemStack stack) {
+		return stack.hasTagCompound() && stack.getTagCompound().getInteger("upgrade") != 0;
+	}
 }

src/main/java/org/squiddev/cctweaks/core/rom/CraftingSetRom.java

@@ -9,6 +9,7 @@
 import net.minecraft.item.Item;
 import net.minecraft.item.ItemStack;
 import net.minecraft.item.crafting.IRecipe;
+import net.minecraft.nbt.NBTTagCompound;
 import net.minecraft.world.World;
 import net.minecraftforge.fml.common.registry.GameRegistry;
 import net.minecraftforge.oredict.RecipeSorter;
@@ -164,4 +165,16 @@ private static ItemStack newDisk(int id) {
 		ItemBase.getTag(result).setInteger("diskID", id);
 		return result;
 	}
+
+	public static ItemStack copyRom(ItemStack to, ItemStack from) {
+		NBTTagCompound fromTag = from.getTagCompound();
+		if (fromTag != null) {
+			NBTTagCompound toTag = to.getTagCompound();
+			if (toTag == null) to.setTagCompound(toTag = new NBTTagCompound());
+
+			if (fromTag.hasKey("rom_id")) toTag.setTag("rom_id", fromTag.getTag("rom_id"));
+		}
+
+		return to;
+	}
 }

src/main/java/org/squiddev/cctweaks/items/CraftingComputerUpgrade.java

@@ -17,6 +17,7 @@
 import net.minecraft.nbt.NBTTagCompound;
 import net.minecraft.world.World;
 import net.minecraftforge.common.ForgeHooks;
+import org.squiddev.cctweaks.core.rom.CraftingSetRom;
 
 /**
  * Handles crafting with ComputerUpgrades
@@ -39,33 +40,33 @@ public ItemStack getCraftingResult(InventoryCrafting inventorycrafting) {
 		int id = computerItem.getComputerID(computerStack);
 		String label = computerItem.getLabel(computerStack);
 
+		ItemStack stack;
 		if (computerItem instanceof ItemComputer) {
-			return ComputerItemFactory.create(id, label, ComputerFamily.Advanced);
+			stack = ComputerItemFactory.create(id, label, ComputerFamily.Advanced);
 		} else if (computerItem instanceof ItemTurtleBase) {
 			ItemTurtleBase turtle = (ItemTurtleBase) computerItem;
 
-			return TurtleItemFactory.create(
+			stack = TurtleItemFactory.create(
 				id, label, turtle.getColour(computerStack), ComputerFamily.Advanced,
 				turtle.getUpgrade(computerStack, TurtleSide.Left),
 				turtle.getUpgrade(computerStack, TurtleSide.Right),
 				turtle.getFuelLevel(computerStack), null
 			);
 		} else if (computerItem instanceof ItemPocketComputer) {
 			ItemPocketComputer pocket = (ItemPocketComputer) computerItem;
+			stack = PocketComputerItemFactory.create(id, label, ComputerFamily.Advanced, pocket.getHasModem(computerStack));
 
-			ItemStack stack = PocketComputerItemFactory.create(id, label, ComputerFamily.Advanced, pocket.getHasModem(computerStack));
-
-			NBTTagCompound toTag = stack.getTagCompound();
-			NBTTagCompound fromTag = computerStack.getTagCompound();
+			NBTTagCompound toTag = ItemBase.getTag(stack);
+			NBTTagCompound fromTag = ItemBase.getTag(computerStack);
 
 			if (fromTag.hasKey("upgrade")) toTag.setTag("upgrade", fromTag.getTag("upgrade"));
 			if (fromTag.hasKey("upgrade_name")) toTag.setTag("upgrade_name", fromTag.getTag("upgrade_name"));
 			if (fromTag.hasKey("upgrade_info")) toTag.setTag("upgrade_info", fromTag.getTag("upgrade_info"));
-
-			return stack;
 		} else {
 			return null;
 		}
+
+		return CraftingSetRom.copyRom(stack, computerStack);
 	}
 
 	/**