Document state storage

src/Compilers/Core/CodeAnalysisTest/Collections/ImmutableSegmentedDictionaryTest.cs

@@ -15,6 +15,8 @@
 using System.Linq;
 using System.Reflection;
 using Microsoft.CodeAnalysis.Collections;
+using Roslyn.Test.Utilities;
+using Roslyn.Utilities;
 using Xunit;
 
 namespace Microsoft.CodeAnalysis.UnitTests.Collections
@@ -299,6 +301,64 @@ public void Indexer_KeyNotFoundException_ContainsKeyInMessage()
             Assert.Contains("'c'", exception.Message);
         }
 
+        [Fact]
+        public void Ordering_NoRemoval()
+        {
+            var map = Empty<int, int>().AddRange(Enumerable.Range(0, 100).Select(i => KeyValuePairUtil.Create(i, i)));
+            AssertEx.Equal(Enumerable.Range(0, 100), map.Select(entry => entry.Value));
+
+            for (int i = 0; i < map.Count; i++)
+            {
+                Assert.Equal(i, map.GetAddedEntry(i).Key);
+                Assert.Equal(i, map.GetAddedEntry(i).Value);
+            }
+        }
+
+        [Fact]
+        public void Ordering_AfterRemoval()
+        {
+            var map = Empty<int, int>()
+                .AddRange(Enumerable.Range(0, 10).Select(i => KeyValuePairUtil.Create(i, i)))
+                .Remove(3);
+
+            Assert.Throws<InvalidOperationException>(() => map.GetAddedEntry(0));
+        }
+
+        [Theory]
+        [InlineData(6, 3)]
+        [InlineData(3, 1)]
+        public void Ordering_AfterRemovalAndCompact(int count, int removeKey)
+        {
+            var builder = Empty<int, int>()
+                .AddRange(Enumerable.Range(0, count).Select(i => KeyValuePairUtil.Create(i, i)))
+                .ToBuilder();
+
+            builder.Remove(removeKey);
+            builder.Compact();
+
+            var map = builder.ToImmutable();
+            for (int i = 0; i < map.Count; i++)
+            {
+                var e = (i < removeKey) ? i : i + 1;
+                Assert.Equal(e, map.GetAddedEntry(i).Key);
+                Assert.Equal(e, map.GetAddedEntry(i).Value);
+            }
+
+            // test some keys that collide
+            for (int i = 0; i < 2 * count; i++)
+            {
+                if (i == removeKey || i >= count)
+                {
+                    Assert.False(map.ContainsKey(i));
+                }
+                else
+                {
+                    Assert.True(map.TryGetValue(i, out var value));
+                    Assert.Equal(i, value);
+                }
+            }
+        }
+
         protected override IImmutableDictionary<TKey, TValue> Empty<TKey, TValue>()
         {
             return ImmutableSegmentedDictionaryTest.Empty<TKey, TValue>();

src/Dependencies/Collections/ImmutableSegmentedDictionary`2+Builder.cs

@@ -181,6 +181,9 @@ public bool Remove(KeyValuePair<TKey, TValue> item)
                 return true;
             }
 
+            public void Compact()
+                => GetOrCreateMutableDictionary().Compact();
+
             public void RemoveRange(IEnumerable<TKey> keys)
             {
                 if (keys is null)

src/Dependencies/Collections/ImmutableSegmentedDictionary`2.cs

@@ -48,6 +48,10 @@ namespace Microsoft.CodeAnalysis.Collections
     /// 
     /// <para>This type is backed by segmented arrays to avoid using the Large Object Heap without impacting algorithmic
     /// complexity.</para>
+    /// 
+    /// <para>The enumerator of the dictionary yields entries in the order they were added to the dictionary provided that 
+    /// no entry has been removed since the dictionary was created.
+    /// </para>
     /// </remarks>
     /// <typeparam name="TKey">The type of the keys in the dictionary.</typeparam>
     /// <typeparam name="TValue">The type of the values in the dictionary.</typeparam>
@@ -129,6 +133,13 @@ TValue IDictionary<TKey, TValue>.this[TKey key]
             set => throw new NotSupportedException();
         }
 
