Wire NullableWalker snapshotting to the public SpeculativeSemanticModel apis (#36563)

Wire NullableWalker snapshotting to the public SpeculativeSemanticModel apis

src/Compilers/CSharp/Portable/Compilation/AttributeSemanticModel.cs

@@ -23,7 +23,7 @@ internal sealed class AttributeSemanticModel : MemberSemanticModel
             SyntaxTreeSemanticModel containingSemanticModelOpt = null,
             SyntaxTreeSemanticModel parentSemanticModelOpt = null,
             int speculatedPosition = 0)
-            : base(syntax, attributeType, new ExecutableCodeBinder(syntax, rootBinder.ContainingMember(), rootBinder), containingSemanticModelOpt, parentSemanticModelOpt, speculatedPosition)
+            : base(syntax, attributeType, new ExecutableCodeBinder(syntax, rootBinder.ContainingMember(), rootBinder), containingSemanticModelOpt, parentSemanticModelOpt, snapshotManagerOpt: null, speculatedPosition)
         {
             Debug.Assert(syntax != null);
             _aliasOpt = aliasOpt;
@@ -98,9 +98,9 @@ internal override BoundNode Bind(Binder binder, CSharpSyntaxNode node, Diagnosti
             }
         }
 
-        protected override BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, out NullableWalker.SnapshotManager snapshotManager)
+        protected override BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, bool createSnapshots, out NullableWalker.SnapshotManager snapshotManager)
         {
-            return NullableWalker.AnalyzeAndRewrite(Compilation, symbol: null, boundRoot, binder, diagnostics, createSnapshots: true, out snapshotManager);
+            return NullableWalker.AnalyzeAndRewrite(Compilation, symbol: null, boundRoot, binder, diagnostics, createSnapshots, out snapshotManager);
         }
 
         internal override bool TryGetSpeculativeSemanticModelCore(SyntaxTreeSemanticModel parentModel, int position, ConstructorInitializerSyntax constructorInitializer, out SemanticModel speculativeModel)

src/Compilers/CSharp/Portable/Compilation/InitializerSemanticModel.cs

@@ -25,7 +25,7 @@ internal sealed class InitializerSemanticModel : MemberSemanticModel
                                      SyntaxTreeSemanticModel containingSemanticModelOpt = null,
                                      SyntaxTreeSemanticModel parentSemanticModelOpt = null,
                                      int speculatedPosition = 0) :
-            base(syntax, symbol, rootBinder, containingSemanticModelOpt, parentSemanticModelOpt, speculatedPosition)
+            base(syntax, symbol, rootBinder, containingSemanticModelOpt, parentSemanticModelOpt, snapshotManagerOpt: null, speculatedPosition)
         {
             Debug.Assert(!(syntax is ConstructorInitializerSyntax));
         }
@@ -258,9 +258,9 @@ internal override bool TryGetSpeculativeSemanticModelForMethodBodyCore(SyntaxTre
             return false;
         }
 
-        protected override BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, out NullableWalker.SnapshotManager snapshotManager)
+        protected override BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, bool createSnapshots, out NullableWalker.SnapshotManager snapshotManager)
         {
-            return NullableWalker.AnalyzeAndRewrite(Compilation, MemberSymbol, boundRoot, binder, diagnostics, createSnapshots: true, out snapshotManager);
+            return NullableWalker.AnalyzeAndRewrite(Compilation, MemberSymbol, boundRoot, binder, diagnostics, createSnapshots, out snapshotManager);
         }
     }
 }

src/Compilers/CSharp/Portable/Compilation/MemberSemanticModel.SpeculativeMemberSemanticModel.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
 using System.Collections.Immutable;
+using System.Diagnostics;
 using Microsoft.CodeAnalysis.CSharp.Symbols;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Roslyn.Utilities;
@@ -19,17 +20,13 @@ private sealed class SpeculativeMemberSemanticModel : MemberSemanticModel
             /// Creates a speculative SemanticModel for a TypeSyntax node at a position within an existing MemberSemanticModel.
             /// </summary>
             public SpeculativeMemberSemanticModel(SyntaxTreeSemanticModel parentSemanticModel, Symbol owner, TypeSyntax root, Binder rootBinder, int position)
-                : base(root, owner, rootBinder, containingSemanticModelOpt: null, parentSemanticModelOpt: parentSemanticModel, speculatedPosition: position)
+                : base(root, owner, rootBinder, containingSemanticModelOpt: null, parentSemanticModelOpt: parentSemanticModel, snapshotManagerOpt: null, speculatedPosition: position)
             {
             }
 
-            protected override BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, out NullableWalker.SnapshotManager snapshotManager)
+            protected override BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, bool createSnapshots, out NullableWalker.SnapshotManager snapshotManager)
             {
-                // https://github.com/dotnet/roslyn/issues/35037: Speculative models are going to have to do something more advanced
-                // here. They will need to run nullable analysis up to the point that is being speculated on, and
-                // then take that state and run analysis on the statement or expression being speculated on.
-                // Currently, it will return incorrect info because it's just running analysis on the speculated
-                // part.
+                Debug.Assert(boundRoot.Syntax is TypeSyntax);
                 return NullableWalker.AnalyzeAndRewrite(Compilation, MemberSymbol as MethodSymbol, boundRoot, binder, diagnostics, createSnapshots: false, out snapshotManager);
             }
 

src/Compilers/CSharp/Portable/Compilation/MemberSemanticModel.cs

@@ -26,6 +26,7 @@ internal abstract partial class MemberSemanticModel : CSharpSemanticModel
         // The bound nodes associated with a syntax node, from highest in the tree to lowest.
         private readonly Dictionary<SyntaxNode, ImmutableArray<BoundNode>> _guardedNodeMap = new Dictionary<SyntaxNode, ImmutableArray<BoundNode>>();
         private Dictionary<SyntaxNode, BoundStatement> _lazyGuardedSynthesizedStatementsMap;
+        private NullableWalker.SnapshotManager _lazySnapshotManager;
 
         internal readonly Binder RootBinder;
 
