Add RegisterDecorator

VContainer/Assets/Tests/DecoratorTest.cs

@@ -0,0 +1,49 @@
+using NUnit.Framework;
+using VContainer.Runtime;
+
+namespace VContainer.Tests
+{
+    interface IDecoratedType
+    {
+    }
+
+    class DecoratedType : IDecoratedType
+    {
+    }
+
+    class Decorator1 : IDecoratedType
+    {
+        readonly IDecoratedType inner;
+
+        public Decorator1(IDecoratedType inner)
+        {
+            this.inner = inner;
+        }
+    }
+
+    class Decorator2 : IDecoratedType
+    {
+        readonly IDecoratedType inner;
+
+        public Decorator2(IDecoratedType inner)
+        {
+            this.inner = inner;
+        }
+    }
+
+    [TestFixture]
+    public class DecoratorTest
+    {
+        [Test]
+        public void Decorate()
+        {
+            var builder = new ContainerBuilder();
+            builder.Register<IDecoratedType, DecoratedType>(Lifetime.Singleton);
+            builder.RegisterDecorator<IDecoratedType, Decorator1>();
+
+            var container = builder.Build();
+            var instance = container.Resolve<IDecoratedType>();
+            Assert.That(instance, Is.TypeOf<Decorator1>());
+        }
+    }
+}

VContainer/Assets/VContainer/Runtime/DecoratorRegistrationBuilder.cs

@@ -0,0 +1,91 @@
+using System;
+using VContainer.Internal;
+
+namespace VContainer.Runtime
+{
+    public static class ConatinerBuilderDecoratorExtensions
+    {
+        public static DecoratorRegistrationBuilder RegisterDecorator<TInterface, TDecorator>(this IContainerBuilder builder)
+        {
+            for (var i = 0; i < builder.Count; i++)
+            {
+                var interfaceTypes = builder[i].InterfaceTypes;
+                if (interfaceTypes != null && interfaceTypes.Contains(typeof(TInterface)))
+                {
+                    return new DecoratorRegistrationBuilder(builder[i], typeof(TDecorator));
+                }
+            }
+            throw new VContainerException(typeof(TInterface), $"No such decorator target: {typeof(TInterface)}");
+        }
+
+        public static DecoratorRegistrationBuilder RegisterDecorator<TInterface, TDecorator>(
+            this IContainerBuilder builder,
+            Func<TInterface, IObjectResolver, TDecorator> factory)
+        {
+            for (var i = 0; i < builder.Count; i++)
+            {
+                var interfaceTypes = builder[i].InterfaceTypes;
+                if (interfaceTypes != null && interfaceTypes.Contains(typeof(TInterface)))
+                {
+                    return new FuncDecoratorRegistrationBuilder<TInterface, TDecorator>(builder[i], factory);
+                }
+            }
+            throw new VContainerException(typeof(TInterface), $"No such decorator target: {typeof(TInterface)}");
+        }
+    }
+
+    public class DecoratorRegistrationBuilder : RegistrationBuilder
+    {
+        readonly RegistrationBuilder inner;
+        readonly Type interfaceType;
+
+        public DecoratorRegistrationBuilder(RegistrationBuilder inner, Type decoratorType)
+            : base(decoratorType, inner.Lifetime)
+        {
+            this.inner = inner;
+            interfaceType = inner.InterfaceTypes != null ? inner.InterfaceTypes[0] : inner.ImplementationType;
+            InterfaceTypes = inner.InterfaceTypes;
+            As(interfaceType);
+        }
+
+        public override Registration Build()
+        {
+            var injector = InjectorCache.GetOrBuild(ImplementationType);
+            var innerRegistration = inner.Build();
+
+            var provider = new FuncInstanceProvider(container =>
+            {
+                var innerInstance = container.Resolve(innerRegistration);
+                var parameters = new IInjectParameter[Parameters == null ? 1 : Parameters.Count];
+                Parameters?.CopyTo(parameters);
+                parameters[parameters.Length - 1] = new TypedParameter(interfaceType, innerInstance);
+                return injector.CreateInstance(container, parameters);
+            });
+            return new Registration(ImplementationType, Lifetime, InterfaceTypes, provider);
+        }
+    }
+
+    public class FuncDecoratorRegistrationBuilder<TInner, TDecorator> : DecoratorRegistrationBuilder
+    {
+        readonly RegistrationBuilder inner;
+        readonly Type interfaceType;
+        readonly Func<TInner, IObjectResolver, TDecorator> factory;
+
+        public FuncDecoratorRegistrationBuilder(RegistrationBuilder inner, Func<TInner, IObjectResolver, TDecorator> factory)
+            : base(inner, typeof(TDecorator))
+        {
+            this.factory = factory;
+        }
+
+        public override Registration Build()
+        {
+            var innerRegistration = inner.Build();
+            var provider = new FuncInstanceProvider(container =>
+            {
+                var innerInstance = container.Resolve(innerRegistration);
+                return factory((TInner)innerInstance, container);
+            });
+            return new Registration(ImplementationType, Lifetime, InterfaceTypes, provider);
+        }
+    }
+}

VContainer/Assets/VContainer/Runtime/RegistrationBuilder.cs

@@ -10,7 +10,7 @@ public class RegistrationBuilder
         protected internal readonly Lifetime Lifetime;
 
         protected internal List<Type> InterfaceTypes;
-        protected internal List<IInjectParameter> Parameters;
+        protected List<IInjectParameter> Parameters;
 
         public RegistrationBuilder(Type implementationType, Lifetime lifetime)
         {