Add contract decoration support

Assets/Reflex.Tests/DecorateTests.cs

@@ -0,0 +1,144 @@
+using System.Collections.Generic;
+using System.Linq;
+using FluentAssertions;
+using NUnit.Framework;
+using Reflex.Core;
+
+namespace Reflex.Tests
+{
+    public class DecorateTests
+    {
+        public interface INumber
+        {
+            int Get();
+        }
+
+        public class Number : INumber
+        {
+            private int _value;
+            public int Get() => _value;
+
+            public static Number FromValue(int value)
+            {
+                return new Number
+                {
+                    _value = value
+                };
+            }
+        }
+
+        public class DoubledNumber : INumber
+        {
+            private readonly INumber _number;
+            public DoubledNumber(INumber number) => _number = number;
+            public int Get() => _number.Get() * 2;
+        }
+
+        public class HalvedNumber : INumber
+        {
+            private readonly INumber _number;
+            public HalvedNumber(INumber number) => _number = number;
+            public int Get() => _number.Get() / 2;
+        }
+
+        public interface IManager
+        {
+        }
+
+        public interface IBundleManager
+        {
+        }
+
+        public class BundleManager : IBundleManager, IManager
+        {
+        }
+
+        public class ResilientBundleManager : IBundleManager, IManager
+        {
+            public ResilientBundleManager(IBundleManager bundleManager)
+            {
+            }
+        }
+
+        [Test]
+        public void Decorate_ShouldDecorateAllMatchingContracts()
+        {
+            var container = new ContainerDescriptor("")
+                .AddSingleton(Number.FromValue(1), contracts: typeof(INumber))
+                .AddSingleton(Number.FromValue(2), contracts: typeof(INumber))
+                .AddSingleton(Number.FromValue(3), contracts: typeof(INumber))
+                .AddDecorator(typeof(DoubledNumber), typeof(INumber))
+                .Build();
+
+            var numbers = container.All<INumber>().Select(n => n.Get()).ToArray();
+            numbers.Length.Should().Be(3);
+            numbers.Should().BeEquivalentTo(new int[] {2, 4, 6});
+        }
+
+        [Test]
+        public void Decorate_ShouldBeAbleToNestDecorations()
+        {
+            var container = new ContainerDescriptor("")
+                .AddSingleton(Number.FromValue(10), contracts: typeof(INumber))
+                .AddDecorator(typeof(DoubledNumber), typeof(INumber))
+                .AddDecorator(typeof(DoubledNumber), typeof(INumber))
+                .AddDecorator(typeof(DoubledNumber), typeof(INumber))
+                .AddDecorator(typeof(DoubledNumber), typeof(INumber))
+                .AddDecorator(typeof(HalvedNumber), typeof(INumber))
+                .AddDecorator(typeof(HalvedNumber), typeof(INumber))
+                .AddDecorator(typeof(DoubledNumber), typeof(INumber))
+                .Build();
+
+            var number = container.Single<INumber>();
+            number.Get().Should().Be(80);
+        }
+
+        [Test]
+        public void DecoratedContract_ShouldReplaceOnlyDecoratedContract()
+        {
+            var container = new ContainerDescriptor("")
+                .AddSingleton(typeof(BundleManager), typeof(IBundleManager), typeof(IManager))
+                .AddDecorator(typeof(ResilientBundleManager), typeof(IBundleManager))
+                .Build();
+
+            container.Resolve<IManager>().GetType().Should().Be<BundleManager>();
+            container.Resolve<IBundleManager>().GetType().Should().Be<ResilientBundleManager>();
+        }
+
+        [Test]
+        public void DecoratedSingleton_ShouldDecorateWithSingleton()
+        {
+            var container = new ContainerDescriptor("")
+                .AddSingleton(typeof(Number), contracts: typeof(INumber))
+                .AddDecorator(typeof(DoubledNumber), typeof(INumber))
+                .Build();
+
+            var distinctInstances = new HashSet<INumber>
+            {
+                container.Resolve<INumber>(),
+                container.Resolve<INumber>(),
+                container.Resolve<INumber>(),
+            };
+
+            distinctInstances.Count.Should().Be(1);
+        }
+
+        [Test]
+        public void DecoratedTransient_ShouldDecorateWithTransient()
+        {
+            var container = new ContainerDescriptor("")
+                .AddTransient(typeof(Number), contracts: typeof(INumber))
+                .AddDecorator(typeof(DoubledNumber), typeof(INumber))
+                .Build();
+
+            var distinctInstances = new HashSet<INumber>
+            {
+                container.Resolve<INumber>(),
+                container.Resolve<INumber>(),
+                container.Resolve<INumber>(),
+            };
+
+            distinctInstances.Count.Should().Be(3);
+        }
+    }
+}

