From 347d9f0d310ac93d566400e73a87e007debe137f Mon Sep 17 00:00:00 2001
From: erri120 <erri120@fossmailer.de>
Date: Mon, 22 Jul 2024 18:46:29 +0200
Subject: [PATCH] Use custom enumerator implementation

Results in less allocations.
---
 .../IndexSegments/Entities.cs                 | 106 ++++++++++++++----
 .../IndexSegments/EntityIds.cs                |  62 +++++++---
 2 files changed, 131 insertions(+), 37 deletions(-)

diff --git a/src/NexusMods.MnemonicDB.Abstractions/IndexSegments/Entities.cs b/src/NexusMods.MnemonicDB.Abstractions/IndexSegments/Entities.cs
index 454d9a72..2d4e9b7c 100644
--- a/src/NexusMods.MnemonicDB.Abstractions/IndexSegments/Entities.cs
+++ b/src/NexusMods.MnemonicDB.Abstractions/IndexSegments/Entities.cs
@@ -1,5 +1,6 @@
 using System.Collections;
 using System.Collections.Generic;
+using JetBrains.Annotations;
 using NexusMods.MnemonicDB.Abstractions.Models;
 
 namespace NexusMods.MnemonicDB.Abstractions.IndexSegments;
@@ -7,6 +8,7 @@ namespace NexusMods.MnemonicDB.Abstractions.IndexSegments;
 /// <summary>
 /// A subview of an IndexSegment that returns a specific entity type
 /// </summary>
+[PublicAPI]
 public readonly struct Entities<TInner, TEntity>(TInner ids, IDb db) :
     IReadOnlyCollection<TEntity>
     where TEntity : IReadOnlyModel<TEntity>
@@ -22,22 +24,49 @@ public readonly struct Entities<TInner, TEntity>(TInner ids, IDb db) :
     /// </summary>
     public int Count => ids.Count;
 
-    /// <inheritdoc />
-    public IEnumerator<TEntity> GetEnumerator()
+    /// <summary>
+    /// Returns an enumerator.
+    /// </summary>
+    public Enumerator GetEnumerator() => new(db, ids);
+    IEnumerator<TEntity> IEnumerable<TEntity>.GetEnumerator() => GetEnumerator();
+    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+
+    /// <summary>
+    /// Enumerator.
+    /// </summary>
+    public struct Enumerator : IEnumerator<TEntity>
     {
-        for (var i = 0; i < Count; i++)
+        private readonly IDb _db;
+        private readonly TInner _ids;
+        private int _index;
+
+        internal Enumerator(IDb db, TInner ids)
         {
-            yield return this[i];
+            _db = db;
+            _ids = ids;
         }
-    }
 
-    IEnumerator IEnumerable.GetEnumerator()
-    {
-        return GetEnumerator();
+        /// <inheritdoc/>
+        public TEntity Current { get; private set; } = default!;
+        object IEnumerator.Current => Current;
+
+        /// <inheritdoc/>
+        public bool MoveNext()
+        {
+            if (_index >= _ids.Count) return false;
+
+            Current = TEntity.Create(_db, _ids[_index++]);
+            return true;
+        }
+
+        /// <inheritdoc/>
+        public void Reset() => _index = 0;
+
+        /// <inheritdoc/>
+        public void Dispose() { }
     }
 }
 
-
 /// <summary>
 /// A wrapper around EntityIds that auto-creates the given ReadModel on-the-fly
 /// </summary>
@@ -58,27 +87,58 @@ public Entities(EntityIds ids, IDb db)
         _db = db;
     }
 
+    /// <inheritdoc />
+    public int Count => _ids.Count;
+
+    /// <summary>
+    /// Returns an enumerator.
+    /// </summary>
+    public Enumerator GetEnumerator() => new(_db, _ids);
+    IEnumerator<TModel> IEnumerable<TModel>.GetEnumerator() => GetEnumerator();
+    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 
     /// <inheritdoc />
-    public IEnumerator<TModel> GetEnumerator()
+    public override string ToString()
     {
-        foreach (var id in _ids)
-        {
-            yield return TModel.Create(_db, id);
-        }
+        return $"Entities<{typeof(TModel).FullName ?? typeof(TModel).Name}>({Count})";
     }
 
