Add module path completions for ARM template files (#3616)

* Add completions for ARM template files

* Remove an extra empty line

* Fix GetDisplayName

* Update test baselines

* Address comments and fix a bug

* Refactoring

Author: shenglol

src/Bicep.Core.Samples/Files/Completions/cwdMCompletions.json

@@ -181,5 +181,24 @@ public void DirExists_should_return_expected_results()
             Directory.CreateDirectory(tempChildDir);
             fileResolver.TryDirExists(PathHelper.FilePathToFileUrl(tempChildDir)).Should().BeTrue();
         }
+
+        [DataTestMethod]
+        [DataRow("", 2, true, "")]
+        [DataRow("a", 2, true, "a")]
+        [DataRow("aa", 2, true, "aa")]
+        [DataRow("aaaa\nbbbbb", 2, true, "aa")]
+        public void TryReadAtMostNCharacters_RegardlessFileContentLength_ReturnsAtMostNCharaters(string fileContents, int n, bool expectedResult, string expectedContents)
+        {
+            var fileResolver = new FileResolver();
+            var tempFile = Path.Combine(Path.GetTempPath(), $"BICEP_TEST_{Guid.NewGuid()}");
+            var tempFileUri = PathHelper.FilePathToFileUrl(tempFile);
+
+            File.WriteAllText(tempFile, fileContents);
+
+            var result = fileResolver.TryReadAtMostNCharaters(tempFileUri, Encoding.UTF8, n, out var readContents);
+
+            result.Should().Be(expectedResult);
+            readContents.Should().Be(expectedContents);
+        }
     }
 }

src/Bicep.Core.Samples/Files/InvalidModules_LF/Completions/cwdCompletions.json

@@ -43,8 +43,6 @@ public bool TryRead(Uri fileUri, [NotNullWhen(true)] out string? fileContents, [
             }
         }
 
-
-
         public bool TryRead(Uri fileUri, [NotNullWhen(true)] out string? fileContents, [NotNullWhen(false)] out DiagnosticBuilder.ErrorBuilderDelegate? failureBuilder, Encoding fileEncoding, int maxCharacters, [NotNullWhen(true)] out Encoding? detectedEncoding)
         {
             if (!fileUri.IsFile)
@@ -152,6 +150,39 @@ public bool TryReadAsBase64(Uri fileUri, [NotNullWhen(true)] out string? fileBas
             }
         }
 
