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

Content.Client/Lobby/UI/HumanoidProfileEditor.xaml.cs

@@ -479,10 +479,10 @@ namespace Content.Client.Lobby.UI
                 return;
             }
 
-            //Setup model
+            // Setup model
-            Dictionary<string, List<string>> model = new();
+            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))
+                if (!_prototypeManager.HasIndex(trait.Category))
-                {
+                    continue;
-                    model.Add(trait.Category, new());
+
-                }
+                var group = traitGroups.GetOrNew(trait.Category);
-                model[trait.Category].Add(trait.ID);
+                group.Add(trait.ID);
             }
 
-            //Create UI view from model
+            // Create UI view from model
-            foreach (var (categoryId, traitId) in 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.
                     };

Content.Shared/Preferences/HumanoidCharacterProfile.cs

@@ -387,22 +387,23 @@ 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>();
+            // null category is assumed to be default.
-            var traitProto = prototypeManager.Index(traitId);
+            if (!protoManager.TryIndex(traitId, out var traitProto))
+                return new(this);
 
-            TraitCategoryPrototype? categoryProto = null;
+            var category = traitProto.Category;
-            if (categoryId != null && categoryId != "default")
-                categoryProto = prototypeManager.Index<TraitCategoryPrototype>(categoryId);
 
-            var list = new HashSet<ProtoId<TraitPrototype>>(_traitPreferences);
+            // Category not found so dump it.
+            TraitCategoryPrototype? traitCategory = null;
 
-            if (pref)
+            if (category != null && !protoManager.TryIndex(category, out traitCategory))
-            {
+                return new(this);
-                list.Add(traitId);
 
-                if (categoryProto == null || categoryProto.MaxTraitPoints < 0)
+            var list = new HashSet<ProtoId<TraitPrototype>>(_traitPreferences) { traitId };
+
+            if (traitCategory == null || traitCategory.MaxTraitPoints < 0)
             {
                 return new(this)
                 {
@@ -413,22 +414,31 @@ namespace Content.Shared.Preferences
             var count = 0;
             foreach (var trait in list)
             {
-                    var traitProtoTemp = prototypeManager.Index(trait);
+                // If trait not found or another category don't count its points.
-                    count += traitProtoTemp.Cost;
+                if (!protoManager.TryIndex<TraitPrototype>(trait, out var otherProto) ||
+                    otherProto.Category != traitCategory)
+                {
+                    continue;
                 }
 
-                if (count > categoryProto.MaxTraitPoints && traitProto.Cost != 0)
+                count += otherProto.Cost;
+            }
+
+            if (count > traitCategory.MaxTraitPoints && traitProto.Cost != 0)
             {
+                return new(this);
+            }
+
             return new(this)
             {
-                        _traitPreferences = _traitPreferences,
+                _traitPreferences = list,
             };
         }
-            }
+
-            else
+        public HumanoidCharacterProfile WithoutTraitPreference(ProtoId<TraitPrototype> traitId, IPrototypeManager protoManager)
         {
+            var list = new HashSet<ProtoId<TraitPrototype>>(_traitPreferences);
             list.Remove(traitId);
-            }
 
             return new(this)
             {
@@ -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);

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!;