Ensure trait groups get validated (#28730)

* Ensure trait groups get validated

The only validation being done was on the UI. I also made the "Default" group match the PascalCase naming schema so might be a slight breaking change but the original PR only got merged a few days ago.

* overwatch

diff --git a/Content.Client/Lobby/UI/HumanoidProfileEditor.xaml.cs b/Content.Client/Lobby/UI/HumanoidProfileEditor.xaml.cs
index e8e619ae6dc6e6ac0ea086bf89791c9c596845a4..b3bc0cafe741397fff6a42d4f4bad9092334ec5e 100644 (file)
--- a/Content.Client/Lobby/UI/HumanoidProfileEditor.xaml.cs
+++ b/Content.Client/Lobby/UI/HumanoidProfileEditor.xaml.cs
@@ -479,10 +479,10 @@ namespace Content.Client.Lobby.UI
                 return;
             }
 
-            //Setup model
-            Dictionary<string, List<string>> model = new();
+            // Setup model
+            Dictionary<string, List<string>> traitGroups = new();
             List<string> defaultTraits = new();
-            model.Add("default", defaultTraits);
+            traitGroups.Add(TraitCategoryPrototype.Default, defaultTraits);
 
             foreach (var trait in traits)
             {
@@ -492,18 +492,19 @@ namespace Content.Client.Lobby.UI
                     continue;
                 }
 
-                if (!model.ContainsKey(trait.Category))
-                {
-                    model.Add(trait.Category, new());
-                }
-                model[trait.Category].Add(trait.ID);
+                if (!_prototypeManager.HasIndex(trait.Category))
+                    continue;
+
+                var group = traitGroups.GetOrNew(trait.Category);
+                group.Add(trait.ID);
             }
 
-            //Create UI view from model
-            foreach (var (categoryId, traitId) in model)
+            // Create UI view from model
+            foreach (var (categoryId, categoryTraits) in traitGroups)
             {
                 TraitCategoryPrototype? category = null;
-                if (categoryId != "default")
+
+                if (categoryId != TraitCategoryPrototype.Default)
                 {
                     category = _prototypeManager.Index<TraitCategoryPrototype>(categoryId);
                     // Label
@@ -518,7 +519,7 @@ namespace Content.Client.Lobby.UI
                 List<TraitPreferenceSelector?> selectors = new();
                 var selectionCount = 0;
 
-                foreach (var traitProto in traitId)
+                foreach (var traitProto in categoryTraits)
                 {
                     var trait = _prototypeManager.Index<TraitPrototype>(traitProto);
                     var selector = new TraitPreferenceSelector(trait);
@@ -529,7 +530,15 @@ namespace Content.Client.Lobby.UI
 
                     selector.PreferenceChanged += preference =>
                     {
-                        Profile = Profile?.WithTraitPreference(trait.ID, categoryId, preference);
+                        if (preference)
+                        {
+                            Profile = Profile?.WithTraitPreference(trait.ID, _prototypeManager);
+                        }
+                        else
+                        {
+                            Profile = Profile?.WithoutTraitPreference(trait.ID, _prototypeManager);
+                        }
+
                         SetDirty();
                         RefreshTraits(); // If too many traits are selected, they will be reset to the real value.
                     };

diff --git a/Content.Shared/Preferences/HumanoidCharacterProfile.cs b/Content.Shared/Preferences/HumanoidCharacterProfile.cs
index 22b4ed98150710b5fbd0842b7d80ff3760b463e1..a823035cd3877d85153c14f0add604c08d3f4eab 100644 (file)
--- a/Content.Shared/Preferences/HumanoidCharacterProfile.cs
+++ b/Content.Shared/Preferences/HumanoidCharacterProfile.cs
@@ -387,47 +387,46 @@ namespace Content.Shared.Preferences
             };
         }
 
-        public HumanoidCharacterProfile WithTraitPreference(ProtoId<TraitPrototype> traitId, string? categoryId, bool pref)
+        public HumanoidCharacterProfile WithTraitPreference(ProtoId<TraitPrototype> traitId, IPrototypeManager protoManager)
         {
-            var prototypeManager = IoCManager.Resolve<IPrototypeManager>();
-            var traitProto = prototypeManager.Index(traitId);
+            // null category is assumed to be default.
+            if (!protoManager.TryIndex(traitId, out var traitProto))
+                return new(this);
 
-            TraitCategoryPrototype? categoryProto = null;
-            if (categoryId != null && categoryId != "default")
-                categoryProto = prototypeManager.Index<TraitCategoryPrototype>(categoryId);
+            var category = traitProto.Category;
 
-            var list = new HashSet<ProtoId<TraitPrototype>>(_traitPreferences);
+            // Category not found so dump it.
+            TraitCategoryPrototype? traitCategory = null;
 
-            if (pref)
-            {
-                list.Add(traitId);
+            if (category != null && !protoManager.TryIndex(category, out traitCategory))
+                return new(this);
 
-                if (categoryProto == null || categoryProto.MaxTraitPoints < 0)
-                {
-                    return new(this)
-                    {
-                        _traitPreferences = list,
-                    };
-                }
+            var list = new HashSet<ProtoId<TraitPrototype>>(_traitPreferences) { traitId };
 
-                var count = 0;
-                foreach (var trait in list)
+            if (traitCategory == null || traitCategory.MaxTraitPoints < 0)
+            {
+                return new(this)
                 {
-                    var traitProtoTemp = prototypeManager.Index(trait);
-                    count += traitProtoTemp.Cost;
-                }
+                    _traitPreferences = list,
+                };
+            }
 
-                if (count > categoryProto.MaxTraitPoints && traitProto.Cost != 0)
+            var count = 0;
+            foreach (var trait in list)
+            {
+                // If trait not found or another category don't count its points.
+                if (!protoManager.TryIndex<TraitPrototype>(trait, out var otherProto) ||
+                    otherProto.Category != traitCategory)
                 {
-                    return new(this)
-                    {
-                        _traitPreferences = _traitPreferences,
-                    };
+                    continue;
                 }
+
+                count += otherProto.Cost;
             }
-            else
+
+            if (count > traitCategory.MaxTraitPoints && traitProto.Cost != 0)
             {
-                list.Remove(traitId);
+                return new(this);
             }
 
             return new(this)
@@ -436,6 +435,17 @@ namespace Content.Shared.Preferences
             };
         }
 
