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)