From 56d32b532db9edb433c5aa69510eb7d46f9eaed8 Mon Sep 17 00:00:00 2001
From: Auxilor <william.favierparsons1@gmail.com>
Date: Thu, 3 Mar 2022 10:12:02 +0000
Subject: [PATCH] Move improvements to CustomGoal

---
 .../com/willfp/eco/core/entities/ai/CustomGoal.java | 13 +++++++++++--
 .../internal/spigot/proxy/common/ai/CustomGoals.kt  |  2 +-
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/eco-api/src/main/java/com/willfp/eco/core/entities/ai/CustomGoal.java b/eco-api/src/main/java/com/willfp/eco/core/entities/ai/CustomGoal.java
index b374ba5f..ab8cba29 100644
--- a/eco-api/src/main/java/com/willfp/eco/core/entities/ai/CustomGoal.java
+++ b/eco-api/src/main/java/com/willfp/eco/core/entities/ai/CustomGoal.java
@@ -3,8 +3,8 @@ package com.willfp.eco.core.entities.ai;
 import org.bukkit.entity.Mob;
 import org.jetbrains.annotations.NotNull;
 
-import java.util.Collection;
 import java.util.EnumSet;
+import java.util.List;
 import java.util.Set;
 
 /**
@@ -98,7 +98,16 @@ public abstract class CustomGoal<T extends Mob> implements EntityGoal<T>, Target
      *
      * @param flags The flags.
      */
-    public void setFlags(@NotNull final Collection<GoalFlag> flags) {
+    public final void setFlags(@NotNull final GoalFlag... flags) {
+        this.setFlags(EnumSet.copyOf(List.of(flags)));
+    }
+
+    /**
+     * Set the flags for the goal.
+     *
+     * @param flags The flags.
+     */
+    public void setFlags(@NotNull final EnumSet<GoalFlag> flags) {
         this.flags.clear();
         this.flags.addAll(flags);
     }
diff --git a/eco-core/core-nms/nms-common/src/main/kotlin/com/willfp/eco/internal/spigot/proxy/common/ai/CustomGoals.kt b/eco-core/core-nms/nms-common/src/main/kotlin/com/willfp/eco/internal/spigot/proxy/common/ai/CustomGoals.kt
index 4ef59e98..98ea0e85 100644
--- a/eco-core/core-nms/nms-common/src/main/kotlin/com/willfp/eco/internal/spigot/proxy/common/ai/CustomGoals.kt
+++ b/eco-core/core-nms/nms-common/src/main/kotlin/com/willfp/eco/internal/spigot/proxy/common/ai/CustomGoals.kt
@@ -72,6 +72,6 @@ private class NMSCustomGoal<T : org.bukkit.entity.Mob>(
 
     override fun setFlags(controls: EnumSet<Flag>) {
         super.setFlags(controls)
-        customEntityGoal.setFlags(controls.toEcoFlags())
+        customEntityGoal.flags = EnumSet.copyOf(controls.toEcoFlags())
     }
 }