From 1652eb6a7bbcac1a65e3e96412372308e4f6b2e9 Mon Sep 17 00:00:00 2001
From: halgari <tbaldridge@gmail.com>
Date: Tue, 16 Jan 2024 17:05:24 -0700
Subject: [PATCH] Implemented secondary indexes as attributes

---
 .../IndexedMultiEntityAttributeDefinition.cs  | 104 ++++++++++++++++++
 .../Events/CreateLoadout.cs                   |   1 +
 .../Model/LoadoutRegistry.cs                  |   4 +
 .../BasicFunctionalityTests.cs                |   5 +
 4 files changed, 114 insertions(+)
 create mode 100644 src/NexusMods.EventSourcing.Abstractions/IndexedMultiEntityAttributeDefinition.cs

diff --git a/src/NexusMods.EventSourcing.Abstractions/IndexedMultiEntityAttributeDefinition.cs b/src/NexusMods.EventSourcing.Abstractions/IndexedMultiEntityAttributeDefinition.cs
new file mode 100644
index 00000000..138655a4
--- /dev/null
+++ b/src/NexusMods.EventSourcing.Abstractions/IndexedMultiEntityAttributeDefinition.cs
@@ -0,0 +1,104 @@
+using System;
+using System.Buffers;
+using System.Buffers.Binary;
+using System.Collections.Generic;
+using NexusMods.EventSourcing.Abstractions.Serialization;
+
+namespace NexusMods.EventSourcing.Abstractions;
+
+/// <summary>
+/// Like <see cref="MultiEntityAttributeDefinition{TOwner,TOther}"/>, but with the data stored in a dictionary
+/// </summary>
+public class IndexedMultiEntityAttributeDefinition<TOwner, TKey, TOther>(string name) : IAttribute<IndexedMultiEntityAccumulator<TKey, TOther>>
+    where TOwner : AEntity<TOwner>, IEntity
+    where TOther : AEntity<TOther>
+    where TKey : notnull
+{
+    /// <inheritdoc />
+    public Type Owner => typeof(TOwner);
+
+    /// <inheritdoc />
+    public string Name => name;
+
+    public IndexedMultiEntityAccumulator<TKey, TOther> CreateAccumulator()
+    {
+        return new IndexedMultiEntityAccumulator<TKey, TOther>();
+    }
+
+    public void Add<TContext>(TContext context, EntityId<TOwner> owner, TKey key, EntityId<TOther> value)
+        where TContext : IEventContext
+    {
+        if (context.GetAccumulator<TOwner, IndexedMultiEntityAttributeDefinition<TOwner, TKey, TOther>,
+                IndexedMultiEntityAccumulator<TKey, TOther>>(owner, this, out var accumulator))
+        {
+            if (accumulator._keys.TryGetValue(value, out var existingKey))
+            {
+                accumulator._values.Remove(existingKey);
+                accumulator._keys.Remove(value);
+            }
+            accumulator._values.Add(key, value);
+            accumulator._keys.Add(value, key);
+        }
+    }
+
+    IAccumulator IAttribute.CreateAccumulator()
+    {
+        return CreateAccumulator();
+    }
+
+    /// <summary>
+    /// Gets the other entities linked to the given entity.
+    /// </summary>
+    /// <param name="entity"></param>
+    /// <returns></returns>
+    public Dictionary<TKey, EntityId<TOther>> Get(TOwner entity)
+    {
+        if (!entity.Context
+                .GetReadOnlyAccumulator<TOwner, IndexedMultiEntityAttributeDefinition<TOwner, TKey, TOther>,
+                    IndexedMultiEntityAccumulator<TKey, TOther>>(entity.Id, this, out var accumulator, true))
+            throw new InvalidOperationException("No accumulator found for entity");
+        return accumulator._values;
+    }
+}
+
+public class IndexedMultiEntityAccumulator<TKey, TOther> : IAccumulator
+    where TOther : AEntity<TOther>
+    where TKey : notnull
+{
+    internal Dictionary<TKey, EntityId<TOther>> _values = new();
+    internal Dictionary<EntityId<TOther>, TKey> _keys = new();
+
+    public void WriteTo(IBufferWriter<byte> writer, ISerializationRegistry registry)
+    {
+        var getSpan = writer.GetSpan(2);
+        BinaryPrimitives.WriteUInt16BigEndian(getSpan, (ushort) _values.Count);
+        writer.Advance(2);
+
+        foreach (var (key, value) in _values)
+        {
+            registry.Serialize(writer, key);
+            registry.Serialize(writer, value);
+        }
+    }
+
+    public int ReadFrom(ref ReadOnlySpan<byte> input, ISerializationRegistry registry)
+    {
+        var originalSize = input.Length;
+        var data = input;
+
+        var count = BinaryPrimitives.ReadUInt16BigEndian(data);
+        data = data.Slice(sizeof(ushort));
+
+        for (var idx = 0; idx < count; idx++)
+        {
+            var written = registry.Deserialize(data, out TKey key);
+            data = data.Slice(written);
+            written = registry.Deserialize(data, out EntityId<TOther> value);
+            data = data.Slice(written);
+            _values.Add(key, value);
+            _keys.Add(value, key);
+        }
+
+        return originalSize - data.Length;
+    }
+}
diff --git a/tests/NexusMods.EventSourcing.TestModel/Events/CreateLoadout.cs b/tests/NexusMods.EventSourcing.TestModel/Events/CreateLoadout.cs
index 52edd223..a5959c99 100644
--- a/tests/NexusMods.EventSourcing.TestModel/Events/CreateLoadout.cs
+++ b/tests/NexusMods.EventSourcing.TestModel/Events/CreateLoadout.cs
@@ -14,6 +14,7 @@ public void Apply<T>(T context) where T : IEventContext
         IEntity.TypeAttribute.New(context, Id);
         Loadout._name.Set(context, Id, Name);
         LoadoutRegistry._loadouts.Add(context, LoadoutRegistry.SingletonId, Id);
