Fix handling of container items that can contain multiple key types (#…

…7666)

For example gauge droppers from Mekanism.

Fixes AppliedEnergistics/Applied-Mekanistics#25.

src/main/java/appeng/api/behaviors/ContainerItemContext.java

@@ -1,5 +1,7 @@
 package appeng.api.behaviors;
 
+import java.util.Map;
+
 import com.google.common.base.Preconditions;
 
 import org.jetbrains.annotations.Nullable;
@@ -11,38 +13,79 @@
 import appeng.api.stacks.AEKeyType;
 import appeng.api.stacks.GenericStack;
 
-public class ContainerItemContext {
-    // slight generic abuse
-    private final ContainerItemStrategy<AEKey, Object> strategy;
-    private final Object context;
-    // used for sanity checking
-    private final AEKeyType type;
+public final class ContainerItemContext {
+    private final Map<AEKeyType, Entry<?>> entries;
 
-    protected ContainerItemContext(ContainerItemStrategy<AEKey, Object> strategy, Object context, AEKeyType type) {
-        this.strategy = strategy;
-        this.context = context;
-        this.type = type;
+    ContainerItemContext(Map<AEKeyType, Entry<?>> entries) {
+        this.entries = entries;
     }
 
     public @Nullable GenericStack getExtractableContent() {
-        return strategy.getExtractableContent(context);
+        for (var entry : entries.entrySet()) {
+            var content = entry.getValue().getExtractableContent();
+            if (content != null) {
+                return content;
+            }
+        }
+        return null;
+    }
+
+    private Entry<?> getEntry(AEKey key) {
+        var keyType = key.getType();
+        Preconditions.checkArgument(entries.containsKey(keyType), "Internal logic error: mismatched key and type");
+        return entries.get(keyType);
     }
 
     public long insert(AEKey key, long amount, Actionable mode) {
-        Preconditions.checkArgument(type.contains(key), "Internal logic error: mismatched key and type");
-        return strategy.insert(context, key, amount, mode);
+        return getEntry(key).insert(key, amount, mode);
     }
 
     public long extract(AEKey key, long amount, Actionable mode) {
-        Preconditions.checkArgument(type.contains(key), "Internal logic error: mismatched key and type");
-        return strategy.extract(context, key, amount, mode);
+        return getEntry(key).extract(key, amount, mode);
     }
 
-    public void playFillSound(Player player, AEKey what) {
-        strategy.playFillSound(player, what);
+    public void playFillSound(Player player, AEKey key) {
+        getEntry(key).playFillSound(player, key);
     }
 
-    public void playEmptySound(Player player, AEKey what) {
-        strategy.playEmptySound(player, what);
+    public void playEmptySound(Player player, AEKey key) {
+        getEntry(key).playEmptySound(player, key);
+    }
+
+    static class Entry<C> {
+        // slight generic abuse
+        private final ContainerItemStrategy<AEKey, C> strategy;
+        private final C context;
+        // used for sanity checking
+        private final AEKeyType type;
+
+        Entry(ContainerItemStrategy<AEKey, C> strategy, C context, AEKeyType type) {
+            this.strategy = strategy;
+            this.context = context;
+            this.type = type;
+        }
+
+        public @Nullable GenericStack getExtractableContent() {
+            return strategy.getExtractableContent(context);
+        }
+
+        public long insert(AEKey key, long amount, Actionable mode) {
+            Preconditions.checkArgument(type.contains(key), "Internal logic error: mismatched key and type");
+            return strategy.insert(context, key, amount, mode);
+        }
+
+        public long extract(AEKey key, long amount, Actionable mode) {
+            Preconditions.checkArgument(type.contains(key), "Internal logic error: mismatched key and type");
+            return strategy.extract(context, key, amount, mode);
+        }
+
+        public void playFillSound(Player player, AEKey what) {
+            strategy.playFillSound(player, what);
+        }
+
+        public void playEmptySound(Player player, AEKey what) {
+            strategy.playEmptySound(player, what);
+        }
+
     }
 }

src/main/java/appeng/api/behaviors/ContainerItemStrategies.java

@@ -1,8 +1,11 @@
 package appeng.api.behaviors;
 
+import java.util.LinkedHashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
+import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 
 import org.jetbrains.annotations.Nullable;
@@ -14,6 +17,7 @@
 import appeng.api.stacks.AEFluidKey;
 import appeng.api.stacks.AEKey;
 import appeng.api.stacks.AEKeyType;
+import appeng.api.stacks.AEKeyTypes;
 import appeng.api.stacks.GenericStack;
 import appeng.util.CowMap;
 
@@ -49,10 +53,13 @@ public static GenericStack getContainedStack(ItemStack stack) {
             return null;
         }
 
-        for (var entry : strategies.getMap().entrySet()) {
-            var content = entry.getValue().getContainedStack(stack);
-            if (content != null) {
-                return content;
+        for (var keyType : AEKeyTypes.getAll()) {
+            var strategy = strategies.getMap().get(keyType);
+            if (strategy != null) {
+                var content = strategy.getContainedStack(stack);
+                if (content != null) {
+                    return content;
+                }
             }
         }
         return null;
@@ -91,24 +98,32 @@ public static ContainerItemContext findCarriedContextForKey(@Nullable AEKey key,
         return findCarriedContext(key == null ? null : key.getType(), player, menu);
     }
 
-    /**
-     * @param keyType Desired key type, or null if any is ok.
-     */
-    @SuppressWarnings("unchecked")
     @Nullable
-    public static ContainerItemContext findCarriedContext(@Nullable AEKeyType keyType, Player player,
-            AbstractContainerMenu menu) {
-        var candidates = keyType == null ? strategies.getMap().keySet() : List.of(keyType);
+    private static ContainerItemContext findContext(@Nullable AEKeyType keyType,
+            Function<ContainerItemStrategy<?, ?>, @Nullable Object> contextFinder) {
+        var candidates = keyType == null ? AEKeyTypes.getAll() : List.of(keyType);
+        Map<AEKeyType, ContainerItemContext.Entry<?>> entries = new LinkedHashMap<>();
         for (var type : candidates) {
             var strategy = strategies.getMap().get(type);
             if (strategy != null) {
-                var context = strategy.findCarriedContext(player, menu);
+                var context = contextFinder.apply(strategy);
                 if (context != null) {
-                    return new ContainerItemContext((ContainerItemStrategy<AEKey, Object>) strategy, context, type);
+                    // noinspection unchecked
+                    entries.put(type, new ContainerItemContext.Entry<>((ContainerItemStrategy<AEKey, Object>) strategy,
+                            context, type));
                 }
             }
         }
-        return null;
+        return entries.isEmpty() ? null : new ContainerItemContext(entries);
+    }
+
+    /**
+     * @param keyType Desired key type, or null if any is ok.
+     */
+    @Nullable
+    public static ContainerItemContext findCarriedContext(@Nullable AEKeyType keyType, Player player,
+            AbstractContainerMenu menu) {
+        return findContext(keyType, strategy -> strategy.findCarriedContext(player, menu));
     }
 
     public static Set<AEKeyType> getSupportedKeyTypes() {
@@ -121,7 +136,6 @@ public static Set<AEKeyType> getSupportedKeyTypes() {
      *
      * @param keyType Desired key type, or null if any is ok.
      */
-    @SuppressWarnings("unchecked")
     @Nullable
     public static ContainerItemContext findOwnedItemContext(@Nullable AEKeyType keyType,
             Player player,
@@ -145,18 +159,8 @@ public static ContainerItemContext findOwnedItemContext(@Nullable AEKeyType keyT
             return null; // Couldn't find the stack in the player inventory
         }
 
-        var candidates = keyType == null ? strategies.getMap().keySet() : List.of(keyType);
-        for (var type : candidates) {
-            var strategy = strategies.getMap().get(type);
-            if (strategy != null) {
-                var context = strategy.findPlayerSlotContext(player, slotIdx);
-                if (context != null) {
-                    return new ContainerItemContext((ContainerItemStrategy<AEKey, Object>) strategy, context, type);
-                }
-            }
-        }
-
-        return null;
+        int slotIdxCopy = slotIdx;
+        return findContext(keyType, strategy -> strategy.findPlayerSlotContext(player, slotIdxCopy));
     }
 
 }