Make EntitySystemManager.DependencyCollection inject EntityQuery, make BUIs inject systems and entity queries (#6394)

* Make EntitySystemManager.DependencyCollection inject EntityQuery

* Make BUIs inject systems and entity queries

* Fix import

* We parallelize those

* RIDER I BEG YOU

* Mocked unit tests are my passion

* Perhaps we do not care about fractional milliseconds

* Forgor to make it debug only

* Use Parallel.For instead of ForEach

* Rider I am going to become the joker

* Fix EntMan resolve

* Now with lazy resolve technology

* Use GetOrAdd

Robust.Shared/GameObjects/Components/UserInterface/BoundUserInterface.cs

@@ -32,7 +32,8 @@ namespace Robust.Shared.GameObjects
 
         protected BoundUserInterface(EntityUid owner, Enum uiKey)
         {
-            IoCManager.InjectDependencies(this);
+            IoCManager.Resolve(ref EntMan);
+            EntMan.EntitySysManager.DependencyCollection.InjectDependencies(this);
             UiSystem = EntMan.System<SharedUserInterfaceSystem>();
 
             Owner = owner;

Robust.Shared/GameObjects/EntitySystemManager.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
 using Prometheus;
 using Robust.Shared.IoC;
 using Robust.Shared.IoC.Exceptions;
@@ -192,6 +193,14 @@ namespace Robust.Shared.GameObjects
                 _systemTypes.Remove(baseType);
             }
 
+            var queryMethod = typeof(EntityManager).GetMethod(nameof(EntityManager.GetEntityQuery), 1, [])!;
+            SystemDependencyCollection.RegisterBaseGenericLazy(
+                typeof(EntityQuery<>),
+                (queryType, dep) => queryMethod
+                    .MakeGenericMethod(queryType.GetGenericArguments()[0])
+                    .Invoke(dep.Resolve<IEntityManager>(), null)!
+            );
+
             SystemDependencyCollection.BuildGraph();
 
             foreach (var systemType in _systemTypes)

Robust.Shared/IoC/DependencyCollection.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Frozen;
 using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
@@ -17,6 +18,11 @@ namespace Robust.Shared.IoC
     public delegate T DependencyFactoryDelegate<out T>()
         where T : class;
 
+    public delegate T DependencyFactoryBaseGenericLazyDelegate<out T>(
+        Type type,
+        IDependencyCollection services)
+        where T : class;
+
     /// <inheritdoc />
     internal sealed class DependencyCollection : IDependencyCollection
     {
@@ -37,6 +43,12 @@ namespace Robust.Shared.IoC
         /// </remarks>
         private FrozenDictionary<Type, object> _services = FrozenDictionary<Type, object>.Empty;
 
+        /// <summary>
+        /// Dictionary that maps the types passed to <see cref="Resolve{T}()"/> to their implementation
+        /// for any types registered through <see cref="RegisterBaseGenericLazy"/>.
+        /// </summary>
+        private readonly ConcurrentDictionary<Type, object> _lazyServices = new();
+
         // Start fields used for building new services.
 
         /// <summary>
@@ -48,6 +60,8 @@ namespace Robust.Shared.IoC
         private readonly Dictionary<Type, DependencyFactoryDelegateInternal<object>> _resolveFactories = new();
         private readonly Queue<Type> _pendingResolves = new();
 
+        private readonly ConcurrentDictionary<Type, DependencyFactoryBaseGenericLazyDelegate<object>> _baseGenericLazyFactories = new();
+
         private readonly object _serviceBuildLock = new();
 
         // End fields for building new services.
@@ -79,8 +93,8 @@ namespace Robust.Shared.IoC
         public IEnumerable<Type> GetRegisteredTypes()
         {
             return _parentCollection != null
-                ? _services.Keys.Concat(_parentCollection.GetRegisteredTypes())
-                : _services.Keys;
+                ? _services.Keys.Concat(_lazyServices.Keys).Concat(_parentCollection.GetRegisteredTypes())
+                : _services.Keys.Concat(_lazyServices.Keys);
         }
 
         public Type[] GetCachedInjectorTypes()
@@ -116,10 +130,7 @@ namespace Robust.Shared.IoC
             FrozenDictionary<Type, object> services,
             [MaybeNullWhen(false)] out object instance)
         {
-            if (!services.TryGetValue(objectType, out instance))
-                return _parentCollection is not null && _parentCollection.TryResolveType(objectType, out instance);
-
-            return true;
+            return TryResolveType(objectType, (IReadOnlyDictionary<Type, object>) services, out instance);
         }
 
         private bool TryResolveType(
@@ -128,7 +139,16 @@ namespace Robust.Shared.IoC
             [MaybeNullWhen(false)] out object instance)
         {
             if (!services.TryGetValue(objectType, out instance))
+            {
+                if (objectType.IsGenericType &&
+                    _baseGenericLazyFactories.TryGetValue(objectType.GetGenericTypeDefinition(), out var factory))
+                {
+                    instance = _lazyServices.GetOrAdd(objectType, type => factory(type, this));
+                    return true;
+                }
+
                 return _parentCollection is not null && _parentCollection.TryResolveType(objectType, out instance);
+            }
 
             return true;
         }
@@ -267,7 +287,7 @@ namespace Robust.Shared.IoC
                 _pendingResolves.Enqueue(interfaceType);
             }
         }
-        
+
         private void CheckRegisterInterface(Type interfaceType, Type implementationType, bool overwrite)
         {
             lock (_serviceBuildLock)
@@ -312,15 +332,24 @@ namespace Robust.Shared.IoC
             Register(type, implementation.GetType(), () => implementation, overwrite);
         }
 
+        public void RegisterBaseGenericLazy(Type interfaceType, DependencyFactoryBaseGenericLazyDelegate<object> factory)
+        {
+            lock (_serviceBuildLock)
+            {
+                _baseGenericLazyFactories[interfaceType] = factory;
+            }
+        }
+
         /// <inheritdoc />
         public void Clear()
         {
-            foreach (var service in _services.Values.OfType<IDisposable>().Distinct())
+            foreach (var service in _services.Values.Concat(_lazyServices.Values).OfType<IDisposable>().Distinct())
             {
                 service.Dispose();
             }
 
             _services = FrozenDictionary<Type, object>.Empty;
+            _lazyServices.Clear();
 
             lock (_serviceBuildLock)
             {

Robust.Shared/IoC/IDependencyCollection.cs

@@ -132,6 +132,15 @@ namespace Robust.Shared.IoC
         /// </param>
         void RegisterInstance(Type type, object implementation, bool overwrite = false);
 
+        /// <summary>
+        ///     Adds a callback to be called when attempting to resolve an unresolved type that matches the specified
+        ///     base generic type, making it accessible to <see cref="IDependencyCollection.Resolve{T}"/>.
+        ///     This instance will only be created the first time that it is attempted to be resolved.
+        /// </summary>
+        /// <param name="genericType">The base generic type of the type that will be resolvable.</param>
+        /// <param name="factory">The callback to call to get an instance of the implementation for that generic type.</param>
+        void RegisterBaseGenericLazy(Type genericType, DependencyFactoryBaseGenericLazyDelegate<object> factory);
+
         /// <summary>
         /// Clear all services and types.
         /// Use this between unit tests and on program shutdown.