change IItemSource to return taken item

src/main/java/com/gtnewhorizon/structurelib/structure/IItemSource.java

@@ -6,20 +6,28 @@ import net.minecraft.entity.player.EntityPlayerMP;
 import net.minecraft.inventory.IInventory;
 import net.minecraft.item.ItemStack;
 
+import javax.annotation.Nonnull;
+import java.util.Map;
 import java.util.function.Predicate;
 
 public interface IItemSource {
-    int take(Predicate<ItemStack> predicate, boolean simulate, int count);
+    /**
+     * We are probably not going to need to get a lot of items from this method
+     * so I don't care about the overhead of storing all fetched items even if downstream code doesn't need it
+     */
+    @Nonnull
+    Map<ItemStack, Integer> take(Predicate<ItemStack> predicate, boolean simulate, int count);
 
-    default boolean takeOne(Predicate<ItemStack> predicate, boolean simulate) {
-        return take(predicate, simulate, 1) == 1;
+    default ItemStack takeOne(Predicate<ItemStack> predicate, boolean simulate) {
+        Map<ItemStack, Integer> take = take(predicate, simulate, 1);
+        return take.isEmpty() ? null : take.keySet().iterator().next();
     }
 
     default boolean takeAll(Predicate<ItemStack> predicate, boolean simulate, int count) {
         // fast path for 1 item take requests
-        if (count == 1) return takeOne(predicate, simulate);
-        int have = take(predicate, true, count);
-        if (have < count) return false;
+        if (count == 1) return takeOne(predicate, simulate) != null;
+        Map<ItemStack, Integer> have = take(predicate, true, count);
+        if (have.values().stream().mapToInt(Integer::intValue).sum() < count) return false;
         take(predicate, simulate, count);
         return true;
     }

src/main/java/com/gtnewhorizon/structurelib/structure/StructureUtility.java

@@ -89,7 +89,7 @@ public class StructureUtility {
             if (check(o, world, x, y, z)) return PlaceResult.ACCEPT;
             // user should place anything here.
             // maybe make this configurable, but for now we try to take some cobble from user
-            if (s.takeOne(ItemStackPredicate.from(Blocks.cobblestone), false)) {
+            if (s.takeOne(ItemStackPredicate.from(Blocks.cobblestone), false) != null) {
                 world.setBlock(x, y, z, Blocks.cobblestone, 0, 2);
             }
             return PlaceResult.REJECT;
@@ -134,7 +134,7 @@ public class StructureUtility {
             return PlaceResult.REJECT;
         Item itemBlock = Item.getItemFromBlock(block);
         int itemMeta = itemBlock instanceof ISpecialItemBlock ? ((ISpecialItemBlock) itemBlock).getItemMetaFromBlockMeta(block, meta) : meta;
-        if (!s.takeOne(ItemStackPredicate.from(itemBlock).withMeta(itemMeta), false))
+        if (s.takeOne(ItemStackPredicate.from(itemBlock).withMeta(itemMeta), false) == null)
             return PlaceResult.REJECT;
         if (block instanceof ICustomBlockSetting) {
             ICustomBlockSetting block2 = (ICustomBlockSetting) block;

src/main/java/com/gtnewhorizon/structurelib/util/InventoryUtility.java

@@ -18,12 +18,11 @@
  */
 package com.gtnewhorizon.structurelib.util;
 
-import com.gtnewhorizon.structurelib.util.ItemStackArrayIterable;
-import com.gtnewhorizon.structurelib.util.SortedRegistry;
 import net.minecraft.entity.player.EntityPlayerMP;
 import net.minecraft.item.ItemStack;
 
 import java.util.Iterator;
+import java.util.Map;
 import java.util.NavigableMap;
 import java.util.TreeMap;
 import java.util.function.Function;
@@ -51,14 +50,15 @@ public class InventoryUtility {
      * @param count     let's hope int size is enough...
      * @return amount taken. never negative nor bigger than count...
      */
-    public static int takeFromInventory(EntityPlayerMP player, Predicate<ItemStack> predicate, boolean simulate, int count) {
-        int found = 0;
+    public static Map<ItemStack, Integer> takeFromInventory(EntityPlayerMP player, Predicate<ItemStack> predicate, boolean simulate, int count) {
+        Map<ItemStack, Integer> store = new ItemStackMap<>();
+        int sum = 0;
         for (Function<EntityPlayerMP, Iterable<ItemStack>> provider : inventoryProviders) {
-            found += takeFromInventory(provider.apply(player), predicate, simulate, count - found);
-            if (found >= count)
-                return found;
+            sum += takeFromInventory(provider.apply(player), predicate, simulate, count - sum, true, store);
+            if (sum >= count)
+                return store;
         }
-        return found;
+        return store;
     }
 
     /**
@@ -71,11 +71,13 @@ public class InventoryUtility {
      * @param count     let's hope int size is enough...
      * @return amount taken. never negative nor bigger than count...
      */
-    public static int takeFromInventory(Iterable<ItemStack> inv, Predicate<ItemStack> predicate, boolean simulate, int count) {
-        return takeFromInventory(inv, predicate, simulate, count, true);
+    public static Map<ItemStack, Integer> takeFromInventory(Iterable<ItemStack> inv, Predicate<ItemStack> predicate, boolean simulate, int count) {
+        Map<ItemStack, Integer> store = new ItemStackMap<>();
+        takeFromInventory(inv, predicate, simulate, count, true, store);
+        return store;
     }
 
-    private static int takeFromInventory(Iterable<ItemStack> inv, Predicate<ItemStack> predicate, boolean simulate, int count, boolean recurse) {
+    private static int takeFromInventory(Iterable<ItemStack> inv, Predicate<ItemStack> predicate, boolean simulate, int count, boolean recurse, Map<ItemStack, Integer> store) {
         int found = 0;
         for (Iterator<ItemStack> iterator = inv.iterator(); iterator.hasNext(); ) {
             ItemStack stack = iterator.next();
@@ -87,10 +89,13 @@ public class InventoryUtility {
                 if (found > count) {
                     if (!simulate) {
                         // leave the surplus in its place
-                        stack.stackSize = found - count;
+                        int surplus = found - count;
+                        store.merge(stack, stack.stackSize - surplus, Integer::sum);
+                        stack.stackSize = surplus;
                     }
                     return count;
                 }
+                store.merge(stack, stack.stackSize, Integer::sum);
                 if (!simulate)
                     iterator.remove();
                 if (found == count)
@@ -100,7 +105,7 @@ public class InventoryUtility {
             for (Function<ItemStack, Iterable<ItemStack>> f : stackExtractors) {
                 Iterable<ItemStack> stacks = f.apply(stack);
                 if (stacks == null) continue;
-                found += takeFromInventory(inv, predicate, simulate, count - found, false);
+                found += takeFromInventory(inv, predicate, simulate, count - found, false, store);
                 if (found > count) return found;
             }
         }

src/main/java/com/gtnewhorizon/structurelib/util/ItemStackMap.java

@@ -0,0 +1,461 @@
+/*
+ * Copyright (c) 2021, glee8e, Code Chicken.
+ *
+ * This file is originally part of NotEnoughtItem by Code Chicken.
+ * It is adapted to implement the standard Map interface by glease
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 2.1 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+package com.gtnewhorizon.structurelib.util;
+import com.google.common.base.Objects;
+import com.google.common.collect.Iterators;
+import net.minecraft.init.Items;
+import net.minecraft.item.Item;
+import net.minecraft.item.ItemStack;
+import net.minecraft.nbt.NBTTagCompound;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.util.AbstractMap;
+import java.util.AbstractSet;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.NoSuchElementException;
+import java.util.Set;
+
+import static net.minecraftforge.oredict.OreDictionary.WILDCARD_VALUE;
+
+/**
+ * A map class for ItemStack keys with wildcard damage/NBT. Optimised for lookup.
+ * <p>
+ * This map does NOT support null values or null keys! null will be silently ignored!
+ * <p>
+ * Originally created by CodeChicken for NotEnoughItems. Adapted to {@code Map<ItemStack, T>} interface by glee8e
+ * @author CodeChiken
+ * @author glee8e
+ */
+public final class ItemStackMap<T> extends AbstractMap<ItemStack, T> {
+	public static final NBTTagCompound WILDCARD_TAG;
+
+	static {
+		WILDCARD_TAG = new NBTTagCompound();
+		WILDCARD_TAG.setBoolean("*", true);
+	}
+
+	private final HashMap<Item, DetailMap> itemMap = new HashMap<>();
+	private int size;
+
+    static ItemStack newItemStack(Item item, int size, int damage, NBTTagCompound tag) {
+        ItemStack stack = new ItemStack(item, size, damage);
+        stack.stackTagCompound = tag;
+        return stack;
+    }
+
+    static int actualDamage(ItemStack stack) {
+        return Items.diamond.getDamage(stack);
+    }
+
+	static int getKeyType(int damage, NBTTagCompound tag) {
+		int i = 0;
+		if (isWildcard(damage)) i = 1;
+		if (isWildcard(tag)) i |= 2;
+		return i;
+	}
+
+	static ItemStack wildcard(Item item) {
+		return newItemStack(item, 1, WILDCARD_VALUE, WILDCARD_TAG);
+	}
+
+	static boolean isWildcard(int damage) {
+		return damage == WILDCARD_VALUE;
+	}
+
+	static boolean isWildcard(NBTTagCompound tag) {
+		return tag != null && tag.getBoolean("*");
+	}
+
+	@Override
+	public T get(Object key) {
+		if (!(key instanceof ItemStack)) return null;
+		ItemStack stack = ((ItemStack) key);
+		if (stack.getItem() == null) return null;
+		DetailMap map = itemMap.get(stack.getItem());
+		return map == null ? null : map.get(stack);
+	}
+
+	@Override
+	public T put(ItemStack key, T value) {
+		if (key == null || key.getItem() == null || value == null)
+			return null;
+		return itemMap.computeIfAbsent(key.getItem(), k -> new DetailMap()).put(key, value);
+	}
+
+	@Override
+	public T remove(Object key) {
+		if (!(key instanceof ItemStack)) return null;
+		ItemStack stack = ((ItemStack) key);
+		if (stack.getItem() == null) return null;
+		DetailMap map = itemMap.get(stack.getItem());
+		return map == null ? null : map.remove(stack);
+	}
+
+	@Override
+	public boolean containsKey(Object key) {
+		if (!(key instanceof ItemStack)) return false;
+		ItemStack stack = ((ItemStack) key);
+		if (stack.getItem() == null) return false;
+		DetailMap map = itemMap.get(stack.getItem());
+		return map != null && map.get(stack) != null;
+	}
+
+	@Override
+	public void clear() {
+		itemMap.clear();
+		size = 0;
+	}
+
+	@Override
+	public Set<Map.Entry<ItemStack, T>> entrySet() {
+		return new SetView();
+	}
+
+	private static class StackMetaKey {
+		public final int damage;
+		public final NBTTagCompound tag;
+
+		public StackMetaKey(int damage, NBTTagCompound tag) {
+			this.damage = damage;
+			this.tag = tag;
+		}
+
+		public StackMetaKey(ItemStack key) {
+			this(actualDamage(key), key.stackTagCompound);
+		}
+
+		public int hashCode() {
+			return Objects.hashCode(damage, tag);
+		}
+
+		public boolean equals(Object o) {
+			if (!(o instanceof StackMetaKey))
+				return false;
+			StackMetaKey t = (StackMetaKey) o;
+			return damage == t.damage && Objects.equal(tag, t.tag);
+		}
+	}
+
+	private static class DetailIter<T> implements Iterator<Map.Entry<ItemStack, T>> {
+		private final Item owner;
+		private final ItemStackMap<T>.DetailMap backing;
+		@Nullable
+		private final Iterator<Map.Entry<Integer, T>> damageIter;
+		@Nullable
+		private final Iterator<Map.Entry<NBTTagCompound, T>> tagIter;
+		@Nullable
+		private final Iterator<Map.Entry<StackMetaKey, T>> metaIter;
+		private DetailIterState state = DetailIterState.NOT_STARTED;
+		private boolean removed = false;
+		private DetailIter(Map.Entry<Item, ItemStackMap<T>.DetailMap> input) {
+			this.owner = input.getKey();
+			this.backing = input.getValue();
+			damageIter = backing.damageMap != null ? backing.damageMap.entrySet().iterator() : null;
+			tagIter = backing.tagMap != null ? backing.tagMap.entrySet().iterator() : null;
+			metaIter = backing.metaMap != null ? backing.metaMap.entrySet().iterator() : null;
+		}
+
+		private DetailIterState nextState() {
+			switch (state) {
+				case NOT_STARTED:
+					if (backing.hasWildcard) return DetailIterState.WILDCARD;
+				case WILDCARD:
+				case DAMAGE:
+					if (damageIter != null && damageIter.hasNext()) return DetailIterState.DAMAGE;
+				case TAG:
+					if (tagIter != null && tagIter.hasNext()) return DetailIterState.TAG;
+				case META:
+					if (metaIter != null && metaIter.hasNext()) return DetailIterState.META;
+				case DONE:
+					return DetailIterState.DONE;
+				default:
+					throw new IllegalStateException("Unexpected value: " + state);
+			}
+		}
+
+		@Override
+		public void remove() {
+			if (removed)
+				throw new IllegalStateException("remove() called twice");
+			state.remove(this);
+			removed = true;
+		}
+
+		@Override
+		public boolean hasNext() {
+			return nextState() != DetailIterState.DONE;
+		}
+
+		@Override
+		public Map.Entry<ItemStack, T> next() {
+			DetailIterState nextState = nextState();
+			if (nextState == DetailIterState.DONE)
+				throw new NoSuchElementException();
+			state = nextState;
+			removed = false;
+			return nextState.get(this);
+		}
+
+		private enum DetailIterState {
+			NOT_STARTED {
+				@Override
+				<T> Map.Entry<ItemStack, T> get(DetailIter<T> iter) {
+					throw new AssertionError("Should not call get on NOT_STARTED");
+				}
+
+				@Override
+				<T> void remove(DetailIter<T> iter) {
+					throw new IllegalStateException("next() never called");
+				}
+			},
+			WILDCARD {
+				@Override
+				<T> Map.Entry<ItemStack, T> get(DetailIter<T> iter) {
+					return new Map.Entry<ItemStack, T>() {
+						private final ItemStack key = wildcard(iter.owner);
+						@Override
+						public ItemStack getKey() {
+							return key;
+						}
+
+						@Override
+						public T getValue() {
+							return iter.backing.wildcard;
+						}
+
+						@Override
+						public T setValue(T value) {
+							return iter.backing.wildcard = value;
+						}
+					};
+				}
+
+				@Override
+				<T> void remove(DetailIter<T> iter) {
+					assert iter.backing.hasWildcard;
+					iter.backing.wildcard = null;
+					iter.backing.hasWildcard = false;
+				}
+			},
+			DAMAGE {
+				@Override
+				<T> Map.Entry<ItemStack, T> get(DetailIter<T> iter) {
+					assert iter.damageIter != null;
+					Map.Entry<Integer, T> entry = iter.damageIter.next();
+					return new ItemStackEntry<>(newItemStack(iter.owner, 1, entry.getKey(), WILDCARD_TAG), entry);
+				}
+
+				@Override
+				<T> void remove(DetailIter<T> iter) {
+					assert iter.damageIter != null;
+					iter.damageIter.remove();
+				}
+			},
+			TAG {
+				@Override
+				<T> Map.Entry<ItemStack, T> get(DetailIter<T> iter) {
+					assert iter.tagIter != null;
+					Map.Entry<NBTTagCompound, T> entry = iter.tagIter.next();
+					return new ItemStackEntry<>(newItemStack(iter.owner, 1, WILDCARD_VALUE, entry.getKey()), entry);
+				}
+
+				@Override
+				<T> void remove(DetailIter<T> iter) {
+					assert iter.tagIter != null;
+					iter.tagIter.remove();
+				}
+			},
+			META {
+				@Override
+				<T> Map.Entry<ItemStack, T> get(DetailIter<T> iter) {
+					assert iter.metaIter != null;
+					Map.Entry<StackMetaKey, T> entry = iter.metaIter.next();
+					return new ItemStackEntry<>(newItemStack(iter.owner, 1, entry.getKey().damage, entry.getKey().tag), entry);
+				}
+
+				@Override
+				<T> void remove(DetailIter<T> iter) {
+					assert iter.metaIter != null;
+					iter.metaIter.remove();
+				}
+			},
+			DONE {
+				@Override
+				<T> Map.Entry<ItemStack, T> get(DetailIter<T> iter) {
+					throw new AssertionError("Should not call get on DONE");
+				}
+
+				@Override
+				<T> void remove(DetailIter<T> iter) {
+					throw new AssertionError("Should not call remove on DONE");
+				}
+			};
+
+			abstract <T> Map.Entry<ItemStack, T> get(DetailIter<T> iter);
+
+			abstract <T> void remove(DetailIter<T> iter);
+
+			private static class ItemStackEntry<T> implements Map.Entry<ItemStack, T> {
+				private final ItemStack key;
+				private final Map.Entry<?, T> entry;
+
+				public ItemStackEntry(ItemStack key, Map.Entry<?, T> entry) {
+					this.key = key;
+					this.entry = entry;
+				}
+
+				@Override
+				public ItemStack getKey() {
+					return key;
+				}
+
+				@Override
+				public T getValue() {
+					return entry.getValue();
+				}
+
+				@Override
+				public T setValue(T value) {
+					return entry.setValue(value);
+				}
+			}
+		}
+	}
+
+	private class DetailMap {
+		private boolean hasWildcard;
+		private T wildcard;
+		private HashMap<Integer, T> damageMap;
+		private HashMap<NBTTagCompound, T> tagMap;
+		private HashMap<StackMetaKey, T> metaMap;
+		private int size;
+
+		public T get(ItemStack key) {
+			if (wildcard != null)
+				return wildcard;
+
+			if (damageMap != null) {
+				final T ret = damageMap.get(actualDamage(key));
+				if (ret != null) return ret;
+			}
+			if (tagMap != null) {
+				final T ret = tagMap.get(key.stackTagCompound);
+				if (ret != null) return ret;
+			}
+			if (metaMap != null)
+				return metaMap.get(new StackMetaKey(key));
+
+			return null;
+		}
+
+		public T put(ItemStack key, T value) {
+			try {
+				switch (getKeyType(actualDamage(key), key.stackTagCompound)) {
+					case 0:
+						if (metaMap == null) metaMap = new HashMap<>();
+						return metaMap.put(new StackMetaKey(key), value);
+					case 1:
+						if (tagMap == null) tagMap = new HashMap<>();
+						return tagMap.put(key.stackTagCompound, value);
+					case 2:
+						if (damageMap == null) damageMap = new HashMap<>();
+						return damageMap.put(actualDamage(key), value);
+					case 3:
+						T ret = wildcard;
+						wildcard = value;
+						hasWildcard = true;
+						return ret;
+				}
+			} finally {
+				updateSize();
+			}
+			return null;
+		}
+
+		public T remove(ItemStack key) {
+			try {
+				switch (getKeyType(actualDamage(key), key.stackTagCompound)) {
+					case 0:
+						return metaMap != null ? metaMap.remove(new StackMetaKey(key)) : null;
+					case 1:
+						return tagMap != null ? tagMap.remove(key.stackTagCompound) : null;
+					case 2:
+						return damageMap != null ? damageMap.remove(actualDamage(key)) : null;
+					case 3:
+						T ret = wildcard;
+						wildcard = null;
+						hasWildcard = false;
+						return ret;
+				}
+			} finally {
+				updateSize();
+			}
+			return null;
+		}
+
+		private void updateSize() {
+			int newSize = (hasWildcard ? 1 : 0) +
+					(metaMap != null ? metaMap.size() : 0) +
+					(tagMap != null ? tagMap.size() : 0) +
+					(damageMap != null ? damageMap.size() : 0);
+
+			if (newSize != size) {
+				ItemStackMap.this.size += newSize - size;
+				size = newSize;
+			}
+		}
+	}
+
+	private class SetView extends AbstractSet<Map.Entry<ItemStack, T>> {
+		@Override
+		public boolean contains(Object o) {
+			if (!(o instanceof Map.Entry)) return false;
+			Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
+			return java.util.Objects.equals(get(entry.getKey()), entry.getValue());
+		}
+
+		@Override
+		public boolean add(Map.Entry<ItemStack, T> itemStackTEntry) {
+			return itemStackTEntry.getKey() != null && itemStackTEntry.getValue() != null &&
+					put(itemStackTEntry.getKey(), itemStackTEntry.getValue()) == null;
+		}
+
+		@Override
+		public boolean remove(Object o) {
+			if (!(o instanceof Map.Entry)) return false;
+			Map.Entry<?, ?> entry = (Map.Entry<?, ?>) o;
+			return ItemStackMap.this.remove(entry.getKey(), entry.getValue());
+		}
+
+		@Nonnull
+		@Override
+		public Iterator<Map.Entry<ItemStack, T>> iterator() {
+			return Iterators.concat(Iterators.transform(itemMap.entrySet().iterator(), DetailIter::new));
+		}
+
+		@Override
+		public int size() {
+			return size;
+		}
+	}
+}