+        public HumanoidCharacterProfile WithoutTraitPreference(ProtoId<TraitPrototype> traitId, IPrototypeManager protoManager)
+        {
+            var list = new HashSet<ProtoId<TraitPrototype>>(_traitPreferences);
+            list.Remove(traitId);
+
+            return new(this)
+            {
+                _traitPreferences = list,
+            };
+        }
+
         public string Summary =>
             Loc.GetString(
                 "humanoid-character-profile-summary",
@@ -605,7 +615,7 @@ namespace Content.Shared.Preferences
             _antagPreferences.UnionWith(antags);
 
             _traitPreferences.Clear();
-            _traitPreferences.UnionWith(traits);
+            _traitPreferences.UnionWith(GetValidTraits(traits, prototypeManager));
 
             // Checks prototypes exist for all loadouts and dump / set to default if not.
             var toRemove = new ValueList<string>();
@@ -627,6 +637,45 @@ namespace Content.Shared.Preferences
             }
         }
 
+        /// <summary>
+        /// Takes in an IEnumerable of traits and returns a List of the valid traits.
+        /// </summary>
+        public List<ProtoId<TraitPrototype>> GetValidTraits(IEnumerable<ProtoId<TraitPrototype>> traits, IPrototypeManager protoManager)
+        {
+            // Track points count for each group.
+            var groups = new Dictionary<string, int>();
+            var result = new List<ProtoId<TraitPrototype>>();
+
+            foreach (var trait in traits)
+            {
+                if (!protoManager.TryIndex(trait, out var traitProto))
+                    continue;
+
+                // Always valid.
+                if (traitProto.Category == null)
+                {
+                    result.Add(trait);
+                    continue;
+                }
+
+                // No category so dump it.
+                if (!protoManager.TryIndex(traitProto.Category, out var category))
+                    continue;
+
+                var existing = groups.GetOrNew(category.ID);
+                existing += traitProto.Cost;
+
+                // Too expensive.
+                if (existing > category.MaxTraitPoints)
+                    continue;
+
+                groups[category.ID] = existing;
+                result.Add(trait);
+            }
+
+            return result;
+        }
+
         public ICharacterProfile Validated(ICommonSession session, IDependencyCollection collection)
         {
             var profile = new HumanoidCharacterProfile(this);

diff --git a/Content.Shared/Traits/TraitCategoryPrototype.cs b/Content.Shared/Traits/TraitCategoryPrototype.cs
index 1da624173a3db4e2a4f4df36d7f7101e49bb3ed6..e5bb919bf07ec532d7b7023681005b35bc039a2a 100644 (file)
--- a/Content.Shared/Traits/TraitCategoryPrototype.cs
+++ b/Content.Shared/Traits/TraitCategoryPrototype.cs
@@ -8,6 +8,8 @@ namespace Content.Shared.Traits;
 [Prototype]
 public sealed partial class TraitCategoryPrototype : IPrototype
 {
+    public const string Default = "Default";
+
     [ViewVariables]
     [IdDataField]
     public string ID { get; private set; } = default!;