Simplify lazy initialization in RazorDocumentServiceProviderWrapper

src/Tools/ExternalAccess/Razor/RazorDocumentServiceProviderWrapper.cs

@@ -1,27 +1,26 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
-#nullable enable annotations
 
 using System;
+using System.Runtime.CompilerServices;
 using Microsoft.CodeAnalysis.Host;
+using Roslyn.Utilities;
 
 namespace Microsoft.CodeAnalysis.ExternalAccess.Razor
 {
     internal sealed class RazorDocumentServiceProviderWrapper : IDocumentServiceProvider, IDocumentOperationService
     {
         private readonly IRazorDocumentServiceProvider _innerDocumentServiceProvider;
-        private readonly object _lock;
 
-        private RazorSpanMappingServiceWrapper? _spanMappingService;
-        private RazorDocumentExcerptServiceWrapper? _excerptService;
-        private RazorDocumentPropertiesServiceWrapper? _documentPropertiesService;
+        // The lazily initialized service fields use StrongBox<T> to explicitly allow null as an initialized value.
+        private StrongBox<ISpanMappingService?>? _lazySpanMappingService;
+        private StrongBox<IDocumentExcerptService?>? _lazyExcerptService;
+        private StrongBox<DocumentPropertiesService?>? _lazyDocumentPropertiesService;
 
         public RazorDocumentServiceProviderWrapper(IRazorDocumentServiceProvider innerDocumentServiceProvider)
         {
             _innerDocumentServiceProvider = innerDocumentServiceProvider ?? throw new ArgumentNullException(nameof(innerDocumentServiceProvider));
-
-            _lock = new object();
         }
 
         public bool CanApplyChange => _innerDocumentServiceProvider.CanApplyChange;
@@ -33,75 +32,44 @@ public RazorDocumentServiceProviderWrapper(IRazorDocumentServiceProvider innerDo
             var serviceType = typeof(TService);
             if (serviceType == typeof(ISpanMappingService))
             {
-                if (_spanMappingService == null)
-                {
-                    lock (_lock)
+                var spanMappingService = LazyInitialization.EnsureInitialized(
+                    ref _lazySpanMappingService,
+                    static documentServiceProvider =>
                     {
-                        if (_spanMappingService == null)
-                        {
-                            var razorMappingService = _innerDocumentServiceProvider.GetService<IRazorSpanMappingService>();
-                            if (razorMappingService != null)
-                            {
-                                _spanMappingService = new RazorSpanMappingServiceWrapper(razorMappingService);
-                            }
-                            else
-                            {
-                                return this as TService;
-                            }
-                        }
-                    }
-                }
+                        var razorMappingService = documentServiceProvider.GetService<IRazorSpanMappingService>();
+                        return razorMappingService != null ? new RazorSpanMappingServiceWrapper(razorMappingService) : null;
+                    },
+                    _innerDocumentServiceProvider);
 
-                return (TService)(object)_spanMappingService;
+                return (TService?)spanMappingService;
             }
 
             if (serviceType == typeof(IDocumentExcerptService))
             {
-                if (_excerptService == null)
-                {
-                    lock (_lock)
+                var excerptService = LazyInitialization.EnsureInitialized(
+                    ref _lazyExcerptService,
+                    static documentServiceProvider =>
                     {
-                        if (_excerptService == null)
-                        {
-                            var excerptService = _innerDocumentServiceProvider.GetService<IRazorDocumentExcerptService>();
-                            if (excerptService != null)
-                            {
-                                _excerptService = new RazorDocumentExcerptServiceWrapper(excerptService);
-                            }
-                            else
-                            {
-                                return this as TService;
-                            }
-                        }
-                    }
-                }
+                        var razorExcerptService = documentServiceProvider.GetService<IRazorDocumentExcerptService>();
+                        return razorExcerptService is not null ? new RazorDocumentExcerptServiceWrapper(razorExcerptService) : null;
+                    },
+                    _innerDocumentServiceProvider);
 
-                return (TService)(object)_excerptService;
+                return (TService?)excerptService;
             }
 
             if (serviceType == typeof(DocumentPropertiesService))
             {
-                if (_documentPropertiesService == null)
-                {
-                    lock (_lock)
+                var documentPropertiesService = LazyInitialization.EnsureInitialized(
+                    ref _lazyDocumentPropertiesService,
+                    static documentServiceProvider =>
                     {
-                        if (_documentPropertiesService == null)
-                        {
-                            var documentPropertiesService = _innerDocumentServiceProvider.GetService<IRazorDocumentPropertiesService>();
-
-                            if (documentPropertiesService != null)
-                            {
-                                _documentPropertiesService = new RazorDocumentPropertiesServiceWrapper(documentPropertiesService);
-                            }
-                            else
-                            {
-                                return this as TService;
-                            }
-                        }
-                    }
-                }
+                        var razorDocumentPropertiesService = documentServiceProvider.GetService<IRazorDocumentPropertiesService>();
+                        return razorDocumentPropertiesService is not null ? new RazorDocumentPropertiesServiceWrapper(razorDocumentPropertiesService) : null;
+                    },
+                    _innerDocumentServiceProvider);
 
-                return (TService)(object)_documentPropertiesService;
+                return (TService?)(object?)documentPropertiesService;
             }
 
             return this as TService;