+        /// <summary>
+        /// Gets <paramref name="index"/>-th entry added to the dictionary.
+        /// </summary>
+        /// <exception cref="InvalidOperationException">An entry has been removed from the dictionary.</exception>
+        public KeyValuePair<TKey, TValue> GetAddedEntry(int index)
+            => _dictionary.GetAddedEntry(index);
+
         public static bool operator ==(ImmutableSegmentedDictionary<TKey, TValue> left, ImmutableSegmentedDictionary<TKey, TValue> right)
             => left.Equals(right);
 
@@ -234,7 +245,7 @@ public ImmutableSegmentedDictionary<TKey, TValue> SetItem(TKey key, TValue value
                 return self;
             }
 
-            var dictionary = new SegmentedDictionary<TKey, TValue>(self._dictionary, self._dictionary.Comparer);
+            var dictionary = new SegmentedDictionary<TKey, TValue>(self._dictionary);
             dictionary[key] = value;
             return new ImmutableSegmentedDictionary<TKey, TValue>(dictionary);
         }

src/Dependencies/Collections/SegmentedArray`1.cs

@@ -121,6 +121,9 @@ T IList<T>.this[int index]
         }
 
         public object Clone()
+            => CloneImpl();
+
+        internal SegmentedArray<T> CloneImpl()
         {
             var items = (T[][])_items.Clone();
             for (var i = 0; i < items.Length; i++)

src/Dependencies/Collections/SegmentedDictionary`2.cs

@@ -24,6 +24,9 @@ namespace Microsoft.CodeAnalysis.Collections
     /// <remarks>
     /// <para>This collection has the same performance characteristics as <see cref="Dictionary{TKey, TValue}"/>, but
     /// uses segmented arrays to avoid allocations in the Large Object Heap.</para>
+    /// 
+    /// <para>The enumerator of the dictionary yields entries in the order they were added to the dictionary provided that 
+    /// no entry has been removed since the dictionary was created.</para>
     /// </remarks>
     /// <typeparam name="TKey">The type of the keys in the dictionary.</typeparam>
     /// <typeparam name="TValue">The type of the values in the dictionary.</typeparam>
@@ -35,15 +38,41 @@ internal sealed class SegmentedDictionary<TKey, TValue> : IDictionary<TKey, TVal
         private SegmentedArray<int> _buckets;
         private SegmentedArray<Entry> _entries;
         private ulong _fastModMultiplier;
+
+        /// <summary>
+        /// Number of used slots.
+        /// </summary>
         private int _count;
+
         private int _freeList;
+
+        /// <summary>
+        /// Number of used slots that represent a removed entry.
+        /// </summary>
         private int _freeCount;
+
         private int _version;
         private readonly IEqualityComparer<TKey>? _comparer;
         private KeyCollection? _keys;
         private ValueCollection? _values;
         private const int StartOfFreeList = -3;
 
+        // copy constructor
+        internal SegmentedDictionary(SegmentedDictionary<TKey, TValue> other)
+        {
+            _comparer = other._comparer;
+
+            if (other._buckets.Length > 0)
+            {
+                _buckets = other._buckets.CloneImpl();
+                _entries = other._entries.CloneImpl();
+                _fastModMultiplier = other._fastModMultiplier;
+                _count = other._count;
+                _freeList = other._freeList;
+                _freeCount = other._freeCount;
+            }
+        }
+
         public SegmentedDictionary()
             : this(0, null)
         {
@@ -176,6 +205,30 @@ public TValue this[TKey key]
             }
         }
 
