File: SolutionExplorer\SourceGeneratedFileItems\SourceGeneratedFileItemSource.cs
Web Access
Project: src\src\VisualStudio\Core\Impl\Microsoft.VisualStudio.LanguageServices.Implementation.csproj (Microsoft.VisualStudio.LanguageServices.Implementation)
// 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;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.Editor.Shared.Utilities;
using Microsoft.CodeAnalysis.Shared.TestHooks;
using Microsoft.CodeAnalysis.Threading;
using Microsoft.Internal.VisualStudio.PlatformUI;
 
namespace Microsoft.VisualStudio.LanguageServices.Implementation.SolutionExplorer;
 
internal sealed class SourceGeneratedFileItemSource(
    SourceGeneratorItem parentGeneratorItem,
    IThreadingContext threadingContext,
    Workspace workspace,
    IAsynchronousOperationListener asyncListener)
    : Shell.IAttachedCollectionSource, ISupportExpansionEvents
{
    private readonly SourceGeneratorItem _parentGeneratorItem = parentGeneratorItem;
    private readonly IThreadingContext _threadingContext = threadingContext;
    private readonly Workspace _workspace = workspace;
    private readonly IAsynchronousOperationListener _asyncListener = asyncListener;
 
    /// <summary>
    /// The returned collection of items. Can only be mutated on the UI thread, as other parts of WPF are subscribed to
    /// the change events and expect that.
    /// </summary>
    private readonly BulkObservableCollectionWithInit<BaseItem> _items = [];
 
    /// <summary>
    /// Gate to guard mutation of <see cref="_resettableDelay"/>.
    /// </summary>
    private readonly object _gate = new();
 
    private readonly CancellationSeries _cancellationSeries = new();
    private ResettableDelay? _resettableDelay;
    private WorkspaceEventRegistration? _workspaceChangedDisposer;
 
    public object SourceItem => _parentGeneratorItem;
 
    // Since we are expensive to compute, always say we have items.
    public bool HasItems => true;
 
    public IEnumerable Items => _items;
 
    private async Task UpdateSourceGeneratedFileItemsAsync(Solution solution, CancellationToken cancellationToken)
    {
        var project = solution.GetProject(_parentGeneratorItem.ProjectId);
 
        if (project == null)
        {
            return;
        }
 
        var sourceGeneratedDocuments = await project.GetSourceGeneratedDocumentsAsync(cancellationToken).ConfigureAwait(false);
        var sourceGeneratedDocumentsForGeneratorById =
            sourceGeneratedDocuments.Where(d => d.Identity.Generator == _parentGeneratorItem.Identity)
            .ToDictionary(d => d.Id);
 
        // We must update the list on the UI thread, since the WPF elements bound to our list expect that
        await _threadingContext.JoinableTaskFactory.SwitchToMainThreadAsync(cancellationToken);
 
        try
        {
            // We're going to incrementally update our items, ensuring we keep the object identity for things we didn't touch.
            // This is because the Solution Explorer itself will use identity to keep track of active items -- if you have an
            // item selected and we were to refresh in the background we don't want to lose that selection. If we just removed
            // and repopulated the list from scratch each time we'd lose the selection.
            _items.BeginBulkOperation();
 
            // Do we already have a "no files" placeholder item?
            if (_items is [NoSourceGeneratedFilesPlaceholderItem])
            {
                // We do -- if we have no items, we're done, since the placeholder is all that needs to be there;
                // otherwise remove it since we have real files now
                if (sourceGeneratedDocumentsForGeneratorById.Count == 0)
                {
                    return;
                }
                else
                {
                    _items.RemoveAt(0);
                }
            }
 
            for (var i = 0; i < _items.Count; i++)
            {
                // If this item that we already have is still a generated document, we'll remove it from our list; the list when we're
                // done is going to have the new items remaining. If it no longer exists, remove it from list.
                if (!sourceGeneratedDocumentsForGeneratorById.Remove(((SourceGeneratedFileItem)_items[i]).DocumentId))
                {
                    _items.RemoveAt(i);
                    i--;
                }
            }
 
            // Whatever is left in sourceGeneratedDocumentsForGeneratorById we should add; if we have nothing to add and nothing
            // in the list after removing anything, then we should add the placeholder.
            if (sourceGeneratedDocumentsForGeneratorById.Count == 0 && _items.Count == 0)
            {
                _items.Add(new NoSourceGeneratedFilesPlaceholderItem());
                return;
            }
 
            foreach (var document in sourceGeneratedDocumentsForGeneratorById.Values)
            {
                // Binary search to figure out where to insert
                var low = 0;
                var high = _items.Count;
 
                while (low < high)
                {
                    var mid = (low + high) / 2;
 
                    if (StringComparer.OrdinalIgnoreCase.Compare(document.HintName, ((SourceGeneratedFileItem)_items[mid]).HintName) < 0)
                    {
                        high = mid;
                    }
                    else
                    {
                        low = mid + 1;
                    }
                }
 
                _items.Insert(low, new SourceGeneratedFileItem(
                    _threadingContext, document.Id, document.HintName, document.Project.Language, _workspace));
            }
        }
        finally
        {
            _items.EndBulkOperation();
            _items.IsInitialized = true;
        }
    }
 
    public void BeforeExpand()
    {
        lock (_gate)
        {
            // We should not have an existing computation active
            Contract.ThrowIfTrue(_cancellationSeries.HasActiveToken);
 
            var cancellationToken = _cancellationSeries.CreateNext();
            var asyncToken = _asyncListener.BeginAsyncOperation(nameof(SourceGeneratedFileItemSource) + "." + nameof(BeforeExpand));
 
            Task.Run(
                async () =>
                {
                    // Since the user just expanded this, we want to do a single population aggressively,
                    // where the only reason we'd cancel is if the user collapsed it again.
                    var solution = _workspace.CurrentSolution;
                    await UpdateSourceGeneratedFileItemsAsync(solution, cancellationToken).ConfigureAwait(false);
 
                    // Now that we've done it the first time, we'll subscribe for future changes
                    lock (_gate)
                    {
                        // It's important we check for cancellation inside our lock: if the user were to collapse
                        // right at this point, we don't want to have a case where we cancelled the work, unsubscribed
                        // in AfterCollapse, and _then_ subscribed here again.
 
                        cancellationToken.ThrowIfCancellationRequested();
                        _workspaceChangedDisposer = _workspace.RegisterWorkspaceChangedHandler(OnWorkspaceChanged);
                        if (_workspace.CurrentSolution != solution)
                        {
                            // The workspace changed while we were doing our initial population, so
                            // refresh it. We'll just call our OnWorkspaceChanged event handler
                            // so this looks like any other change.
                            OnWorkspaceChanged(new WorkspaceChangeEventArgs(WorkspaceChangeKind.SolutionChanged, solution, _workspace.CurrentSolution));
                        }
                    }
                },
                cancellationToken).CompletesAsyncOperation(asyncToken);
        }
    }
 
    public void AfterCollapse()
    {
        StopUpdating();
    }
 
    private void StopUpdating()
    {
        lock (_gate)
        {
            _cancellationSeries.CreateNext(new CancellationToken(canceled: true));
            _workspaceChangedDisposer?.Dispose();
            _workspaceChangedDisposer = null;
            _resettableDelay = null;
        }
    }
 
    private void OnWorkspaceChanged(WorkspaceChangeEventArgs e)
    {
        if (!e.NewSolution.ContainsProject(_parentGeneratorItem.ProjectId))
        {
            StopUpdating();
        }
 
        lock (_gate)
        {
            // If we already have a ResettableDelay, just delay it further; otherwise we either have no delay
            // or the actual processing began, and we need to start over
            if (_resettableDelay != null)
            {
                _resettableDelay.Reset();
            }
            else
            {
                // Time to start the work all over again. We'll ensure any previous work is cancelled
                var cancellationToken = _cancellationSeries.CreateNext();
                var asyncToken = _asyncListener.BeginAsyncOperation(nameof(SourceGeneratedFileItemSource) + "." + nameof(OnWorkspaceChanged));
 
                // We're going to go with a really long delay: once the user expands this we will keep it updated, but it's fairly
                // unlikely to change in a lot of cases if a generator only produces a stable set of names.
                _resettableDelay = new ResettableDelay(delayInMilliseconds: 5000, _asyncListener, cancellationToken);
                _resettableDelay.Task.ContinueWith(_ =>
                {
                    lock (_gate)
                    {
                        // We've started off this work, so if another change comes in we need to start a delay all over again
                        _resettableDelay = null;
                    }
 
                    cancellationToken.ThrowIfCancellationRequested();
 
                    return UpdateSourceGeneratedFileItemsAsync(_workspace.CurrentSolution, cancellationToken);
                }, cancellationToken, TaskContinuationOptions.OnlyOnRanToCompletion, TaskScheduler.Default).Unwrap().CompletesAsyncOperation(asyncToken);
            }
        }
    }
}