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.IntegrationTests/GameObjects/EntitySystemManagerOrderTest.cs

@@ -1,6 +1,3 @@
-using System;
-using System.Collections.Generic;
-using System.Linq;
 using Moq;
 using NUnit.Framework;
 using Robust.Server.Configuration;
@@ -88,7 +85,6 @@ namespace Robust.UnitTesting.Shared.GameObjects
             deps.Register<IDynamicTypeFactoryInternal, DynamicTypeFactory>();
             deps.RegisterInstance<IModLoader>(new Mock<IModLoader>().Object);
             deps.Register<IEntitySystemManager, EntitySystemManager>();
-            deps.RegisterInstance<IEntityManager>(new Mock<IEntityManager>().Object);
             // WHEN WILL THE SUFFERING END
             deps.RegisterInstance<IReplayRecordingManager>(new Mock<IReplayRecordingManager>().Object);
 
@@ -104,6 +100,15 @@ namespace Robust.UnitTesting.Shared.GameObjects
 
             deps.RegisterInstance<IReflectionManager>(reflectionMock.Object);
 
+            // Never
+            var componentFactoryMock = new Mock<IComponentFactory>();
+            componentFactoryMock.Setup(p => p.AllRegisteredTypes).Returns(Enumerable.Empty<Type>());
+            deps.RegisterInstance<IComponentFactory>(componentFactoryMock.Object);
+
+            var entityManagerMock = new Mock<IEntityManager>();
+            entityManagerMock.Setup(p => p.ComponentFactory).Returns(componentFactoryMock.Object);
+            deps.RegisterInstance<IEntityManager>(entityManagerMock.Object);
+
             deps.BuildGraph();
 
             IoCManager.InitThread(deps, true);

Robust.Shared.IntegrationTests/GameObjects/EntitySystemManager_Tests.cs

@@ -1,11 +1,10 @@
-using System;
-using System.Collections.Generic;
-using System.Linq;
 using NUnit.Framework;
 using Robust.Shared.Analyzers;
+using Robust.Shared.Configuration;
 using Robust.Shared.GameObjects;
 using Robust.Shared.IoC;
 using Robust.Shared.IoC.Exceptions;
+using Robust.Shared.Physics.Components;
 
 namespace Robust.UnitTesting.Shared.GameObjects
 {
@@ -39,6 +38,14 @@ namespace Robust.UnitTesting.Shared.GameObjects
             [Dependency] public readonly ESystemDepA ESystemDepA = default!;
         }
 
+        internal sealed class ESystemDepAll : EntitySystem
+        {
+            [Dependency] public readonly ESystemDepA ESystemDepA = default!;
+            [Dependency] public readonly IConfigurationManager Config = default!;
+            [Dependency] public readonly EntityQuery<TransformComponent> TransformQuery = default!;
+            [Dependency] public readonly EntityQuery<PhysicsComponent> PhysicsQuery = default!;
+        }
+
         /*
          ESystemBase (Abstract)
            - ESystemA
@@ -58,6 +65,7 @@ namespace Robust.UnitTesting.Shared.GameObjects
             syssy.LoadExtraSystemType<ESystemC>();
             syssy.LoadExtraSystemType<ESystemDepA>();
             syssy.LoadExtraSystemType<ESystemDepB>();
+            syssy.LoadExtraSystemType<ESystemDepAll>();
             syssy.Initialize(false);
         }
 
@@ -103,5 +111,16 @@ namespace Robust.UnitTesting.Shared.GameObjects
             Assert.That(sysB.ESystemDepA, Is.EqualTo(sysA));
         }
 
+        [Test]
+        public void DependencyInjectionTest()
+        {
+            var esm = IoCManager.Resolve<IEntitySystemManager>();
+            var sys = esm.GetEntitySystem<ESystemDepAll>();
+
+            Assert.That(sys.ESystemDepA, Is.Not.Null);
+            Assert.That(sys.Config, Is.Not.Null);
+            Assert.That(sys.TransformQuery, Is.Not.Default);
+            Assert.That(sys.PhysicsQuery, Is.Not.Default);
+        }
     }
 }

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.Concat(_lazyServices.Keys).Concat(_parentCollection.GetRegisteredTypes())
-                : _services.Keys;
+                : _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 TryResolveType(objectType, (IReadOnlyDictionary<Type, object>) services, out instance);
-                return _parentCollection is not null && _parentCollection.TryResolveType(objectType, out instance);
-
-            return true;
         }
 
         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;
         }
@@ -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.