Merge pull request #73237 from CyrusNajmabadi/treeSource2

src/Workspaces/Core/Portable/Workspace/Solution/DocumentState.cs

@@ -31,7 +31,7 @@ internal partial class DocumentState : TextDocumentState
     private readonly ParseOptions? _options;
 
     // null if the document doesn't support syntax trees:
-    private readonly AsyncLazy<TreeAndVersion>? _treeSource;
+    private readonly ITreeAndVersionSource? _treeSource;
 
     protected DocumentState(
         LanguageServices languageServices,
@@ -40,7 +40,7 @@ protected DocumentState(
         ParseOptions? options,
         ITextAndVersionSource textSource,
         LoadTextOptions loadTextOptions,
-        AsyncLazy<TreeAndVersion>? treeSource)
+        ITreeAndVersionSource? treeSource)
         : base(languageServices.SolutionServices, documentServiceProvider, attributes, textSource, loadTextOptions)
     {
         Contract.ThrowIfFalse(_options is null == _treeSource is null);
@@ -79,7 +79,7 @@ public DocumentState(
         }
     }
 
-    public AsyncLazy<TreeAndVersion>? TreeSource => _treeSource;
+    public ITreeAndVersionSource? TreeSource => _treeSource;
 
     [MemberNotNullWhen(true, nameof(_treeSource))]
     [MemberNotNullWhen(true, nameof(TreeSource))]
@@ -97,15 +97,15 @@ public SourceCodeKind SourceCodeKind
     public bool IsGenerated
         => Attributes.IsGenerated;
 
