Added RegistrationBuilder variations with callback support

VContainer/Assets/VContainer/Runtime/ContainerBuilderExtensions.cs

@@ -12,33 +12,72 @@ public static RegistrationBuilder Register(
             Type type,
             Lifetime lifetime)
             => builder.Register(new RegistrationBuilder(type, lifetime));
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static RegistrationBuilder Register(
+            this IContainerBuilder builder,
+            Type type,
+            Lifetime lifetime,
+            Action<object, IObjectResolver> callback)
+            => builder.Register(new RegistrationBuilderWithCallback(type, lifetime, callback));
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static RegistrationBuilder Register<T>(
             this IContainerBuilder builder,
             Lifetime lifetime)
             => builder.Register(typeof(T), lifetime);
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static RegistrationBuilder Register<T>(
+            this IContainerBuilder builder,
+            Lifetime lifetime,
+            Action<T, IObjectResolver> callback)
+            => builder.Register(typeof(T), lifetime, (instance, resolver) => callback((T)instance, resolver));
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static RegistrationBuilder Register<TInterface, TImplement>(
             this IContainerBuilder builder,
             Lifetime lifetime)
             where TImplement : TInterface
             => builder.Register<TImplement>(lifetime).As<TInterface>();
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static RegistrationBuilder Register<TInterface, TImplement>(
+            this IContainerBuilder builder,
+            Lifetime lifetime,
+            Action<TImplement, IObjectResolver> callback)
+            where TImplement : TInterface
+            => builder.Register(lifetime, callback).As<TInterface>();
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static RegistrationBuilder Register<TInterface1, TInterface2, TImplement>(
             this IContainerBuilder builder,
             Lifetime lifetime)
             where TImplement : TInterface1, TInterface2
             => builder.Register<TImplement>(lifetime).As(typeof(TInterface1), typeof(TInterface2));
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static RegistrationBuilder Register<TInterface1, TInterface2, TImplement>(
+            this IContainerBuilder builder,
+            Lifetime lifetime,
+            Action<TImplement, IObjectResolver> callback)
+            where TImplement : TInterface1, TInterface2
+            => builder.Register(lifetime, callback).As(typeof(TInterface1), typeof(TInterface2));
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static RegistrationBuilder Register<TInterface1, TInterface2, TInterface3, TImplement>(
             this IContainerBuilder builder,
             Lifetime lifetime)
             where TImplement : TInterface1, TInterface2, TInterface3
             => builder.Register<TImplement>(lifetime).As(typeof(TInterface1), typeof(TInterface2), typeof(TInterface3));
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static RegistrationBuilder Register<TInterface1, TInterface2, TInterface3, TImplement>(
+            this IContainerBuilder builder,
+            Lifetime lifetime,
+            Action<TImplement, IObjectResolver> callback)
+            where TImplement : TInterface1, TInterface2, TInterface3
+            => builder.Register(lifetime, callback).As(typeof(TInterface1), typeof(TInterface2), typeof(TInterface3));
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static RegistrationBuilder Register<TInterface>(
@@ -47,6 +86,15 @@ public static RegistrationBuilder Register<TInterface>(
             Lifetime lifetime)
             where TInterface : class
             => builder.Register(new FuncRegistrationBuilder(implementationConfiguration, typeof(TInterface), lifetime));
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public static RegistrationBuilder Register<TInterface>(
+            this IContainerBuilder builder,
+            Func<IObjectResolver, TInterface> implementationConfiguration,
+            Lifetime lifetime,
+            Action<TInterface, IObjectResolver> callback)
+            where TInterface : class
+            => builder.Register(new FuncRegistrationBuilderWithCallback(implementationConfiguration, typeof(TInterface), lifetime, (instance, resolver) => callback((TInterface)instance, resolver)));
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public static RegistrationBuilder RegisterInstance<TInterface>(

VContainer/Assets/VContainer/Runtime/Internal/FuncRegistrationBuilderWithCallback.cs

@@ -0,0 +1,26 @@
+using System;
+
+namespace VContainer.Internal
+{
+    sealed class FuncRegistrationBuilderWithCallback : RegistrationBuilder
+    {
+        readonly Func<IObjectResolver, object> implementationProvider;
+        readonly Action<object, IObjectResolver> callback;
+
+        public FuncRegistrationBuilderWithCallback(
+            Func<IObjectResolver, object> implementationProvider,
+            Type implementationType,
+            Lifetime lifetime, 
+            Action<object, IObjectResolver> callback) : base(implementationType, lifetime)
+        {
+            this.implementationProvider = implementationProvider;
+            this.callback = callback;
+        }
+
+        public override Registration Build()
+        {
+            var spawner = new InstanceProviderCallbackDecorator(new FuncInstanceProvider(implementationProvider), callback);
+            return new Registration(ImplementationType, Lifetime, InterfaceTypes, spawner);
+        }
+    }
+}

VContainer/Assets/VContainer/Runtime/Internal/InstanceProviders/InstanceProviderCallbackDecorator.cs

@@ -0,0 +1,23 @@
+using System;
+
+namespace VContainer.Internal
+{
+    sealed class InstanceProviderCallbackDecorator : IInstanceProvider
+    {
+        readonly Action<object, IObjectResolver> callback;
+        readonly IInstanceProvider instanceProvider;
+
+        public InstanceProviderCallbackDecorator(IInstanceProvider instanceProvider, Action<object, IObjectResolver> callback)
+        {
+            this.instanceProvider = instanceProvider;
+            this.callback = callback;
+        }
+
+        public object SpawnInstance(IObjectResolver resolver)
+        {
+            var instance = instanceProvider.SpawnInstance(resolver);
+            callback(instance, resolver);
+            return instance;
+        }
+    }
+}

VContainer/Assets/VContainer/Runtime/Internal/RegistrationBuilderWithCallback.cs

@@ -0,0 +1,26 @@
+using System;
+
+namespace VContainer.Internal
+{
+    sealed class RegistrationBuilderWithCallback : RegistrationBuilder
+    {
+        readonly Action<object, IObjectResolver> callback;
+
+        public RegistrationBuilderWithCallback(Type implementationType, Lifetime lifetime, Action<object, IObjectResolver> callback) 
+            : base(implementationType, lifetime)
+        {
+            this.callback = callback;
+        }
+
+        public override Registration Build()
+        {
+            var injector = InjectorCache.GetOrBuild(ImplementationType);
+            var spawner = new InstanceProviderCallbackDecorator(new InstanceProvider(injector, Parameters), callback);
+            return new Registration(
+                ImplementationType,
+                Lifetime,
+                InterfaceTypes,
+                spawner);
+        }
+    }
+}

VContainer/Assets/VContainer/Tests/ContainerTest.cs

@@ -518,5 +518,35 @@ public void Inject()
             var ctorInjectable = new ServiceA(new NoDependencyServiceA());
             Assert.DoesNotThrow(() => container.Inject(ctorInjectable));
         }
+
+        [Test]
+        public void RegisterWithCallback()
+        {
+            {
+                NoDependencyServiceA fromCallback = null;
+
+                var builder = new ContainerBuilder();
+                builder.Register<NoDependencyServiceA>(Lifetime.Scoped, (service, _) => fromCallback = service);
+
+                var container = builder.Build();
+
+                var fromResolve = container.Resolve<NoDependencyServiceA>();
+
+                Assert.That(fromResolve, Is.Not.Null.And.EqualTo(fromCallback));
+            }
+
+            {
+                NoDependencyServiceA fromCallback = null;
+
+                var builder = new ContainerBuilder();
+                builder.Register(_ => new NoDependencyServiceA(), Lifetime.Scoped, (service, _) => fromCallback = service);
+
+                var container = builder.Build();
+
+                var fromResolve = container.Resolve<NoDependencyServiceA>();
+
+                Assert.That(fromResolve, Is.Not.Null.And.EqualTo(fromCallback));
+            }
+        }
     }
 }