[1.21.4] Add Datagen Helper for Creating Modded Datapack Registry Obj…

…ects (#1882)

src/main/java/net/neoforged/neoforge/data/event/GatherDataEvent.java

@@ -15,9 +15,11 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -26,11 +28,13 @@
 import net.minecraft.client.resources.ClientPackSource;
 import net.minecraft.client.resources.IndexedAssetSource;
 import net.minecraft.core.HolderLookup;
+import net.minecraft.core.RegistrySetBuilder;
 import net.minecraft.data.DataGenerator;
 import net.minecraft.data.DataProvider;
 import net.minecraft.data.PackOutput;
 import net.minecraft.data.tags.TagsProvider;
 import net.minecraft.network.chat.Component;
+import net.minecraft.resources.ResourceKey;
 import net.minecraft.server.packs.PackLocationInfo;
 import net.minecraft.server.packs.PackResources;
 import net.minecraft.server.packs.PackType;
@@ -46,7 +50,10 @@
 import net.neoforged.fml.ModList;
 import net.neoforged.fml.event.IModBusEvent;
 import net.neoforged.fml.loading.FMLEnvironment;
+import net.neoforged.neoforge.common.conditions.ICondition;
+import net.neoforged.neoforge.common.data.DatapackBuiltinEntriesProvider;
 import net.neoforged.neoforge.resource.ResourcePackLoader;
+import org.apache.commons.lang3.function.Consumers;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Nullable;
 
@@ -55,6 +62,9 @@ public abstract class GatherDataEvent extends Event implements IModBusEvent {
     private final DataGeneratorConfig config;
     private final ModContainer modContainer;
 
+    @Nullable
+    private CompletableFuture<HolderLookup.Provider> registriesWithModdedEntries = null;
+
     public GatherDataEvent(final ModContainer mc, final DataGenerator dataGenerator, final DataGeneratorConfig dataGeneratorConfig) {
         this.modContainer = mc;
         this.dataGenerator = dataGenerator;
@@ -81,7 +91,7 @@ public DataGenerator getGenerator() {
     }
 
     public CompletableFuture<HolderLookup.Provider> getLookupProvider() {
-        return this.config.lookupProvider;
+        return Objects.requireNonNullElse(this.registriesWithModdedEntries, this.config.lookupProvider);
     }
 
     public boolean includeDev() {
@@ -216,14 +226,40 @@ public <T extends DataProvider> T createProvider(DataProviderFromOutput<T> build
     }
 
     public <T extends DataProvider> T createProvider(DataProviderFromOutputLookup<T> builder) {
-        return addProvider(builder.create(dataGenerator.getPackOutput(), config.lookupProvider));
+        return addProvider(builder.create(dataGenerator.getPackOutput(), this.getLookupProvider()));
     }
 
     public void createBlockAndItemTags(DataProviderFromOutputLookup<TagsProvider<Block>> blockTagsProvider, ItemTagsProvider itemTagsProvider) {
         var blockTags = createProvider(blockTagsProvider);
         addProvider(itemTagsProvider.create(this.getGenerator().getPackOutput(), this.getLookupProvider(), blockTags.contentsGetter()));
     }
 
+    public void createDatapackRegistryObjects(RegistrySetBuilder datapackEntriesBuilder) {
+        this.createDatapackRegistryObjects(datapackEntriesBuilder, Set.of(this.modContainer.getModId()));
+    }
+
+    public void createDatapackRegistryObjects(RegistrySetBuilder datapackEntriesBuilder, Set<String> modIds) {
+        this.createDatapackRegistryObjects(datapackEntriesBuilder, Consumers.nop(), modIds);
+    }
+
+    public void createDatapackRegistryObjects(RegistrySetBuilder datapackEntriesBuilder, Map<ResourceKey<?>, List<ICondition>> conditions) {
+        this.createDatapackRegistryObjects(datapackEntriesBuilder, conditions, Set.of(this.modContainer.getModId()));
+    }
+
+    public void createDatapackRegistryObjects(RegistrySetBuilder datapackEntriesBuilder, Map<ResourceKey<?>, List<ICondition>> conditions, Set<String> modIds) {
+        var registries = this.createProvider((output, lookupProvider) -> new DatapackBuiltinEntriesProvider(output, lookupProvider, datapackEntriesBuilder, conditions, modIds));
+        this.registriesWithModdedEntries = registries.getRegistryProvider();
+    }
+
+    public void createDatapackRegistryObjects(RegistrySetBuilder datapackEntriesBuilder, Consumer<BiConsumer<ResourceKey<?>, ICondition>> conditionsBuilder) {
+        this.createDatapackRegistryObjects(datapackEntriesBuilder, conditionsBuilder, Set.of(this.modContainer.getModId()));
+    }
+
+    public void createDatapackRegistryObjects(RegistrySetBuilder datapackEntriesBuilder, Consumer<BiConsumer<ResourceKey<?>, ICondition>> conditionsBuilder, Set<String> modIds) {
+        var registries = this.createProvider((output, lookupProvider) -> new DatapackBuiltinEntriesProvider(output, lookupProvider, datapackEntriesBuilder, conditionsBuilder, modIds));
+        this.registriesWithModdedEntries = registries.getRegistryProvider();
+    }
+
     @FunctionalInterface
     public interface DataProviderFromOutput<T extends DataProvider> {
         T create(PackOutput output);