From 54c13985e07d13eddd2c3840e23c04f293274913 Mon Sep 17 00:00:00 2001
From: Sam Harwell <Sam.Harwell@microsoft.com>
Date: Fri, 26 Feb 2021 15:17:13 -0800
Subject: [PATCH] Use ImmutableSegmentedDictionary for faster lookup in
 AnalyzerDriver

---
 .../DiagnosticAnalyzer/AnalyzerDriver.cs        | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalyzerDriver.cs b/src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalyzerDriver.cs
index 14f2380c9dfdc..04b651cedfa92 100644
--- a/src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalyzerDriver.cs
+++ b/src/Compilers/Core/Portable/DiagnosticAnalyzer/AnalyzerDriver.cs
@@ -10,6 +10,7 @@
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.CodeAnalysis.Collections;
 using Microsoft.CodeAnalysis.Diagnostics.Telemetry;
 using Microsoft.CodeAnalysis.Operations;
 using Microsoft.CodeAnalysis.PooledObjects;
@@ -158,16 +159,16 @@ private ImmutableDictionary<DiagnosticAnalyzer, SemaphoreSlim> AnalyzerGateMap
             }
         }
 
-        private ImmutableDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags>? _lazyGeneratedCodeAnalysisFlagsMap;
+        private ImmutableSegmentedDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> _lazyGeneratedCodeAnalysisFlagsMap;
 
         /// <summary>
         /// Map from analyzers to their <see cref="GeneratedCodeAnalysisFlags"/> setting. 
         /// </summary>
-        private ImmutableDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> GeneratedCodeAnalysisFlagsMap
+        private ImmutableSegmentedDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> GeneratedCodeAnalysisFlagsMap
         {
             get
             {
-                Debug.Assert(_lazyGeneratedCodeAnalysisFlagsMap != null);
+                Debug.Assert(!_lazyGeneratedCodeAnalysisFlagsMap.IsDefault);
                 return _lazyGeneratedCodeAnalysisFlagsMap;
             }
         }
@@ -532,7 +533,7 @@ private ImmutableHashSet<DiagnosticAnalyzer> ComputeSymbolStartAnalyzers(Immutab
 
         private static bool ComputeShouldSkipAnalysisOnGeneratedCode(
             ImmutableHashSet<DiagnosticAnalyzer> analyzers,
-            ImmutableDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> generatedCodeAnalysisFlagsMap,
+            ImmutableSegmentedDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> generatedCodeAnalysisFlagsMap,
             bool treatAllCodeAsNonGeneratedCode)
         {
             foreach (var analyzer in analyzers)
@@ -549,7 +550,7 @@ private static bool ComputeShouldSkipAnalysisOnGeneratedCode(
         /// <summary>
         /// Returns true if all analyzers need to analyze and report diagnostics in generated code - we can assume all code to be non-generated code.
         /// </summary>
-        private static bool ComputeShouldTreatAllCodeAsNonGeneratedCode(ImmutableHashSet<DiagnosticAnalyzer> analyzers, ImmutableDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> generatedCodeAnalysisFlagsMap)
+        private static bool ComputeShouldTreatAllCodeAsNonGeneratedCode(ImmutableHashSet<DiagnosticAnalyzer> analyzers, ImmutableSegmentedDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> generatedCodeAnalysisFlagsMap)
         {
             foreach (var analyzer in analyzers)
             {
@@ -570,7 +571,7 @@ private bool ShouldSkipAnalysisOnGeneratedCode(DiagnosticAnalyzer analyzer)
 
         private static bool ShouldSkipAnalysisOnGeneratedCode(
             DiagnosticAnalyzer analyzer,
-            ImmutableDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> generatedCodeAnalysisFlagsMap,
+            ImmutableSegmentedDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags> generatedCodeAnalysisFlagsMap,
             bool treatAllCodeAsNonGeneratedCode)
         {
             if (treatAllCodeAsNonGeneratedCode)
@@ -2163,13 +2164,13 @@ private static async Task<ImmutableDictionary<DiagnosticAnalyzer, SemaphoreSlim>
             return builder.ToImmutable();
         }
 
-        private static async Task<ImmutableDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags>> CreateGeneratedCodeAnalysisFlagsMapAsync(
+        private static async Task<ImmutableSegmentedDictionary<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags>> CreateGeneratedCodeAnalysisFlagsMapAsync(
             ImmutableHashSet<DiagnosticAnalyzer> analyzers,
             AnalyzerManager analyzerManager,
             AnalyzerExecutor analyzerExecutor,
             SeverityFilter severityFilter)
         {
-            var builder = ImmutableDictionary.CreateBuilder<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags>();
+            var builder = ImmutableSegmentedDictionary.CreateBuilder<DiagnosticAnalyzer, GeneratedCodeAnalysisFlags>();
             foreach (var analyzer in analyzers)
             {
                 Debug.Assert(!IsDiagnosticAnalyzerSuppressed(analyzer, analyzerExecutor.Compilation.Options, analyzerManager, analyzerExecutor, severityFilter));