Replace Exit Sub when converting Subs to Functions

This assigns the return value and then exists the function instead of exiting the procedure that gets converted. These two statements are inserted separated by a statement separator instead of doing it on multiple lines to avoid having a separate code path if the Exit Sub is in a single line if statement.

Author: MDoerner

Rubberduck.CodeAnalysis/QuickFixes/Concrete/ChangeProcedureToFunctionQuickFix.cs

@@ -59,23 +59,31 @@ public override void Fix(IInspectionResult result, IRewriteSession rewriteSessio
             var arg = parameterizedDeclaration.Parameters.First(p => p.IsByRef || p.IsImplicitByRef);
             var argIndex = parameterizedDeclaration.Parameters.IndexOf(arg);
 
-            UpdateSignature(result.Target, arg, rewriteSession);
+            UpdateProcedure(result.Target, arg, rewriteSession);
             foreach (var reference in result.Target.References.Where(reference => !reference.IsDefaultMemberAccess))
             {
                 UpdateCall(reference, argIndex, rewriteSession);
             }
         }
 
-        private void UpdateSignature(Declaration target, ParameterDeclaration arg, IRewriteSession rewriteSession)
+        private void UpdateProcedure(Declaration target, ParameterDeclaration arg, IRewriteSession rewriteSession)
         {
             var subStmt = (VBAParser.SubStmtContext) target.Context;
             var argContext = (VBAParser.ArgContext)arg.Context;
-
+            var argName = argContext.unrestrictedIdentifier().GetText();
             var rewriter = rewriteSession.CheckOutModuleRewriter(target.QualifiedModuleName);
 
+            UpdateSignature(subStmt, arg, rewriter);
+            AddReturnStatement(subStmt, argName, rewriter);
+            ReplaceExitSubs(subStmt, argName, rewriter);
+        }
+
+        private void UpdateSignature(VBAParser.SubStmtContext subStmt, ParameterDeclaration arg, IModuleRewriter rewriter)
+        {
             rewriter.Replace(subStmt.SUB(), Tokens.Function);
             rewriter.Replace(subStmt.END_SUB(), "End Function");
 
+            var argContext = (VBAParser.ArgContext)arg.Context;
             rewriter.InsertAfter(subStmt.argList().Stop.TokenIndex, $" As {arg.AsTypeName}");
 
             if (arg.IsByRef)
@@ -86,11 +94,26 @@ private void UpdateSignature(Declaration target, ParameterDeclaration arg, IRewr
             {
                 rewriter.InsertBefore(argContext.unrestrictedIdentifier().Start.TokenIndex, Tokens.ByVal);
             }
+        }
 
-            var returnStmt = $"    {subStmt.subroutineName().GetText()} = {argContext.unrestrictedIdentifier().GetText()}{Environment.NewLine}";
+        private void AddReturnStatement(VBAParser.SubStmtContext subStmt, string argName, IModuleRewriter rewriter)
+        {
+            var returnStmt = $"    {subStmt.subroutineName().GetText()} = {argName}{Environment.NewLine}";
+            // This exploits that the VBE will realign the End Function statement automatically. 
             rewriter.InsertBefore(subStmt.END_SUB().Symbol.TokenIndex, returnStmt);
         }
 
+        private void ReplaceExitSubs(VBAParser.SubStmtContext subStmt, string argName, IModuleRewriter rewriter)
+        {
+            // We use a statement separator here to be able to deal with single line if statments without too much issues.
+            var exitFunctionCode = $"{subStmt.subroutineName().GetText()} = {argName}: Exit Function";
+            foreach (var exitSub in subStmt.GetDescendents<VBAParser.ExitStmtContext>())
+            {
+                rewriter.Replace(exitSub, exitFunctionCode);
+            }
+        }
+
+
         private void UpdateCall(IdentifierReference reference, int argIndex, IRewriteSession rewriteSession)
         {
             var rewriter = rewriteSession.CheckOutModuleRewriter(reference.QualifiedModuleName);

RubberduckTests/QuickFixes/ChangeProcedureToFunctionQuickFixTests.cs

@@ -192,6 +192,120 @@ Foo fizz
     Foo = arg1
 End Function
 
+Sub Goo(ByVal a As Integer)
+    Dim fizz As Integer
+    fizz = Foo(fizz)
+End Sub";
+
+            var actualCode = ApplyQuickFixToFirstInspectionResult(inputCode, state => new ProcedureCanBeWrittenAsFunctionInspection(state));
+            Assert.AreEqual(expectedCode, actualCode);
+        }
+
+        // Based on issue #6139 at https://github.com/rubberduck-vba/Rubberduck/issues/6139
+        [Test]
+        [Category("QuickFixes")]
+        public void ProcedureShouldBeFunction_QuickFixWorks_ExitSub()
+        {
+            const string inputCode =
+                @"Private Sub Foo(ByRef arg1 As Integer)
+    If condition Then 
+        Exit Sub
+    End If
+
+    arg1 = 42
+End Sub
+
+Sub Goo(ByVal a As Integer)
+    Dim fizz As Integer
+    Foo fizz
+End Sub";
+
+            const string expectedCode =
+                @"Private Function Foo(ByVal arg1 As Integer) As Integer
+    If condition Then 
+        Foo = arg1: Exit Function
+    End If
+
+    arg1 = 42
+    Foo = arg1
+End Function
+
+Sub Goo(ByVal a As Integer)
+    Dim fizz As Integer
+    fizz = Foo(fizz)
+End Sub";
+
+            var actualCode = ApplyQuickFixToFirstInspectionResult(inputCode, state => new ProcedureCanBeWrittenAsFunctionInspection(state));
+            Assert.AreEqual(expectedCode, actualCode);
+        }
+
+        [Test]
+        [Category("QuickFixes")]
+        public void ProcedureShouldBeFunction_QuickFixWorks_ExitSubIndentationRespected()
+        {
+            const string inputCode =
+                @"Private Sub Foo(ByRef arg1 As Integer)
+    If condition Then
+        If otherCondition Then
+             Exit Sub
+        End If
+    End If
+
+    arg1 = 42
+End Sub
+
+Sub Goo(ByVal a As Integer)
+    Dim fizz As Integer
+    Foo fizz
+End Sub";
+
+            const string expectedCode =
+                @"Private Function Foo(ByVal arg1 As Integer) As Integer
+    If condition Then
+        If otherCondition Then
+             Foo = arg1: Exit Function
+        End If
+    End If
+
+    arg1 = 42
+    Foo = arg1
+End Function
+
+Sub Goo(ByVal a As Integer)
+    Dim fizz As Integer
+    fizz = Foo(fizz)
+End Sub";
+
+            var actualCode = ApplyQuickFixToFirstInspectionResult(inputCode, state => new ProcedureCanBeWrittenAsFunctionInspection(state));
+            Assert.AreEqual(expectedCode, actualCode);
+        }
+
+        [Test]
+        [Category("QuickFixes")]
+        public void ProcedureShouldBeFunction_QuickFixWorks_ExitSub_InSingleLineIf()
+        {
+            const string inputCode =
+                @"Private Sub Foo(ByRef arg1 As Integer)
+    If condition Then Exit Sub
+    If otherContirion Then: Else Exit Sub 
+
+    arg1 = 42
+End Sub
+
+Sub Goo(ByVal a As Integer)
+    Dim fizz As Integer
+    Foo fizz
+End Sub";
+
+            const string expectedCode =
+                @"Private Function Foo(ByVal arg1 As Integer) As Integer
+    If condition Then Foo = arg1: Exit Function
+    If otherContirion Then: Else Foo = arg1: Exit Function 
+
+    arg1 = 42
+    Foo = arg1
+End Function
+
 Sub Goo(ByVal a As Integer)
     Dim fizz As Integer
     fizz = Foo(fizz)