From 0d0c7f4e06084eee6e439a1ba11aab4ce95ce9db Mon Sep 17 00:00:00 2001
From: "Christian G. Warden" <cwarden@xerus.org>
Date: Tue, 23 Jul 2024 10:43:20 -0500
Subject: [PATCH] Fix Formatting of `while` and `for` Loops

Fix bug in formatting of empty `while` loops.

Fix formatting of `for` loops.
---
 formatter/format_test.go | 11 ++++++++++-
 formatter/visitors.go    |  4 ++--
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/formatter/format_test.go b/formatter/format_test.go
index 681f27e..b85d8fb 100644
--- a/formatter/format_test.go
+++ b/formatter/format_test.go
@@ -253,6 +253,15 @@ OneDayDischargeFollowUp.twoHoursDelay,
 	OneDayDischargeFollowUp.twoHoursDelay,
 	1
 );`},
+			{
+				`while(true);`,
+				`while (true);`},
+			{
+				`for (Campaign campaign : results) resultList.add(campaign.Group__c);`,
+				`for (Campaign campaign : results) {
+	resultList.add(campaign.Group__c);
+}
+`},
 		}
 	for _, tt := range tests {
 		input := antlr.NewInputStream(tt.input)
@@ -269,7 +278,7 @@ OneDayDischargeFollowUp.twoHoursDelay,
 			t.Errorf("Unexpected result parsing apex")
 		}
 		if out != tt.output {
-			t.Errorf("unexpected format.  expected:\n%s\ngot:\n%s\n", tt.output, out)
+			t.Errorf("unexpected format.  expected:\n%q\ngot:\n%q\n", tt.output, out)
 		}
 	}
 
diff --git a/formatter/visitors.go b/formatter/visitors.go
index f71b2a3..f58f23a 100644
--- a/formatter/visitors.go
+++ b/formatter/visitors.go
@@ -235,7 +235,7 @@ func (v *FormatVisitor) VisitIfStatement(ctx *parser.IfStatementContext) interfa
 }
 
 func (v *FormatVisitor) VisitWhileStatement(ctx *parser.WhileStatementContext) interface{} {
-	if s := ctx.Statement; s == nil {
+	if s := ctx.Statement(); s == nil {
 		return fmt.Sprintf("while %s;", v.visitRule(ctx.ParExpression()))
 	}
 	if block := ctx.Statement().Block(); block != nil {
@@ -250,7 +250,7 @@ func (v *FormatVisitor) VisitForStatement(ctx *parser.ForStatementContext) inter
 		if statement.Block() != nil {
 			return fmt.Sprintf("for (%s) %s", v.visitRule(ctx.ForControl()), v.visitRule(ctx.Statement()))
 		} else {
-			return fmt.Sprintf("for (%s) {\n%s}\n", v.visitRule(ctx.ForControl()), v.indent(v.visitRule(ctx.Statement()).(string)))
+			return fmt.Sprintf("for (%s) {\n%s\n}\n", v.visitRule(ctx.ForControl()), v.indent(v.visitRule(ctx.Statement()).(string)))
 		}
 	} else {
 		return fmt.Sprintf("for (%s);", v.visitRule(ctx.ForControl()))