File: Snippets\AbstractSnippetInfoService.cs
Web Access
Project: src\src\VisualStudio\Core\Def\Microsoft.VisualStudio.LanguageServices_pxr0p0dn_wpftmp.csproj (Microsoft.VisualStudio.LanguageServices)
// 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.
 
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;
using System.Runtime.InteropServices;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
using Microsoft.CodeAnalysis.Shared.TestHooks;
using Microsoft.CodeAnalysis.Snippets;
using Microsoft.VisualStudio.Editor;
using Microsoft.VisualStudio.TextManager.Interop;
using Microsoft.VisualStudio.Threading;
using Roslyn.Utilities;
 
namespace Microsoft.VisualStudio.LanguageServices.Implementation.Snippets;
 
/// <summary>
/// This service is created on the UI thread during package initialization, but it must not
/// block the initialization process.
/// </summary>
internal abstract class AbstractSnippetInfoService : ISnippetInfoService, IVsExpansionEvents
{
    private readonly Guid _languageGuidForSnippets;
    private IVsExpansionManager? _expansionManager;
 
    /// <summary>
    /// Initialize these to empty values. When returning from <see cref="GetSnippetsIfAvailable "/> 
    /// and <see cref="SnippetShortcutExists_NonBlocking"/>, we return the current set of known 
    /// snippets rather than waiting for initial results.
    /// </summary>
    protected ImmutableArray<SnippetInfo> snippets = [];
    protected IImmutableSet<string> snippetShortcuts = ImmutableHashSet.Create<string>();
 
    // Guard the snippets and snippetShortcut fields so that returned result sets are always
    // complete.
    protected object cacheGuard = new();
 
    private readonly IAsynchronousOperationListener _waiter;
    private readonly IThreadingContext _threadingContext;
 
    public AbstractSnippetInfoService(
        IThreadingContext threadingContext,
        Shell.IAsyncServiceProvider serviceProvider,
        Guid languageGuidForSnippets,
        IAsynchronousOperationListenerProvider listenerProvider)
    {
        _waiter = listenerProvider.GetListener(FeatureAttribute.Snippets);
        _languageGuidForSnippets = languageGuidForSnippets;
        _threadingContext = threadingContext;
 
        _threadingContext.RunWithShutdownBlockAsync((_) => InitializeAndPopulateSnippetsCacheAsync(serviceProvider));
    }
 
    private async Task InitializeAndPopulateSnippetsCacheAsync(Shell.IAsyncServiceProvider asyncServiceProvider)
    {
        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync();
        var textManager = (IVsTextManager2?)await asyncServiceProvider.GetServiceAsync(typeof(SVsTextManager)).ConfigureAwait(true);
        Assumes.Present(textManager);
 
        if (textManager.GetExpansionManager(out _expansionManager) == VSConstants.S_OK)
        {
            ComEventSink.Advise<IVsExpansionEvents>(_expansionManager, this);
            await PopulateSnippetCacheAsync().ConfigureAwait(false);
        }
    }
 
    public int OnAfterSnippetsUpdate()
    {
        _threadingContext.ThrowIfNotOnUIThread();
 
        if (_expansionManager != null)
        {
            _threadingContext.RunWithShutdownBlockAsync((_) => PopulateSnippetCacheAsync());
        }
 
        return VSConstants.S_OK;
    }
 
    public int OnAfterSnippetsKeyBindingChange([ComAliasName("Microsoft.VisualStudio.OLE.Interop.DWORD")] uint dwCmdGuid, [ComAliasName("Microsoft.VisualStudio.OLE.Interop.DWORD")] uint dwCmdId, [ComAliasName("Microsoft.VisualStudio.OLE.Interop.BOOL")] int fBound)
        => VSConstants.S_OK;
 
    public IEnumerable<SnippetInfo> GetSnippetsIfAvailable()
    {
        // Immediately return the known set of snippets, even if we're still in the process
        // of calculating a more up-to-date list.
        lock (cacheGuard)
        {
            return snippets;
        }
    }
 
    public bool SnippetShortcutExists_NonBlocking(string shortcut)
    {
        if (shortcut == null)
        {
            return false;
        }
 
        // Check against the known set of snippets, even if we're still in the process of
        // calculating a more up-to-date list.
        lock (cacheGuard)
        {
            return snippetShortcuts.Contains(shortcut);
        }
    }
 
    public virtual bool ShouldFormatSnippet(SnippetInfo snippetInfo)
        => false;
 
    private async Task PopulateSnippetCacheAsync()
    {
        using var token = _waiter.BeginAsyncOperation(GetType().Name + ".Start");
        RoslynDebug.Assert(_expansionManager != null);
 
        // In Dev14 Update2+ the platform always provides an IExpansion Manager
        var expansionManager = (IExpansionManager)_expansionManager;
        // Call the asynchronous IExpansionManager API from a background thread
        await TaskScheduler.Default;
        var expansionEnumerator = await expansionManager.EnumerateExpansionsAsync(
            _languageGuidForSnippets,
            0, // shortCutOnly
            [], // types
            0, // countTypes
            1, // includeNULLTypes
            1 // includeDulicates: Allows snippets with the same title but different shortcuts
            ).ConfigureAwait(false);
 
        // The rest of the process requires being on the UI thread, see the explanation on
        // PopulateSnippetCacheFromExpansionEnumeration for details
        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync();
        PopulateSnippetCacheFromExpansionEnumeration(expansionEnumerator);
    }
 
