From 3f14624d46b5bf423dda5f315cd4a4fd26e641d7 Mon Sep 17 00:00:00 2001
From: Camotoy <20743703+Camotoy@users.noreply.github.com>
Date: Mon, 8 Mar 2021 15:53:47 -0500
Subject: [PATCH] Fix some villager block trades being unable to stack

---
 .../java/world/JavaTradeListTranslator.java   |  7 +++++
 .../world/block/BlockTranslator.java          | 27 ++++++++++++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/connector/src/main/java/org/geysermc/connector/network/translators/java/world/JavaTradeListTranslator.java b/connector/src/main/java/org/geysermc/connector/network/translators/java/world/JavaTradeListTranslator.java
index 9ecffe1d1e1..ac194eca729 100644
--- a/connector/src/main/java/org/geysermc/connector/network/translators/java/world/JavaTradeListTranslator.java
+++ b/connector/src/main/java/org/geysermc/connector/network/translators/java/world/JavaTradeListTranslator.java
@@ -44,6 +44,7 @@
 import org.geysermc.connector.network.translators.item.ItemEntry;
 import org.geysermc.connector.network.translators.item.ItemRegistry;
 import org.geysermc.connector.network.translators.item.ItemTranslator;
+import org.geysermc.connector.network.translators.world.block.BlockTranslator;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -140,6 +141,12 @@ private NbtMap getItemTag(GeyserSession session, ItemStack stack, int specialPri
             NbtMap tag = itemData.getTag().toBuilder().build();
             builder.put("tag", tag);
         }
+        NbtMap blockTag = BlockTranslator.getBedrockBlockNbt(itemEntry.getJavaIdentifier());
+        if (blockTag != null) {
+            // This fixes certain blocks being unable to stack after grabbing one
+            builder.putCompound("Block", blockTag);
+            builder.putShort("Damage", (short) 0);
+        }
         return builder.build();
     }
 }
diff --git a/connector/src/main/java/org/geysermc/connector/network/translators/world/block/BlockTranslator.java b/connector/src/main/java/org/geysermc/connector/network/translators/world/block/BlockTranslator.java
index 1d7c86a53e3..e23c56ee5d0 100644
--- a/connector/src/main/java/org/geysermc/connector/network/translators/world/block/BlockTranslator.java
+++ b/connector/src/main/java/org/geysermc/connector/network/translators/world/block/BlockTranslator.java
@@ -98,6 +98,12 @@ public class BlockTranslator {
 
     public static final int JAVA_RUNTIME_SPAWNER_ID;
 
+    /**
+     * Contains a map of Java blocks to their respective Bedrock block tag, if the Java identifier is different from Bedrock.
+     * Required to fix villager trades with these blocks.
+     */
+    private static final Map<String, NbtMap> JAVA_IDENTIFIER_TO_BEDROCK_TAG;
+
     private static final int BLOCK_STATE_VERSION = 17825808;
 
     static {
@@ -112,6 +118,8 @@ public class BlockTranslator {
             throw new AssertionError("Unable to get blocks from runtime block states", e);
         }
 
+        JAVA_IDENTIFIER_TO_BEDROCK_TAG = new Object2ObjectOpenHashMap<>(blocksTag.size());
+
         // New since 1.16.100 - find the block runtime ID by the order given to us in the block palette,
         // as we no longer send a block palette
         Object2IntMap<NbtMap> blockStateOrderedMap = new Object2IntOpenHashMap<>(blocksTag.size());
@@ -188,15 +196,19 @@ public class BlockTranslator {
             BlockStateValues.storeBlockStateValues(entry.getKey(), javaRuntimeId, entry.getValue());
 
             String cleanJavaIdentifier = entry.getKey().split("\\[")[0];
+            String bedrockIdentifier = entry.getValue().get("bedrock_identifier").asText();
+
+            boolean javaIdentifierSameAsBedrock = cleanJavaIdentifier.equals(bedrockIdentifier);
 
             if (!JAVA_ID_TO_JAVA_IDENTIFIER_MAP.containsValue(cleanJavaIdentifier)) {
                 uniqueJavaId++;
                 JAVA_ID_TO_JAVA_IDENTIFIER_MAP.put(uniqueJavaId, cleanJavaIdentifier);
+                if (!javaIdentifierSameAsBedrock) {
+                    JAVA_IDENTIFIER_TO_BEDROCK_TAG.put(cleanJavaIdentifier, blockTag);
+                }
             }
 
-            String bedrockIdentifier = entry.getValue().get("bedrock_identifier").asText();
-
-            if (!cleanJavaIdentifier.equals(bedrockIdentifier)) {
+            if (!javaIdentifierSameAsBedrock) {
                 JAVA_TO_BEDROCK_IDENTIFIERS.put(cleanJavaIdentifier, bedrockIdentifier);
             }
 
@@ -393,4 +405,13 @@ public static String getPickItem(int javaId) {
     public static String[] getAllBlockIdentifiers() {
         return JAVA_ID_TO_JAVA_IDENTIFIER_MAP.values().toArray(new String[0]);
     }
+
+    /**
+     * @param cleanJavaIdentifier the clean Java identifier of the block to look up
+     *
+     * @return the block tag of the block name mapped from Java to Bedrock.
+     */
+    public static NbtMap getBedrockBlockNbt(String cleanJavaIdentifier) {
+        return JAVA_IDENTIFIER_TO_BEDROCK_TAG.get(cleanJavaIdentifier);
+    }
 }