Fix dropping comments after handling unexpectedCommentError (#8553)

Signed-off-by: Sebastian Spaink <sebastianspaink@gmail.com>

v1/format/format.go

@@ -953,7 +953,7 @@ func (w *writer) writeExpr(expr *ast.Expr, comments []*ast.Comment) ([]*ast.Comm
 	// Print on same row if already there, otherwise increase indent a print remaining
 	if withs[0].Location.Row == lastRow {
 		if comments, err = w.writeWith(withs[0], comments, false); err != nil {
-			return nil, err
+			return comments, err
 		}
 		lastRow, withs = withs[0].Location.Row, withs[1:]
 	}
@@ -1150,7 +1150,7 @@ func (w *writer) writeWith(with *ast.With, comments []*ast.Comment, indented boo
 	w.write(" as ")
 	comments, err = w.writeTerm(with.Value, comments)
 	if err != nil {
-		return nil, err
+		return comments, err
 	}
 	return comments, nil
 }

v1/format/testfiles/v1/test_issue_nukedcomment.rego

@@ -0,0 +1,23 @@
+package p
+
+rule_that_formats_ok if {
+	r := other with input.x as {
+		"key": {"value˚": {{
+			"another": [
+				[["b", "c"], "3:1:3:8"], # this
+				[["b"], "4:1:4:8"], # is
+				[["c"], "5:1:5:8"], # fine
+			],
+		}}},
+	}
+
+	r == {{}}
+}
+
+# this is a comment
+# there are many like this
+# but this one is mine
+this_comment_above_gets_nuked_and_that_is_not_good if {
+	r := set()
+	r == set()
+}

v1/format/testfiles/v1/test_issue_nukedcomment.rego.formatted

@@ -0,0 +1,23 @@
+package p
+
+rule_that_formats_ok if {
+	r := other with input.x as {
+		"key": {"value˚": {{
+			"another": [
+				[["b", "c"], "3:1:3:8"], # this
+				[["b"], "4:1:4:8"], # is
+				[["c"], "5:1:5:8"], # fine
+			],
+		}}},
+	}
+
+	r == {{}}
+}
+
+# this is a comment
+# there are many like this
+# but this one is mine
+this_comment_above_gets_nuked_and_that_is_not_good if {
+	r := set()
+	r == set()
+}