Skip branches in generated MoveNext() for singleton iterators (#813)

Skip branches in generated `MoveNext()` for singleton iterators

src/coverlet.core/Symbols/CecilSymbolHelper.cs

@@ -412,7 +412,13 @@ public static List<BranchPoint> GetBranchPoints(MethodDefinition methodDefinitio
 
             bool isAsyncStateMachineMoveNext = IsMoveNextInsideAsyncStateMachine(methodDefinition);
             bool isMoveNextInsideAsyncStateMachineProlog = isAsyncStateMachineMoveNext && IsMoveNextInsideAsyncStateMachineProlog(methodDefinition);
+
+            // State machine for enumerator uses `brfalse.s`/`beq` or `switch` opcode depending on how many `yield` we have in the method body.
+            // For more than one `yield` a `switch` is emitted so we should only skip the first branch. In case of a single `yield` we need to
+            // skip the first two branches to avoid reporting a phantom branch. The first branch (`brfalse.s`) jumps to the `yield`ed value,
+            // the second one (`beq`) exits the enumeration.
             bool skipFirstBranch = IsMoveNextInsideEnumerator(methodDefinition);
+            bool skipSecondBranch = false;
 
             foreach (Instruction instruction in instructions.Where(instruction => instruction.OpCode.FlowControl == FlowControl.Cond_Branch))
             {
@@ -421,6 +427,13 @@ public static List<BranchPoint> GetBranchPoints(MethodDefinition methodDefinitio
                     if (skipFirstBranch)
                     {
                         skipFirstBranch = false;
+                        skipSecondBranch = instruction.OpCode.Code != Code.Switch;
+                        continue;
+                    }
+
+                    if (skipSecondBranch)
+                    {
+                        skipSecondBranch = false;
                         continue;
                     }
 

test/coverlet.core.tests/Coverage/CoverageTest.Yield.cs

@@ -0,0 +1,164 @@
+using System.IO;
+using System.Threading.Tasks;
+
+using Coverlet.Core.Samples.Tests;
+using Tmds.Utils;
+using Xunit;
+
+namespace Coverlet.Core.Tests
+{
+    public partial class CoverageTests : ExternalProcessExecutionTest
+    {
+        [Fact]
+        public void Yield_Single()
+        {
+            string path = Path.GetTempFileName();
+            try
+            {
+                FunctionExecutor.Run(async (string[] pathSerialize) =>
+                {
+                    CoveragePrepareResult coveragePrepareResult = await TestInstrumentationHelper.Run<Yield>(instance =>
+                    {
+                        foreach (var _ in instance.One()) ;
+
+                        return Task.CompletedTask;
+                    }, persistPrepareResultToFile: pathSerialize[0]);
+
+                    return 0;
+                }, new string[] { path });
+
+                CoverageResult result = TestInstrumentationHelper.GetCoverageResult(path);
+
+                result.Document("Instrumentation.Yield.cs")
+                      .Method("System.Boolean Coverlet.Core.Samples.Tests.Yield/<One>d__0::MoveNext()")
+                      .AssertLinesCovered((9, 1))
+                      .ExpectedTotalNumberOfBranches(0);
+            }
+            finally
+            {
+                File.Delete(path);
+            }
+        }
+
+        [Fact]
+        public void Yield_Two()
+        {
+            string path = Path.GetTempFileName();
+            try
+            {
+                FunctionExecutor.Run(async (string[] pathSerialize) =>
+                {
+                    CoveragePrepareResult coveragePrepareResult = await TestInstrumentationHelper.Run<Yield>(instance =>
+                    {
+                        foreach (var _ in instance.Two()) ;
+
+                        return Task.CompletedTask;
+                    }, persistPrepareResultToFile: pathSerialize[0]);
+                    return 0;
+                }, new string[] { path });
+
+                CoverageResult result = TestInstrumentationHelper.GetCoverageResult(path);
+
+                result.Document("Instrumentation.Yield.cs")
+                      .Method("System.Boolean Coverlet.Core.Samples.Tests.Yield/<Two>d__1::MoveNext()")
+                      .AssertLinesCovered((14, 1), (15, 1))
+                      .ExpectedTotalNumberOfBranches(0);
+            }
+            finally
+            {
+                File.Delete(path);
+            }
+        }
+
+        [Fact]
+        public void Yield_SingleWithSwitch()
+        {
+            string path = Path.GetTempFileName();
+            try
+            {
+                FunctionExecutor.Run(async (string[] pathSerialize) =>
+                {
+                    CoveragePrepareResult coveragePrepareResult = await TestInstrumentationHelper.Run<Yield>(instance =>
+                    {
+                        foreach (var _ in instance.OneWithSwitch(2)) ;
+
+                        return Task.CompletedTask;
+                    }, persistPrepareResultToFile: pathSerialize[0]);
+
+                    return 0;
+                }, new string[] { path });
+
+                CoverageResult result = TestInstrumentationHelper.GetCoverageResult(path);
+
+                result.Document("Instrumentation.Yield.cs")
+                      .Method("System.Boolean Coverlet.Core.Samples.Tests.Yield/<OneWithSwitch>d__2::MoveNext()")
+                      .AssertLinesCovered(BuildConfiguration.Debug, (21, 1), (30, 1), (31, 1), (37, 1))
+                      .ExpectedTotalNumberOfBranches(1);
+            }
+            finally
+            {
+                File.Delete(path);
+            }
+        }
+
+        [Fact]
+        public void Yield_Three()
+        {
+            string path = Path.GetTempFileName();
+            try
+            {
+                FunctionExecutor.Run(async (string[] pathSerialize) =>
+                {
+                    CoveragePrepareResult coveragePrepareResult = await TestInstrumentationHelper.Run<Yield>(instance =>
+                    {
+                        foreach (var _ in instance.Three()) ;
+
+                        return Task.CompletedTask;
+                    }, persistPrepareResultToFile: pathSerialize[0]);
+                    return 0;
+                }, new string[] { path });
+
+                CoverageResult result = TestInstrumentationHelper.GetCoverageResult(path);
+
+                result.Document("Instrumentation.Yield.cs")
+                      .Method("System.Boolean Coverlet.Core.Samples.Tests.Yield/<Three>d__3::MoveNext()")
+                      .AssertLinesCovered((42, 1), (43, 1), (44, 1))
+                      .ExpectedTotalNumberOfBranches(0);
+            }
+            finally
+            {
+                File.Delete(path);
+            }
+        }
+
+        [Fact]
+        public void Yield_Enumerable()
+        {
+            string path = Path.GetTempFileName();
+            try
+            {
+                FunctionExecutor.Run(async (string[] pathSerialize) =>
+                {
+                    CoveragePrepareResult coveragePrepareResult = await TestInstrumentationHelper.Run<Yield>(instance =>
+                    {
+                        foreach (var _ in instance.Enumerable(new[] { "one", "two", "three", "four" })) ;
+
+                        return Task.CompletedTask;
+                    }, persistPrepareResultToFile: pathSerialize[0]);
+                    return 0;
+                }, new string[] { path });
+
+                CoverageResult result = TestInstrumentationHelper.GetCoverageResult(path);
+
+                result.Document("Instrumentation.Yield.cs")
+                      .Method("System.Boolean Coverlet.Core.Samples.Tests.Yield/<Enumerable>d__4::MoveNext()")
+                      .AssertLinesCovered(BuildConfiguration.Debug, (48, 1), (49, 1), (50, 4), (51, 5), (52, 1), (54, 4), (55, 4), (56, 4), (57, 1))
+                      .ExpectedTotalNumberOfBranches(1);
+            }
+            finally
+            {
+                File.Delete(path);
+            }
+        }
+    }
+}

test/coverlet.core.tests/Coverage/InstrumenterHelper.Assertions.cs

@@ -79,11 +79,46 @@ public static Document Document(this CoverageResult coverageResult, string docNa
             throw new XunitException($"Document not found '{docName}'");
         }
 
+        public static Document Method(this Document document, string methodName)
+        {
+            var methodDoc = new Document { Path = document.Path, Index = document.Index };
+
+            if (!document.Lines.Any() && !document.Branches.Any())
+            {
+                return methodDoc;
+            }
+
+            if (document.Lines.Values.All(l => l.Method != methodName) && document.Branches.Values.All(l => l.Method != methodName))
+            {
+                var methods = document.Lines.Values.Select(l => $"'{l.Method}'")
+                    .Concat(document.Branches.Values.Select(b => $"'{b.Method}'"))
+                    .Distinct();
+                throw new XunitException($"Method '{methodName}' not found. Methods in document: {string.Join(", ", methods)}");
+            }
+
+            foreach (var line in document.Lines.Where(l => l.Value.Method == methodName))
+            {
+                methodDoc.Lines[line.Key] = line.Value;
+            }
+
+            foreach (var branch in document.Branches.Where(b => b.Value.Method == methodName))
+            {
+                methodDoc.Branches[branch.Key] = branch.Value;
+            }
+
+            return methodDoc;
+        }
+
         public static Document AssertBranchesCovered(this Document document, params (int line, int ordinal, int hits)[] lines)
         {
             return AssertBranchesCovered(document, BuildConfiguration.Debug | BuildConfiguration.Release, lines);
         }
 
+        public static Document ExpectedTotalNumberOfBranches(this Document document, int totalExpectedBranch)
+        {
+            return ExpectedTotalNumberOfBranches(document, BuildConfiguration.Debug | BuildConfiguration.Release, totalExpectedBranch);
+        }
+
         public static Document ExpectedTotalNumberOfBranches(this Document document, BuildConfiguration configuration, int totalExpectedBranch)
         {
             if (document is null)

test/coverlet.core.tests/Samples/Instrumentation.Yield.cs

@@ -0,0 +1,59 @@
+// Remember to use full name because adding new using directives change line numbers
+
+namespace Coverlet.Core.Samples.Tests
+{
+    public class Yield
+    {
+        public System.Collections.Generic.IEnumerable<int> One()
+        {
+            yield return 1;
+        }
+
+        public System.Collections.Generic.IEnumerable<int> Two()
+        {
+            yield return 1;
+            yield return 2;
+        }
+
+        public System.Collections.Generic.IEnumerable<int> OneWithSwitch(int n)
+        {
+            int result;
+            switch (n)
+            {
+                case 0:
+                    result = 10;
+                    break;
+                case 1:
+                    result = 11;
+                    break;
+                case 2:
+                    result = 12;
+                    break;
+                default:
+                    result = -1;
+                    break;
+            }
+
+            yield return result;
+        }
+
+        public System.Collections.Generic.IEnumerable<int> Three()
+        {
+            yield return 1;
+            yield return 2;
+            yield return 3;
+        }
+
+        public System.Collections.Generic.IEnumerable<string> Enumerable(System.Collections.Generic.IList<string> ls)
+        {
+            foreach (
+                    var item
+                    in
+                    ls
+                    )
+            {
+                yield return item;
+            }
+        }
+    }
+}

test/coverlet.core.tests/Samples/Samples.cs

@@ -172,6 +172,14 @@ public IEnumerable<string> Fetch()
         }
     }
 
+    public class SingletonIterator
+    {
+        public IEnumerable<string> Fetch()
+        {
+            yield return "one";
+        }
+    }
+
     public class AsyncAwaitStateMachine
     {
         async public Task AsyncAwait()

test/coverlet.core.tests/Symbols/CecilSymbolHelperTests.cs

@@ -259,7 +259,22 @@ public void GetBranchPoints_IgnoresSwitchIn_GeneratedMoveNext()
 
             // assert
             Assert.Empty(points);
+        }
+
+        [Fact]
+        public void GetBranchPoints_IgnoresBranchesIn_GeneratedMoveNextForSingletonIterator()
+        {
+            // arrange
+            var nestedName = typeof(SingletonIterator).GetNestedTypes(BindingFlags.NonPublic).First().Name;
+            var type = _module.Types.FirstOrDefault(x => x.FullName == typeof(SingletonIterator).FullName);
+            var nestedType = type.NestedTypes.FirstOrDefault(x => x.FullName.EndsWith(nestedName));
+            var method = nestedType.Methods.First(x => x.FullName.EndsWith("::MoveNext()"));
+
+            // act
+            var points = CecilSymbolHelper.GetBranchPoints(method);
 
+            // assert
+            Assert.Empty(points);
         }
 
         [Fact]