Add ability to define default attributes and spawn restirctions within the entity type builder. (#828)

* Add ability to define default attributes within the entity type builder.

* Update fabric-object-builder-api-v1/src/main/java/net/fabricmc/fabric/api/object/builder/v1/entity/FabricDefaultAttributeRegistry.java

Co-authored-by: Pyrofab <redstoneinfire@gmail.com>

* Add ability to specify spawn restriction. Add a Living and Mob builder to handle attributes and spawn restrictions seperately from default builder.

* Let's not cast like a maniac

* Add misc overloads for living builder

* Add a few null checks

* Allow factory to be defined in builder

* Make spawn group a builder parameter

Co-authored-by: Pyrofab <redstoneinfire@gmail.com>

fabric-object-builder-api-v1/src/main/java/net/fabricmc/fabric/api/object/builder/v1/entity/FabricDefaultAttributeRegistry.java

@@ -16,6 +16,8 @@
 
 package net.fabricmc.fabric.api.object.builder.v1.entity;
 
+import java.util.function.Supplier;
+
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -55,8 +57,11 @@ public final class FabricDefaultAttributeRegistry {
 	 * <p>If a registration overrides another, a debug log message will be emitted. Existing registrations
 	 * can be checked at {@link net.minecraft.entity.attribute.DefaultAttributeRegistry#hasDefinitionFor(EntityType)}.</p>
 	 *
+	 * <p>For convenience, this can also be done on the {@link FabricEntityTypeBuilder} to simplify the building process.
+	 *
 	 * @param type    the entity type
 	 * @param builder the builder that creates the default attribute
+	 * @see	FabricEntityTypeBuilder.Living#defaultAttributes(Supplier)
 	 */
 	public static void register(EntityType<? extends LivingEntity> type, DefaultAttributeContainer.Builder builder) {
 		if (DefaultAttributeRegistryAccessor.getRegistry().put(type, builder.build()) != null) {

fabric-object-builder-api-v1/src/main/java/net/fabricmc/fabric/api/object/builder/v1/entity/FabricEntityTypeBuilder.java

@@ -16,17 +16,27 @@
 
 package net.fabricmc.fabric.api.object.builder.v1.entity;
 
+import java.util.Objects;
+import java.util.function.Supplier;
+
 import com.google.common.collect.ImmutableSet;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
 import net.minecraft.entity.Entity;
+import net.minecraft.entity.LivingEntity;
 import net.minecraft.entity.SpawnGroup;
 import net.minecraft.entity.EntityDimensions;
 import net.minecraft.entity.EntityType;
 import net.minecraft.block.Block;
+import net.minecraft.entity.SpawnRestriction;
+import net.minecraft.entity.attribute.DefaultAttributeContainer;
+import net.minecraft.entity.mob.MobEntity;
+import net.minecraft.world.Heightmap;
+import net.minecraft.world.World;
 
 import net.fabricmc.fabric.impl.object.builder.FabricEntityType;
+import net.fabricmc.fabric.mixin.object.builder.SpawnRestrictionAccessor;
 
 /**
  * Extended version of {@link EntityType.Builder} with added registration for
@@ -36,8 +46,8 @@ import net.fabricmc.fabric.impl.object.builder.FabricEntityType;
  */
 public class FabricEntityTypeBuilder<T extends Entity> {
 	private static final Logger LOGGER = LogManager.getLogger();
-	private final SpawnGroup spawnGroup;
+	private SpawnGroup spawnGroup;
-	private final EntityType.EntityFactory<T> function;
+	private EntityType.EntityFactory<T> factory;
 	private boolean saveable = true;
 	private boolean summonable = true;
 	private int trackingDistance = 5;
@@ -48,12 +58,25 @@ public class FabricEntityTypeBuilder<T extends Entity> {
 	private EntityDimensions dimensions = EntityDimensions.changing(-1.0f, -1.0f);
 	private ImmutableSet<Block> specificSpawnBlocks = ImmutableSet.of();
 
-	protected FabricEntityTypeBuilder(SpawnGroup spawnGroup, EntityType.EntityFactory<T> function) {
+	protected FabricEntityTypeBuilder(SpawnGroup spawnGroup, EntityType.EntityFactory<T> factory) {
 		this.spawnGroup = spawnGroup;
-		this.function = function;
+		this.factory = factory;
 		this.spawnableFarFromPlayer = spawnGroup == SpawnGroup.CREATURE || spawnGroup == SpawnGroup.MISC;
 	}
 
+	/**
+	 * Creates an entity type builder.
+	 *
+	 * <p>This entity's spawn group will automatically be set to {@link SpawnGroup#MISC}.
+	 *
+	 * @param <T> the type of entity
+	 *
+	 * @return a new entity type builder
+	 */
+	public static <T extends Entity> FabricEntityTypeBuilder<T> create() {
+		return create(SpawnGroup.MISC);
+	}
+
 	/**
 	 * Creates an entity type builder.
 	 *
@@ -63,20 +86,60 @@ public class FabricEntityTypeBuilder<T extends Entity> {
 	 * @return a new entity type builder
 	 */
 	public static <T extends Entity> FabricEntityTypeBuilder<T> create(SpawnGroup spawnGroup) {
-		return new FabricEntityTypeBuilder<>(spawnGroup, (t, w) -> null);
+		return create(spawnGroup, FabricEntityTypeBuilder::emptyFactory);
 	}
 
 	/**
 	 * Creates an entity type builder.
 	 *
 	 * @param spawnGroup the entity spawn group
-	 * @param function the entity function used to create this entity
+	 * @param factory the entity factory used to create this entity
 	 * @param <T> the type of entity
 	 *
 	 * @return a new entity type builder
 	 */
-	public static <T extends Entity> FabricEntityTypeBuilder<T> create(SpawnGroup spawnGroup, EntityType.EntityFactory<T> function) {
+	public static <T extends Entity> FabricEntityTypeBuilder<T> create(SpawnGroup spawnGroup, EntityType.EntityFactory<T> factory) {
-		return new FabricEntityTypeBuilder<>(spawnGroup, function);
+		return new FabricEntityTypeBuilder<>(spawnGroup, factory);
+	}
+
+	/**
+	 * Creates an entity type builder for a living entity.
+	 *
+	 * <p>This entity's spawn group will automatically be set to {@link SpawnGroup#MISC}.
+	 *
+	 * @param <T> the type of entity
+	 *
+	 * @return a new living entity type builder
+	 */
+	public static <T extends LivingEntity> FabricEntityTypeBuilder.Living<T> createLiving() {
+		return new FabricEntityTypeBuilder.Living<>(SpawnGroup.MISC, FabricEntityTypeBuilder::emptyFactory);
+	}
+
+	/**
+	 * Creates an entity type builder for a mob entity.
+	 *
+	 * @param <T> the type of entity
+	 *
+	 * @return a new mob entity type builder
+	 */
+	public static <T extends MobEntity> FabricEntityTypeBuilder.Mob<T> createMob() {
+		return new FabricEntityTypeBuilder.Mob<>(SpawnGroup.MISC, FabricEntityTypeBuilder::emptyFactory);
+	}
+
+	private static <T extends Entity> T emptyFactory(EntityType<T> type, World world) {
+		return null;
+	}
+
+	public FabricEntityTypeBuilder<T> spawnGroup(SpawnGroup group) {
+		Objects.requireNonNull(group, "Spawn group cannot be null");
+		this.spawnGroup = group;
+		return this;
+	}
+
+	public FabricEntityTypeBuilder<T> entityFactory(EntityType.EntityFactory<T> factory) {
+		Objects.requireNonNull(factory, "Entity Factory cannot be null");
+		this.factory = factory;
+		return this;
 	}
 
 	/**
@@ -122,6 +185,7 @@ public class FabricEntityTypeBuilder<T extends Entity> {
 	 * @return this builder for chaining
 	 */
 	public FabricEntityTypeBuilder<T> dimensions(EntityDimensions dimensions) {
+		Objects.requireNonNull(dimensions, "Cannot set null dimensions");
 		this.dimensions = dimensions;
 		return this;
 	}
@@ -159,8 +223,221 @@ public class FabricEntityTypeBuilder<T extends Entity> {
 			// TODO: Flesh out once modded datafixers exist.
 		}
 
-		EntityType<T> type = new FabricEntityType<T>(this.function, this.spawnGroup, this.saveable, this.summonable, this.fireImmune, this.spawnableFarFromPlayer, this.specificSpawnBlocks, dimensions, trackingDistance, updateIntervalTicks, alwaysUpdateVelocity);
+		EntityType<T> type = new FabricEntityType<>(this.factory, this.spawnGroup, this.saveable, this.summonable, this.fireImmune, this.spawnableFarFromPlayer, this.specificSpawnBlocks, dimensions, trackingDistance, updateIntervalTicks, alwaysUpdateVelocity);
 
 		return type;
 	}
+
+	/**
+	 * An extended version of {@link FabricEntityTypeBuilder} with support for features on present on {@link LivingEntity living entities}, such as default attributes.
+	 *
+	 * @param <T> Entity class.
+	 */
+	public static class Living<T extends LivingEntity> extends FabricEntityTypeBuilder<T> {
+		/* @Nullable */
+		private Supplier<DefaultAttributeContainer.Builder> defaultAttributeBuilder;
+
+		protected Living(SpawnGroup spawnGroup, EntityType.EntityFactory<T> function) {
+			super(spawnGroup, function);
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> spawnGroup(SpawnGroup group) {
+			super.spawnGroup(group);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> entityFactory(EntityType.EntityFactory<T> factory) {
+			super.entityFactory(factory);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> disableSummon() {
+			super.disableSummon();
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> disableSaving() {
+			super.disableSaving();
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> fireImmune() {
+			super.fireImmune();
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> spawnableFarFromPlayer() {
+			super.spawnableFarFromPlayer();
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> dimensions(EntityDimensions dimensions) {
+			super.dimensions(dimensions);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> trackable(int trackingDistanceBlocks, int updateIntervalTicks) {
+			super.trackable(trackingDistanceBlocks, updateIntervalTicks);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> trackable(int trackingDistanceBlocks, int updateIntervalTicks, boolean alwaysUpdateVelocity) {
+			super.trackable(trackingDistanceBlocks, updateIntervalTicks, alwaysUpdateVelocity);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Living<T> specificSpawnBlocks(Block... blocks) {
+			super.specificSpawnBlocks(blocks);
+			return this;
+		}
+
+		/**
+		 * Sets the default attributes for a type of living entity.
+		 *
+		 * <p>This can be used in a fashion similar to this:
+		 * <blockquote><pre>
+		 * FabricEntityTypeBuilder.createLiving()
+		 * 	.spawnGroup(SpawnGroup.CREATURE)
+		 * 	.entityFactory(MyCreature::new)
+		 * 	.defaultAttributes(LivingEntity::createLivingAttributes)
+		 * 	...
+		 * 	.build();
+		 * </pre></blockquote>
+		 *
+		 * @param defaultAttributeBuilder a function to generate the default attribute builder from the entity type
+		 * @return this builder for chaining
+		 */
+		public FabricEntityTypeBuilder.Living<T> defaultAttributes(Supplier<DefaultAttributeContainer.Builder> defaultAttributeBuilder) {
+			Objects.requireNonNull(defaultAttributeBuilder, "Cannot set null attribute builder");
+			this.defaultAttributeBuilder = defaultAttributeBuilder;
+			return this;
+		}
+
+		@Override
+		public EntityType<T> build() {
+			final EntityType<T> type = super.build();
+
+			if (this.defaultAttributeBuilder != null) {
+				FabricDefaultAttributeRegistry.register(type, this.defaultAttributeBuilder.get());
+			}
+
+			return type;
+		}
+	}
+
+	/**
+	 * An extended version of {@link FabricEntityTypeBuilder} with support for features on present on {@link MobEntity mob entities}, such as spawn restrictions.
+	 *
+	 * @param <T> Entity class.
+	 */
+	public static class Mob<T extends MobEntity> extends FabricEntityTypeBuilder.Living<T> {
+		private SpawnRestriction.Location restrictionLocation;
+		private Heightmap.Type restrictionHeightmap;
+		private SpawnRestriction.SpawnPredicate<T> spawnPredicate;
+
+		protected Mob(SpawnGroup spawnGroup, EntityType.EntityFactory<T> function) {
+			super(spawnGroup, function);
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> spawnGroup(SpawnGroup group) {
+			super.spawnGroup(group);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> entityFactory(EntityType.EntityFactory<T> factory) {
+			super.entityFactory(factory);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> disableSummon() {
+			super.disableSummon();
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> disableSaving() {
+			super.disableSaving();
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> fireImmune() {
+			super.fireImmune();
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> spawnableFarFromPlayer() {
+			super.spawnableFarFromPlayer();
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> dimensions(EntityDimensions dimensions) {
+			super.dimensions(dimensions);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> trackable(int trackingDistanceBlocks, int updateIntervalTicks) {
+			super.trackable(trackingDistanceBlocks, updateIntervalTicks);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> trackable(int trackingDistanceBlocks, int updateIntervalTicks, boolean alwaysUpdateVelocity) {
+			super.trackable(trackingDistanceBlocks, updateIntervalTicks, alwaysUpdateVelocity);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> specificSpawnBlocks(Block... blocks) {
+			super.specificSpawnBlocks(blocks);
+			return this;
+		}
+
+		@Override
+		public FabricEntityTypeBuilder.Mob<T> defaultAttributes(Supplier<DefaultAttributeContainer.Builder> defaultAttributeBuilder) {
+			super.defaultAttributes(defaultAttributeBuilder);
+			return this;
+		}
+
+		/**
+		 * Registers a spawn restriction for this entity.
+		 *
+		 * <p>This is used by mobs to determine whether Minecraft should spawn an entity within a certain context.
+		 *
+		 * @return this builder for chaining.
+		 */
+		public FabricEntityTypeBuilder.Mob<T> spawnRestriction(SpawnRestriction.Location location, Heightmap.Type heightmap, SpawnRestriction.SpawnPredicate<T> spawnPredicate) {
+			this.restrictionLocation = Objects.requireNonNull(location, "Location cannot be null.");
+			this.restrictionHeightmap = Objects.requireNonNull(heightmap, "Heightmap type cannot be null.");
+			this.spawnPredicate = Objects.requireNonNull(spawnPredicate, "Spawn predicate cannot be null.");
+			return this;
+		}
+
+		@Override
+		public EntityType<T> build() {
+			EntityType<T> type = super.build();
+
+			if (this.spawnPredicate != null) {
+				SpawnRestrictionAccessor.callRegister(type, this.restrictionLocation, this.restrictionHeightmap, this.spawnPredicate);
+			}
+
+			return type;
+		}
+	}
 }

fabric-object-builder-api-v1/src/main/java/net/fabricmc/fabric/mixin/object/builder/SpawnRestrictionAccessor.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) 2016, 2017, 2018, 2019 FabricMC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package net.fabricmc.fabric.mixin.object.builder;
+
+import org.spongepowered.asm.mixin.Mixin;
+import org.spongepowered.asm.mixin.gen.Invoker;
+
+import net.minecraft.entity.EntityType;
+import net.minecraft.entity.SpawnRestriction;
+import net.minecraft.entity.mob.MobEntity;
+import net.minecraft.world.Heightmap;
+
+@Mixin(SpawnRestriction.class)
+public interface SpawnRestrictionAccessor {
+	@Invoker
+	static <T extends MobEntity> void callRegister(EntityType<T> type, SpawnRestriction.Location location, Heightmap.Type heightmap, SpawnRestriction.SpawnPredicate<T> spawnPredicate) {
+		throw new AssertionError("This should not occur!");
+	}
+}

fabric-object-builder-api-v1/src/main/resources/fabric-object-builder-v1.mixins.json

@@ -11,6 +11,7 @@
     "MaterialBuilderAccessor",
     "MixinBlock",
     "PointOfInterestTypeAccessor",
+    "SpawnRestrictionAccessor",
     "VillagerProfessionAccessor"
   ],
   "client": [