Merge pull request #6171 from MDoerner/FixSubToFunctionQuickFixForExi…

…tSub

Replace Exit Sub when converting Subs to Functions

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)