    /// <remarks>
    /// This method must be called on the UI thread because it eventually calls into
    /// IVsExpansionEnumeration.Next, which must be called on the UI thread due to an issue
    /// with how the call is marshalled.
    /// 
    /// The second parameter for IVsExpansionEnumeration.Next is defined like this:
    ///    [ComAliasName("Microsoft.VisualStudio.TextManager.Interop.VsExpansion")] IntPtr[] rgelt
    ///
    /// We pass a pointer for rgelt that we expect to be populated as the result. This
    /// eventually calls into the native CExpansionEnumeratorShim::Next method, which has the
    /// same contract of expecting a non-null rgelt that it can drop expansion data into. When
    /// we call from the UI thread, this transition from managed code to the
    /// CExpansionEnumeratorShim goes smoothly and everything works.
    ///
    /// When we call from a background thread, the COM marshaller has to move execution to the
    /// UI thread, and as part of this process it uses the interface as defined in the idl to
    /// set up the appropriate arguments to pass. The same parameter from the idl is defined as
    ///    [out, size_is(celt), length_is(*pceltFetched)] VsExpansion **rgelt
    ///
    /// Because rgelt is specified as an <c>out</c> parameter, the marshaller is discarding the
    /// pointer we passed and substituting the null reference. This then causes a null
    /// reference exception in the shim. Calling from the UI thread avoids this marshaller.
    /// </remarks>
    private void PopulateSnippetCacheFromExpansionEnumeration(IVsExpansionEnumeration expansionEnumerator)
    {
        _threadingContext.ThrowIfNotOnUIThread();
 
        var updatedSnippets = ExtractSnippetInfo(expansionEnumerator);
        var updatedSnippetShortcuts = GetShortcutsHashFromSnippets(updatedSnippets);
 
        lock (cacheGuard)
        {
            snippets = updatedSnippets;
            snippetShortcuts = updatedSnippetShortcuts;
        }
    }
 
    private ImmutableArray<SnippetInfo> ExtractSnippetInfo(IVsExpansionEnumeration expansionEnumerator)
    {
        _threadingContext.ThrowIfNotOnUIThread();
 
        var snippetListBuilder = ImmutableArray.CreateBuilder<SnippetInfo>();
        var snippetInfo = new VsExpansion();
        var pSnippetInfo = new IntPtr[1];
 
        try
        {
            // Allocate enough memory for one VSExpansion structure. This memory is filled in by the Next method.
            pSnippetInfo[0] = Marshal.AllocCoTaskMem(Marshal.SizeOf(snippetInfo));
 
            expansionEnumerator.GetCount(out var count);
 
            for (uint i = 0; i < count; i++)
            {
                expansionEnumerator.Next(1, pSnippetInfo, out var fetched);
                if (fetched > 0)
                {
                    // Convert the returned blob of data into a structure that can be read in managed code.
                    snippetInfo = ConvertToVsExpansionAndFree(pSnippetInfo[0]);
 
                    if (!string.IsNullOrEmpty(snippetInfo.shortcut))
                    {
                        snippetListBuilder.Add(new SnippetInfo(snippetInfo.shortcut, snippetInfo.title, snippetInfo.description, snippetInfo.path));
                    }
                }
            }
        }
        finally
        {
            Marshal.FreeCoTaskMem(pSnippetInfo[0]);
        }
 
        return snippetListBuilder.ToImmutableAndClear();
    }
 
    protected static IImmutableSet<string> GetShortcutsHashFromSnippets(ImmutableArray<SnippetInfo> updatedSnippets)
    {
        return new HashSet<string>(updatedSnippets.Select(s => s.Shortcut), StringComparer.OrdinalIgnoreCase)
            .ToImmutableHashSet(StringComparer.OrdinalIgnoreCase);
    }
 
    private static VsExpansion ConvertToVsExpansionAndFree(IntPtr expansionPtr)
    {
        var buffer = (VsExpansionWithIntPtrs)Marshal.PtrToStructure(expansionPtr, typeof(VsExpansionWithIntPtrs));
        var expansion = new VsExpansion();
 
        ConvertToStringAndFree(ref buffer.DescriptionPtr, ref expansion.description);
        ConvertToStringAndFree(ref buffer.PathPtr, ref expansion.path);
        ConvertToStringAndFree(ref buffer.ShortcutPtr, ref expansion.shortcut);
        ConvertToStringAndFree(ref buffer.TitlePtr, ref expansion.title);
 
        return expansion;
    }
 
    private static void ConvertToStringAndFree(ref IntPtr ptr, ref string? str)
    {
        if (ptr != IntPtr.Zero)
        {
            str = Marshal.PtrToStringBSTR(ptr);
            Marshal.FreeBSTR(ptr);
            ptr = IntPtr.Zero;
        }
    }
 
    /// <summary>
    /// This structure is used to facilitate the interop calls with IVsExpansionEnumeration.
    /// </summary>
    [StructLayout(LayoutKind.Sequential)]
    private struct VsExpansionWithIntPtrs
    {
        public IntPtr PathPtr;
        public IntPtr TitlePtr;
        public IntPtr ShortcutPtr;
        public IntPtr DescriptionPtr;
    }
}