Follow up feedback, optimization to reuse semantic models

src/Bicep.Core/Semantics/Compilation.cs

@@ -6,6 +6,7 @@
 using System.Linq;
 using Bicep.Core.Configuration;
 using Bicep.Core.Diagnostics;
+using Bicep.Core.Extensions;
 using Bicep.Core.TypeSystem;
 using Bicep.Core.Workspaces;
 
@@ -15,18 +16,24 @@ public class Compilation
     {
         private readonly ImmutableDictionary<ISourceFile, Lazy<ISemanticModel>> lazySemanticModelLookup;
 
-        public Compilation(IResourceTypeProvider resourceTypeProvider, SourceFileGrouping sourceFileGrouping)
+        public Compilation(IResourceTypeProvider resourceTypeProvider, SourceFileGrouping sourceFileGrouping, Compilation? previousCompilation = null)
         {
+            var parentFiles = (previousCompilation is not null) ? 
+                previousCompilation.SourceFileGrouping.GetFilesDependingOn(sourceFileGrouping.EntryPoint) :
+                ImmutableHashSet.Create<ISourceFile>();
+
             this.SourceFileGrouping = sourceFileGrouping;
             this.ResourceTypeProvider = resourceTypeProvider;
             this.lazySemanticModelLookup = sourceFileGrouping.SourceFiles.ToImmutableDictionary(
                 sourceFile => sourceFile,
-                sourceFile => new Lazy<ISemanticModel>(() => sourceFile switch
-                {
-                    BicepFile bicepFile => new SemanticModel(this, bicepFile, SourceFileGrouping.FileResolver),
-                    ArmTemplateFile armTemplateFile => new ArmTemplateSemanticModel(armTemplateFile),
-                    _ => throw new ArgumentOutOfRangeException(nameof(sourceFile)),
-                }));
+                sourceFile => 
+                    (!parentFiles.Contains(sourceFile) ? previousCompilation?.lazySemanticModelLookup.TryGetValue(sourceFile) : null) ??
+                    new(() => sourceFile switch
+                    {
+                        BicepFile bicepFile => new SemanticModel(this, bicepFile, SourceFileGrouping.FileResolver),
+                        ArmTemplateFile armTemplateFile => new ArmTemplateSemanticModel(armTemplateFile),
+                        _ => throw new ArgumentOutOfRangeException(nameof(sourceFile)),
+                    }));
         }
 
         public SourceFileGrouping SourceFileGrouping { get; }

src/Bicep.Core/Semantics/SemanticModel.cs

@@ -3,6 +3,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.Immutable;
+using System.Diagnostics;
 using System.Linq;
 using Bicep.Core.Analyzers.Linter;
 using Bicep.Core.Configuration;
