Add ItemVariant#withComponentChanges and FluidVariant#withComponentChanges (#4082)

* add ItemVariant#withChanges and FluidVariant#withChanges

* withChanges -> withComponentChanges

* add TransferVariant#withComponentChanges

* make TransferVariant#withComponentChanges throw

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/fluid/FluidVariant.java

@@ -84,4 +84,12 @@ public interface FluidVariant extends TransferVariant<Fluid> {
 	default RegistryEntry<Fluid> getRegistryEntry() {
 		return getFluid().getRegistryEntry();
 	}
+
+	/**
+	 * Creates a copy of this FluidVariant with the provided component changes applied.
+	 * @param changes the changes to apply
+	 * @return the new variant with the changes applied
+	 */
+	@Override
+	FluidVariant withComponentChanges(ComponentChanges changes);
 }

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/item/ItemVariant.java

@@ -106,4 +106,14 @@ public interface ItemVariant extends TransferVariant<Item> {
 		if (isBlank()) return ItemStack.EMPTY;
 		return new ItemStack(getRegistryEntry(), count, getComponents());
 	}
+
+	/**
+	 * Creates a copy of this ItemVariant with the provided component changes applied.
+	 * @param changes the changes to apply
+	 * @return the new variant with the changes applied
+	 *
+	 * @see ItemStack#applyUnvalidatedChanges(ComponentChanges)
+	 */
+	@Override
+	ItemVariant withComponentChanges(ComponentChanges changes);
 }

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/api/transfer/v1/storage/TransferVariant.java

@@ -75,4 +75,13 @@ public interface TransferVariant<O> {
 	default boolean isOf(O object) {
 		return getObject() == object;
 	}
+
+	/**
+	 * Creates a copy of this TransferVariant with the provided component changes applied.
+	 * @param changes the changes to apply
+	 * @return the new variant with the changes applied
+	 */
+	default TransferVariant<O> withComponentChanges(ComponentChanges changes) {
+		throw new UnsupportedOperationException("withComponentChanges is not supported by this TransferVariant");
+	}
 }

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/TransferApiImpl.java

@@ -20,12 +20,17 @@ import java.util.AbstractList;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.NoSuchElementException;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import net.minecraft.component.ComponentChanges;
+import net.minecraft.component.ComponentType;
+
 import net.fabricmc.fabric.api.transfer.v1.storage.SlottedStorage;
 import net.fabricmc.fabric.api.transfer.v1.storage.Storage;
 import net.fabricmc.fabric.api.transfer.v1.storage.StorageView;
@@ -107,4 +112,24 @@ public class TransferApiImpl {
 			}
 		};
 	}
+
+	public static ComponentChanges mergeChanges(ComponentChanges base, ComponentChanges applied) {
+		ComponentChanges.Builder builder = ComponentChanges.builder();
+
+		writeChangesTo(base, builder);
+		writeChangesTo(applied, builder);
+
+		return builder.build();
+	}
+
+	@SuppressWarnings("unchecked")
+	private static void writeChangesTo(ComponentChanges changes, ComponentChanges.Builder builder) {
+		for (Map.Entry<ComponentType<?>, Optional<?>> entry : changes.entrySet()) {
+			if (entry.getValue().isPresent()) {
+				builder.add((ComponentType<Object>) entry.getKey(), entry.getValue().get());
+			} else {
+				builder.remove(entry.getKey());
+			}
+		}
+	}
 }

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/fluid/FluidVariantImpl.java

@@ -31,6 +31,7 @@ import net.minecraft.registry.entry.RegistryEntry;
 import net.minecraft.util.Identifier;
 
 import net.fabricmc.fabric.api.transfer.v1.fluid.FluidVariant;
+import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
 
 public class FluidVariantImpl implements FluidVariant {
 	public static FluidVariant of(Fluid fluid, ComponentChanges components) {
@@ -95,6 +96,11 @@ public class FluidVariantImpl implements FluidVariant {
 		return componentMap;
 	}
 
+	@Override
+	public FluidVariant withComponentChanges(ComponentChanges changes) {
+		return of(fluid, TransferApiImpl.mergeChanges(getComponents(), changes));
+	}
+
 	@Override
 	public String toString() {
 		return "FluidVariant{fluid=" + fluid + ", components=" + components + '}';

fabric-transfer-api-v1/src/main/java/net/fabricmc/fabric/impl/transfer/item/ItemVariantImpl.java

@@ -28,6 +28,7 @@ import net.minecraft.item.Items;
 import net.minecraft.registry.entry.RegistryEntry;
 
 import net.fabricmc.fabric.api.transfer.v1.item.ItemVariant;
+import net.fabricmc.fabric.impl.transfer.TransferApiImpl;
 
 public class ItemVariantImpl implements ItemVariant {
 	public static ItemVariant of(Item item, ComponentChanges components) {
@@ -76,6 +77,11 @@ public class ItemVariantImpl implements ItemVariant {
 		return getCachedStack().getComponents();
 	}
 
+	@Override
+	public ItemVariant withComponentChanges(ComponentChanges changes) {
+		return of(item, TransferApiImpl.mergeChanges(getComponents(), changes));
+	}
+
 	@Override
 	public boolean isBlank() {
 		return item == Items.AIR;

fabric-transfer-api-v1/src/test/java/net/fabricmc/fabric/test/transfer/unittests/FluidVariantTests.java

@@ -18,11 +18,15 @@ package net.fabricmc.fabric.test.transfer.unittests;
 
 import static net.fabricmc.fabric.test.transfer.TestUtil.assertEquals;
 
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 
+import net.minecraft.component.ComponentChanges;
+import net.minecraft.component.DataComponentTypes;
 import net.minecraft.fluid.Fluid;
 import net.minecraft.fluid.Fluids;
+import net.minecraft.util.Unit;
 
 import net.fabricmc.fabric.api.transfer.v1.fluid.FluidVariant;
 
@@ -40,6 +44,28 @@ class FluidVariantTests extends AbstractTransferApiTest {
 		assertEquals(FluidVariant.of(Fluids.LAVA), FluidVariant.of(Fluids.FLOWING_LAVA));
 	}
 
+	@Test
+	public void testWithComponentChanges() {
+		FluidVariant variant = FluidVariant.of(Fluids.WATER, ComponentChanges.builder()
+				.add(DataComponentTypes.HIDE_TOOLTIP, Unit.INSTANCE)
+				.build());
+
+		FluidVariant newVariant = variant.withComponentChanges(ComponentChanges.builder()
+				.remove(DataComponentTypes.HIDE_TOOLTIP)
+				.add(DataComponentTypes.GLIDER, Unit.INSTANCE)
+				.build());
+
+		Assertions.assertFalse(
+				newVariant.getComponentMap().contains(DataComponentTypes.HIDE_TOOLTIP),
+				"New variant's HIDE_TOOLTIP component was removed, but is still present"
+		);
+
+		Assertions.assertTrue(
+				newVariant.getComponentMap().contains(DataComponentTypes.GLIDER),
+				"New variant's GLIDER component was added, but is not present"
+		);
+	}
+
 	private static void assertFluidEquals(Fluid fluid, FluidVariant... variants) {
 		for (FluidVariant variant : variants) {
 			if (variant.getFluid() != fluid) {