fixed out of bounds in line mapping

src/OmniSharp.Roslyn.CSharp/Helpers/LocationExtensions.cs

@@ -1,8 +1,10 @@
 using System;
+using System.Collections.Generic;
 using System.IO;
 using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.Text;
 using OmniSharp.Models;
 
 namespace OmniSharp.Helpers
@@ -17,22 +19,55 @@ public static QuickFix GetQuickFix(this Location location, OmniSharpWorkspace wo
             var lineSpan = Path.GetExtension(location.SourceTree.FilePath).Equals(".cake", StringComparison.OrdinalIgnoreCase)
                 ? location.GetLineSpan()
                 : location.GetMappedLineSpan();
-            var path = lineSpan.Path;
-            var documents = workspace.GetDocuments(path);
 
-            var line = lineSpan.StartLinePosition.Line;
-            var text = location.SourceTree.GetText().Lines[line].ToString();
+            var documents = workspace.GetDocuments(lineSpan.Path);
+            var sourceText = GetSourceText(location, documents, lineSpan.HasMappedPath);
+            var text = GetText(location, sourceText, lineSpan.StartLinePosition.Line);
 
             return new QuickFix
             {
-                Text = text.Trim(),
-                FileName = path,
-                Line = line,
+                Text = text,
+                FileName = lineSpan.Path,
+                Line = lineSpan.StartLinePosition.Line,
                 Column = lineSpan.HasMappedPath ? 0 : lineSpan.StartLinePosition.Character, // when a #line directive maps into a separate file, assume columns (0,0)
                 EndLine = lineSpan.EndLinePosition.Line,
                 EndColumn = lineSpan.HasMappedPath ? 0 : lineSpan.EndLinePosition.Character,
                 Projects = documents.Select(document => document.Project.Name).ToArray()
             };
+
+            static SourceText GetSourceText(Location location, IEnumerable<Document> documents, bool hasMappedPath)
+            {
+                // if we have a mapped linespan and we found a corresponding document, pick that one
+                // otherwise use the SourceText of current location
+                if (hasMappedPath)
+                {
+                    if (documents != null && documents.Any() && documents.First().TryGetText(out SourceText sourceText))
+                    {
+                        // we have a mapped document that exists in workspace
+                        return sourceText;
+                    }
+
+                    // we have a mapped document that doesn't exist in workspace
+                    // in that case we have to always fall back to original linespan
+                    return null;
+                }
+
+                // unmapped document so just continue with current SourceText
+                return location.SourceTree.GetText();
+            }
+
+            static string GetText(Location location, SourceText sourceText, int startLine)
+            {
+                // bounds check in case the mapping is incorrect, since user can put whatever line they want
+                if (sourceText != null && sourceText.Lines.Count > startLine)
+                {
+                    return sourceText.Lines[startLine].ToString();
+                }
+
+                // in case we fall out of bounds, we shouldn't crash, fallback to text from the unmapped span and the current file
+                var fallBackLineSpan = location.GetLineSpan();
+                return location.SourceTree.GetText().Lines[fallBackLineSpan.StartLinePosition.Line].ToString();
+            }
         }
     }
 }

tests/OmniSharp.Roslyn.CSharp.Tests/FindReferencesFacts.cs

@@ -1,4 +1,5 @@
-using System.Linq;
+using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
 using OmniSharp.Models;
 using OmniSharp.Models.FindUsages;
@@ -116,73 +117,102 @@ public FooConsumer()
             Assert.Equal(2, usages.QuickFixes.Count());
         }
 