@@ -29,6 +30,8 @@ public class SemanticModel : ISemanticModel
 
         public SemanticModel(Compilation compilation, BicepFile sourceFile, IFileResolver fileResolver)
         {
+            Trace.Write($"Reloading {sourceFile.FileUri}");
+
             Compilation = compilation;
             SourceFile = sourceFile;
             FileResolver = fileResolver;

src/Bicep.Core/Workspaces/SourceFileGrouping.cs

@@ -1,26 +1,36 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
+using System.Collections.Generic;
 using System.Collections.Immutable;
 using System.Diagnostics.CodeAnalysis;
 using Bicep.Core.Diagnostics;
 using Bicep.Core.FileSystem;
 using Bicep.Core.Syntax;
+using System.Linq;
 using static Bicep.Core.Diagnostics.DiagnosticBuilder;
 
 namespace Bicep.Core.Workspaces
 {
     public class SourceFileGrouping
     {
         private readonly ImmutableDictionary<ModuleDeclarationSyntax, ISourceFile> sourceFilesByModuleDeclaration;
+        private readonly ImmutableDictionary<ISourceFile, ImmutableHashSet<ISourceFile>> sourceFileParentLookup;
 
         private readonly ImmutableDictionary<ModuleDeclarationSyntax, ErrorBuilderDelegate> errorBuildersByModuleDeclaration;
 
-        public SourceFileGrouping(IFileResolver fileResolver, BicepFile entryPoint, ImmutableHashSet<ISourceFile> sourceFiles, ImmutableDictionary<ModuleDeclarationSyntax, ISourceFile> sourceFilesByModuleDeclaration, ImmutableDictionary<ModuleDeclarationSyntax, ErrorBuilderDelegate> errorBuildersByModuleDeclaration)
+        public SourceFileGrouping(
+            IFileResolver fileResolver,
+            BicepFile entryPoint,
+            ImmutableHashSet<ISourceFile> sourceFiles,
+            ImmutableDictionary<ModuleDeclarationSyntax, ISourceFile> sourceFilesByModuleDeclaration,
+            ImmutableDictionary<ISourceFile, ImmutableHashSet<ISourceFile>> sourceFileParentLookup,
+            ImmutableDictionary<ModuleDeclarationSyntax, ErrorBuilderDelegate> errorBuildersByModuleDeclaration)
         {
             this.FileResolver = fileResolver;
             this.EntryPoint = entryPoint;
             this.SourceFiles = sourceFiles;
             this.sourceFilesByModuleDeclaration = sourceFilesByModuleDeclaration;
+            this.sourceFileParentLookup = sourceFileParentLookup;
             this.errorBuildersByModuleDeclaration = errorBuildersByModuleDeclaration;
         }
         public IFileResolver FileResolver { get; }
@@ -32,6 +42,26 @@ public SourceFileGrouping(IFileResolver fileResolver, BicepFile entryPoint, Immu
         public ISourceFile LookUpModuleSourceFile(ModuleDeclarationSyntax moduleDeclaration) =>
             this.sourceFilesByModuleDeclaration[moduleDeclaration];
 
+        public ImmutableHashSet<ISourceFile> GetFilesDependingOn(ISourceFile sourceFile)
+        {
+            var filesToCheck = new Queue<ISourceFile>(new [] { sourceFile });
+            var knownFiles = new HashSet<ISourceFile>(new [] { sourceFile });
+
+            while (filesToCheck.TryDequeue(out var current))
+            {
+                if (!knownFiles.Contains(current))
+                {
+                    foreach (var parent in sourceFileParentLookup[current].Where(x => !knownFiles.Contains(x)))
+                    {
+                        knownFiles.Add(parent);
+                        filesToCheck.Enqueue(parent);
+                    }
+                }
+            }
+
+            return knownFiles.ToImmutableHashSet();
+        }
+
         public bool TryLookUpModuleErrorDiagnostic(ModuleDeclarationSyntax moduleDeclaration, [NotNullWhen(true)] out ErrorDiagnostic? errorDiagnostic)
         {
             if (this.errorBuildersByModuleDeclaration.TryGetValue(moduleDeclaration, out var errorBuilder))

src/Bicep.Core/Workspaces/SourceFileGroupingBuilder.cs

@@ -59,13 +59,14 @@ private SourceFileGrouping Build(Uri entryFileUri)
                 throw new InvalidOperationException($"Expected {nameof(PopulateRecursive)} to return a BicepFile.");
             }
 
-            this.ReportFailuresForCycles();
+            var sourceFileParentLookup = this.ReportFailuresForCycles();
 
             return new SourceFileGrouping(
                 fileResolver,
                 entryPoint,
                 sourceFilesByUri.Values.ToImmutableHashSet(),
                 sourceFilesByModuleDeclaration.ToImmutableDictionary(),
+                sourceFileParentLookup.ToImmutableDictionary(x => x.Key, x => x.ToImmutableHashSet()),
                 errorBuildersByModuleDeclaration.ToImmutableDictionary());
         }
 
@@ -245,31 +246,33 @@ public static bool ValidateFilePath(string pathName, [NotNullWhen(false)] out Er
             return moduleUri;
         }
 
-        private void ReportFailuresForCycles()
+        private ILookup<ISourceFile, ISourceFile> ReportFailuresForCycles()
         {
             var sourceFileGraph = this.sourceFilesByUri.Values
                 .SelectMany(sourceFile => GetModuleDeclarations(sourceFile)
                     .Where(this.sourceFilesByModuleDeclaration.ContainsKey)
                     .Select(moduleDeclaration => this.sourceFilesByModuleDeclaration[moduleDeclaration])
                     .Distinct()
                     .Select(referencedFile => (sourceFile, referencedFile)))
-                .ToLookup(x => x.sourceFile, x => x.referencedFile);
+                .ToLookup(x => x.referencedFile, x => x.sourceFile);
 
             var cycles = CycleDetector<ISourceFile>.FindCycles(sourceFileGraph);
-            foreach (var kvp in sourceFilesByModuleDeclaration)
+            foreach (var (moduleDeclaration, moduleSourceFile) in sourceFilesByModuleDeclaration)
             {
-                if (cycles.TryGetValue(kvp.Value, out var cycle))
+                if (cycles.TryGetValue(moduleSourceFile, out var cycle))
                 {
                     if (cycle.Length == 1)
                     {
-                        errorBuildersByModuleDeclaration[kvp.Key] = x => x.CyclicModuleSelfReference();
+                        errorBuildersByModuleDeclaration[moduleDeclaration] = x => x.CyclicModuleSelfReference();
                     }
                     else
                     {
-                        errorBuildersByModuleDeclaration[kvp.Key] = x => x.CyclicModule(cycle.Select(x => x.FileUri.LocalPath));
+                        errorBuildersByModuleDeclaration[moduleDeclaration] = x => x.CyclicModule(cycle.Reverse().Select(x => x.FileUri.LocalPath));
                     }
                 }
             }
+
+            return sourceFileGraph;
         }
 
         private static IEnumerable<ModuleDeclarationSyntax> GetModuleDeclarations(ISourceFile sourceFile) => sourceFile is BicepFile bicepFile

src/Bicep.LangServer/BicepCompilationManager.cs

@@ -7,6 +7,7 @@
 using System.Linq;
 using Bicep.Core;
 using Bicep.Core.Extensions;
+using Bicep.Core.FileSystem;
 using Bicep.Core.Workspaces;
 using Bicep.LanguageServer.CompilationManager;
 using Bicep.LanguageServer.Extensions;
@@ -24,15 +25,17 @@ public class BicepCompilationManager : ICompilationManager
         private readonly IWorkspace workspace;
         private readonly ILanguageServerFacade server;
         private readonly ICompilationProvider provider;
+        private readonly IFileResolver fileResolver;
 
         // represents compilations of open bicep files
         private readonly ConcurrentDictionary<DocumentUri, CompilationContext> activeContexts = new ConcurrentDictionary<DocumentUri, CompilationContext>();
 
-        public BicepCompilationManager(ILanguageServerFacade server, ICompilationProvider provider, IWorkspace workspace)
+        public BicepCompilationManager(ILanguageServerFacade server, ICompilationProvider provider, IWorkspace workspace, IFileResolver fileResolver)
         {
             this.server = server;
             this.provider = provider;
             this.workspace = workspace;
+            this.fileResolver = fileResolver;
         }
 
         public void UpsertCompilation(DocumentUri documentUri, int? version, string fileContents, string? languageId = null)
@@ -144,12 +147,14 @@ private ImmutableArray<ISourceFile> CloseCompilationInternal(DocumentUri documen
         {
             try
             {
-                var context = this.provider.Create(workspace, documentUri);
+                var sourceFileGrouping = SourceFileGroupingBuilder.Build(fileResolver, workspace, documentUri.ToUri());
 
-                var output = workspace.UpsertSourceFiles(context.Compilation.SourceFileGrouping.SourceFiles);
+                var context = this.activeContexts.AddOrUpdate(
+                    documentUri,
+                    (documentUri) => this.provider.Create(sourceFileGrouping, documentUri),
+                    (documentUri, prevContext) => this.provider.Update(sourceFileGrouping, documentUri, prevContext));
 
-                // there shouldn't be concurrent upsert requests (famous last words...), so a simple overwrite should be sufficient
-                this.activeContexts[documentUri] = context;
+                var output = workspace.UpsertSourceFiles(sourceFileGrouping.SourceFiles);
 
                 // convert all the diagnostics to LSP diagnostics
                 var diagnostics = GetDiagnosticsFromContext(context).ToDiagnostics(context.LineStarts);

src/Bicep.LangServer/Providers/BicepCompilationProvider.cs

@@ -24,12 +24,24 @@ public BicepCompilationProvider(IResourceTypeProvider resourceTypeProvider, IFil
             this.fileResolver = fileResolver;
         }
 
-        public CompilationContext Create(IReadOnlyWorkspace workspace, DocumentUri documentUri)
+        public CompilationContext Create(SourceFileGrouping sourceFileGrouping, DocumentUri documentUri)
         {
-            var sourceFileGrouping = SourceFileGroupingBuilder.Build(fileResolver, workspace, documentUri.ToUri());
             var compilation = new Compilation(resourceTypeProvider, sourceFileGrouping);
 
             return new CompilationContext(compilation);
         }
+
+        public CompilationContext Update(SourceFileGrouping sourceFileGrouping, DocumentUri documentUri, CompilationContext prevContext)
+        {
+            if (prevContext.Compilation is { SourceFileGrouping: {} previousGrouping } &&
+                sourceFileGrouping.SourceFiles.Count == sourceFileGrouping.SourceFiles.Intersect(previousGrouping.SourceFiles).Count)
+            {
+                return prevContext;
+            }
+
+            var compilation = new Compilation(resourceTypeProvider, sourceFileGrouping, prevContext.Compilation);
+
+            return new CompilationContext(compilation);
+        }
     }
 }

src/Bicep.LangServer/Providers/ICompilationProvider.cs

@@ -9,6 +9,8 @@ namespace Bicep.LanguageServer.Providers
 {
     public interface ICompilationProvider
     {
-        CompilationContext Create(IReadOnlyWorkspace workspace, DocumentUri documentUri);
+        CompilationContext Create(SourceFileGrouping sourceFileGrouping, DocumentUri documentUri);
+
+        CompilationContext Update(SourceFileGrouping sourceFileGrouping, DocumentUri documentUri, CompilationContext prevContext);
     }
 }

src/Bicep.LangServer/Server.cs

@@ -5,6 +5,7 @@
 using System.IO.Pipelines;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Diagnostics;
 using Bicep.Core.Emit;
 using Bicep.Core.FileSystem;
 using Bicep.Core.TypeSystem;
@@ -17,8 +18,11 @@
 using Bicep.LanguageServer.Snippets;
 using Bicep.LanguageServer.Telemetry;
 using Microsoft.Extensions.DependencyInjection;
+using OmniSharp.Extensions.LanguageServer.Protocol.Window;
 using OmniSharp.Extensions.LanguageServer.Server;
 using OmnisharpLanguageServer = OmniSharp.Extensions.LanguageServer.Server.LanguageServer;
+using OmniSharp.Extensions.LanguageServer.Protocol.Server;
+using Bicep.LanguageServer.Utils;
 
 namespace Bicep.LanguageServer
 {
@@ -71,12 +75,16 @@ private Server(CreationOptions creationOptions, Action<LanguageServerOptions> on
 
                 onOptionsFunc(options);
             });
+
+            Trace.Listeners.Add(new ServerLogTraceListener(server));
         }
 
         public async Task RunAsync(CancellationToken cancellationToken)
         {
             await server.Initialize(cancellationToken);
 
+            server.LogInfo($"Running on processId {Environment.ProcessId}");
+
             await server.WaitForExit;
         }
 

src/Bicep.LangServer/Snippets/SnippetsProvider.cs

@@ -212,6 +212,7 @@ private ImmutableDictionary<DeclaredSymbol, ImmutableHashSet<ResourceDependency>
                 bicepFile,
                 ImmutableHashSet.Create<ISourceFile>(bicepFile),
                 ImmutableDictionary.Create<ModuleDeclarationSyntax, ISourceFile>(),
+                ImmutableDictionary.Create<ISourceFile, ImmutableHashSet<ISourceFile>>(),
                 ImmutableDictionary.Create<ModuleDeclarationSyntax, DiagnosticBuilder.ErrorBuilderDelegate>());
 
             Compilation compilation = new Compilation(AzResourceTypeProvider.CreateWithAzTypes(), sourceFileGrouping);

src/Bicep.LangServer/Utils/ServerLogTraceListener.cs

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+using System.Diagnostics;
+using OmniSharp.Extensions.LanguageServer.Protocol.Models;
+using OmniSharp.Extensions.LanguageServer.Protocol.Window;
+using OmniSharp.Extensions.LanguageServer.Protocol.Server;
+
+namespace Bicep.LanguageServer.Utils
+{
+    class ServerLogTraceListener : TraceListener
+    {
+        private readonly ILanguageServer server;
+
+        public ServerLogTraceListener(ILanguageServer server)
+        {
+            this.server = server;
+        }
+
+        public override void Write(string? message)
+        {
+            if (message is not null)
+            {
+                server.Log(new LogMessageParams { Type = MessageType.Log, Message = message});
+            }
+        }
+
+        public override void WriteLine(string? message)
+        {
+            if (message is not null)
+            {
+                server.Log(new LogMessageParams { Type = MessageType.Log, Message = message});
+            }
+        }
+    }
+}