@@ -46,6 +47,7 @@ internal abstract partial class MemberSemanticModel : CSharpSemanticModel
             Binder rootBinder,
             SyntaxTreeSemanticModel containingSemanticModelOpt,
             SyntaxTreeSemanticModel parentSemanticModelOpt,
+            NullableWalker.SnapshotManager snapshotManagerOpt,
             int speculatedPosition)
         {
             Debug.Assert(root != null);
@@ -53,6 +55,7 @@ internal abstract partial class MemberSemanticModel : CSharpSemanticModel
             Debug.Assert(parentSemanticModelOpt == null ^ containingSemanticModelOpt == null);
             Debug.Assert(containingSemanticModelOpt == null || !containingSemanticModelOpt.IsSpeculativeSemanticModel);
             Debug.Assert(parentSemanticModelOpt == null || !parentSemanticModelOpt.IsSpeculativeSemanticModel, CSharpResources.ChainingSpeculativeModelIsNotSupported);
+            Debug.Assert(snapshotManagerOpt == null || parentSemanticModelOpt != null);
 
             _root = root;
             _memberSymbol = memberSymbol;
@@ -60,6 +63,7 @@ internal abstract partial class MemberSemanticModel : CSharpSemanticModel
             this.RootBinder = rootBinder.WithAdditionalFlags(GetSemanticModelBinderFlags());
             _containingSemanticModelOpt = containingSemanticModelOpt;
             _parentSemanticModelOpt = parentSemanticModelOpt;
+            _lazySnapshotManager = snapshotManagerOpt;
             _speculatedPosition = speculatedPosition;
 
             _operationFactory = new Lazy<CSharpOperationFactory>(() => new CSharpOperationFactory(this));
@@ -131,6 +135,15 @@ internal override MemberSemanticModel GetMemberModel(SyntaxNode node)
             return IsInTree(node) ? this : null;
         }
 
+        protected NullableWalker.SnapshotManager SnapshotManager
+        {
+            get
+            {
+                EnsureRootBoundForNullabilityIfNecessary();
+                return _lazySnapshotManager;
+            }
+        }
+
         internal sealed override bool TryGetSpeculativeSemanticModelCore(SyntaxTreeSemanticModel parentModel, int position, TypeSyntax type, SpeculativeBindingOption bindingOption, out SemanticModel speculativeModel)
         {
             var expression = SyntaxFactory.GetStandaloneExpression(type);
@@ -1382,15 +1395,15 @@ private ImmutableArray<BoundNode> GuardedAddBoundTreeAndGetBoundNodeFromMap(CSha
             return _guardedNodeMap.TryGetValue(syntax, out result) ? result : default(ImmutableArray<BoundNode>);
         }
 
-        protected void UnguardedAddBoundTreeForStandaloneSyntax(SyntaxNode syntax, BoundNode bound)
+        protected void UnguardedAddBoundTreeForStandaloneSyntax(SyntaxNode syntax, BoundNode bound, NullableWalker.SnapshotManager manager = null)
         {
             using (_nodeMapLock.DisposableWrite())
             {
-                GuardedAddBoundTreeForStandaloneSyntax(syntax, bound);
+                GuardedAddBoundTreeForStandaloneSyntax(syntax, bound, manager);
             }
         }
 
-        protected void GuardedAddBoundTreeForStandaloneSyntax(SyntaxNode syntax, BoundNode bound)
+        protected void GuardedAddBoundTreeForStandaloneSyntax(SyntaxNode syntax, BoundNode bound, NullableWalker.SnapshotManager manager = null)
         {
             Debug.Assert(_nodeMapLock.IsWriteLockHeld);
             bool alreadyInTree = false;
@@ -1417,6 +1430,25 @@ protected void GuardedAddBoundTreeForStandaloneSyntax(SyntaxNode syntax, BoundNo
                     // expressions can be added individually.
                     NodeMapBuilder.AddToMap(bound, _guardedNodeMap, syntax);
                 }
+
+                // If we're a speculative model, then we should never be passed a new manager, as we should not be recording
+                // new snapshot info, only using what was given when the model was created. If we're not a speculative model,
+                // then there are 3 possibilities:
+                // 1. Nullable analysis wasn't enabled, and we should never be passed a manager.
+                // 2. Nullable analysis is enabled, but we're not adding info for a root node and we shouldn't be passed
+                //    new snapshots. This can occur when we're asked to bind a standalone node that wasn't in the original
+                //    tree, such as aliases. These nodes do not affect analysis, so we should not be attempting to save
+                //    snapshot information for use in speculative models created off this one.
+                // 3. Nullable analysis is enabled, and we were passed a new manager. If this is the case, we must be adding
+                //    cached nodes for the root syntax node, and the existing snapshot manager must be null.
+                Debug.Assert((IsSpeculativeSemanticModel && manager is null) ||
+                             (!IsSpeculativeSemanticModel &&
+                              (manager is null && (!Compilation.NullableAnalysisEnabled || syntax != Root)) ||
+                              (manager is object && syntax == Root && Compilation.NullableAnalysisEnabled && _lazySnapshotManager is null)));
+                if (!IsSpeculativeSemanticModel && manager is object)
+                {
+                    _lazySnapshotManager = manager;
+                }
             }
         }
 
@@ -1625,7 +1657,7 @@ private BoundNode GetBoundLambdaOrQuery(CSharpSyntaxNode lambdaOrQuery)
                 // https://github.com/dotnet/roslyn/issues/35038: We need to do a rewrite here, and create a test that can hit this.
 #if DEBUG
                 var diagnostics = new DiagnosticBag();
-                _ = RewriteNullableBoundNodesWithSnapshots(boundOuterExpression, incrementalBinder, diagnostics, out _);
+                _ = RewriteNullableBoundNodesWithSnapshots(boundOuterExpression, incrementalBinder, diagnostics, takeSnapshots: false, out _);
 #endif
 
                 nodes = GuardedAddBoundTreeAndGetBoundNodeFromMap(lambdaOrQuery, boundOuterExpression);