-    protected static AsyncLazy<TreeAndVersion> CreateLazyFullyParsedTree(
+    protected static ITreeAndVersionSource CreateLazyFullyParsedTree(
         ITextAndVersionSource newTextSource,
         LoadTextOptions loadTextOptions,
         string? filePath,
         ParseOptions options,
         LanguageServices languageServices,
         PreservationMode mode = PreservationMode.PreserveValue)
     {
-        return AsyncLazy.Create(
+        return SimpleTreeAndVersionSource.Create(
             static (arg, c) => FullyParseTreeAsync(arg.newTextSource, arg.loadTextOptions, arg.filePath, arg.options, arg.languageServices, arg.mode, c),
             static (arg, c) => FullyParseTree(arg.newTextSource, arg.loadTextOptions, arg.filePath, arg.options, arg.languageServices, arg.mode, c),
             arg: (newTextSource, loadTextOptions, filePath, options, languageServices, mode));
@@ -163,19 +163,19 @@ private static TreeAndVersion CreateTreeAndVersion(
         return new TreeAndVersion(tree, textAndVersion.Version);
     }
 
-    private static AsyncLazy<TreeAndVersion> CreateLazyIncrementallyParsedTree(
-        AsyncLazy<TreeAndVersion> oldTreeSource,
+    private static ITreeAndVersionSource CreateLazyIncrementallyParsedTree(
+        ITreeAndVersionSource oldTreeSource,
         ITextAndVersionSource newTextSource,
         LoadTextOptions loadTextOptions)
     {
-        return AsyncLazy.Create(
+        return SimpleTreeAndVersionSource.Create(
             static (arg, c) => IncrementallyParseTreeAsync(arg.oldTreeSource, arg.newTextSource, arg.loadTextOptions, c),
             static (arg, c) => IncrementallyParseTree(arg.oldTreeSource, arg.newTextSource, arg.loadTextOptions, c),
             arg: (oldTreeSource, newTextSource, loadTextOptions));
     }
 
     private static async Task<TreeAndVersion> IncrementallyParseTreeAsync(
-        AsyncLazy<TreeAndVersion> oldTreeSource,
+        ITreeAndVersionSource oldTreeSource,
         ITextAndVersionSource newTextSource,
         LoadTextOptions loadTextOptions,
         CancellationToken cancellationToken)
@@ -197,7 +197,7 @@ private static async Task<TreeAndVersion> IncrementallyParseTreeAsync(
     }
 
     private static TreeAndVersion IncrementallyParseTree(
-        AsyncLazy<TreeAndVersion> oldTreeSource,
+        ITreeAndVersionSource oldTreeSource,
         ITextAndVersionSource newTextSource,
         LoadTextOptions loadTextOptions,
         CancellationToken cancellationToken)
@@ -346,7 +346,7 @@ private DocumentState SetParseOptions(ParseOptions options, bool onlyPreprocesso
             throw new InvalidOperationException();
         }
 
-        AsyncLazy<TreeAndVersion>? newTreeSource = null;
+        ITreeAndVersionSource? newTreeSource = null;
 
         // Optimization: if we are only changing preprocessor directives, and we've already parsed the existing tree
         // and it didn't have any, we can avoid a reparse since the tree will be parsed the same.
@@ -368,7 +368,7 @@ private DocumentState SetParseOptions(ParseOptions options, bool onlyPreprocesso
             }
 
             if (newTree is not null)
-                newTreeSource = AsyncLazy.Create(new TreeAndVersion(newTree, existingTreeAndVersion.Version));
+                newTreeSource = SimpleTreeAndVersionSource.Create(new TreeAndVersion(newTree, existingTreeAndVersion.Version));
         }
 
         // If we weren't able to reuse in a smart way, just reparse
@@ -457,7 +457,7 @@ public DocumentState UpdateFilePath(string? filePath)
 
     protected override TextDocumentState UpdateText(ITextAndVersionSource newTextSource, PreservationMode mode, bool incremental)
     {
-        AsyncLazy<TreeAndVersion>? newTreeSource;
+        ITreeAndVersionSource? newTreeSource;
 
         if (!SupportsSyntaxTree)
         {
@@ -528,7 +528,7 @@ internal DocumentState UpdateTree(SyntaxNode newRoot, PreservationMode mode)
             _options,
             textSource: text,
             LoadTextOptions,
-            treeSource: AsyncLazy.Create(treeAndVersion));
+            treeSource: SimpleTreeAndVersionSource.Create(treeAndVersion));
 
         // use static method so we don't capture references to this
         static (ITextAndVersionSource, TreeAndVersion) CreateTreeWithLazyText(

src/Workspaces/Core/Portable/Workspace/Solution/DocumentState_LinkedFileReuse.cs

@@ -13,6 +13,29 @@ namespace Microsoft.CodeAnalysis;
 
 internal partial class DocumentState
 {
+    /// <summary>
+    /// <see cref="ITreeAndVersionSource"/> when we're linked to another file (a 'sibling') and will attempt to reuse
+    /// that sibling's tree as our own. Note: we won't know if we can actually use the contents of that sibling file
+    /// until we actually go and realize it, as it may contains constructs (like pp-directives) that prevent use.  In
+    /// that case, we'll fall back to a normal incremental parse between our original <paramref
+    /// name="originalTreeSource"/> and the latest <em>text</em> contents of our sibling's file.
+    /// </summary>
+    private sealed class LinkedFileReuseTreeAndVersionSource(
+        ITreeAndVersionSource originalTreeSource,
+        AsyncLazy<TreeAndVersion> lazyComputation) : ITreeAndVersionSource
+    {
+        public readonly ITreeAndVersionSource OriginalTreeSource = originalTreeSource;
+
+        public Task<TreeAndVersion> GetValueAsync(CancellationToken cancellationToken)
+            => lazyComputation.GetValueAsync(cancellationToken);
+
+        public TreeAndVersion GetValue(CancellationToken cancellationToken)
+            => lazyComputation.GetValue(cancellationToken);
+
+        public bool TryGetValue([NotNullWhen(true)] out TreeAndVersion? value)
+            => lazyComputation.TryGetValue(out value);
+    }
+
     /// <summary>
     /// Returns a new instance of this document state that points to <paramref name="siblingTextSource"/> as the
     /// text contents of the document, and which will produce a syntax tree that reuses from <paramref
@@ -21,7 +44,7 @@ internal partial class DocumentState
     /// </summary>
     public DocumentState UpdateTextAndTreeContents(
         ITextAndVersionSource siblingTextSource,
-        AsyncLazy<TreeAndVersion>? siblingTreeSource,
+        ITreeAndVersionSource? siblingTreeSource,
         bool forceEvenIfTreesWouldDiffer)
     {
         if (!SupportsSyntaxTree)
@@ -38,50 +61,52 @@ public DocumentState UpdateTextAndTreeContents(
 
         Contract.ThrowIfNull(siblingTreeSource);
 
+        // We don't want to point at a long chain of transformations as our sibling files change, deferring to each next
+        // link of the chain to potentially do the work (or potentially failing out). So, if we're about to do this,
+        // instead return our original tree-source so that in the case we are unable to use the sibling file's root, we
+        // can do a single step incremental parse between our original tree and the final sibling text.
+        //
+        // We only need to look one deep here as we'll pull that tree source forward to our level.  If another link is
+        // later added to us, it will do the same thing.
+        var originalTreeSource = this.TreeSource;
+        if (originalTreeSource is LinkedFileReuseTreeAndVersionSource linkedFileTreeAndVersionSource)
+            originalTreeSource = linkedFileTreeAndVersionSource.OriginalTreeSource;
+
         // Always pass along the sibling text.  We will always be in sync with that.
 
-        // if a sibling tree source is provided, then we'll want to attempt to use the tree it creates, to share as
-        // much memory as possible with linked files.  However, we can't point at that source directly.  If we did,
-        // we'd produce the *exact* same tree-reference as another file.  That would be bad as it would break the
-        // invariant that each document gets a unique SyntaxTree.  So, instead, we produce a ValueSource that defers
-        // to the provided source, gets the tree from it, and then wraps its root in a new tree for us.
+        // Defer to static helper to make sure we don't accidentally capture anything else we don't want off of 'this'
+        // (like "this.TreeSource").
+        return UpdateTextAndTreeContentsWorker(
+            this.Attributes, this.LanguageServices, this.Services, this.LoadTextOptions, this.ParseOptions,
+            originalTreeSource, siblingTextSource, siblingTreeSource, forceEvenIfTreesWouldDiffer);
+    }
 
-        // copy data from this entity, and pass to static helper, so we don't keep this green node alive.
+    private static DocumentState UpdateTextAndTreeContentsWorker(
+        DocumentInfo.DocumentAttributes attributes,
+        LanguageServices languageServices,
+        IDocumentServiceProvider services,
+        LoadTextOptions loadTextOptions,
+        ParseOptions parseOptions,
+        ITreeAndVersionSource originalTreeSource,
+        ITextAndVersionSource siblingTextSource,
+        ITreeAndVersionSource siblingTreeSource,
+        bool forceEvenIfTreesWouldDiffer)
+    {
+        // if a sibling tree source is provided, then we'll want to attempt to use the tree it creates, to share as much
+        // memory as possible with linked files.  However, we can't point at that source directly.  If we did, we'd
+        // produce the *exact* same tree-reference as another file.  That would be bad as it would break the invariant
+        // that each document gets a unique SyntaxTree.  So, instead, we produce a tree-source that defers to the
+        // provided source, gets the tree from it, and then wraps its root in a new tree for us.
 
-        var filePath = this.Attributes.SyntaxTreeFilePath;
-        var languageServices = this.LanguageServices;
-        var loadTextOptions = this.LoadTextOptions;
-        var parseOptions = this.ParseOptions;
-        var textAndVersionSource = this.TextAndVersionSource;
-        var treeSource = this.TreeSource;
+        var lazyComputation = AsyncLazy.Create(
+            static (arg, cancellationToken) => TryReuseSiblingTreeAsync(arg.filePath, arg.languageServices, arg.loadTextOptions, arg.parseOptions, arg.originalTreeSource, arg.siblingTextSource, arg.siblingTreeSource, arg.forceEvenIfTreesWouldDiffer, cancellationToken),
+            static (arg, cancellationToken) => TryReuseSiblingTree(arg.filePath, arg.languageServices, arg.loadTextOptions, arg.parseOptions, arg.originalTreeSource, arg.siblingTextSource, arg.siblingTreeSource, arg.forceEvenIfTreesWouldDiffer, cancellationToken),
+            arg: (filePath: attributes.SyntaxTreeFilePath, languageServices, loadTextOptions, parseOptions, originalTreeSource, siblingTextSource, siblingTreeSource, forceEvenIfTreesWouldDiffer));
 
-        var newTreeSource = GetReuseTreeSource(
-            filePath, languageServices, loadTextOptions, parseOptions, treeSource, siblingTextSource, siblingTreeSource, forceEvenIfTreesWouldDiffer);
+        var newTreeSource = new LinkedFileReuseTreeAndVersionSource(originalTreeSource, lazyComputation);
 
         return new DocumentState(
-            languageServices,
-            Services,
-            Attributes,
-            _options,
-            siblingTextSource,
-            LoadTextOptions,
-            newTreeSource);
-
-        static AsyncLazy<TreeAndVersion> GetReuseTreeSource(
-            string filePath,
-            LanguageServices languageServices,
-            LoadTextOptions loadTextOptions,
-            ParseOptions parseOptions,
-            AsyncLazy<TreeAndVersion> treeSource,
-            ITextAndVersionSource siblingTextSource,
-            AsyncLazy<TreeAndVersion> siblingTreeSource,
-            bool forceEvenIfTreesWouldDiffer)
-        {
-            return AsyncLazy.Create(
-                static (arg, cancellationToken) => TryReuseSiblingTreeAsync(arg.filePath, arg.languageServices, arg.loadTextOptions, arg.parseOptions, arg.treeSource, arg.siblingTextSource, arg.siblingTreeSource, arg.forceEvenIfTreesWouldDiffer, cancellationToken),
-                static (arg, cancellationToken) => TryReuseSiblingTree(arg.filePath, arg.languageServices, arg.loadTextOptions, arg.parseOptions, arg.treeSource, arg.siblingTextSource, arg.siblingTreeSource, arg.forceEvenIfTreesWouldDiffer, cancellationToken),
-                arg: (filePath, languageServices, loadTextOptions, parseOptions, treeSource, siblingTextSource, siblingTreeSource, forceEvenIfTreesWouldDiffer));
-        }
+            languageServices, services, attributes, parseOptions, siblingTextSource, loadTextOptions, newTreeSource);
 
         static bool TryReuseSiblingRoot(
             string filePath,
@@ -186,9 +211,9 @@ static async Task<TreeAndVersion> TryReuseSiblingTreeAsync(
             LanguageServices languageServices,
             LoadTextOptions loadTextOptions,
             ParseOptions parseOptions,
-            AsyncLazy<TreeAndVersion> treeSource,
+            ITreeAndVersionSource treeSource,
             ITextAndVersionSource siblingTextSource,
-            AsyncLazy<TreeAndVersion> siblingTreeSource,
+            ITreeAndVersionSource siblingTreeSource,
             bool forceEvenIfTreesWouldDiffer,
             CancellationToken cancellationToken)
         {
@@ -209,9 +234,9 @@ static TreeAndVersion TryReuseSiblingTree(
             LanguageServices languageServices,
             LoadTextOptions loadTextOptions,
             ParseOptions parseOptions,
-            AsyncLazy<TreeAndVersion> treeSource,
+            ITreeAndVersionSource treeSource,
             ITextAndVersionSource siblingTextSource,
-            AsyncLazy<TreeAndVersion> siblingTreeSource,
+            ITreeAndVersionSource siblingTreeSource,
             bool forceEvenIfTreesWouldDiffer,
             CancellationToken cancellationToken)
         {

src/Workspaces/Core/Portable/Workspace/Solution/SourceGeneratedDocumentState.cs

@@ -105,7 +105,7 @@ private SourceGeneratedDocumentState(
         ITextAndVersionSource textSource,
         SourceText text,
         LoadTextOptions loadTextOptions,
-        AsyncLazy<TreeAndVersion> treeSource,
+        ITreeAndVersionSource treeSource,
         Lazy<Checksum> lazyContentHash,
         DateTime generationDateTime)
         : base(languageServices, documentServiceProvider, attributes, options, textSource, loadTextOptions, treeSource)

src/Workspaces/Core/Portable/Workspace/Solution/VersionSource/ITreeAndVersionSource.cs

@@ -0,0 +1,20 @@
+// 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.Diagnostics.CodeAnalysis;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.CodeAnalysis;
+
+/// <summary>
+/// Similar to <see cref="ITextAndVersionSource"/>, but for trees.  Allows hiding (or introspecting) the details of how
+/// a tree is created for a particular document.
+/// </summary>
+internal interface ITreeAndVersionSource
+{
+    Task<TreeAndVersion> GetValueAsync(CancellationToken cancellationToken);
+    TreeAndVersion GetValue(CancellationToken cancellationToken);
+    bool TryGetValue([NotNullWhen(true)] out TreeAndVersion? value);
+}

src/Workspaces/Core/Portable/Workspace/Solution/VersionSource/SimpleTreeAndVersionSource.cs

@@ -0,0 +1,44 @@
+// 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.Diagnostics.CodeAnalysis;
+using System.Threading;
+using System.Threading.Tasks;
+using Roslyn.Utilities;
+
+namespace Microsoft.CodeAnalysis;
+
+/// <summary>
+/// Simple implementation of <see cref="ITreeAndVersionSource"/> backed by an opaque <see
+/// cref="AsyncLazy{TreeAndVersion}"/>."/>
+/// </summary>
+internal sealed class SimpleTreeAndVersionSource : ITreeAndVersionSource
+{
+    private readonly AsyncLazy<TreeAndVersion> _source;
+
+    private SimpleTreeAndVersionSource(AsyncLazy<TreeAndVersion> source)
+    {
+        _source = source;
+    }
+
+    public Task<TreeAndVersion> GetValueAsync(CancellationToken cancellationToken)
+        => _source.GetValueAsync(cancellationToken);
+
+    public TreeAndVersion GetValue(CancellationToken cancellationToken)
+        => _source.GetValue(cancellationToken);
+
+    public bool TryGetValue([NotNullWhen(true)] out TreeAndVersion? value)
+        => _source.TryGetValue(out value);
+
+    public static SimpleTreeAndVersionSource Create<TArg>(
+        Func<TArg, CancellationToken, Task<TreeAndVersion>> asynchronousComputeFunction,
+        Func<TArg, CancellationToken, TreeAndVersion>? synchronousComputeFunction, TArg arg)
+    {
+        return new(AsyncLazy<TreeAndVersion>.Create(asynchronousComputeFunction, synchronousComputeFunction, arg));
+    }
+
+    public static SimpleTreeAndVersionSource Create(TreeAndVersion source)
+        => new(AsyncLazy.Create(source));
+}