diff --git a/PkmnLib.Dynamic/ScriptHandling/ScriptIterator.cs b/PkmnLib.Dynamic/ScriptHandling/ScriptIterator.cs index 21dd414..96646ee 100644 --- a/PkmnLib.Dynamic/ScriptHandling/ScriptIterator.cs +++ b/PkmnLib.Dynamic/ScriptHandling/ScriptIterator.cs @@ -8,6 +8,7 @@ namespace PkmnLib.Dynamic.ScriptHandling; public class ScriptIterator : IEnumerable { private readonly IReadOnlyList> _scripts; + private int _currentIndex = 0; /// public ScriptIterator(IReadOnlyList> scripts) @@ -15,39 +16,41 @@ public class ScriptIterator : IEnumerable _scripts = scripts; } - private IEnumerable GetAsEnumerable() + /// + public IEnumerator GetEnumerator() { - foreach (var enumerable in _scripts) + while (_currentIndex < _scripts.Count) { - if (enumerable is IScriptSet set) + var enumerable = _scripts[_currentIndex++]; + switch (enumerable) { - // We can't assume that the set remains the same size, so we need to iterate through it - // with a for loop. - for (var j = 0; j < set.Count; j++) + case IScriptSet set: { - var script = set.At(j); - // We can ignore empty scripts. - if (script.IsEmpty) - continue; - yield return script; - } - } - else - { - foreach (var script in enumerable) - { - // We can ignore empty scripts. - if (script.IsEmpty) - continue; - yield return script; + // We can't assume that the set remains the same size, so we need to iterate through it + // with a for loop. + for (var j = 0; j < set.Count; j++) + { + var script = set.At(j); + // We can ignore empty scripts. + if (script.IsEmpty) + continue; + yield return script; + } + break; } + case ScriptContainer { IsEmpty: true }: + continue; + case ScriptContainer container: + yield return container; + break; + default: + throw new InvalidOperationException( + $"Unexpected script type: {enumerable.GetType().Name}. Expected IScriptSet or ScriptContainer."); } } + _currentIndex = 0; } - /// - public IEnumerator GetEnumerator() => GetAsEnumerable().GetEnumerator(); - /// IEnumerator IEnumerable.GetEnumerator() => GetEnumerator(); } \ No newline at end of file