@@ -1807,15 +1839,14 @@ private static Binder GetLambdaEnclosingBinder(int position, CSharpSyntaxNode st
         /// analysis does not run a subset of code, so we need to fully bind the entire member
         /// first
         /// </summary>
-        private void EnsureRootBoundForNullabilityIfNecessary()
+        protected void EnsureRootBoundForNullabilityIfNecessary()
         {
             // In DEBUG without nullable analysis enabled, we want to use a temp diagnosticbag
             // that can't produce any observable side effects
             DiagnosticBag diagnostics = _ignoredDiagnostics;
 
             // If we're in DEBUG mode, always enable the analysis, but throw away the results
-            // https://github.com/dotnet/roslyn/issues/35037: Disable speculative semantic models for now.
-            if (!Compilation.NullableAnalysisEnabled || IsSpeculativeSemanticModel)
+            if (!Compilation.NullableAnalysisEnabled)
             {
 #if DEBUG
                 diagnostics = new DiagnosticBag();
@@ -1824,6 +1855,14 @@ private void EnsureRootBoundForNullabilityIfNecessary()
 #endif
             }
 
+            // If this isn't a speculative model and we have a snapshot manager,
+            // then we've already done all the work necessary and we should avoid
+            // taking an unnecessary read lock.
+            if (!IsSpeculativeSemanticModel && _lazySnapshotManager is object)
+            {
+                return;
+            }
+
             var bindableRoot = GetBindableSyntaxNode(Root);
             using var upgradeableLock = _nodeMapLock.DisposableUpgradeableRead();
 
@@ -1843,22 +1882,45 @@ private void EnsureRootBoundForNullabilityIfNecessary()
 
             Debug.Assert(Root == GetBindableSyntaxNode(Root));
 
-            // https://github.com/dotnet/roslyn/issues/35037: Speculative models are going to have to do something more advanced
-            // here. They will need to run nullable analysis up to the point that is being speculated on, and
-            // then take that state and run analysis on the statement or expression being speculated on.
-            // Currently, it will return incorrect info because it's just running analysis on the speculated
-            // part.
             var binder = GetEnclosingBinder(GetAdjustedNodePosition(bindableRoot));
             var boundRoot = Bind(binder, bindableRoot, diagnostics);
-            boundRoot = RewriteNullableBoundNodesWithSnapshots(boundRoot, binder, diagnostics, out _);
+            if (IsSpeculativeSemanticModel)
+            {
+                ensureSpeculativeNodeBound();
+            }
+            else
+            {
+                bindAndRewrite(takeSnapshots: true);
+            }
 
-            if (Compilation.NullableAnalysisEnabled && !IsSpeculativeSemanticModel)
+            void ensureSpeculativeNodeBound()
             {
+                // Not all speculative models are created with existing snapshots. Attributes,
+                // TypeSyntaxes, and MethodBodies do not depend on existing state in a member,
+                // and so the SnapshotManager can be null in these cases.
+                if (_lazySnapshotManager is null)
+                {
+                    bindAndRewrite(takeSnapshots: false);
+                    return;
+                }
+
+                var analyzedNullabilitiesBuilder = ImmutableDictionary.CreateBuilder<BoundExpression, (NullabilityInfo, TypeSymbol)>();
+                boundRoot = NullableWalker.AnalyzeAndRewriteSpeculation(_speculatedPosition, boundRoot, binder, _lazySnapshotManager);
                 GuardedAddBoundTreeForStandaloneSyntax(bindableRoot, boundRoot);
             }
+
+            void bindAndRewrite(bool takeSnapshots)
+            {
+                boundRoot = RewriteNullableBoundNodesWithSnapshots(boundRoot, binder, diagnostics, takeSnapshots, out var snapshotManager);
+#if DEBUG
+                // Don't actually cache the results if the nullable analysis is not enabled in debug mode.
+                if (!Compilation.NullableAnalysisEnabled) return;
+#endif
+                GuardedAddBoundTreeForStandaloneSyntax(bindableRoot, boundRoot, snapshotManager);
+            }
         }
 
-        protected abstract BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, out NullableWalker.SnapshotManager snapshotManager);
+        protected abstract BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, bool takeSnapshots, out NullableWalker.SnapshotManager snapshotManager);
 
         /// <summary>
         /// Get all bounds nodes associated with a node, ordered from highest to lowest in the bound tree.

src/Compilers/CSharp/Portable/Compilation/MethodBodySemanticModel.cs

@@ -15,8 +15,9 @@ internal sealed class MethodBodySemanticModel : MemberSemanticModel
             CSharpSyntaxNode syntax,
             SyntaxTreeSemanticModel containingSemanticModelOpt = null,
             SyntaxTreeSemanticModel parentSemanticModelOpt = null,
+            NullableWalker.SnapshotManager snapshotManagerOpt = null,
             int speculatedPosition = 0)
