mfgames-nitride-cil/src/MfGames.Nitride/Entities/EntityScanner.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;

using FluentValidation;

using MfGames.Gallium;

namespace MfGames.Nitride.Entities;

/// <summary>
/// Implements a Nitride operation that scans the entities as they are
/// passed through the `Run` method and gathers information into a
/// `Dictionary&lt;string, List&lt;Entity&gt;&gt;` which then can be
/// queries by later operations. This handles making sure the entire
/// input has been processed before operating.
/// </summary>
[WithProperties]
public partial class EntityScanner : OperationBase
{
    private readonly object locker;

    private readonly ConcurrentDictionary<string, List<Entity>> results;

    private readonly IValidator<EntityScanner> validator;

    private bool done;

    public EntityScanner(IValidator<EntityScanner> validator)
    {
        this.validator = validator;
        this.locker = new object();
        this.results = new ConcurrentDictionary<string, List<Entity>>();
    }

    /// <summary>
    /// Gets or sets a callback function that gets the keys associated with
    /// the given entity.
    /// </summary>
    public Func<Entity, IEnumerable<string>?> GetKeysFromEntity { get; set; } =
        null!;

    /// <summary>
    /// Gets the list of entities associated with the given key. if the key has not
    /// been
    /// seen, this returns an empty collection.
    /// </summary>
    /// <param name="key">The key to search for.</param>
    /// <returns>A list of entities associated with the given key.</returns>
    /// <exception cref="InvalidOperationException">
    /// If the input has not been
    /// completely processed, this exception is thrown.
    /// </exception>
    public IEnumerable<Entity> GetScannedEntities(string key)
    {
        // Make sure we're done processing.
        this.CheckDone();

        // We have the list, so return it or an empty list.
        if (this.results.TryGetValue(key, out List<Entity>? list))
        {
            return list.AsReadOnly();
        }

        // We didn't have the list but we always return something.
        return Array.Empty<Entity>();
    }

    /// <summary>
    /// Gets a list of all known keys from the scanner.
    /// </summary>
    /// <returns></returns>
    public IEnumerable<string> GetScannedKeys()
    {
        this.CheckDone();

        return this.results.Keys.ToImmutableList();
    }

    /// <summary>
    /// Gets a dictionary of all the results from the scanner.
    /// </summary>
    /// <returns></returns>
    public ImmutableDictionary<string, List<Entity>> GetScannedResults()
    {
        this.CheckDone();

        return this.results.ToImmutableDictionary();
    }

    /// <inheritdoc />
    public override IEnumerable<Entity> Run(IEnumerable<Entity> input)
    {
        // Make sure we have sane data.
        this.validator.ValidateAndThrow(this);

        // Reset our done flag to handle re-entrant calls.
        lock (this.locker)
        {
            this.done = false;
        }

        // Loop through the entities and process each one.
        foreach (Entity? entity in input)
        {
            // Scan the given entity and see where it needs to be included.
            // The entity is added to each of the keys returned by this class.
            IEnumerable<string>? keysFromEntity =
                this.GetKeysFromEntity(entity);

            if (keysFromEntity != null)
            {
                foreach (string key in keysFromEntity)
                {
                    this.results.AddOrUpdate(
                        key,
                        _ => new List<Entity> { entity },
                        (
                            _,
                            list) => list.Union(new[] { entity })
                            .ToList());
                }
            }

            // Finish processing this entity.
            yield return entity;
        }

        // We are done, so flip our flag and we're done processing.
        lock (this.locker)
        {
            this.done = true;
        }
    }

    private void CheckDone()
    {
        lock (this.locker)
        {
            // Make sure we are done, otherwise give a useful message.
            if (!this.done)
            {
                throw new InvalidOperationException(
                    "Have not finished processing through the input for scanning. To finish"
                    + " processing, have an operation that causes the enumerable to resolve. This can"
                    + " can be as simple as a `.ToList()` operator or using another resolving operation"
                    + " such as one implementing `IResolvingOperation` before any calls to `GetEntities`.");
            }
        }
    }
}