+        public bool TryReadAtMostNCharaters(Uri fileUri, Encoding fileEncoding, int n, [NotNullWhen(true)] out string? fileContents)
+        {
+            if (!fileUri.IsFile || n <= 0)
+            {
+                fileContents = null;
+                return false;
+            }
+
+            try
+            {
+                if (Directory.Exists(fileUri.LocalPath))
+                {
+                    // Docs suggest this is the error to throw when we give a directory. 
+                    // A trailing backslash causes windows not to throw this exception.
+                    throw new UnauthorizedAccessException($"Access to the path '{fileUri.LocalPath}' is denied.");
+                }
+
+                using var fileStream = File.OpenRead(fileUri.LocalPath);
+                using var sr = new StreamReader(fileStream, fileEncoding, true);
+                
+                var buffer = new char[n];
+                n = sr.ReadBlock(buffer, 0, n);
+
+                fileContents = new string(buffer.Take(n).ToArray());
+                return true;
+            }
+            catch (Exception)
+            {
+                fileContents = null;
+                return false;
+            }
+        }
+
         public Uri? TryResolveFilePath(Uri parentFileUri, string childFilePath)
         {
             if (!Uri.TryCreate(parentFileUri, childFilePath, out var relativeUri))

src/Bicep.Core.Samples/Files/InvalidModules_LF/Completions/cwdFileCompletions.json

@@ -21,14 +21,16 @@ public interface IFileResolver
         bool TryRead(Uri fileUri, [NotNullWhen(true)] out string? fileContents, [NotNullWhen(false)] out DiagnosticBuilder.ErrorBuilderDelegate? failureBuilder);
 
         bool TryRead(Uri fileUri, [NotNullWhen(true)] out string? fileContents, [NotNullWhen(false)] out DiagnosticBuilder.ErrorBuilderDelegate? failureBuilder, Encoding fileEncoding, int maxCharacters, [NotNullWhen(true)] out Encoding? detectedEncoding);
+
+        bool TryReadAtMostNCharaters(Uri fileUri, Encoding fileEncoding, int n, [NotNullWhen(true)] out string? fileContents);
+
         /// <summary>
         /// Tries to resolve a child file path relative to a parent module file path.
         /// </summary>
         /// <param name="parentFileUri">The file URI of the parent.</param>
         /// <param name="childFilePath">The file path of the child.</param>
         Uri? TryResolveFilePath(Uri parentFileUri, string childFilePath);
 
-
         /// <summary>
         /// Tries to get Directories given a uri and pattern. Both argument and returned URIs MUST have a trailing '/'
         /// </summary>

src/Bicep.Core.UnitTests/FileSystem/FileResolverTests.cs

@@ -57,6 +57,18 @@ public bool TryRead(Uri fileUri, [NotNullWhen(true)] out string? fileContents, [
             return true;
         }
 
+        public bool TryReadAtMostNCharaters(Uri fileUri, Encoding fileEncoding, int n, [NotNullWhen(true)] out string? fileContents)
+        {
+            if (!fileLookup.TryGetValue(fileUri, out fileContents))
+            {
+                fileContents = null;
+                return false;
+            }
+
+            fileContents = new string(fileContents.Take(n).ToArray());
+            return true;
+        }
+
         public Uri? TryResolveFilePath(Uri parentFileUri, string childFilePath)
         {
             if (!Uri.TryCreate(parentFileUri, childFilePath, out var relativeUri))

src/Bicep.Core/FileSystem/FileResolver.cs

@@ -5,6 +5,7 @@
 using System.Collections.Immutable;
 using System.Linq;
 using System.Text;
+using System.Text.RegularExpressions;
 using Bicep.Core.Parsing;
 using Bicep.Core.Resources;
 using Bicep.Core.TypeSystem;
@@ -48,9 +49,11 @@ public static class LanguageConstants
         public const string TargetScopeTypeSubscription = "subscription";
         public const string TargetScopeTypeResourceGroup = "resourceGroup";
 
-        public static ImmutableSortedSet<string> DeclarationKeywords = new[] { ParameterKeyword, VariableKeyword, ResourceKeyword, OutputKeyword, ModuleKeyword }.ToImmutableSortedSet(StringComparer.Ordinal);
+        public static readonly Regex ArmTemplateSchemaRegex = new(@"https?:\/\/schema\.management\.azure\.com\/schemas\/([^""\/]+\/[a-zA-Z]*[dD]eploymentTemplate\.json)#?");
 
-        public static ImmutableSortedSet<string> ContextualKeywords = DeclarationKeywords
+        public static readonly ImmutableSortedSet<string> DeclarationKeywords = new[] { ParameterKeyword, VariableKeyword, ResourceKeyword, OutputKeyword, ModuleKeyword }.ToImmutableSortedSet(StringComparer.Ordinal);
+
+        public static readonly ImmutableSortedSet<string> ContextualKeywords = DeclarationKeywords
             .Add(TargetScopeKeyword)
             .Add(IfKeyword)
             .Add(ForKeyword)

src/Bicep.Core/FileSystem/IFileResolver.cs

@@ -45,6 +45,16 @@ public class CompletionTests
         [NotNull]
         public TestContext? TestContext { get; set; }
 
+        public static string GetDisplayName(MethodInfo info, object[] row)
+        {
+            row.Should().HaveCount(3);
+            row[0].Should().BeOfType<DataSet>();
+            row[1].Should().BeOfType<string>();
+            row[2].Should().BeAssignableTo<IList<Position>>();
+
+            return $"{info.Name}_{((DataSet)row[0]).Name}_{row[1]}";
+        }
+
         [TestMethod]
         public async Task EmptyFileShouldProduceDeclarationCompletions()
         {
@@ -786,14 +796,64 @@ public async Task RequestCompletions_MatchingNodeIsBooleanOrIntegerOrNullLiteral
             await RunCompletionScenarioTest(this.TestContext, fileWithCursors, AssertAllCompletionsEmpty);
         }
 
-        private static async Task RunCompletionScenarioTest(TestContext testContext, string fileWithCursors, Action<IEnumerable<CompletionList?>> assertAction)
+        [TestMethod]
+        public async Task RequestModulePathCompletions_ArmTemplateFilesInDir_ReturnsCompletionsIncludingArmTemplatePaths()
         {
-            var (file, cursors) = ParserHelper.GetFileWithCursors(fileWithCursors);
-            var bicepFile = SourceFileFactory.CreateBicepFile(new Uri("file:///path/to/main.bicep"), file);
-            var client = await IntegrationTestHelper.StartServerWithTextAsync(testContext, file, bicepFile.FileUri, resourceTypeProvider: BuiltInTestTypes.Create());
-            var completions = await RequestCompletions(client, bicepFile, cursors);
+            var mainUri = DocumentUri.FromFileSystemPath("/path/to/main.bicep");
+            var armTemplateUri1 = DocumentUri.FromFileSystemPath("/path/to/template1.arm");
+            var armTemplateUri2 = DocumentUri.FromFileSystemPath("/path/to/template2.json");
+            var armTemplateUri3 = DocumentUri.FromFileSystemPath("/path/to/template3.jsonc");
+            var armTemplateUri4 = DocumentUri.FromFileSystemPath("/path/to/template4.json");
+            var armTemplateUri5 = DocumentUri.FromFileSystemPath("/path/to/template5.json");
+            var jsonUri1 = DocumentUri.FromFileSystemPath("/path/to/json1.json");
+            var jsonUri2 = DocumentUri.FromFileSystemPath("/path/to/json2.json");
+            var bicepModuleUri1 = DocumentUri.FromFileSystemPath("/path/to/module1.txt");
+            var bicepModuleUri2 = DocumentUri.FromFileSystemPath("/path/to/module2.bicep");
+            var bicepModuleUri3 = DocumentUri.FromFileSystemPath("/path/to/module3.bicep");
+
+            var (mainFileText, cursors) = ParserHelper.GetFileWithCursors(@"
+module mod1 './module1.txt' = {}
+module mod2 './template3.jsonc' = {}
+module mod2 './|' = {}
+");
+            var mainFile = SourceFileFactory.CreateBicepFile(mainUri.ToUri(), mainFileText);
+            var schema = "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#";
+
+            var fileTextsByUri = new Dictionary<Uri, string>
+            {
+                [mainUri.ToUri()] = mainFileText,
+                [armTemplateUri1.ToUri()] = "",
+                [armTemplateUri2.ToUri()] = @$"{{ ""schema"": ""{schema}"" }}",
+                [armTemplateUri3.ToUri()] = @"{}",
+                [armTemplateUri4.ToUri()] = new string('x', 2000 - schema.Length) + schema,
+                [armTemplateUri5.ToUri()] = new string('x', 2002 - schema.Length) + schema,
+                [jsonUri1.ToUri()] = "{}",
+                [jsonUri2.ToUri()] = @"[{ ""name"": ""value"" }]",
+                [bicepModuleUri1.ToUri()] = "param foo string",
+                [bicepModuleUri2.ToUri()] = "param bar bool",
+                [bicepModuleUri3.ToUri()] = "",
+            };
 
-            assertAction(completions);
+            var fileResolver = new InMemoryFileResolver(fileTextsByUri);
+
+            var client = await IntegrationTestHelper.StartServerWithTextAsync(
+                TestContext,
+                mainFileText,
+                mainUri,
+                resourceTypeProvider: BuiltInTestTypes.Create(),
+                fileResolver: fileResolver);
+
+            var completionLists = await RequestCompletions(client, mainFile, cursors);
+            completionLists.Should().HaveCount(1);
+
+            var completionItems = completionLists.Single()!.Items;
+            completionItems.Should().SatisfyRespectively(
+                x => x.Label.Should().Be("module2.bicep"),
+                x => x.Label.Should().Be("module3.bicep"),
+                x => x.Label.Should().Be("template1.arm"),
+                x => x.Label.Should().Be("template2.json"),
+                x => x.Label.Should().Be("template3.jsonc"),
+                x => x.Label.Should().Be("template4.json"));
         }
 
         [TestMethod]
@@ -970,14 +1030,14 @@ private void ValidateCompletions(DataSet dataSet, string setName, List<(Position
 
         private static string GetGlobalCompletionSetPath(string setName) => Path.Combine("src", "Bicep.Core.Samples", "Files", DataSet.TestCompletionsDirectory, GetFullSetName(setName));
 
-        public static string GetDisplayName(MethodInfo info, object[] row)
+        private static async Task RunCompletionScenarioTest(TestContext testContext, string fileWithCursors, Action<IEnumerable<CompletionList?>> assertAction)
         {
-            row.Should().HaveCount(3);
-            row[0].Should().BeOfType<DataSet>();
-            row[1].Should().BeOfType<string>();
-            row[2].Should().BeAssignableTo<IList<Position>>();
+            var (file, cursors) = ParserHelper.GetFileWithCursors(fileWithCursors);
+            var bicepFile = SourceFileFactory.CreateBicepFile(new Uri("file:///path/to/main.bicep"), file);
+            var client = await IntegrationTestHelper.StartServerWithTextAsync(testContext, file, bicepFile.FileUri, resourceTypeProvider: BuiltInTestTypes.Create());
+            var completions = await RequestCompletions(client, bicepFile, cursors);
 
-            return $"{info.Name}_{((DataSet)row[0]).Name}_{row[1]}";
+            assertAction(completions);
         }
 
         private static string FormatPosition(Position position) => $"({position.Line}, {position.Character})";

src/Bicep.Core/FileSystem/InMemoryFileResolver.cs

@@ -291,31 +291,66 @@ private IEnumerable<CompletionItem> GetModulePathCompletions(SemanticModel model
                 files = FileResolver.GetFiles(queryParent, "");
                 dirs = FileResolver.GetDirectories(queryParent, "");
             }
-            // "./" will not be preserved when making relative Uris. We have to go and manually add it.
-            // Prioritize .bicep files higher than other files.
-            var fileItems = files
-                .Where(file => file != model.SourceFile.FileUri)
-                .Where(file => file.Segments.Last().EndsWith(LanguageConstants.LanguageFileExtension))
-                .Select(file => CreateModulePathCompletionBuilder(
-                    file.Segments.Last(),
-                    (entered.StartsWith("./") ? "./" : "") + cwdUri.MakeRelativeUri(file).ToString(),
-                    context.ReplacementRange,
-                    CompletionItemKind.File,
-                    file.Segments.Last().EndsWith(LanguageConstants.LanguageId) ? CompletionPriority.High : CompletionPriority.Medium)
-                .Build())
-                .ToList();
 
+            // Prioritize .bicep files higher than other files.
+            var bicepFileItems = CreateFileCompletionItems(files, cwdUri, IsBicepFile, CompletionPriority.High);
+            var armTemplateFileItems = CreateFileCompletionItems(files, cwdUri, IsArmTemplateFileLike, CompletionPriority.Medium);
             var dirItems = dirs
-                .Select(dir => CreateModulePathCompletionBuilder(
-                    dir.Segments.Last(),
-                    (entered.StartsWith("./") ? "./" : "") + cwdUri.MakeRelativeUri(dir).ToString(),
-                    context.ReplacementRange,
-                    CompletionItemKind.Folder,
-                    CompletionPriority.Medium)
-                .WithCommand(new Command { Name = EditorCommands.RequestCompletions })
-                .Build())
-                .ToList();
-            return fileItems.Concat(dirItems);
+                .Select(dir =>
+                    CreateModulePathCompletionBuilder(
+                        dir.Segments.Last(),
+                        // "./" will not be preserved when making relative Uris. We have to go and manually add it.
+                        (entered.StartsWith("./") ? "./" : "") + cwdUri.MakeRelativeUri(dir).ToString(),
+                        context.ReplacementRange,
+                        CompletionItemKind.Folder,
+                        CompletionPriority.Low)
+                    .WithCommand(new Command { Name = EditorCommands.RequestCompletions })
+                    .Build());
+
+            return bicepFileItems.Concat(armTemplateFileItems).Concat(dirItems);
+
+            // Local functions.
+            IEnumerable<CompletionItem> CreateFileCompletionItems(IEnumerable<Uri> fileUris, Uri cwdUri, Predicate<Uri> predicate, CompletionPriority priority) => fileUris
+                .Where(fileUri => fileUri != model.SourceFile.FileUri && predicate(fileUri))
+                .Select(fileUri =>
+                    CreateModulePathCompletionBuilder(
+                        fileUri.Segments.Last(),
+                        (entered.StartsWith("./") ? "./" : "") + cwdUri.MakeRelativeUri(fileUri).ToString(),
+                        context.ReplacementRange,
+                        CompletionItemKind.File,
+                        priority)
+                    .Build());
+
+            bool IsBicepFile(Uri fileUri) => PathHelper.HasBicepExtension(fileUri);
+
+            bool IsArmTemplateFileLike(Uri fileUri)
+            {
+                if (PathHelper.HasExtension(fileUri, LanguageConstants.ArmTemplateFileExtension))
+                {
+                    return true;
+                }
+
+                if (model.Compilation.SourceFileGrouping.SourceFiles.Any(sourceFile =>
+                        sourceFile is ArmTemplateFile &&
+                        sourceFile.FileUri.LocalPath.Equals(fileUri.LocalPath, PathHelper.PathComparison)))
+                {
+                    return true;
+                }
+
+                if (!PathHelper.HasExtension(fileUri, LanguageConstants.JsonFileExtension) &&
+                    !PathHelper.HasExtension(fileUri, LanguageConstants.JsoncFileExtension))
+                {
+                    return false;
+                }
+
+                if (FileResolver.TryReadAtMostNCharaters(fileUri, Encoding.UTF8, 2000, out var fileContents) &&
+                    LanguageConstants.ArmTemplateSchemaRegex.IsMatch(fileContents))
+                {
+                    return true;
+                }
+
+                return false;
+            }
         }
 
         private static IEnumerable<CompletionItem> GetParameterTypeSnippets(Compilation compitation, BicepCompletionContext context)