-            : base(syntax, owner, rootBinder, containingSemanticModelOpt, parentSemanticModelOpt, speculatedPosition)
+            : base(syntax, owner, rootBinder, containingSemanticModelOpt, parentSemanticModelOpt, snapshotManagerOpt, speculatedPosition)
         {
             Debug.Assert((object)owner != null);
             Debug.Assert(owner.Kind == SymbolKind.Method);
@@ -28,14 +29,14 @@ internal sealed class MethodBodySemanticModel : MemberSemanticModel
         /// Creates a SemanticModel for the method.
         /// </summary>
         internal static MethodBodySemanticModel Create(SyntaxTreeSemanticModel containingSemanticModel, MethodSymbol owner, ExecutableCodeBinder executableCodeBinder,
-                                                       CSharpSyntaxNode syntax, BoundNode boundNode = null)
+                                                       CSharpSyntaxNode syntax, BoundNode boundNode = null, NullableWalker.SnapshotManager snapshotManager = null)
         {
             Debug.Assert(containingSemanticModel != null);
             var result = new MethodBodySemanticModel(owner, executableCodeBinder, syntax, containingSemanticModel);
 
             if (boundNode != null)
             {
-                result.UnguardedAddBoundTreeForStandaloneSyntax(syntax, boundNode);
+                result.UnguardedAddBoundTreeForStandaloneSyntax(syntax, boundNode, snapshotManager);
             }
 
             return result;
@@ -70,14 +71,14 @@ internal override BoundNode Bind(Binder binder, CSharpSyntaxNode node, Diagnosti
         /// <summary>
         /// Creates a speculative SemanticModel for a method body that did not appear in the original source code.
         /// </summary>
-        internal static MethodBodySemanticModel CreateSpeculative(SyntaxTreeSemanticModel parentSemanticModel, MethodSymbol owner, StatementSyntax syntax, Binder rootBinder, int position)
+        internal static MethodBodySemanticModel CreateSpeculative(SyntaxTreeSemanticModel parentSemanticModel, MethodSymbol owner, StatementSyntax syntax, Binder rootBinder, NullableWalker.SnapshotManager snapshotManagerOpt, int position)
         {
             Debug.Assert(parentSemanticModel != null);
             Debug.Assert(syntax != null);
             Debug.Assert(rootBinder != null);
             Debug.Assert(rootBinder.IsSemanticModelBinder);
 
-            return new MethodBodySemanticModel(owner, rootBinder, syntax, parentSemanticModelOpt: parentSemanticModel, speculatedPosition: position);
+            return new MethodBodySemanticModel(owner, rootBinder, syntax, parentSemanticModelOpt: parentSemanticModel, snapshotManagerOpt: snapshotManagerOpt, speculatedPosition: position);
         }
 
         /// <summary>
@@ -137,7 +138,9 @@ private bool GetSpeculativeSemanticModelForMethodBody(SyntaxTreeSemanticModel pa
 
             var executablebinder = new ExecutableCodeBinder(body, methodSymbol, binder ?? this.RootBinder);
             var blockBinder = executablebinder.GetBinder(body).WithAdditionalFlags(GetSemanticModelBinderFlags());
-            speculativeModel = CreateSpeculative(parentModel, methodSymbol, body, blockBinder, position);
+            // We don't pass the snapshot manager along here, because we're speculating about an entirely new body and it should not
+            // be influenced by any existing code in the body.
+            speculativeModel = CreateSpeculative(parentModel, methodSymbol, body, blockBinder, snapshotManagerOpt: null, position);
             return true;
         }
 
@@ -159,7 +162,7 @@ internal override bool TryGetSpeculativeSemanticModelCore(SyntaxTreeSemanticMode
 
             var methodSymbol = (MethodSymbol)this.MemberSymbol;
             binder = new ExecutableCodeBinder(statement, methodSymbol, binder);
-            speculativeModel = CreateSpeculative(parentModel, methodSymbol, statement, binder, position);
+            speculativeModel = CreateSpeculative(parentModel, methodSymbol, statement, binder, SnapshotManager, position);
             return true;
         }
 
@@ -205,9 +208,9 @@ internal override bool TryGetSpeculativeSemanticModelCore(SyntaxTreeSemanticMode
             return false;
         }
 
-        protected override BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, out NullableWalker.SnapshotManager snapshotManager)
+        protected override BoundNode RewriteNullableBoundNodesWithSnapshots(BoundNode boundRoot, Binder binder, DiagnosticBag diagnostics, bool createSnapshots, out NullableWalker.SnapshotManager snapshotManager)
         {
-            return NullableWalker.AnalyzeAndRewrite(Compilation, MemberSymbol, boundRoot, binder, diagnostics, createSnapshots: true, out snapshotManager);
+            return NullableWalker.AnalyzeAndRewrite(Compilation, MemberSymbol, boundRoot, binder, diagnostics, createSnapshots, out snapshotManager);
         }
     }
 }

src/Compilers/CSharp/Portable/Compiler/MethodCompiler.cs