+        /// <summary>
+        /// Gets <paramref name="index"/>-th entry added to the dictionary.
+        /// </summary>
+        /// <remarks>
+        /// Requires that no entry has been removed from the dictionary since the dictionary was created or
+        /// <see cref="Compact()"/> was called.
+        /// </remarks>
+        /// <exception cref="InvalidOperationException">An entry has been removed from the dictionary.</exception>
+        public KeyValuePair<TKey, TValue> GetAddedEntry(int index)
+        {
+            if (_freeCount > 0)
+            {
+                throw new InvalidOperationException();
+            }
+
+            if (index < 0 || index >= _count)
+            {
+                ThrowHelper.ThrowArgumentOutOfRange_IndexException();
+            }
+
+            ref var entry = ref _entries[index];
+            return new(entry._key, entry._value);
+        }
+
         public void Add(TKey key, TValue value)
         {
             var modified = TryInsert(key, value, InsertionBehavior.ThrowOnExisting);
@@ -945,6 +998,15 @@ public void TrimExcess()
         /// once it is known that no new elements will be added.
         /// </remarks>
         public void TrimExcess(int capacity)
+            => TrimExcess(capacity, forceCompact: false);
+
+        /// <summary>
+        /// Compacts the underlying storage.
+        /// </summary>
+        public void Compact()
+            => TrimExcess(Count, forceCompact: true);
+
+        private void TrimExcess(int capacity, bool forceCompact)
         {
             if (capacity < Count)
             {
@@ -954,24 +1016,31 @@ public void TrimExcess(int capacity)
             var newSize = HashHelpers.GetPrime(capacity);
             var oldEntries = _entries;
             var currentCapacity = oldEntries.Length;
-            if (newSize >= currentCapacity)
+            if (newSize >= currentCapacity && !forceCompact)
             {
                 return;
             }
 
             var oldCount = _count;
             _version++;
-            Initialize(newSize);
+            if (newSize < currentCapacity)
+            {
+                Initialize(newSize);
+            }
+            else
+            {
+                ((IList)_buckets).Clear();
+            }
+
             var entries = _entries;
             var count = 0;
             for (var i = 0; i < oldCount; i++)
             {
-                var hashCode = oldEntries[i]._hashCode; // At this point, we know we have entries.
                 if (oldEntries[i]._next >= -1)
                 {
                     ref var entry = ref entries[count];
                     entry = oldEntries[i];
-                    ref var bucket = ref GetBucket(hashCode);
+                    ref var bucket = ref GetBucket(entry._hashCode);
                     entry._next = bucket - 1; // Value in _buckets is 1-based
                     bucket = count + 1;
                     count++;
@@ -980,6 +1049,7 @@ public void TrimExcess(int capacity)
 
             _count = count;
             _freeCount = 0;
+            _freeList = -1;
         }
 
         bool ICollection.IsSynchronized => false;

src/EditorFeatures/CSharpTest/Workspaces/WorkspaceTests.cs

@@ -220,14 +220,14 @@ public async Task TestAddedSubmissionParseTreeHasEmptyFilePath()
             // Check that a parse tree for a submission has an empty file path.
             var tree1 = await workspace.CurrentSolution
                 .GetProjectState(project1.Id)
-                .GetDocumentState(document1.Id)
+                .DocumentStates.GetState(document1.Id)
                 .GetSyntaxTreeAsync(CancellationToken.None);
             Assert.Equal("", tree1.FilePath);
 
             // Check that a parse tree for a script does not have an empty file path.
             var tree2 = await workspace.CurrentSolution
                 .GetProjectState(project2.Id)
-                .GetDocumentState(document2.Id)
+                .DocumentStates.GetState(document2.Id)
                 .GetSyntaxTreeAsync(CancellationToken.None);
             Assert.Equal("a.csx", tree2.FilePath);
         }
@@ -821,7 +821,7 @@ public async Task TestAnalyzerConfigFile_Properties()
 
             Assert.Equal(1, project.Documents.Count());
             Assert.Equal(1, project.AnalyzerConfigDocuments.Count());
-            Assert.Equal(1, project.State.AnalyzerConfigDocumentIds.Count());
+            Assert.Equal(1, project.State.AnalyzerConfigDocumentStates.Count);
 
             var doc = project.GetDocument(analyzerConfigDoc.Id);
             Assert.Null(doc);

src/Features/Core/Portable/SolutionCrawler/WorkCoordinator.cs

@@ -792,7 +792,7 @@ public int GetDocumentCount(Solution solution)
                 {
                     foreach (var projectState in solution.State.ProjectStates)
                     {
-                        count += projectState.Value.DocumentIds.Count;
+                        count += projectState.Value.DocumentStates.Count;
                     }
 
                     return count;

src/Tools/IdeCoreBenchmarks/ProjectOperationBenchmarks.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Linq;
 using System.Text;
 using BenchmarkDotNet.Attributes;
 using Microsoft.CodeAnalysis;
@@ -13,6 +14,8 @@ namespace IdeCoreBenchmarks
 {
     public class ProjectOperationBenchmarks
     {
+        private static readonly SourceText s_newText = SourceText.From("text");
+
         [MemoryDiagnoser]
         public class IterateDocuments
         {
@@ -21,6 +24,7 @@ public class IterateDocuments
             private Project _hundredProject;
             private Project _thousandsProject;
 
+
             public IterateDocuments()
             {
                 // These fields are initialized in GlobalSetup
@@ -97,6 +101,17 @@ public int Documents()
 
                 return count;
             }
+
+            [Benchmark(Description = "Solution.WithDocumentText")]
+            public void WithDocumentText()
+            {
+                var solution = Project.Solution;
+                var documentId = Project.DocumentIds.FirstOrDefault();
+                if (documentId != null)
+                {
+                    var _ = solution.WithDocumentText(documentId, s_newText);
+                }
+            }
         }
     }
 }

src/VisualStudio/Core/Def/Implementation/ProjectSystem/VisualStudioProject.cs

@@ -234,7 +234,7 @@ private static void TryReportCompilationThrownAway(SolutionState solutionState,
             // We log the number of syntax trees that have been parsed even if there was no compilation created yet
             var projectState = solutionState.GetRequiredProjectState(projectId);
             var parsedTrees = 0;
-            foreach (var documentState in projectState.DocumentStates.Values)
+            foreach (var documentState in projectState.DocumentStates.States)
             {
                 if (documentState.TryGetSyntaxTree(out _))
                 {

src/VisualStudio/Core/Def/Implementation/ProjectSystem/VisualStudioWorkspaceImpl.RemoveAnalyzerConfigDocumentUndoUnit.cs

@@ -19,7 +19,7 @@ public RemoveAnalyzerConfigDocumentUndoUnit(
             }
 
             protected override IReadOnlyList<DocumentId> GetDocumentIds(Project fromProject)
-                => fromProject.State.AnalyzerConfigDocumentIds.AsImmutable();
+                => fromProject.State.AnalyzerConfigDocumentStates.Ids;
 
             protected override TextDocument? GetDocument(Solution currentSolution)
                 => currentSolution.GetAnalyzerConfigDocument(this.DocumentId);

src/Workspaces/Core/Portable/FindSymbols/SymbolTree/SymbolTreeInfo_Source.cs

@@ -72,11 +72,9 @@ private static async Task<Checksum> ComputeSourceSymbolsChecksumAsync(ProjectSta
             var serializer = projectState.LanguageServices.WorkspaceServices.GetService<ISerializerService>();
             var projectStateChecksums = await projectState.GetStateChecksumsAsync(cancellationToken).ConfigureAwait(false);
 
-            // Order the documents by FilePath.  Default ordering in the RemoteWorkspace is
-            // to be ordered by Guid (which is not consistent across VS sessions).
-            var textChecksumsTasks = projectState.DocumentStates.OrderBy(d => d.Value.FilePath, StringComparer.Ordinal).Select(async d =>
+            var textChecksumsTasks = projectState.DocumentStates.States.Select(async state =>
             {
-                var documentStateChecksum = await d.Value.GetStateChecksumsAsync(cancellationToken).ConfigureAwait(false);
+                var documentStateChecksum = await state.GetStateChecksumsAsync(cancellationToken).ConfigureAwait(false);
                 return documentStateChecksum.Text;
             });
 

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

@@ -36,13 +36,13 @@ public IEnumerator<Checksum> GetEnumerator()
         IEnumerator IEnumerable.GetEnumerator()
             => GetEnumerator();
 
-        internal static async Task FindAsync<TKey, TValue>(
-            ImmutableSortedDictionary<TKey, TValue> documentStates,
+        internal static async Task FindAsync<TState>(
+            TextDocumentStates<TState> documentStates,
             HashSet<Checksum> searchingChecksumsLeft,
             Dictionary<Checksum, object> result,
-            CancellationToken cancellationToken) where TValue : TextDocumentState
+            CancellationToken cancellationToken) where TState : TextDocumentState
         {
-            foreach (var (_, state) in documentStates)
+            foreach (var state in documentStates.States)
             {
                 Contract.ThrowIfFalse(state.TryGetStateChecksums(out var stateChecksums));