-        [Fact]
-        public async Task CanFindReferencesWithLineMapping()
+        [Theory]
+        [InlineData(9, "public FooConsumer()")]
+        [InlineData(100, "new Foo().bar();")]
+        public async Task CanFindReferencesWithLineMapping(int mappingLine, string expectedMappingText)
         {
-            const string code = @"
+            var code = @"
                 public class Foo
                 {
-                    public void b$$ar() { }
+public void b$$ar() { }
                 }
 
                 public class FooConsumer
                 {
-                    public FooConsumer()
+public FooConsumer()
                     {
-#line 1
-                        new Foo().bar();
+#line "+mappingLine+@"
+new Foo().bar();
 #line default
                     }
                 }";
 
             var usages = await FindUsagesAsync(code);
             Assert.Equal(2, usages.QuickFixes.Count());
 
-            var mappedResult = usages.QuickFixes.FirstOrDefault(x => x.Line == 0);
-            var regularResult = usages.QuickFixes.FirstOrDefault(x => x.Line == 3);
-            Assert.NotNull(mappedResult);
-            Assert.NotNull(regularResult);
+            var quickFixes = usages.QuickFixes.OrderBy(x => x.Line);
+            var regularResult = quickFixes.ElementAt(0);
+            var mappedResult = quickFixes.ElementAt(1);
+
+            Assert.Equal("dummy.cs", regularResult.FileName);
+            Assert.Equal("dummy.cs", mappedResult.FileName);
+
+            Assert.Equal("public void bar() { }", regularResult.Text);
+            Assert.Equal(expectedMappingText, mappedResult.Text);
+
+            Assert.Equal(3, regularResult.Line);
+            Assert.Equal(mappingLine-1, mappedResult.Line);
 
             // regular result has regular postition
-            Assert.Equal(32, regularResult.Column);
-            Assert.Equal(35, regularResult.EndColumn);
+            Assert.Equal(12, regularResult.Column);
+            Assert.Equal(15, regularResult.EndColumn);
+
+            // mapped result has column 0,0
+            Assert.Equal(10, mappedResult.Column);
+            Assert.Equal(13, mappedResult.EndColumn);
         }
 
-        [Fact]
-        public async Task CanFindReferencesWithLineMappingAcrossFiles()
+        [Theory]
+        [InlineData(1, "// hello", true)] // everything correct
+        [InlineData(100, "new Foo().bar();", true)] // file exists in workspace but mapping incorrect
+        [InlineData(1, "new Foo().bar();", false)] // file doesn't exist in workspace but mapping correct
+        public async Task CanFindReferencesWithLineMappingAcrossFiles(int mappingLine, string expectedMappingText, bool mappedFileExistsInWorkspace)
         {
-            var testFiles = new[]
+            var testFiles = new List<TestFile>()
             {
                 new TestFile("a.cs", @"
                 public class Foo
                 {
-                    public void b$$ar() { }
+public void b$$ar() { }
                 }
 
                 public class FooConsumer
                 {
                     public FooConsumer()
                     {
-#line 1 ""b.cs""
-                        new Foo().bar();
+#line "+mappingLine+@" ""b.cs""
+new Foo().bar();
 #line default
                     }
                 }"),
-                new TestFile("b.cs",
-                    @"// hello")
+
             };
 
-            var usages = await FindUsagesAsync(testFiles, onlyThisFile: false);
+            if (mappedFileExistsInWorkspace)
+            {
+                testFiles.Add(new TestFile("b.cs",
+                    @"// hello"));
+            }
+
+            var usages = await FindUsagesAsync(testFiles.ToArray(), onlyThisFile: false);
             Assert.Equal(2, usages.QuickFixes.Count());
 
-            var mappedResult = usages.QuickFixes.FirstOrDefault(x => x.Line == 0 && x.FileName == "b.cs");
-            var regularResult = usages.QuickFixes.FirstOrDefault(x => x.Line == 3 && x.FileName == "a.cs");
-            Assert.NotNull(mappedResult);
-            Assert.NotNull(regularResult);
+            var regularResult = usages.QuickFixes.ElementAt(0);
+            var mappedResult = usages.QuickFixes.ElementAt(1);
+
+            Assert.Equal("a.cs", regularResult.FileName);
+            Assert.Equal("b.cs", mappedResult.FileName);
+
+            Assert.Equal(3, regularResult.Line);
+            Assert.Equal(mappingLine-1, mappedResult.Line);
+
+            Assert.Equal("public void bar() { }", regularResult.Text);
+            Assert.Equal(expectedMappingText, mappedResult.Text);
 
             // regular result has regular postition
-            Assert.Equal(32, regularResult.Column);
-            Assert.Equal(35, regularResult.EndColumn);
+            Assert.Equal(12, regularResult.Column);
+            Assert.Equal(15, regularResult.EndColumn);
 
             // mapped result has column 0,0
             Assert.Equal(0, mappedResult.Column);