@@ -925,7 +925,7 @@ public override object VisitField(FieldSymbol symbol, TypeCompilationState argum
 
                 // initializers that have been analyzed but not yet lowered.
                 BoundStatementList analyzedInitializers = null;
-                (SyntaxNode Syntax, BoundNode Body, ExecutableCodeBinder Binder) forSemanticModel = default;
+                (SyntaxNode Syntax, BoundNode Body, ExecutableCodeBinder Binder, NullableWalker.SnapshotManager SnapshotManager) forSemanticModel = default;
                 ImportChain importChain = null;
                 var hasTrailingExpression = false;
 
@@ -1045,7 +1045,7 @@ public override object VisitField(FieldSymbol symbol, TypeCompilationState argum
 
                     if (body != null)
                     {
-                        (SyntaxNode Syntax, BoundNode Body, ExecutableCodeBinder Binder) forSemanticModelToUseInLambda = forSemanticModel;
+                        (SyntaxNode Syntax, BoundNode Body, ExecutableCodeBinder Binder, NullableWalker.SnapshotManager SnapshotManager) forSemanticModelToUseInLambda = forSemanticModel;
 
                         lazySemanticModel = new Lazy<SemanticModel>(() =>
                         {
@@ -1058,9 +1058,12 @@ public override object VisitField(FieldSymbol symbol, TypeCompilationState argum
                                                             (rootSyntax) =>
                                                             {
                                                                 Debug.Assert(rootSyntax == forSemanticModelToUseInLambda.Syntax);
-                                                                return MethodBodySemanticModel.Create(semanticModel, methodSymbol,
-                                                                                                      forSemanticModelToUseInLambda.Binder, rootSyntax,
-                                                                                                      forSemanticModelToUseInLambda.Body);
+                                                                return MethodBodySemanticModel.Create(semanticModel,
+                                                                                                      methodSymbol,
+                                                                                                      forSemanticModelToUseInLambda.Binder,
+                                                                                                      rootSyntax,
+                                                                                                      forSemanticModelToUseInLambda.Body,
+                                                                                                      forSemanticModelToUseInLambda.SnapshotManager);
                                                             });
                             }
 
@@ -1598,13 +1601,14 @@ internal static BoundBlock BindMethodBody(MethodSymbol method, TypeCompilationSt
         // NOTE: can return null if the method has no body.
         private static BoundBlock BindMethodBody(MethodSymbol method, TypeCompilationState compilationState, DiagnosticBag diagnostics,
                                                  out ImportChain importChain, out bool originalBodyNested,
-                                                 out (SyntaxNode Syntax, BoundNode Body, ExecutableCodeBinder Binder) forSemanticModel)
+                                                 out (SyntaxNode Syntax, BoundNode Body, ExecutableCodeBinder Binder, NullableWalker.SnapshotManager snapshotManager) forSemanticModel)
         {
             originalBodyNested = false;
             importChain = null;
             forSemanticModel = default;
 
             BoundBlock body;
+            NullableWalker.SnapshotManager snapshotManager = null;
 
             var sourceMethod = method as SourceMemberMethodSymbol;
             if ((object)sourceMethod != null)
@@ -1653,9 +1657,9 @@ internal static BoundBlock BindMethodBody(MethodSymbol method, TypeCompilationSt
                         // also run the nullable walker, and issue duplicate warnings. We should try to only run the pass
                         // once.
                         // https://github.com/dotnet/roslyn/issues/35041
-                        methodBody = NullableWalker.AnalyzeAndRewrite(bodyBinder.Compilation, method, methodBody, bodyBinder, new DiagnosticBag(), createSnapshots: false, out _);
+                        methodBody = NullableWalker.AnalyzeAndRewrite(bodyBinder.Compilation, method, methodBody, bodyBinder, new DiagnosticBag(), createSnapshots: true, out snapshotManager);
                     }
-                    forSemanticModel = (syntaxNode, methodBody, bodyBinder);
+                    forSemanticModel = (syntaxNode, methodBody, bodyBinder, snapshotManager);
 
                     switch (methodBody.Kind)
                     {

src/Compilers/CSharp/Portable/FlowAnalysis/NullableWalker.SnapshotManager.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
+using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics;
@@ -18,56 +19,64 @@ internal sealed class SnapshotManager
             /// </summary>
             private readonly ImmutableArray<SharedWalkerState> _walkerGlobalStates;
             /// <summary>
-            /// The dictionary key corresponds to Syntax position. The dictionary should be sorted in descending order.
+            /// The snapshot array should be sorted in ascending order by the position tuple element in order for the binary search algorithm to
+            /// function correctly.
             /// </summary>
-            private readonly ImmutableSortedDictionary<int, Snapshot> _incrementalSnapshots;
+            private readonly ImmutableArray<(int position, Snapshot snapshot)> _incrementalSnapshots;
 
-            private SnapshotManager(ImmutableArray<SharedWalkerState> walkerGlobalStates, ImmutableSortedDictionary<int, Snapshot> incrementalSnapshots)
+            private static readonly Func<(int position, Snapshot snapshot), int, int> BinarySearchComparer = (current, target) => current.position.CompareTo(target);
+
+            private SnapshotManager(ImmutableArray<SharedWalkerState> walkerGlobalStates, ImmutableArray<(int position, Snapshot snapshot)> incrementalSnapshots)
             {
                 _walkerGlobalStates = walkerGlobalStates;
                 _incrementalSnapshots = incrementalSnapshots;
+
+#if DEBUG
+                Debug.Assert(!incrementalSnapshots.IsDefaultOrEmpty);
+                int previousPosition = incrementalSnapshots[0].position;
+                for (int i = 1; i < incrementalSnapshots.Length; i++)
+                {
+                    int currentPosition = incrementalSnapshots[i].position;
+                    Debug.Assert(currentPosition > previousPosition);
+                    previousPosition = currentPosition;
+                }
+#endif
             }
 
-            internal NullableWalker RestoreWalkerToAnalyzeNewNode(
+            internal (NullableWalker, VariableState) RestoreWalkerToAnalyzeNewNode(
                 int position,
                 BoundNode nodeToAnalyze,
                 Binder binder,
                 ImmutableDictionary<BoundExpression, (NullabilityInfo, TypeSymbol)>.Builder analyzedNullabilityMap)
             {
-                // TODO: Use a more efficient storage and lookup mechanism like an AVL tree here.
-                // https://github.com/dotnet/roslyn/issues/35037
-                Snapshot incrementalSnapshot = default;
-                bool foundSnapshot = false;
-                foreach (var (currentPosition, currentSnapshot) in _incrementalSnapshots)
-                {
-                    if (currentPosition <= position)
-                    {
-                        incrementalSnapshot = currentSnapshot;
-                        foundSnapshot = true;
-                        break;
-                    }
-                }
+                var snapshotPosition = _incrementalSnapshots.BinarySearch(position, BinarySearchComparer);
 
-                if (!foundSnapshot)
+                if (snapshotPosition < 0)
                 {
-                    return null;
+                    // BinarySearch returns the next higher position. Always take the one most close but behind the requested position
+                    snapshotPosition = (~snapshotPosition) - 1;
+
+                    // If there was none in the snapshots before the target position, just take index 0
+                    if (snapshotPosition < 0) snapshotPosition = 0;
                 }
 
+                (_, Snapshot incrementalSnapshot) = _incrementalSnapshots[snapshotPosition];
                 var globalState = _walkerGlobalStates[incrementalSnapshot.SharedStateIndex];
                 var variableState = new VariableState(globalState.VariableSlot, globalState.VariableBySlot, globalState.VariableTypes, incrementalSnapshot.VariableState.Clone());
                 var method = globalState.Symbol as MethodSymbol;
-                return new NullableWalker(
-                    binder.Compilation,
-                    globalState.Symbol,
-                    useMethodSignatureParameterTypes: !(method is null),
-                    method,
-                    nodeToAnalyze,
-                    binder,
-                    binder.Conversions,
-                    variableState,
-                    returnTypesOpt: null,
-                    analyzedNullabilityMap,
-                    snapshotBuilderOpt: null);
+                return (new NullableWalker(binder.Compilation,
+                                           globalState.Symbol,
+                                           useMethodSignatureParameterTypes: !(method is null),
+                                           method,
+                                           nodeToAnalyze,
+                                           binder,
+                                           binder.Conversions,
+                                           variableState,
+                                           returnTypesOpt: null,
+                                           analyzedNullabilityMap,
+                                           snapshotBuilderOpt: null,
+                                           isSpeculative: true),
+                        variableState);
             }
 
 #if DEBUG
@@ -78,20 +87,16 @@ internal void VerifyNode(BoundNode node)
                     return;
                 }
 
-                Debug.Assert(_incrementalSnapshots.ContainsKey(node.Syntax.SpanStart), $"Did not find a snapshot for {node} `{node.Syntax}.`");
-                Debug.Assert(_walkerGlobalStates.Length > _incrementalSnapshots[node.Syntax.SpanStart].SharedStateIndex, $"Did not find global state for {node} `{node.Syntax}`.");
-            }
-#endif
+                int nodePosition = node.Syntax.SpanStart;
+                int position = _incrementalSnapshots.BinarySearch(nodePosition, BinarySearchComparer);
 
-            /// <summary>
-            /// Simple int comparer that orders in descending order, so that searching for a snapshot will find the
-            /// closest position before a given position.
-            /// </summary>
-            private sealed class DescendingIntComparer : IComparer<int>
-            {
-                internal static readonly DescendingIntComparer Singleton = new DescendingIntComparer();
-                public int Compare(int x, int y) => y.CompareTo(x);
+                if (position < 0)
+                {
+                    Debug.Fail($"Did not find a snapshot for {node} `{node.Syntax}.`");
+                }
+                Debug.Assert(_walkerGlobalStates.Length > _incrementalSnapshots[position].snapshot.SharedStateIndex, $"Did not find global state for {node} `{node.Syntax}`.");
             }
+#endif
 
             internal sealed class Builder
             {
@@ -100,7 +105,7 @@ internal sealed class Builder
                 /// <summary>
                 /// Snapshots are kept in a dictionary of position -> snapshot at that position. These are stored in descending order.
                 /// </summary>
-                private readonly ImmutableSortedDictionary<int, Snapshot>.Builder _incrementalSnapshots = ImmutableSortedDictionary.CreateBuilder<int, Snapshot>(DescendingIntComparer.Singleton);
+                private readonly SortedDictionary<int, Snapshot> _incrementalSnapshots = new SortedDictionary<int, Snapshot>();
                 /// <summary>
                 /// Every walker is walking a specific symbol, and can potentially walk each symbol multiple times
                 /// to get to a stable state. Each of these symbols gets a single global state slot, which this
@@ -115,9 +120,9 @@ internal SnapshotManager ToManagerAndFree()
                     Debug.Assert(_symbolToSlot.Count == _walkerStates.Count);
                     Debug.Assert(_symbolToSlot.Count > 0);
                     _symbolToSlot.Free();
-                    return new SnapshotManager(
-                        _walkerStates.ToImmutableAndFree(),
-                        _incrementalSnapshots.ToImmutable());
+                    var snapshotsArray = Roslyn.Utilities.EnumerableExtensions.SelectAsArray<KeyValuePair<int, Snapshot>, (int, Snapshot)>(_incrementalSnapshots, (kvp) => (kvp.Key, kvp.Value));
+
+                    return new SnapshotManager(_walkerStates.ToImmutableAndFree(), snapshotsArray);
                 }
 
                 internal int EnterNewWalker(Symbol symbol)

src/Compilers/CSharp/Portable/FlowAnalysis/NullableWalker.cs

@@ -167,6 +167,12 @@ public VisitArgumentResult(VisitResult visitResult, Optional<LocalState> stateFo
         /// </summary>
         private PooledDictionary<BoundExpression, TypeWithState> _methodGroupReceiverMapOpt;
 
+        /// <summary>
+        /// True if we're analyzing speculative code. This turns off some initialization steps
+        /// that would otherwise be taken.
+        /// </summary>
+        private readonly bool _isSpeculative;
+
 #if DEBUG
         /// <summary>
         /// Contains the expressions that should not be inserted into <see cref="_analyzedNullabilityMapOpt"/>.
@@ -330,7 +336,8 @@ protected override void Free()
             VariableState initialState,
             ArrayBuilder<(BoundReturnStatement, TypeWithAnnotations)> returnTypesOpt,
             ImmutableDictionary<BoundExpression, (NullabilityInfo, TypeSymbol)>.Builder analyzedNullabilityMapOpt,
-            SnapshotManager.Builder snapshotBuilderOpt)
+            SnapshotManager.Builder snapshotBuilderOpt,
+            bool isSpeculative = false)
             // Members of variables are tracked up to a fixed depth, to avoid cycles. The
             // maxSlotDepth value is arbitrary but large enough to allow most scenarios.
             : base(compilation, symbol, node, EmptyStructTypeCache.CreatePrecise(), trackUnassignments: true, maxSlotDepth: 5)
@@ -342,6 +349,7 @@ protected override void Free()
             _analyzedNullabilityMapOpt = analyzedNullabilityMapOpt;
             _returnTypesOpt = returnTypesOpt;
             _snapshotBuilderOpt = snapshotBuilderOpt;
+            _isSpeculative = isSpeculative;
 
             if (initialState != null)
             {
@@ -388,16 +396,19 @@ protected override ImmutableArray<PendingBranch> Scan(ref bool badRegion)
                 _returnTypesOpt.Clear();
             }
             this.Diagnostics.Clear();
-            ParameterSymbol methodThisParameter = MethodThisParameter;
             this.regionPlace = RegionPlace.Before;
-            EnterParameters(); // assign parameters
-            if (!(methodThisParameter is null))
+            if (!_isSpeculative)
             {
-                EnterParameter(methodThisParameter, methodThisParameter.TypeWithAnnotations);
-            }
+                ParameterSymbol methodThisParameter = MethodThisParameter;
+                EnterParameters(); // assign parameters
+                if (!(methodThisParameter is null))
+                {
+                    EnterParameter(methodThisParameter, methodThisParameter.TypeWithAnnotations);
+                }
 
-            // We need to create a snapshot even of the first node, because we want to have the state of the initial parameters.
-            _snapshotBuilderOpt?.TakeIncrementalSnapshot(methodMainNode, State);
+                // We need to create a snapshot even of the first node, because we want to have the state of the initial parameters.
+                _snapshotBuilderOpt?.TakeIncrementalSnapshot(methodMainNode, State);
+            }
 
             ImmutableArray<PendingBranch> pendingReturns = base.Scan(ref badRegion);
             return pendingReturns;
@@ -472,6 +483,28 @@ protected override ImmutableArray<PendingBranch> Scan(ref bool badRegion)
             return new NullabilityRewriter(analyzedNullabilitiesMap).Visit(node);
         }
 
+        internal static BoundNode AnalyzeAndRewriteSpeculation(
+            int position,
+            BoundNode node,
+            Binder binder,
+            SnapshotManager originalSnapshots)
+        {
+            var analyzedNullabilities = ImmutableDictionary.CreateBuilder<BoundExpression, (NullabilityInfo, TypeSymbol)>();
+            var (walker, initialState) = originalSnapshots.RestoreWalkerToAnalyzeNewNode(position, node, binder, analyzedNullabilities);
+            Analyze(walker, symbol: null, diagnostics: null, initialState, snapshotBuilderOpt: null);
+
+            var analyzedNullabilitiesMap = analyzedNullabilities.ToImmutable();
+
+#if DEBUG
+            if (binder.Compilation.NullableAnalysisEnabled)
+            {
+                DebugVerifier.Verify(analyzedNullabilitiesMap, snapshotManagerOpt: null, node);
+            }
+#endif
+
+            return new NullabilityRewriter(analyzedNullabilitiesMap).Visit(node);
+        }
+
         internal static void AnalyzeIfNeeded(
             Binder binder,
             BoundAttribute attribute,
@@ -539,19 +572,28 @@ protected override ImmutableArray<PendingBranch> Scan(ref bool badRegion)
             ArrayBuilder<(BoundReturnStatement, TypeWithAnnotations)> returnTypesOpt)
         {
             Debug.Assert(diagnostics != null);
-            var walker = new NullableWalker(
-                compilation,
-                symbol,
-                useMethodSignatureParameterTypes,
-                methodSignatureOpt,
-                node,
-                binder,
-                conversions,
-                initialState,
-                returnTypesOpt,
-                analyzedNullabilityMapOpt,
-                snapshotBuilderOpt);
+            var walker = new NullableWalker(compilation,
+                                                    symbol,
+                                                    useMethodSignatureParameterTypes,
+                                                    methodSignatureOpt,
+                                                    node,
+                                                    binder,
+                                                    conversions,
+                                                    initialState,
+                                                    returnTypesOpt,
+                                                    analyzedNullabilityMapOpt,
+                                                    snapshotBuilderOpt);
+
+            Analyze(walker, symbol, diagnostics, initialState, snapshotBuilderOpt);
+        }
 
+        private static void Analyze(
+            NullableWalker walker,
+            Symbol symbol,
+            DiagnosticBag diagnostics,
+            VariableState initialState,
+            SnapshotManager.Builder snapshotBuilderOpt)
+        {
             try
             {
                 bool badRegion = false;
@@ -559,7 +601,7 @@ protected override ImmutableArray<PendingBranch> Scan(ref bool badRegion)
                 var previousSlot = snapshotBuilderOpt?.EnterNewWalker(symbol) ?? -1;
                 ImmutableArray<PendingBranch> returns = walker.Analyze(ref badRegion, initialLocalState);
                 snapshotBuilderOpt?.ExitWalker(walker.SaveSharedState(), previousSlot);
-                diagnostics.AddRange(walker.Diagnostics);
+                diagnostics?.AddRange(walker.Diagnostics);
                 Debug.Assert(!badRegion);
             }
             catch (CancelledByStackGuardException ex) when (diagnostics != null)
@@ -575,7 +617,7 @@ protected override ImmutableArray<PendingBranch> Scan(ref bool badRegion)
         private SharedWalkerState SaveSharedState() =>
             new SharedWalkerState(
                 _variableSlot.ToImmutableDictionary(),
-                variableBySlot.ToImmutableArray(),
+                ImmutableArray.Create(variableBySlot, start: 0, length: nextVariableSlot),
                 _variableTypes.ToImmutableDictionary(),
                 _symbol);
 

src/Compilers/CSharp/Test/Symbol/Symbols/Source/NullablePublicAPITests.cs

@@ -706,24 +706,29 @@ void M()
 
             comp.VerifyDiagnostics();
             comp.VerifyAnalyzerDiagnostics(new[] { new NullabilityPrinter() }, null, null, true,
-                Diagnostic("CA9999_NullabilityPrinter", "o").WithArguments("o", "MaybeNull", "Annotated").WithLocation(7, 13),
-                Diagnostic("CA9999_NullabilityPrinter", "o").WithArguments("o", "MaybeNull", "Annotated").WithLocation(8, 13),
-                Diagnostic("CA9999_NullabilityPrinter", "o").WithArguments("o", "NotNull", "NotAnnotated").WithLocation(9, 13));
+                Diagnostic("CA9999_NullabilityPrinter", "o").WithArguments("o", "MaybeNull", "Annotated", "MaybeNull").WithLocation(7, 13),
+                Diagnostic("CA9999_NullabilityPrinter", "o").WithArguments("o", "MaybeNull", "Annotated", "MaybeNull").WithLocation(8, 13),
+                Diagnostic("CA9999_NullabilityPrinter", "o").WithArguments("o", "NotNull", "NotAnnotated", "NotNull").WithLocation(9, 13));
         }
 
         private class NullabilityPrinter : DiagnosticAnalyzer
         {
             public override ImmutableArray<DiagnosticDescriptor> SupportedDiagnostics => ImmutableArray.Create(s_descriptor);
 
-            private static DiagnosticDescriptor s_descriptor = new DiagnosticDescriptor(id: "CA9999_NullabilityPrinter", title: "CA9999_NullabilityPrinter", messageFormat: "Nullability of '{0}' is '{1}':'{2}'", category: "Test", defaultSeverity: DiagnosticSeverity.Warning, isEnabledByDefault: true);
+            private static DiagnosticDescriptor s_descriptor = new DiagnosticDescriptor(id: "CA9999_NullabilityPrinter", title: "CA9999_NullabilityPrinter", messageFormat: "Nullability of '{0}' is '{1}':'{2}'. Speculative flow state is '{3}'", category: "Test", defaultSeverity: DiagnosticSeverity.Warning, isEnabledByDefault: true);
 
             public override void Initialize(AnalysisContext context)
             {
+                var newSource = (ExpressionStatementSyntax)SyntaxFactory.ParseStatement("_ = o;");
+                var oReference = ((AssignmentExpressionSyntax)newSource.Expression).Right;
+
                 context.RegisterSyntaxNodeAction(syntaxContext =>
                 {
                     if (syntaxContext.Node.ToString() == "_") return;
                     var info = syntaxContext.SemanticModel.GetTypeInfo(syntaxContext.Node);
-                    syntaxContext.ReportDiagnostic(CodeAnalysis.Diagnostic.Create(s_descriptor, syntaxContext.Node.GetLocation(), syntaxContext.Node, info.Nullability.FlowState, info.Nullability.Annotation));
+                    Assert.True(syntaxContext.SemanticModel.TryGetSpeculativeSemanticModel(syntaxContext.Node.SpanStart, newSource, out var specModel));
+                    var specInfo = specModel.GetTypeInfo(oReference);
+                    syntaxContext.ReportDiagnostic(CodeAnalysis.Diagnostic.Create(s_descriptor, syntaxContext.Node.GetLocation(), syntaxContext.Node, info.Nullability.FlowState, info.Nullability.Annotation, specInfo.Nullability.FlowState));
                 }, SyntaxKind.IdentifierName);
             }
         }
@@ -876,5 +881,103 @@ void M2(List<C?> l1, List<C> l2)
                 );
             comp.VerifyTypes();
         }
