Support trade rebalance experiment

fabric-object-builder-api-v1/src/main/java/net/fabricmc/fabric/api/object/builder/v1/trade/TradeOfferHelper.java

@@ -16,9 +16,13 @@
 
 package net.fabricmc.fabric.api.object.builder.v1.trade;
 
+import java.util.Collection;
 import java.util.List;
 import java.util.function.Consumer;
 
+import org.jetbrains.annotations.ApiStatus;
+
+import net.minecraft.util.Identifier;
 import net.minecraft.village.TradeOffers;
 import net.minecraft.village.VillagerProfession;
 
@@ -30,6 +34,9 @@
 public final class TradeOfferHelper {
 	/**
 	 * Registers trade offer factories for use by villagers.
+	 * This adds the same trade offers to current and rebalanced trades.
+	 * To add separate offers for the rebalanced trade experiment, use
+	 * {@link #registerVillagerOffers(VillagerProfession, int, VillagerOffersAdder)}.
 	 *
 	 * <p>Below is an example, of registering a trade offer factory to be added a blacksmith with a profession level of 3:
 	 * <blockquote><pre>
@@ -43,11 +50,38 @@ public final class TradeOfferHelper {
 	 * @param factories a consumer to provide the factories
 	 */
 	public static void registerVillagerOffers(VillagerProfession profession, int level, Consumer<List<TradeOffers.Factory>> factories) {
+		TradeOfferInternals.registerVillagerOffers(profession, level, (trades, rebalanced) -> factories.accept(trades));
+	}
+
+	/**
+	 * Registers trade offer factories for use by villagers.
+	 * This method allows separate offers to be added depending on whether the rebalanced
+	 * trade experiment is enabled.
+	 * If a particular profession's rebalanced trade offers are not added at all, it falls back
+	 * to the regular trade offers.
+	 *
+	 * <p>Below is an example, of registering a trade offer factory to be added a blacksmith with a profession level of 3:
+	 * <blockquote><pre>
+	 * TradeOfferHelper.registerVillagerOffers(VillagerProfession.BLACKSMITH, 3, (factories, rebalanced) -> {
+	 * 	factories.add(new CustomTradeFactory(...);
+	 * });
+	 * </pre></blockquote>
+	 *
+	 * <p><strong>Experimental feature</strong>. This API may receive changes as necessary to adapt to further experiment changes.
+	 *
+	 * @param profession the villager profession to assign the trades to
+	 * @param level the profession level the villager must be to offer the trades
+	 * @param factories a consumer to provide the factories
+	 */
+	@ApiStatus.Experimental
+	public static void registerVillagerOffers(VillagerProfession profession, int level, VillagerOffersAdder factories) {
 		TradeOfferInternals.registerVillagerOffers(profession, level, factories);
 	}
 
 	/**
 	 * Registers trade offer factories for use by wandering trades.
+	 * This does not add offers for the rebalanced trade experiment.
+	 * To add rebalanced trades, use {@link #registerRebalancedWanderingTraderOffers}.
 	 *
 	 * @param level the level the trades
 	 * @param factory a consumer to provide the factories
@@ -56,6 +90,20 @@ public static void registerWanderingTraderOffers(int level, Consumer<List<TradeO
 		TradeOfferInternals.registerWanderingTraderOffers(level, factory);
 	}
 
+	/**
+	 * Registers trade offer factories for use by wandering trades.
+	 * This only adds offers for the rebalanced trade experiment.
+	 * To add regular trades, use {@link #registerWanderingTraderOffers(int, Consumer)}.
+	 *
+	 * <p><strong>Experimental feature</strong>. This API may receive changes as necessary to adapt to further experiment changes.
+	 *
+	 * @param factory a consumer to add trade offers
+	 */
+	@ApiStatus.Experimental
+	public static synchronized void registerRebalancedWanderingTraderOffers(Consumer<WanderingTraderOffersBuilder> factory) {
+		factory.accept(new TradeOfferInternals.WanderingTraderOffersBuilderImpl());
+	}
+
 	/**
 	 * @deprecated This never did anything useful.
 	 */
@@ -66,4 +114,115 @@ public static void refreshOffers() {
 
 	private TradeOfferHelper() {
 	}
+
+	@FunctionalInterface
+	public interface VillagerOffersAdder {
+		void onRegister(List<TradeOffers.Factory> factories, boolean rebalanced);
+	}
+
+	/**
+	 * A builder for rebalanced wandering trader offers.
+	 *
+	 * <p><strong>Experimental feature</strong>. This API may receive changes as necessary to adapt to further experiment changes.
+	 *
+	 * @see #registerRebalancedWanderingTraderOffers(Consumer)
+	 */
+	@ApiStatus.NonExtendable
+	@ApiStatus.Experimental
+	public interface WanderingTraderOffersBuilder {
+		/**
+		 * The pool ID for the "buy items" pool.
+		 * Two trade offers are picked from this pool.
+		 *
+		 * <p>In vanilla, this pool contains offers to buy water buckets, baked potatoes, etc.
+		 * for emeralds.
+		 */
+		Identifier BUY_ITEMS_POOL = new Identifier("minecraft", "buy_items");
+		/**
+		 * The pool ID for the "sell special items" pool.
+		 * Two trade offers are picked from this pool.
+		 *
+		 * <p>In vanilla, this pool contains offers to sell logs, enchanted iron pickaxes, etc.
+		 */
+		Identifier SELL_SPECIAL_ITEMS_POOL = new Identifier("minecraft", "sell_special_items");
+		/**
+		 * The pool ID for the "sell common items" pool.
+		 * Five trade offers are picked from this pool.
+		 *
+		 * <p>In vanilla, this pool contains offers to sell flowers, saplings, etc.
+		 */
+		Identifier SELL_COMMON_ITEMS_POOL = new Identifier("minecraft", "sell_common_items");
+
+		/**
+		 * Adds a new pool to the offer list. Exactly {@code count} offers are picked from
+		 * {@code factories} and offered to customers.
+		 * @param id the ID to be assigned to this pool, to allow further modification
+		 * @param count the number of offers to be picked from {@code factories}
+		 * @param factories the trade offer factories
+		 * @return this builder, for chaining
+		 * @throws IllegalArgumentException if {@code count} is not positive or if {@code factories} is empty
+		 */
+		WanderingTraderOffersBuilder pool(Identifier id, int count, TradeOffers.Factory... factories);
+
+		/**
+		 * Adds a new pool to the offer list. Exactly {@code count} offers are picked from
+		 * {@code factories} and offered to customers.
+		 * @param id the ID to be assigned to this pool, to allow further modification
+		 * @param count the number of offers to be picked from {@code factories}
+		 * @param factories the trade offer factories
+		 * @return this builder, for chaining
+		 * @throws IllegalArgumentException if {@code count} is not positive or if {@code factories} is empty
+		 */
+		default WanderingTraderOffersBuilder pool(Identifier id, int count, Collection<? extends TradeOffers.Factory> factories) {
+			return pool(id, count, factories.toArray(TradeOffers.Factory[]::new));
+		}
+
+		/**
+		 * Adds trade offers to the offer list. All offers from {@code factories} are
+		 * offered to each customer.
+		 * @param id the ID to be assigned to this pool, to allow further modification
+		 * @param factories the trade offer factories
+		 * @return this builder, for chaining
+		 * @throws IllegalArgumentException if {@code factories} is empty
+		 */
+		default WanderingTraderOffersBuilder addAll(Identifier id, Collection<? extends TradeOffers.Factory> factories) {
+			return pool(id, factories.size(), factories);
+		}
+
+		/**
+		 * Adds trade offers to the offer list. All offers from {@code factories} are
+		 * offered to each customer.
+		 * @param id the ID to be assigned to this pool, to allow further modification
+		 * @param factories the trade offer factories
+		 * @return this builder, for chaining
+		 * @throws IllegalArgumentException if {@code factories} is empty
+		 */
+		default WanderingTraderOffersBuilder addAll(Identifier id, TradeOffers.Factory... factories) {
+			return pool(id, factories.length, factories);
+		}
+
+		/**
+		 * Adds trade offers to an existing pool identified by an ID.
+		 *
+		 * <p>See the constants for vanilla trade offer pool IDs that are always available.
+		 * @param pool the pool ID
+		 * @param factories the trade offer factories
+		 * @return this builder, for chaining
+		 * @throws IndexOutOfBoundsException if {@code pool} is out of bounds
+		 */
+		WanderingTraderOffersBuilder addOffersToPool(Identifier pool, TradeOffers.Factory... factories);
+
+		/**
+		 * Adds trade offers to an existing pool identified by an ID.
+		 *
+		 * <p>See the constants for vanilla trade offer pool IDs that are always available.
+		 * @param pool the pool ID
+		 * @param factories the trade offer factories
+		 * @return this builder, for chaining
+		 * @throws IndexOutOfBoundsException if {@code pool} is out of bounds
+		 */
+		default WanderingTraderOffersBuilder addOffersToPool(Identifier pool, Collection<TradeOffers.Factory> factories) {
+			return addOffersToPool(pool, factories.toArray(TradeOffers.Factory[]::new));
+		}
+	}
 }

fabric-object-builder-api-v1/src/main/java/net/fabricmc/fabric/impl/object/builder/TradeOfferInternals.java

@@ -17,30 +17,57 @@
 package net.fabricmc.fabric.impl.object.builder;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
 
 import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
 import it.unimi.dsi.fastutil.ints.Int2ObjectOpenHashMap;
+import it.unimi.dsi.fastutil.objects.Object2IntMap;
+import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap;
 import org.apache.commons.lang3.ArrayUtils;
-import org.slf4j.LoggerFactory;
+import org.apache.commons.lang3.tuple.Pair;
 import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import net.minecraft.util.Identifier;
+import net.minecraft.util.Util;
 import net.minecraft.village.TradeOffers;
 import net.minecraft.village.VillagerProfession;
 
+import net.fabricmc.fabric.api.object.builder.v1.trade.TradeOfferHelper;
+
 public final class TradeOfferInternals {
 	private static final Logger LOGGER = LoggerFactory.getLogger("fabric-object-builder-api-v1");
 
 	private TradeOfferInternals() {
 	}
 
+	/**
+	 * Make the rebalanced profession map modifiable, then copy all vanilla
+	 * professions' trades to prevent modifications from propagating to the rebalanced one.
+	 */
+	private static void initVillagerTrades() {
+		if (!(TradeOffers.REBALANCED_PROFESSION_TO_LEVELED_TRADE instanceof HashMap)) {
+			Map<VillagerProfession, Int2ObjectMap<TradeOffers.Factory[]>> map = new HashMap<>(TradeOffers.REBALANCED_PROFESSION_TO_LEVELED_TRADE);
+
+			for (Map.Entry<VillagerProfession, Int2ObjectMap<TradeOffers.Factory[]>> trade : TradeOffers.PROFESSION_TO_LEVELED_TRADE.entrySet()) {
+				if (!map.containsKey(trade.getKey())) map.put(trade.getKey(), trade.getValue());
+			}
+
+			TradeOffers.REBALANCED_PROFESSION_TO_LEVELED_TRADE = map;
+		}
+	}
+
 	// synchronized guards against concurrent modifications - Vanilla does not mutate the underlying arrays (as of 1.16),
 	// so reads will be fine without locking.
-	public static synchronized void registerVillagerOffers(VillagerProfession profession, int level, Consumer<List<TradeOffers.Factory>> factory) {
+	public static synchronized void registerVillagerOffers(VillagerProfession profession, int level, TradeOfferHelper.VillagerOffersAdder factory) {
 		Objects.requireNonNull(profession, "VillagerProfession may not be null.");
-		registerOffers(TradeOffers.PROFESSION_TO_LEVELED_TRADE.computeIfAbsent(profession, key -> new Int2ObjectOpenHashMap<>()), level, factory);
+		initVillagerTrades();
+		registerOffers(TradeOffers.PROFESSION_TO_LEVELED_TRADE.computeIfAbsent(profession, key -> new Int2ObjectOpenHashMap<>()), level, trades -> factory.onRegister(trades, false));
+		registerOffers(TradeOffers.REBALANCED_PROFESSION_TO_LEVELED_TRADE.computeIfAbsent(profession, key -> new Int2ObjectOpenHashMap<>()), level, trades -> factory.onRegister(trades, true));
 	}
 
 	public static synchronized void registerWanderingTraderOffers(int level, Consumer<List<TradeOffers.Factory>> factory) {
@@ -63,4 +90,62 @@ public static void printRefreshOffersWarning() {
 		Throwable loggingThrowable = new Throwable();
 		LOGGER.warn("TradeOfferHelper#refreshOffers does not do anything, yet it was called! Stack trace:", loggingThrowable);
 	}
+
+	public static class WanderingTraderOffersBuilderImpl implements TradeOfferHelper.WanderingTraderOffersBuilder {
+		private static final Object2IntMap<Identifier> ID_TO_INDEX = Util.make(new Object2IntOpenHashMap<>(), idToIndex -> {
+			idToIndex.put(BUY_ITEMS_POOL, 0);
+			idToIndex.put(SELL_SPECIAL_ITEMS_POOL, 1);
+			idToIndex.put(SELL_COMMON_ITEMS_POOL, 2);
+		});
+
+		private static final Map<Identifier, TradeOffers.Factory[]> DELAYED_MODIFICATIONS = new HashMap<>();
+
+		/**
+		 * Make the trade list modifiable.
+		 */
+		static void initWanderingTraderTrades() {
+			if (!(TradeOffers.REBALANCED_WANDERING_TRADER_TRADES instanceof ArrayList)) {
+				TradeOffers.REBALANCED_WANDERING_TRADER_TRADES = new ArrayList<>(TradeOffers.REBALANCED_WANDERING_TRADER_TRADES);
+			}
+		}
+
+		@Override
+		public TradeOfferHelper.WanderingTraderOffersBuilder pool(Identifier id, int count, TradeOffers.Factory... factories) {
+			if (factories.length == 0) throw new IllegalArgumentException("cannot add empty pool");
+			if (count <= 0) throw new IllegalArgumentException("count must be positive");
+
+			Objects.requireNonNull(id, "id cannot be null");
+
+			if (ID_TO_INDEX.containsKey(id)) throw new IllegalArgumentException("pool id %s is already registered".formatted(id));
+
+			Pair<TradeOffers.Factory[], Integer> pool = Pair.of(factories, count);
+			initWanderingTraderTrades();
+			ID_TO_INDEX.put(id, TradeOffers.REBALANCED_WANDERING_TRADER_TRADES.size());
+			TradeOffers.REBALANCED_WANDERING_TRADER_TRADES.add(pool);
+			TradeOffers.Factory[] delayedModifications = DELAYED_MODIFICATIONS.remove(id);
+
+			if (delayedModifications != null) addOffersToPool(id, delayedModifications);
+
+			return this;
+		}
+
+		@Override
+		public TradeOfferHelper.WanderingTraderOffersBuilder addOffersToPool(Identifier pool, TradeOffers.Factory... factories) {
+			if (!ID_TO_INDEX.containsKey(pool)) {
+				DELAYED_MODIFICATIONS.compute(pool, (id, current) -> {
+					if (current == null) return factories;
+
+					return ArrayUtils.addAll(current, factories);
+				});
+				return this;
+			}
+
+			int poolIndex = ID_TO_INDEX.getInt(pool);
+			initWanderingTraderTrades();
+			Pair<TradeOffers.Factory[], Integer> poolPair = TradeOffers.REBALANCED_WANDERING_TRADER_TRADES.get(poolIndex);
+			TradeOffers.Factory[] modified = ArrayUtils.addAll(poolPair.getLeft(), factories);
+			TradeOffers.REBALANCED_WANDERING_TRADER_TRADES.set(poolIndex, Pair.of(modified, poolPair.getRight()));
+			return this;
+		}
+	}
 }

fabric-object-builder-api-v1/src/main/resources/fabric-object-builder-api-v1.accesswidener

@@ -5,6 +5,8 @@ accessible    method    net/minecraft/world/poi/PointOfInterestTypes    register
 
 extendable class net/minecraft/block/entity/BlockEntityType$BlockEntityFactory
 accessible class net/minecraft/village/TradeOffers$TypeAwareBuyForOneEmeraldFactory
+mutable field net/minecraft/village/TradeOffers REBALANCED_PROFESSION_TO_LEVELED_TRADE Ljava/util/Map;
+mutable field net/minecraft/village/TradeOffers REBALANCED_WANDERING_TRADER_TRADES Ljava/util/List;
 
 accessible method net/minecraft/entity/SpawnRestriction register (Lnet/minecraft/entity/EntityType;Lnet/minecraft/entity/SpawnRestriction$Location;Lnet/minecraft/world/Heightmap$Type;Lnet/minecraft/entity/SpawnRestriction$SpawnPredicate;)V