-    IEnumerator IEnumerable.GetEnumerator()
+    /// <summary>
+    /// Enumerator.
+    /// </summary>
+    public struct Enumerator : IEnumerator<TModel>
     {
-        return GetEnumerator();
-    }
+        private readonly EntityIds _ids;
+        private readonly IDb _db;
+        private int _index;
 
-    /// <inheritdoc />
-    public int Count => _ids.Count;
+        internal Enumerator(IDb db, EntityIds ids)
+        {
+            _db = db;
+            _ids = ids;
+        }
 
-    /// <inheritdoc />
-    public override string ToString()
-    {
-        return $"Entities<{typeof(TModel).FullName ?? typeof(TModel).Name}>({Count})";
+        /// <inheritdoc/>
+        public TModel Current { get; private set; } = default!;
+
+        object IEnumerator.Current => Current;
+
+        /// <inheritdoc/>
+        public bool MoveNext()
+        {
+            if (_index >= _ids.Count) return false;
+            var id = _ids[_index++];
+
+            Current = TModel.Create(_db, id);
+            return true;
+        }
+
+        /// <inheritdoc/>
+        public void Reset() => _index = 0;
+
+        /// <inheritdoc/>
+        public void Dispose()
+        {
+        }
     }
 }
diff --git a/src/NexusMods.MnemonicDB.Abstractions/IndexSegments/EntityIds.cs b/src/NexusMods.MnemonicDB.Abstractions/IndexSegments/EntityIds.cs
index fcbd0df7..993aa1f2 100644
--- a/src/NexusMods.MnemonicDB.Abstractions/IndexSegments/EntityIds.cs
+++ b/src/NexusMods.MnemonicDB.Abstractions/IndexSegments/EntityIds.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using JetBrains.Annotations;
 using NexusMods.MnemonicDB.Abstractions.Models;
 
 namespace NexusMods.MnemonicDB.Abstractions.IndexSegments;
@@ -8,7 +9,8 @@ namespace NexusMods.MnemonicDB.Abstractions.IndexSegments;
 /// <summary>
 /// A subview of an IndexSegment that returns the entity ids of the segment
 /// </summary>
-public struct EntityIds(IndexSegment segment, int start, int end) :
+[PublicAPI]
+public readonly struct EntityIds(IndexSegment segment, int start, int end) :
     IReadOnlyCollection<EntityId>, IIndexSegment<EntityId>
 {
     /// <summary>
@@ -34,19 +36,12 @@ public EntityId[] ToArray()
         return arr;
     }
 
-    /// <inheritdoc />
-    public IEnumerator<EntityId> GetEnumerator()
-    {
-        for (var i = start; i < end; i++)
-        {
-            yield return this[i];
-        }
-    }
-
-    IEnumerator IEnumerable.GetEnumerator()
-    {
-        return GetEnumerator();
-    }
+    /// <summary>
+    /// Returns an enumerator.
+    /// </summary>
+    public Enumerator GetEnumerator() => new(segment, start, end);
+    IEnumerator<EntityId> IEnumerable<EntityId>.GetEnumerator() => GetEnumerator();
+    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 
     /// <summary>
     /// Creates a view of these Ids that auto-casts every Id into a model of the given model type
@@ -56,4 +51,43 @@ public Entities<TModel> AsModels<TModel>(IDb db)
     {
         return new Entities<TModel>(this, db);
     }
+
+    /// <summary>
+    /// Enumerator.
+    /// </summary>
+    public struct Enumerator : IEnumerator<EntityId>
+    {
+        private readonly IndexSegment _segment;
+        private readonly int _start;
+        private readonly int _end;
+        private int _index;
+
+        internal Enumerator(IndexSegment segment, int start, int end)
+        {
+            _segment = segment;
+            _start = start;
+            _end = end;
+
+            _index = start;
+        }
+
+        /// <inheritdoc/>
+        public EntityId Current { get; private set; } = default!;
+        object IEnumerator.Current => Current;
+
+        /// <inheritdoc/>
+        public bool MoveNext()
+        {
+            if (_index >= _end) return false;
+            Current = _segment[_index + _start].E;
+            _index += 1;
+            return true;
+        }
+
+        /// <inheritdoc/>
+        public void Reset() => _index = 0;
+
+        /// <inheritdoc/>
+        public void Dispose() { }
+    }
 }