#nullable enable using System; using System.Collections; using System.Collections.Generic; using System.Collections.Immutable; using System.Linq; namespace Barotrauma { public class PrefabSelector : IEnumerable where T : notnull, Prefab { public T? BasePrefab { get { lock (overrides) { return basePrefabInternal; } } } public T? ActivePrefab { get { lock (overrides) { return activePrefabInternal; } } } public void Add(T prefab, bool isOverride) { lock (overrides) { AddInternal(prefab, isOverride); } } public void RemoveIfContains(T prefab) { lock (overrides) { RemoveIfContainsInternal(prefab); } } public void Remove(T prefab) { lock (overrides) { RemoveInternal(prefab); } } public void RemoveByFile(ContentFile file, Action? callback = null) { lock (overrides) { RemoveByFileInternal(file, callback); } } public void Sort() { lock (overrides) { SortInternal(); } } public bool IsEmpty { get { lock (overrides) { return isEmptyInternal; } } } public bool Contains(T prefab) { lock (overrides) { return ContainsInternal(prefab); } } public bool IsOverride(T prefab) { lock (overrides) { return IsOverrideInternal(prefab); } } #region Underlying implementations of the public methods, done separately to avoid nested locking private T? basePrefabInternal; private readonly List overrides = new List(); private T? activePrefabInternal => overrides.Any() ? overrides.First() : basePrefabInternal; private void AddInternal(T prefab, bool isOverride) { if (isOverride) { if (overrides.Contains(prefab)) { throw new InvalidOperationException($"Duplicate prefab in PrefabSelector ({typeof(T)}, {prefab.Identifier}, {prefab.ContentFile.ContentPackage.Name})"); } overrides.Add(prefab); } else { if (BasePrefab != null) { string prefabName = prefab is MapEntityPrefab mapEntityPrefab ? $"\"{mapEntityPrefab.OriginalName}\", \"{prefab.Identifier}\"" : $"\"{prefab.Identifier}\""; throw new InvalidOperationException( $"Failed to add the prefab {prefabName} ({prefab.GetType()}) from \"{prefab.ContentPackage?.Name ?? "[NULL]"}\" ({prefab.ContentPackage?.Dir ?? ""}): " + $"a prefab with the same identifier from \"{ActivePrefab!.ContentPackage?.Name ?? "[NULL]"}\" ({ActivePrefab!.ContentPackage?.Dir ?? ""}) already exists; try overriding"); } basePrefabInternal = prefab; } SortInternal(); } private void RemoveIfContainsInternal(T prefab) { if (!ContainsInternal(prefab)) { return; } RemoveInternal(prefab); } private void RemoveInternal(T prefab) { if (basePrefabInternal == prefab) { basePrefabInternal = null; } else if (overrides.Contains(prefab)) { overrides.Remove(prefab); } else { throw new InvalidOperationException($"Can't remove prefab from PrefabSelector ({typeof(T)}, {prefab.Identifier}, {prefab.ContentFile.ContentPackage.Name})"); } prefab.Dispose(); SortInternal(); } private void RemoveByFileInternal(ContentFile file, Action? callback) { for (int i = overrides.Count-1; i >= 0; i--) { var prefab = overrides[i]; if (prefab.ContentFile == file) { RemoveInternal(prefab); callback?.Invoke(prefab); } } if (basePrefabInternal is { ContentFile: var baseFile } p && baseFile == file) { RemoveInternal(basePrefabInternal); callback?.Invoke(p); } } private void SortInternal() { overrides.Sort((p1, p2) => (p1.ContentPackage?.Index ?? int.MaxValue) - (p2.ContentPackage?.Index ?? int.MaxValue)); } private bool isEmptyInternal => basePrefabInternal is null && !overrides.Any(); private bool ContainsInternal(T prefab) => basePrefabInternal == prefab || overrides.Contains(prefab); private int IndexOfInternal(T prefab) => basePrefabInternal == prefab ? overrides.Count : overrides.IndexOf(prefab); private bool IsOverrideInternal(T prefab) => IndexOfInternal(prefab) > 0; #endregion public IEnumerator GetEnumerator() { T? basePrefab; ImmutableArray overrideClone; lock (overrides) { basePrefab = basePrefabInternal; overrideClone = overrides.ToImmutableArray(); } if (basePrefab != null) { yield return basePrefab; } foreach (T prefab in overrideClone) { yield return prefab; } } IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); } }