1 files changed, 51 insertions, 77 deletions

diff --git a/src/WixToolset.Core/WixToolsetServiceProvider.cs b/src/WixToolset.Core/WixToolsetServiceProvider.cs
index a018b8dc..0c40c56d 100644
--- a/src/WixToolset.Core/WixToolsetServiceProvider.cs
+++ b/src/WixToolset.Core/WixToolsetServiceProvider.cs
@@ -3,6 +3,7 @@
 namespace WixToolset.Core
 {
     using System;
+    using System.Collections.Generic;
     using WixToolset.Core.CommandLine;
     using WixToolset.Core.ExtensibilityServices;
     using WixToolset.Data;
@@ -11,95 +12,68 @@ namespace WixToolset.Core
 
     public class WixToolsetServiceProvider : IServiceProvider
     {
-        private ExtensionManager extensionManager;
-        private Messaging messaging;
-        private ParseHelper parseHelper;
-        private PreprocessHelper preprocessHelper;
-        private TupleDefinitionCreator tupleDefinitionCreator;
-        private WindowsInstallerBackendHelper windowsInstallerBackendHelper;
-
-        public object GetService(Type serviceType)
+        public WixToolsetServiceProvider()
         {
-            if (serviceType == null) throw new ArgumentNullException(nameof(serviceType));
-
-            // Transients.
-            if (serviceType == typeof(IPreprocessContext))
-            {
-                return new PreprocessContext(this);
-            }
-
-            if (serviceType == typeof(ICompileContext))
-            {
-                return new CompileContext(this);
-            }
-
-            if (serviceType == typeof(ILinkContext))
-            {
-                return new LinkContext(this);
-            }
-
-            if (serviceType == typeof(IBindContext))
-            {
-                return new BindContext(this);
-            }
-
-            if (serviceType == typeof(ILayoutContext))
-            {
-                return new LayoutContext(this);
-            }
-
-            if (serviceType == typeof(IResolveContext))
-            {
-                return new ResolveContext(this);
-            }
-
-            if (serviceType == typeof(IInscribeContext))
-            {
-                return new InscribeContext(this);
-            }
-
-            if (serviceType == typeof(ICommandLineContext))
+            this.CreationFunctions = new Dictionary<Type, Func<IServiceProvider, Dictionary<Type, object>, object>>
             {
-                return new CommandLineContext(this);
-            }
+            // Singletons.
+                { typeof(IExtensionManager), (provider, singletons) => AddSingleton(singletons, typeof(IExtensionManager), new ExtensionManager()) },
+                { typeof(IMessaging), (provider, singletons) => AddSingleton(singletons, typeof(IMessaging), new Messaging()) },
+                { typeof(ITupleDefinitionCreator), (provider, singletons) => AddSingleton(singletons, typeof(ITupleDefinitionCreator), new TupleDefinitionCreator(provider)) },
+                { typeof(IParseHelper), (provider, singletons) => AddSingleton(singletons, typeof(IParseHelper), new ParseHelper(provider)) },
+                { typeof(IPreprocessHelper), (provider, singletons) => AddSingleton(singletons, typeof(IPreprocessHelper), new PreprocessHelper(provider)) },
+                { typeof(IWindowsInstallerBackendHelper), (provider, singletons) => AddSingleton(singletons, typeof(IWindowsInstallerBackendHelper), new WindowsInstallerBackendHelper(provider)) },
 
-            if (serviceType == typeof(ICommandLine))
-            {
-                return new CommandLineParser();
-            }
+            // Transients.
+                { typeof(ICommandLineContext), (provider, singletons) => new CommandLineContext(provider) },
+                { typeof(ICommandLine), (provider, singletons) => new CommandLineParser() },
+                { typeof(IPreprocessContext), (provider, singletons) => new PreprocessContext(provider) },
+                { typeof(ICompileContext), (provider, singletons) => new CompileContext(provider) },
+                { typeof(ILinkContext), (provider, singletons) => new LinkContext(provider) },
+                { typeof(IResolveContext), (provider, singletons) => new ResolveContext(provider) },
+                { typeof(IBindContext), (provider, singletons) => new BindContext(provider) },
+                { typeof(ILayoutContext), (provider, singletons) => new LayoutContext(provider) },
+                { typeof(IInscribeContext), (provider, singletons) => new InscribeContext(provider) },
+            };
+
+            this.Singletons = new Dictionary<Type, object>();
+        }
 
-            // Singletons.
-            if (serviceType == typeof(IExtensionManager))
-            {
-                return this.extensionManager = this.extensionManager ?? new ExtensionManager();
-            }
+        private Dictionary<Type, Func<IServiceProvider, Dictionary<Type, object>, object>> CreationFunctions { get; }
 
-            if (serviceType == typeof(IMessaging))
-            {
-                return this.messaging = this.messaging ?? new Messaging();
-            }
+        private Dictionary<Type, object> Singletons { get; }
 
-            if (serviceType == typeof(ITupleDefinitionCreator))
-            {
-                return this.tupleDefinitionCreator = this.tupleDefinitionCreator ?? new TupleDefinitionCreator(this);
-            }
+        public object GetService(Type serviceType)
+        {
+            if (serviceType == null) throw new ArgumentNullException(nameof(serviceType));
 
-            if (serviceType == typeof(IParseHelper))
+            if (!this.Singletons.TryGetValue(serviceType, out var service))
             {
-                return this.parseHelper = this.parseHelper ?? new ParseHelper(this);
+                if (this.CreationFunctions.TryGetValue(serviceType, out var creationFunction))
+                {
+                    service = creationFunction(this, this.Singletons);
+
+#if DEBUG
+                    if (!serviceType.IsAssignableFrom(service?.GetType()))
+                    {
+                        throw new InvalidOperationException($"Creation function for service type: {serviceType.Name} created incompatible service with type: {service?.GetType()}");
+                    }
+#endif
+                }
             }
 
-            if (serviceType == typeof(IPreprocessHelper))
-            {
-                return this.preprocessHelper = this.preprocessHelper ?? new PreprocessHelper(this);
-            }
+            return service ?? throw new ArgumentException($"Unknown service type: {serviceType.Name}", nameof(serviceType));
+        }
 
-            if (serviceType == typeof(IWindowsInstallerBackendHelper))
-            {
-                return this.windowsInstallerBackendHelper = this.windowsInstallerBackendHelper ?? new WindowsInstallerBackendHelper(this);
-            }
+        public void AddService(Type serviceType, Func<IServiceProvider, Dictionary<Type, object>, object> creationFunction)
+        {
+            this.CreationFunctions[serviceType] = creationFunction;
+        }
 
-            throw new ArgumentException($"Unknown service type: {serviceType.Name}", nameof(serviceType));
+        private static object AddSingleton(Dictionary<Type, object> singletons, Type type, object service)
+        {
+            singletons.Add(type, service);
+            return service;
         }
     }
 }