#nullable enable using System; using System.Collections; using System.Collections.Generic; using System.Collections.Immutable; using System.Linq; using System.Threading; using Barotrauma.Threading; namespace Barotrauma { public class PrefabSelector : IEnumerable where T : notnull, Prefab { private readonly ReaderWriterLockSlim rwl = new ReaderWriterLockSlim(); public T? BasePrefab { get { using (new ReadLock(rwl)) { return basePrefabInternal; } } } public T? ActivePrefab { get { using (new ReadLock(rwl)) { return activePrefabInternal; } } } public void Add(T prefab, bool isOverride) { using (new WriteLock(rwl)) { AddInternal(prefab, isOverride); } } public void RemoveIfContains(T prefab) { using (new WriteLock(rwl)) { RemoveIfContainsInternal(prefab); } } public void Remove(T prefab) { using (new WriteLock(rwl)) { RemoveInternal(prefab); } } public void RemoveByFile(ContentFile file, Action? callback = null) { var removed = new List(); using (new WriteLock(rwl)) { for (int i = overrides.Count-1; i >= 0; i--) { var prefab = overrides[i]; if (prefab.ContentFile == file) { RemoveInternal(prefab); removed.Add(prefab); } } if (basePrefabInternal is { ContentFile: var baseFile } p && baseFile == file) { RemoveInternal(basePrefabInternal); removed.Add(p); } } if (callback != null) { removed.ForEach(callback); } } public void Sort() { using (new WriteLock(rwl)) { SortInternal(); } } public bool IsEmpty { get { using (new ReadLock(rwl)) { return isEmptyInternal; } } } public bool Contains(T prefab) { using (new ReadLock(rwl)) { return ContainsInternal(prefab); } } public bool IsOverride(T prefab) { using (new ReadLock(rwl)) { 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.Count > 0 ? 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 (basePrefabInternal != 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 \"{activePrefabInternal!.ContentPackage?.Name ?? "[NULL]"}\" ({activePrefabInternal!.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 SortInternal() { overrides.Sort((p1, p2) => (p1.ContentPackage?.Index ?? int.MaxValue) - (p2.ContentPackage?.Index ?? int.MaxValue)); } private bool isEmptyInternal => basePrefabInternal is null && overrides.Count == 0; 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; using (new ReadLock(rwl)) { basePrefab = basePrefabInternal; overrideClone = overrides.ToImmutableArray(); } if (basePrefab != null) { yield return basePrefab; } foreach (T prefab in overrideClone) { yield return prefab; } } IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); } }