+
+        [Fact]
+        public void SpeculativeSemanticModel_BasicTest()
+        {
+            var source = @"
+class C
+{
+    void M(string? s1)
+    {
+        if (s1 != null)
+        {
+            s1.ToString();
+        }
+
+        s1?.ToString();
+
+        s1 = """";
+        var s2 = s1 == null ? """" : s1;
+    }
+}";
+
+            var comp = CreateCompilation(source, options: WithNonNullTypesTrue(), parseOptions: TestOptions.Regular8WithNullableAnalysis);
+            comp.VerifyDiagnostics();
+
+            var syntaxTree = comp.SyntaxTrees[0];
+            var root = syntaxTree.GetRoot();
+            var model = comp.GetSemanticModel(syntaxTree);
+
+            var ifStatement = root.DescendantNodes().OfType<IfStatementSyntax>().Single();
+            var conditionalAccessExpression = root.DescendantNodes().OfType<ConditionalAccessExpressionSyntax>().Single();
+            var ternary = root.DescendantNodes().OfType<ConditionalExpressionSyntax>().Single();
+
+            var newSource = ((ExpressionStatementSyntax)SyntaxFactory.ParseStatement(@"_ = s1 == """" ? s1 : s1;"));
+            var newTernary = (ConditionalExpressionSyntax)((AssignmentExpressionSyntax)newSource.Expression).Right;
+            var inCondition = ((BinaryExpressionSyntax)newTernary.Condition).Left;
+            var whenTrue = newTernary.WhenTrue;
+            var whenFalse = newTernary.WhenFalse;
+
+            // Before the if statement
+            verifySpeculativeModel(ifStatement.SpanStart, PublicNullableFlowState.MaybeNull);
+
+            // In if statement consequence
+            verifySpeculativeModel(ifStatement.Statement.SpanStart, PublicNullableFlowState.NotNull);
+
+            // Before the conditional access
+            verifySpeculativeModel(conditionalAccessExpression.SpanStart, PublicNullableFlowState.MaybeNull);
+
+            // After the conditional access
+            verifySpeculativeModel(conditionalAccessExpression.WhenNotNull.SpanStart, PublicNullableFlowState.NotNull);
+
+            // In the conditional whenTrue
+            verifySpeculativeModel(ternary.WhenTrue.SpanStart, PublicNullableFlowState.MaybeNull);
+
+            // In the conditional whenFalse
+            verifySpeculativeModel(ternary.WhenFalse.SpanStart, PublicNullableFlowState.NotNull);
+
+            void verifySpeculativeModel(int spanStart, PublicNullableFlowState conditionFlowState)
+            {
+                Assert.True(model.TryGetSpeculativeSemanticModel(spanStart, newSource, out var speculativeModel));
+                var speculativeTypeInfo = speculativeModel.GetTypeInfo(inCondition);
+                Assert.Equal(conditionFlowState, speculativeTypeInfo.Nullability.FlowState);
+                speculativeTypeInfo = speculativeModel.GetTypeInfo(whenTrue);
+                Assert.Equal(PublicNullableFlowState.NotNull, speculativeTypeInfo.Nullability.FlowState);
+                speculativeTypeInfo = speculativeModel.GetTypeInfo(whenFalse);
+                Assert.Equal(conditionFlowState, speculativeTypeInfo.Nullability.FlowState);
+            }
+        }
+
+        [Fact]
+        public void SpeculativeModel_Properties()
+        {
+            var source = @"
+class C
+{
+    object? Foo
+    {
+        get
+        {
+            object? x = null;
+            return x;
+        }
+    }
+}";
+
+            var comp = CreateCompilation(source, options: WithNonNullTypesTrue(), parseOptions: TestOptions.Regular8WithNullableAnalysis);
+            comp.VerifyDiagnostics();
+
+            var syntaxTree = comp.SyntaxTrees[0];
+            var root = syntaxTree.GetRoot();
+            var model = comp.GetSemanticModel(syntaxTree);
+
+            var returnStatement = root.DescendantNodes().OfType<ReturnStatementSyntax>().Single();
+            var newSource = (BlockSyntax)SyntaxFactory.ParseStatement("{ var y = x ?? new object(); y.ToString(); }");
+            var yReference = ((MemberAccessExpressionSyntax)newSource.DescendantNodes().OfType<InvocationExpressionSyntax>().Single().Expression).Expression;
+            Assert.True(model.TryGetSpeculativeSemanticModel(returnStatement.SpanStart, newSource, out var specModel));
+            var speculativeTypeInfo = specModel.GetTypeInfo(yReference);
+            Assert.Equal(PublicNullableFlowState.NotNull, speculativeTypeInfo.Nullability.FlowState);
+        }
     }
 }