Assets/Reflex/Core/ContainerDescriptor.cs

@@ -2,9 +2,11 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Linq;
+using Reflex.Enums;
 using Reflex.Exceptions;
 using Reflex.Extensions;
 using Reflex.Generics;
+using Reflex.Injectors;
 using Reflex.Resolvers;
 
 namespace Reflex.Core
@@ -15,6 +17,7 @@ public class ContainerDescriptor
         private Container _parent;
         private List<ResolverDescriptor> _descriptors = new();
         public event Action<Container> OnContainerBuilt;
+
         public ContainerDescriptor(string name, Container parent = null)
         {
             _name = name;
@@ -26,18 +29,18 @@ public Container Build()
             Build(out var disposables, out var resolversByContract, out var toStart);
             var container = new Container(_name, resolversByContract, disposables);
             container.SetParent(_parent);
-            
+
             // Clear references
             _name = null;
             _parent = null;
             _descriptors = null;
-            
+
             // Call initializers
             foreach (var startable in toStart.Select(r => (IStartable) r.Resolve(container)))
             {
                 startable.Start();
             }
-            
+
             OnContainerBuilt?.Invoke(container);
             return container;
         }
@@ -51,7 +54,7 @@ public ContainerDescriptor AddSingleton(Type concrete)
         {
             return AddSingleton(concrete, concrete);
         }
-        
+
         public ContainerDescriptor AddSingleton(object instance, params Type[] contracts)
         {
             return Add(instance.GetType(), contracts, new SingletonValueResolver(instance));
@@ -72,7 +75,7 @@ object Proxy(Container container)
                 return factory.Invoke(container);
             }
         }
-        
+
         public ContainerDescriptor AddSingleton<T>(Func<Container, T> factory)
         {
             return AddSingleton(factory, typeof(T));
@@ -87,7 +90,7 @@ public ContainerDescriptor AddTransient(Type concrete)
         {
             return AddTransient(concrete, concrete);
         }
-        
+
         public ContainerDescriptor AddTransient<T>(Func<Container, T> factory, params Type[] contracts)
         {
             var resolver = new TransientFactoryResolver(Proxy, typeof(T));
@@ -98,12 +101,46 @@ object Proxy(Container container)
                 return factory.Invoke(container);
             }
         }
-        
+
         public ContainerDescriptor AddTransient<T>(Func<Container, T> factory)
         {
             return AddTransient(factory, typeof(T));
         }
 
+        public ContainerDescriptor AddDecorator(Type concrete, Type contract)
+        {
+            var count = _descriptors.Count;
+
+            for (int i = 0; i < count; i++)
+            {
+                var descriptor = _descriptors[i];
+
+                if (!descriptor.Contracts.Contains(contract))
+                {
+                    continue;
+                }
+
+                var decoratedContract = new DecoratedType(contract);
+
+                object Factory(Container container)
+                {
+                    var dependency = container.Resolve(decoratedContract);
+                    return ConstructorInjector.Construct(concrete, new[] {dependency});
+                }
+
+                var resolver = descriptor.Resolver.Lifetime == Lifetime.Singleton
+                    ? (Resolver) new SingletonFactoryResolver(Factory, concrete)
+                    : (Resolver) new TransientFactoryResolver(Factory, concrete);
+
+                AddWithoutValidation(new[] {contract}, resolver);
+
+                var originalDescriptorNewContracts = descriptor.Contracts.Except(new[] {contract}).Concat(new[] {decoratedContract}).ToArray();
+                _descriptors[i] = new ResolverDescriptor(descriptor.Resolver, originalDescriptorNewContracts);
+            }
+
+            return this;
+        }
+
         private void Build(out DisposableCollection disposables, out Dictionary<Type, List<Resolver>> resolversByContract, out IEnumerable<Resolver> toStart)
         {
             disposables = new DisposableCollection();
@@ -146,6 +183,13 @@ private ContainerDescriptor Add(Type concrete, Type[] contracts, Resolver resolv
             return this;
         }
 
+        private ContainerDescriptor AddWithoutValidation(Type[] contracts, Resolver resolver)
+        {
+            var resolverDescriptor = new ResolverDescriptor(resolver, contracts);
+            _descriptors.Add(resolverDescriptor);
+            return this;
+        }
+
         [Conditional("UNITY_EDITOR")]
         private static void ValidateContracts(Type concrete, params Type[] contracts)
         {