+        LoadoutRegistry._loadoutNames.Add(context, LoadoutRegistry.SingletonId, Name, Id);
     }
     public static EntityId<Loadout> Create(ITransaction tx, string name)
     {
diff --git a/tests/NexusMods.EventSourcing.TestModel/Model/LoadoutRegistry.cs b/tests/NexusMods.EventSourcing.TestModel/Model/LoadoutRegistry.cs
index d98072c7..1cc257fe 100644
--- a/tests/NexusMods.EventSourcing.TestModel/Model/LoadoutRegistry.cs
+++ b/tests/NexusMods.EventSourcing.TestModel/Model/LoadoutRegistry.cs
@@ -18,4 +18,8 @@ public class LoadoutRegistry(IEntityContext context) : AEntity<LoadoutRegistry>(
     public ReadOnlyObservableCollection<Loadout> Loadouts => _loadouts.Get(this);
     internal static readonly MultiEntityAttributeDefinition<LoadoutRegistry, Loadout> _loadouts = new(nameof(Loadouts));
 
+
+    public Dictionary<string, EntityId<Loadout>> LoadoutNames => _loadoutNames.Get(this);
+    internal static readonly IndexedMultiEntityAttributeDefinition<LoadoutRegistry, string, Loadout> _loadoutNames = new(nameof(LoadoutNames));
+
 }
diff --git a/tests/NexusMods.EventSourcing.Tests/BasicFunctionalityTests.cs b/tests/NexusMods.EventSourcing.Tests/BasicFunctionalityTests.cs
index b5a183c5..84d96787 100644
--- a/tests/NexusMods.EventSourcing.Tests/BasicFunctionalityTests.cs
+++ b/tests/NexusMods.EventSourcing.Tests/BasicFunctionalityTests.cs
@@ -31,6 +31,11 @@ public void CanSetupBasicLoadout()
         var loadout = _ctx.Get(loadoutId);
         loadout.Should().NotBeNull();
         loadout.Name.Should().Be("Test");
+
+        var loadoutRegistry = _ctx.Get<LoadoutRegistry>();
+        loadoutRegistry.Loadouts.Should().NotBeEmpty();
+        loadoutRegistry.Loadouts.First().Name.Should().Be("Test");
+        loadoutRegistry.LoadoutNames["Test"].Should().Be(loadoutId);
     }
 
     [Fact]