src/Workspaces/SharedUtilitiesAndExtensions/Compiler/Core/Utilities/LazyInitialization.cs

@@ -4,6 +4,7 @@
 
 using System;
 using System.Diagnostics.CodeAnalysis;
+using System.Runtime.CompilerServices;
 using System.Threading;
 
 namespace Roslyn.Utilities
@@ -22,7 +23,7 @@ internal static T InterlockedStore<T>([NotNull] ref T? target, T value) where T
         /// more than once by multiple threads, but only one of those values will successfully be written to the target.</param>
         /// <returns>The target value.</returns>
         public static T EnsureInitialized<T>([NotNull] ref T? target, Func<T> valueFactory) where T : class
-            => Volatile.Read(ref target!) ?? InterlockedStore(ref target!, valueFactory());
+            => Volatile.Read(ref target!) ?? InterlockedStore(ref target, valueFactory());
 
         /// <summary>
         /// Ensure that the given target value is initialized (not null) in a thread-safe manner.
@@ -37,7 +38,41 @@ public static T EnsureInitialized<T>([NotNull] ref T? target, Func<T> valueFacto
         public static T EnsureInitialized<T, U>([NotNull] ref T? target, Func<U, T> valueFactory, U state)
             where T : class
         {
-            return Volatile.Read(ref target!) ?? InterlockedStore(ref target!, valueFactory(state));
+            return Volatile.Read(ref target!) ?? InterlockedStore(ref target, valueFactory(state));
+        }
+
+        /// <summary>
+        /// Ensure that the given target value is initialized in a thread-safe manner. This overload supports the
+        /// initialization of value types, and reference type fields where <see langword="null"/> is considered an
+        /// initialized value.
+        /// </summary>
+        /// <typeparam name="T">The type of the target value.</typeparam>
+        /// <param name="target">A target value box to initialize.</param>
+        /// <param name="valueFactory">A factory delegate to create a new instance of the target value. Note that this delegate may be called
+        /// more than once by multiple threads, but only one of those values will successfully be written to the target.</param>
+        /// <returns>The target value.</returns>
+        public static T? EnsureInitialized<T>([NotNull] ref StrongBox<T?>? target, Func<T?> valueFactory)
+        {
+            var box = Volatile.Read(ref target!) ?? InterlockedStore(ref target, new StrongBox<T?>(valueFactory()));
+            return box.Value;
+        }
+
+        /// <summary>
+        /// Ensure that the given target value is initialized in a thread-safe manner. This overload supports the
+        /// initialization of value types, and reference type fields where <see langword="null"/> is considered an
+        /// initialized value.
+        /// </summary>
+        /// <typeparam name="T">The type of the target value.</typeparam>
+        /// <param name="target">A target value box to initialize.</param>
+        /// <typeparam name="U">The type of the <paramref name="state"/> argument passed to the value factory.</typeparam>
+        /// <param name="valueFactory">A factory delegate to create a new instance of the target value. Note that this delegate may be called
+        /// more than once by multiple threads, but only one of those values will successfully be written to the target.</param>
+        /// <param name="state">An argument passed to the value factory.</param>
+        /// <returns>The target value.</returns>
+        public static T? EnsureInitialized<T, U>([NotNull] ref StrongBox<T?>? target, Func<U, T?> valueFactory, U state)
+        {
+            var box = Volatile.Read(ref target!) ?? InterlockedStore(ref target, new StrongBox<T?>(valueFactory(state)));
+            return box.Value;
         }
     }
 }