From 6cf1536224eadeabbd646f6714a438d47ea86229 Mon Sep 17 00:00:00 2001 From: Max Doerner Date: Mon, 23 Oct 2023 23:04:44 +0200 Subject: [PATCH] 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. --- .../ChangeProcedureToFunctionQuickFix.cs | 31 ++++- .../ChangeProcedureToFunctionQuickFixTests.cs | 114 ++++++++++++++++++ 2 files changed, 141 insertions(+), 4 deletions(-) diff --git a/Rubberduck.CodeAnalysis/QuickFixes/Concrete/ChangeProcedureToFunctionQuickFix.cs b/Rubberduck.CodeAnalysis/QuickFixes/Concrete/ChangeProcedureToFunctionQuickFix.cs index d355ac16f2..bce7c03a0f 100644 --- a/Rubberduck.CodeAnalysis/QuickFixes/Concrete/ChangeProcedureToFunctionQuickFix.cs +++ b/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()) + { + rewriter.Replace(exitSub, exitFunctionCode); + } + } + + private void UpdateCall(IdentifierReference reference, int argIndex, IRewriteSession rewriteSession) { var rewriter = rewriteSession.CheckOutModuleRewriter(reference.QualifiedModuleName); diff --git a/RubberduckTests/QuickFixes/ChangeProcedureToFunctionQuickFixTests.cs b/RubberduckTests/QuickFixes/ChangeProcedureToFunctionQuickFixTests.cs index 87ba5aca50..5760184105 100644 --- a/RubberduckTests/QuickFixes/ChangeProcedureToFunctionQuickFixTests.cs +++ b/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)