diff --git a/v1/format/format.go b/v1/format/format.go
index 5b4382321c..f4ac287eeb 100644
--- a/v1/format/format.go
+++ b/v1/format/format.go
@@ -7,6 +7,7 @@ package format
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"regexp"
 	"slices"
@@ -84,9 +85,9 @@ func SourceWithOpts(filename string, src []byte, opts Opts) ([]byte, error) {
 		checkOpts.RequireIfKeyword = false
 		checkOpts.RequireContainsKeyword = false
 		checkOpts.RequireRuleBodyOrValue = false
-		errors := ast.CheckRegoV1WithOptions(module, checkOpts)
-		if len(errors) > 0 {
-			return nil, errors
+		errs := ast.CheckRegoV1WithOptions(module, checkOpts)
+		if len(errs) > 0 {
+			return nil, errs
 		}
 	}
 
@@ -99,7 +100,7 @@ func SourceWithOpts(filename string, src []byte, opts Opts) ([]byte, error) {
 }
 
 // MustAst is a helper function to format a Rego AST element. If any errors
-// occurs this function will panic. This is mostly used for test
+// occur this function will panic. This is mostly used for test
 func MustAst(x interface{}) []byte {
 	bs, err := Ast(x)
 	if err != nil {
@@ -109,7 +110,7 @@ func MustAst(x interface{}) []byte {
 }
 
 // MustAstWithOpts is a helper function to format a Rego AST element. If any errors
-// occurs this function will panic. This is mostly used for test
+// occur this function will panic. This is mostly used for test
 func MustAstWithOpts(x interface{}, opts Opts) []byte {
 	bs, err := AstWithOpts(x, opts)
 	if err != nil {
@@ -262,30 +263,63 @@ func AstWithOpts(x interface{}, opts Opts) ([]byte, error) {
 				x.Imports = ensureFutureKeywordImport(x.Imports, kw)
 			}
 		}
-		w.writeModule(x)
+		err := w.writeModule(x)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case *ast.Package:
-		w.writePackage(x, nil)
+		_, err := w.writePackage(x, nil)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case *ast.Import:
-		w.writeImports([]*ast.Import{x}, nil)
+		_, err := w.writeImports([]*ast.Import{x}, nil)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case *ast.Rule:
-		w.writeRule(x, false /* isElse */, nil)
+		_, err := w.writeRule(x, false /* isElse */, nil)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case *ast.Head:
-		w.writeHead(x,
+		_, err := w.writeHead(x,
 			false, // isDefault
 			false, // isExpandedConst
 			nil)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case ast.Body:
-		w.writeBody(x, nil)
+		_, err := w.writeBody(x, nil)
+		if err != nil {
+			return nil, err
+		}
 	case *ast.Expr:
-		w.writeExpr(x, nil)
+		_, err := w.writeExpr(x, nil)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case *ast.With:
-		w.writeWith(x, nil, false)
+		_, err := w.writeWith(x, nil, false)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case *ast.Term:
-		w.writeTerm(x, nil)
+		_, err := w.writeTerm(x, nil)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case ast.Value:
-		w.writeTerm(&ast.Term{Value: x, Location: &ast.Location{}}, nil)
+		_, err := w.writeTerm(&ast.Term{Value: x, Location: &ast.Location{}}, nil)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	case *ast.Comment:
-		w.writeComments([]*ast.Comment{x})
+		err := w.writeComments([]*ast.Comment{x})
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 	default:
 		return nil, fmt.Errorf("not an ast element: %v", x)
 	}
@@ -338,16 +372,17 @@ func defaultLocation(x ast.Node) *ast.Location {
 type writer struct {
 	buf bytes.Buffer
 
-	indent    string
-	level     int
-	inline    bool
-	beforeEnd *ast.Comment
-	delay     bool
-	errs      ast.Errors
-	fmtOpts   fmtOpts
+	indent                  string
+	level                   int
+	inline                  bool
+	beforeEnd               *ast.Comment
+	delay                   bool
+	errs                    ast.Errors
+	fmtOpts                 fmtOpts
+	writeCommentOnFinalLine bool
 }
 
-func (w *writer) writeModule(module *ast.Module) {
+func (w *writer) writeModule(module *ast.Module) error {
 	var pkg *ast.Package
 	var others []interface{}
 	var comments []*ast.Comment
@@ -369,23 +404,41 @@ func (w *writer) writeModule(module *ast.Module) {
 	visitor.Walk(module)
 
 	sort.Slice(comments, func(i, j int) bool {
-		return locLess(comments[i], comments[j])
+		l, err := locLess(comments[i], comments[j])
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
+		return l
 	})
 
 	sort.Slice(others, func(i, j int) bool {
-		return locLess(others[i], others[j])
+		l, err := locLess(others[i], others[j])
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
+		return l
 	})
 
 	comments = trimTrailingWhitespaceInComments(comments)
 
-	comments = w.writePackage(pkg, comments)
+	var err error
+	comments, err = w.writePackage(pkg, comments)
+	if err != nil {
+		return err
+	}
 	var imports []*ast.Import
 	var rules []*ast.Rule
 	for len(others) > 0 {
 		imports, others = gatherImports(others)
-		comments = w.writeImports(imports, comments)
+		comments, err = w.writeImports(imports, comments)
+		if err != nil {
+			return err
+		}
 		rules, others = gatherRules(others)
-		comments = w.writeRules(rules, comments)
+		comments, err = w.writeRules(rules, comments)
+		if err != nil {
+			return err
+		}
 	}
 
 	for i, c := range comments {
@@ -394,6 +447,8 @@ func (w *writer) writeModule(module *ast.Module) {
 			w.write("\n")
 		}
 	}
+
+	return nil
 }
 
 func trimTrailingWhitespaceInComments(comments []*ast.Comment) []*ast.Comment {
@@ -404,8 +459,12 @@ func trimTrailingWhitespaceInComments(comments []*ast.Comment) []*ast.Comment {
 	return comments
 }
 
-func (w *writer) writePackage(pkg *ast.Package, comments []*ast.Comment) []*ast.Comment {
-	comments = w.insertComments(comments, pkg.Location)
+func (w *writer) writePackage(pkg *ast.Package, comments []*ast.Comment) ([]*ast.Comment, error) {
+	var err error
+	comments, err = w.insertComments(comments, pkg.Location)
+	if err != nil {
+		return nil, err
+	}
 
 	w.startLine()
 
@@ -413,33 +472,53 @@ func (w *writer) writePackage(pkg *ast.Package, comments []*ast.Comment) []*ast.
 	path := make(ast.Ref, len(pkg.Path)-1)
 	if len(path) == 0 {
 		w.errs = append(w.errs, ast.NewError(ast.FormatErr, pkg.Location, "invalid package path: %s", pkg.Path))
-		return comments
+		return comments, nil
 	}
 
 	path[0] = ast.VarTerm(string(pkg.Path[1].Value.(ast.String)))
 	copy(path[1:], pkg.Path[2:])
 
 	w.write("package ")
-	w.writeRef(path, nil)
+	_, err = w.writeRef(path, nil)
+	if err != nil {
+		return nil, err
+	}
 
 	w.blankLine()
 
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeComments(comments []*ast.Comment) {
+func (w *writer) writeComments(comments []*ast.Comment) error {
 	for i := range comments {
-		if i > 0 && locCmp(comments[i], comments[i-1]) > 1 {
-			w.blankLine()
+		if i > 0 {
+			l, err := locCmp(comments[i], comments[i-1])
+			if err != nil {
+				return err
+			}
+			if l > 1 {
+				w.blankLine()
+			}
 		}
+
 		w.writeLine(comments[i].String())
 	}
+
+	return nil
 }
 
-func (w *writer) writeRules(rules []*ast.Rule, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeRules(rules []*ast.Rule, comments []*ast.Comment) ([]*ast.Comment, error) {
 	for i, rule := range rules {
-		comments = w.insertComments(comments, rule.Location)
-		comments = w.writeRule(rule, false, comments)
+		var err error
+		comments, err = w.insertComments(comments, rule.Location)
+		if err != nil && !errors.As(err, &unexpectedCommentError{}) {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
+
+		comments, err = w.writeRule(rule, false, comments)
+		if err != nil && !errors.As(err, &unexpectedCommentError{}) {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 
 		if i < len(rules)-1 && w.groupableOneLiner(rule) {
 			next := rules[i+1]
@@ -449,10 +528,9 @@ func (w *writer) writeRules(rules []*ast.Rule, comments []*ast.Comment) []*ast.C
 				continue
 			}
 		}
-
 		w.blankLine()
 	}
-	return comments
+	return comments, nil
 }
 
 var expandedConst = ast.NewBody(ast.NewExpr(ast.InternedBooleanTerm(true)))
@@ -469,9 +547,9 @@ func (w *writer) groupableOneLiner(rule *ast.Rule) bool {
 	return (w.fmtOpts.regoV1 || w.fmtOpts.ifs) && partialSetException
 }
 
-func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment) ([]*ast.Comment, error) {
 	if rule == nil {
-		return comments
+		return comments, nil
 	}
 
 	if !isElse {
@@ -487,14 +565,26 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment)
 	// in the formatted code instead of expanding the bodies into rules, so we
 	// pretend that the rule has no body in this case.
 	isExpandedConst := rule.Body.Equal(expandedConst) && rule.Else == nil
+	w.writeCommentOnFinalLine = isExpandedConst
 
-	comments = w.writeHead(rule.Head, rule.Default, isExpandedConst, comments)
+	var err error
+	var unexpectedComment bool
+	comments, err = w.writeHead(rule.Head, rule.Default, isExpandedConst, comments)
+	if err != nil {
+		if errors.As(err, &unexpectedCommentError{}) {
+			unexpectedComment = true
+		} else {
+			return nil, err
+		}
+	}
 
 	if len(rule.Body) == 0 || isExpandedConst {
 		w.endLine()
-		return comments
+		return comments, nil
 	}
 
+	w.writeCommentOnFinalLine = true
+
 	// this excludes partial sets UNLESS `contains` is used
 	partialSetException := w.fmtOpts.contains || rule.Head.Value != nil
 
@@ -503,20 +593,38 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment)
 		if len(rule.Body) == 1 {
 			if rule.Body[0].Location.Row == rule.Head.Location.Row {
 				w.write(" ")
-				comments = w.writeExpr(rule.Body[0], comments)
+				var err error
+				comments, err = w.writeExpr(rule.Body[0], comments)
+				if err != nil {
+					return nil, err
+				}
 				w.endLine()
 				if rule.Else != nil {
-					comments = w.writeElse(rule, comments)
+					comments, err = w.writeElse(rule, comments)
+					if err != nil {
+						return nil, err
+					}
 				}
-				return comments
+				return comments, nil
 			}
 		}
 	}
-	w.write(" {")
-	w.endLine()
+	if unexpectedComment && len(comments) > 0 {
+		w.write(" { ")
+	} else {
+		w.write(" {")
+		w.endLine()
+	}
+
 	w.up()
 
-	comments = w.writeBody(rule.Body, comments)
+	comments, err = w.writeBody(rule.Body, comments)
+	if err != nil {
+		// the unexpected comment error is passed up to be handled by writeHead
+		if !errors.As(err, &unexpectedCommentError{}) {
+			return nil, err
+		}
+	}
 
 	var closeLoc *ast.Location
 
@@ -528,20 +636,28 @@ func (w *writer) writeRule(rule *ast.Rule, isElse bool, comments []*ast.Comment)
 		closeLoc = closingLoc(0, 0, '{', '}', rule.Location)
 	}
 
-	comments = w.insertComments(comments, closeLoc)
+	comments, err = w.insertComments(comments, closeLoc)
+	if err != nil {
+		return nil, err
+	}
 
-	w.down()
+	if err := w.down(); err != nil {
+		return nil, err
+	}
 	w.startLine()
 	w.write("}")
 	if rule.Else != nil {
-		comments = w.writeElse(rule, comments)
+		comments, err = w.writeElse(rule, comments)
+		if err != nil {
+			return nil, err
+		}
 	}
-	return comments
+	return comments, nil
 }
 
 var elseVar ast.Value = ast.Var("else")
 
-func (w *writer) writeElse(rule *ast.Rule, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeElse(rule *ast.Rule, comments []*ast.Comment) ([]*ast.Comment, error) {
 	// If there was nothing else on the line before the "else" starts
 	// then preserve this style of else block, otherwise it will be
 	// started as an "inline" else eg:
@@ -592,7 +708,11 @@ func (w *writer) writeElse(rule *ast.Rule, comments []*ast.Comment) []*ast.Comme
 
 	rule.Else.Head.Reference = ast.Ref{elseHeadReference}
 	rule.Else.Head.Args = nil
-	comments = w.insertComments(comments, rule.Else.Head.Location)
+	var err error
+	comments, err = w.insertComments(comments, rule.Else.Head.Location)
+	if err != nil {
+		return nil, err
+	}
 
 	if hasCommentAbove && !wasInline {
 		// The comments would have ended the line, be sure to start one again
@@ -610,14 +730,27 @@ func (w *writer) writeElse(rule *ast.Rule, comments []*ast.Comment) []*ast.Comme
 	return w.writeRule(rule.Else, true, comments)
 }
 
-func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeHead(head *ast.Head, isDefault bool, isExpandedConst bool, comments []*ast.Comment) ([]*ast.Comment, error) {
 	ref := head.Ref()
 	if head.Key != nil && head.Value == nil && !head.HasDynamicRef() {
 		ref = ref.GroundPrefix()
 	}
 	if w.fmtOpts.refHeads || len(ref) == 1 {
-		w.writeRef(ref, comments)
+		var err error
+		comments, err = w.writeRef(ref, comments)
+		if err != nil {
+			return nil, err
+		}
 	} else {
+		// if there are comments within the object in the rule head, don't format it
+		if len(comments) > 0 && ref[1].Location.Row == comments[0].Location.Row {
+			comments, err := w.writeUnformatted(head.Location, comments)
+			if err != nil {
+				return nil, err
+			}
+			return comments, nil
+		}
+
 		w.write(ref[0].String())
 		w.write("[")
 		w.write(ref[1].String())
@@ -630,16 +763,28 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, comm
 		for _, arg := range head.Args {
 			args = append(args, arg)
 		}
-		comments = w.writeIterable(args, head.Location, closingLoc(0, 0, '(', ')', head.Location), comments, w.listWriter())
+		var err error
+		comments, err = w.writeIterable(args, head.Location, closingLoc(0, 0, '(', ')', head.Location), comments, w.listWriter())
 		w.write(")")
+		if err != nil {
+			return comments, err
+		}
 	}
 	if head.Key != nil {
 		if w.fmtOpts.contains && head.Value == nil {
 			w.write(" contains ")
-			comments = w.writeTerm(head.Key, comments)
+			var err error
+			comments, err = w.writeTerm(head.Key, comments)
+			if err != nil {
+				return comments, err
+			}
 		} else if head.Value == nil { // no `if` for p[x] notation
 			w.write("[")
-			comments = w.writeTerm(head.Key, comments)
+			var err error
+			comments, err = w.writeTerm(head.Key, comments)
+			if err != nil {
+				return comments, err
+			}
 			w.write("]")
 		}
 	}
@@ -661,7 +806,7 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, comm
 			// If the value location is the same as the location of the head,
 			// we know that the value is generated, i.e. f(1)
 			// Don't print the value (` = true`) as it is implied.
-			return comments
+			return comments, nil
 		}
 
 		if head.Assign || w.fmtOpts.regoV1 {
@@ -670,24 +815,35 @@ func (w *writer) writeHead(head *ast.Head, isDefault, isExpandedConst bool, comm
 		} else {
 			w.write(" = ")
 		}
-		comments = w.writeTerm(head.Value, comments)
+		var err error
+		comments, err = w.writeTerm(head.Value, comments)
+		if err != nil {
+			return comments, err
+		}
 	}
-	return comments
+	return comments, nil
 }
 
-func (w *writer) insertComments(comments []*ast.Comment, loc *ast.Location) []*ast.Comment {
+func (w *writer) insertComments(comments []*ast.Comment, loc *ast.Location) ([]*ast.Comment, error) {
 	before, at, comments := partitionComments(comments, loc)
-	w.writeComments(before)
+
+	err := w.writeComments(before)
+	if err != nil {
+		return nil, err
+	}
 	if len(before) > 0 && loc.Row-before[len(before)-1].Location.Row > 1 {
 		w.blankLine()
 	}
 
-	w.beforeLineEnd(at)
-	return comments
+	return comments, w.beforeLineEnd(at)
 }
 
-func (w *writer) writeBody(body ast.Body, comments []*ast.Comment) []*ast.Comment {
-	comments = w.insertComments(comments, body.Loc())
+func (w *writer) writeBody(body ast.Body, comments []*ast.Comment) ([]*ast.Comment, error) {
+	var err error
+	comments, err = w.insertComments(comments, body.Loc())
+	if err != nil {
+		return comments, err
+	}
 	for i, expr := range body {
 		// Insert a blank line in before the expression if it was not right
 		// after the previous expression.
@@ -704,14 +860,21 @@ func (w *writer) writeBody(body ast.Body, comments []*ast.Comment) []*ast.Commen
 		}
 		w.startLine()
 
-		comments = w.writeExpr(expr, comments)
+		comments, err = w.writeExpr(expr, comments)
+		if err != nil && !errors.As(err, &unexpectedCommentError{}) {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
 		w.endLine()
 	}
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeExpr(expr *ast.Expr, comments []*ast.Comment) []*ast.Comment {
-	comments = w.insertComments(comments, expr.Location)
+func (w *writer) writeExpr(expr *ast.Expr, comments []*ast.Comment) ([]*ast.Comment, error) {
+	var err error
+	comments, err = w.insertComments(comments, expr.Location)
+	if err != nil {
+		return comments, err
+	}
 	if !w.inline {
 		w.startLine()
 	}
@@ -722,37 +885,65 @@ func (w *writer) writeExpr(expr *ast.Expr, comments []*ast.Comment) []*ast.Comme
 
 	switch t := expr.Terms.(type) {
 	case *ast.SomeDecl:
-		comments = w.writeSomeDecl(t, comments)
+		comments, err = w.writeSomeDecl(t, comments)
+		if err != nil {
+			return nil, err
+		}
 	case *ast.Every:
-		comments = w.writeEvery(t, comments)
+		comments, err = w.writeEvery(t, comments)
+		if err != nil {
+			return nil, err
+		}
 	case []*ast.Term:
-		comments = w.writeFunctionCall(expr, comments)
+		comments, err = w.writeFunctionCall(expr, comments)
+		if err != nil {
+			return comments, err
+		}
 	case *ast.Term:
-		comments = w.writeTerm(t, comments)
+		comments, err = w.writeTerm(t, comments)
+		if err != nil {
+			return comments, err
+		}
 	}
 
-	var indented bool
+	var indented, down bool
 	for i, with := range expr.With {
 		if i == 0 || with.Location.Row == expr.With[i-1].Location.Row { // we're on the same line
-			comments = w.writeWith(with, comments, false)
+			comments, err = w.writeWith(with, comments, false)
+			if err != nil {
+				return nil, err
+			}
 		} else { // we're on a new line
 			if !indented {
 				indented = true
 
 				w.up()
-				defer w.down()
+				down = true
 			}
 			w.endLine()
 			w.startLine()
-			comments = w.writeWith(with, comments, true)
+			comments, err = w.writeWith(with, comments, true)
+			if err != nil {
+				return nil, err
+			}
 		}
 	}
 
-	return comments
+	if down {
+		if err := w.down(); err != nil {
+			return nil, err
+		}
+	}
+
+	return comments, nil
 }
 
-func (w *writer) writeSomeDecl(decl *ast.SomeDecl, comments []*ast.Comment) []*ast.Comment {
-	comments = w.insertComments(comments, decl.Location)
+func (w *writer) writeSomeDecl(decl *ast.SomeDecl, comments []*ast.Comment) ([]*ast.Comment, error) {
+	var err error
+	comments, err = w.insertComments(comments, decl.Location)
+	if err != nil {
+		return nil, err
+	}
 	w.write("some ")
 
 	row := decl.Location.Row
@@ -769,41 +960,66 @@ func (w *writer) writeSomeDecl(decl *ast.SomeDecl, comments []*ast.Comment) []*a
 				w.write(" ")
 			}
 
-			comments = w.writeTerm(term, comments)
+			comments, err = w.writeTerm(term, comments)
+			if err != nil {
+				return nil, err
+			}
 
 			if i < len(decl.Symbols)-1 {
 				w.write(",")
 			}
 		case ast.Call:
-			comments = w.writeInOperator(false, val[1:], comments, decl.Location, ast.BuiltinMap[val[0].String()].Decl)
+			comments, err = w.writeInOperator(false, val[1:], comments, decl.Location, ast.BuiltinMap[val[0].String()].Decl)
+			if err != nil {
+				return nil, err
+			}
 		}
 	}
 
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeEvery(every *ast.Every, comments []*ast.Comment) []*ast.Comment {
-	comments = w.insertComments(comments, every.Location)
+func (w *writer) writeEvery(every *ast.Every, comments []*ast.Comment) ([]*ast.Comment, error) {
+	var err error
+	comments, err = w.insertComments(comments, every.Location)
+	if err != nil {
+		return nil, err
+	}
 	w.write("every ")
 	if every.Key != nil {
-		comments = w.writeTerm(every.Key, comments)
+		comments, err = w.writeTerm(every.Key, comments)
+		if err != nil {
+			return nil, err
+		}
 		w.write(", ")
 	}
-	comments = w.writeTerm(every.Value, comments)
+	comments, err = w.writeTerm(every.Value, comments)
+	if err != nil {
+		return nil, err
+	}
 	w.write(" in ")
-	comments = w.writeTerm(every.Domain, comments)
+	comments, err = w.writeTerm(every.Domain, comments)
+	if err != nil {
+		return nil, err
+	}
 	w.write(" {")
-	comments = w.writeComprehensionBody('{', '}', every.Body, every.Loc(), every.Loc(), comments)
+	comments, err = w.writeComprehensionBody('{', '}', every.Body, every.Loc(), every.Loc(), comments)
+	if err != nil {
+		// the unexpected comment error is passed up to be handled by writeHead
+		if !errors.As(err, &unexpectedCommentError{}) {
+			return nil, err
+		}
+	}
 
 	if len(every.Body) == 1 &&
 		every.Body[0].Location.Row == every.Location.Row {
 		w.write(" ")
 	}
 	w.write("}")
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeFunctionCall(expr *ast.Expr, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeFunctionCall(expr *ast.Expr, comments []*ast.Comment) ([]*ast.Comment, error) {
 
 	terms := expr.Terms.([]*ast.Term)
 	operator := terms[0].Value.String()
@@ -821,19 +1037,31 @@ func (w *writer) writeFunctionCall(expr *ast.Expr, comments []*ast.Comment) []*a
 	numDeclArgs := bi.Decl.Arity()
 	numCallArgs := len(terms) - 1
 
+	var err error
 	switch numCallArgs {
 	case numDeclArgs: // Print infix where result is unassigned (e.g., x != y)
-		comments = w.writeTerm(terms[1], comments)
+		comments, err = w.writeTerm(terms[1], comments)
+		if err != nil {
+			return nil, err
+		}
 		w.write(" " + bi.Infix + " ")
 		return w.writeTerm(terms[2], comments)
-
 	case numDeclArgs + 1: // Print infix where result is assigned (e.g., z = x + y)
-		comments = w.writeTerm(terms[3], comments)
+		comments, err = w.writeTerm(terms[3], comments)
+		if err != nil {
+			return nil, err
+		}
 		w.write(" " + ast.Equality.Infix + " ")
-		comments = w.writeTerm(terms[1], comments)
+		comments, err = w.writeTerm(terms[1], comments)
+		if err != nil {
+			return nil, err
+		}
 		w.write(" " + bi.Infix + " ")
-		comments = w.writeTerm(terms[2], comments)
-		return comments
+		comments, err = w.writeTerm(terms[2], comments)
+		if err != nil {
+			return nil, err
+		}
+		return comments, nil
 	}
 	// NOTE(Trolloldem): in this point we are operating with a built-in function with the
 	// wrong arity even when the assignment notation is used
@@ -841,7 +1069,7 @@ func (w *writer) writeFunctionCall(expr *ast.Expr, comments []*ast.Comment) []*a
 	return w.writeFunctionCallPlain(terms, comments)
 }
 
-func (w *writer) writeFunctionCallPlain(terms []*ast.Term, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeFunctionCallPlain(terms []*ast.Term, comments []*ast.Comment) ([]*ast.Comment, error) {
 	w.write(terms[0].String() + "(")
 	defer w.write(")")
 	args := make([]interface{}, len(terms)-1)
@@ -849,45 +1077,133 @@ func (w *writer) writeFunctionCallPlain(terms []*ast.Term, comments []*ast.Comme
 		args[i] = t
 	}
 	loc := terms[0].Location
-	return w.writeIterable(args, loc, closingLoc(0, 0, '(', ')', loc), comments, w.listWriter())
+	var err error
+	comments, err = w.writeIterable(args, loc, closingLoc(0, 0, '(', ')', loc), comments, w.listWriter())
+	if err != nil {
+		return nil, err
+	}
+	return comments, nil
 }
 
-func (w *writer) writeWith(with *ast.With, comments []*ast.Comment, indented bool) []*ast.Comment {
-	comments = w.insertComments(comments, with.Location)
+func (w *writer) writeWith(with *ast.With, comments []*ast.Comment, indented bool) ([]*ast.Comment, error) {
+	var err error
+	comments, err = w.insertComments(comments, with.Location)
+	if err != nil {
+		return nil, err
+	}
 	if !indented {
 		w.write(" ")
 	}
 	w.write("with ")
-	comments = w.writeTerm(with.Target, comments)
+	comments, err = w.writeTerm(with.Target, comments)
+	if err != nil {
+		return nil, err
+	}
 	w.write(" as ")
-	return w.writeTerm(with.Value, comments)
+	comments, err = w.writeTerm(with.Value, comments)
+	if err != nil {
+		return nil, err
+	}
+	return comments, nil
 }
 
-func (w *writer) writeTerm(term *ast.Term, comments []*ast.Comment) []*ast.Comment {
-	return w.writeTermParens(false, term, comments)
+func (w *writer) writeTerm(term *ast.Term, comments []*ast.Comment) ([]*ast.Comment, error) {
+	currentComments := make([]*ast.Comment, len(comments))
+	copy(currentComments, comments)
+
+	currentLen := w.buf.Len()
+
+	comments, err := w.writeTermParens(false, term, comments)
+	if err != nil {
+		if errors.As(err, &unexpectedCommentError{}) {
+			w.buf.Truncate(currentLen)
+
+			comments, uErr := w.writeUnformatted(term.Location, currentComments)
+			if uErr != nil {
+				return nil, uErr
+			}
+			return comments, err
+		}
+		return nil, err
+	}
+
+	return comments, nil
 }
 
-func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Comment) []*ast.Comment {
-	comments = w.insertComments(comments, term.Location)
+// writeUnformatted writes the unformatted text instead and updates the comment state
+func (w *writer) writeUnformatted(location *ast.Location, currentComments []*ast.Comment) ([]*ast.Comment, error) {
+	if len(location.Text) == 0 {
+		return nil, errors.New("original unformatted text is empty")
+	}
+
+	rawRule := string(location.Text)
+	rowNum := len(strings.Split(rawRule, "\n"))
+
+	w.write(string(location.Text))
+
+	comments := make([]*ast.Comment, 0, len(currentComments))
+	for _, c := range currentComments {
+		// if there is a body then wait to write the last comment
+		if w.writeCommentOnFinalLine && c.Location.Row == location.Row+rowNum-1 {
+			w.write(" " + string(c.Location.Text))
+			continue
+		}
+
+		// drop comments that occur within the rule raw text
+		if c.Location.Row < location.Row+rowNum-1 {
+			continue
+		}
+		comments = append(comments, c)
+	}
+	return comments, nil
+}
+
+func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Comment) ([]*ast.Comment, error) {
+	var err error
+	comments, err = w.insertComments(comments, term.Location)
+	if err != nil {
+		return nil, err
+	}
 	if !w.inline {
 		w.startLine()
 	}
 
 	switch x := term.Value.(type) {
 	case ast.Ref:
-		comments = w.writeRef(x, comments)
+		comments, err = w.writeRef(x, comments)
+		if err != nil {
+			return nil, err
+		}
 	case ast.Object:
-		comments = w.writeObject(x, term.Location, comments)
+		comments, err = w.writeObject(x, term.Location, comments)
+		if err != nil {
+			return nil, err
+		}
 	case *ast.Array:
-		comments = w.writeArray(x, term.Location, comments)
+		comments, err = w.writeArray(x, term.Location, comments)
+		if err != nil {
+			return nil, err
+		}
 	case ast.Set:
-		comments = w.writeSet(x, term.Location, comments)
+		comments, err = w.writeSet(x, term.Location, comments)
+		if err != nil {
+			return nil, err
+		}
 	case *ast.ArrayComprehension:
-		comments = w.writeArrayComprehension(x, term.Location, comments)
+		comments, err = w.writeArrayComprehension(x, term.Location, comments)
+		if err != nil {
+			return nil, err
+		}
 	case *ast.ObjectComprehension:
-		comments = w.writeObjectComprehension(x, term.Location, comments)
+		comments, err = w.writeObjectComprehension(x, term.Location, comments)
+		if err != nil {
+			return nil, err
+		}
 	case *ast.SetComprehension:
-		comments = w.writeSetComprehension(x, term.Location, comments)
+		comments, err = w.writeSetComprehension(x, term.Location, comments)
+		if err != nil {
+			return nil, err
+		}
 	case ast.String:
 		if term.Location.Text[0] == '`' {
 			// To preserve raw strings, we need to output the original text,
@@ -899,7 +1215,10 @@ func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Co
 	case ast.Var:
 		w.write(w.formatVar(x))
 	case ast.Call:
-		comments = w.writeCall(parens, x, term.Location, comments)
+		comments, err = w.writeCall(parens, x, term.Location, comments)
+		if err != nil {
+			return nil, err
+		}
 	case fmt.Stringer:
 		w.write(x.String())
 	}
@@ -907,17 +1226,21 @@ func (w *writer) writeTermParens(parens bool, term *ast.Term, comments []*ast.Co
 	if !w.inline {
 		w.startLine()
 	}
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeRef(x ast.Ref, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeRef(x ast.Ref, comments []*ast.Comment) ([]*ast.Comment, error) {
 	if len(x) > 0 {
 		parens := false
 		_, ok := x[0].Value.(ast.Call)
 		if ok {
 			parens = x[0].Location.Text[0] == 40 // Starts with "("
 		}
-		comments = w.writeTermParens(parens, x[0], comments)
+		var err error
+		comments, err = w.writeTermParens(parens, x[0], comments)
+		if err != nil {
+			return nil, err
+		}
 		path := x[1:]
 		for _, t := range path {
 			switch p := t.Value.(type) {
@@ -927,13 +1250,21 @@ func (w *writer) writeRef(x ast.Ref, comments []*ast.Comment) []*ast.Comment {
 				w.writeBracketed(w.formatVar(p))
 			default:
 				w.write("[")
-				comments = w.writeTerm(t, comments)
+				comments, err = w.writeTerm(t, comments)
+				if err != nil {
+					if errors.As(err, &unexpectedCommentError{}) {
+						// add a new line so that the closing bracket isn't part of the unexpected comment
+						w.write("\n")
+					} else {
+						return nil, err
+					}
+				}
 				w.write("]")
 			}
 		}
 	}
 
-	return comments
+	return comments, nil
 }
 
 func (w *writer) writeBracketed(str string) {
@@ -958,7 +1289,7 @@ func (*writer) formatVar(v ast.Var) string {
 	return v.String()
 }
 
-func (w *writer) writeCall(parens bool, x ast.Call, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeCall(parens bool, x ast.Call, loc *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 	bi, ok := ast.BuiltinMap[x[0].String()]
 	if !ok || bi.Infix == "" {
 		return w.writeFunctionCallPlain(x, comments)
@@ -980,20 +1311,27 @@ func (w *writer) writeCall(parens bool, x ast.Call, loc *ast.Location, comments
 	// built-in function
 	if bi.Decl.Arity() != len(x)-1 {
 		w.errs = append(w.errs, ArityFormatMismatchError(x[1:], x[0].String(), loc, bi.Decl))
-		return comments
+		return comments, nil
 	}
 
-	comments = w.writeTermParens(true, x[1], comments)
+	var err error
+	comments, err = w.writeTermParens(true, x[1], comments)
+	if err != nil {
+		return nil, err
+	}
 	w.write(" " + bi.Infix + " ")
-	comments = w.writeTermParens(true, x[2], comments)
+	comments, err = w.writeTermParens(true, x[2], comments)
+	if err != nil {
+		return nil, err
+	}
 	if parens {
 		w.write(")")
 	}
 
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeInOperator(parens bool, operands []*ast.Term, comments []*ast.Comment, loc *ast.Location, f *types.Function) []*ast.Comment {
+func (w *writer) writeInOperator(parens bool, operands []*ast.Term, comments []*ast.Comment, loc *ast.Location, f *types.Function) ([]*ast.Comment, error) {
 
 	if len(operands) != f.Arity() {
 		// The number of operands does not math the arity of the `in` operator
@@ -1002,33 +1340,49 @@ func (w *writer) writeInOperator(parens bool, operands []*ast.Term, comments []*
 			operator = ast.MemberWithKey.Name
 		}
 		w.errs = append(w.errs, ArityFormatMismatchError(operands, operator, loc, f))
-		return comments
+		return comments, nil
 	}
 	kw := "in"
+	var err error
 	switch len(operands) {
 	case 2:
-		comments = w.writeTermParens(true, operands[0], comments)
+		comments, err = w.writeTermParens(true, operands[0], comments)
+		if err != nil {
+			return nil, err
+		}
 		w.write(" ")
 		w.write(kw)
 		w.write(" ")
-		comments = w.writeTermParens(true, operands[1], comments)
+		comments, err = w.writeTermParens(true, operands[1], comments)
+		if err != nil {
+			return nil, err
+		}
 	case 3:
 		if parens {
 			w.write("(")
 			defer w.write(")")
 		}
-		comments = w.writeTermParens(true, operands[0], comments)
+		comments, err = w.writeTermParens(true, operands[0], comments)
+		if err != nil {
+			return nil, err
+		}
 		w.write(", ")
-		comments = w.writeTermParens(true, operands[1], comments)
+		comments, err = w.writeTermParens(true, operands[1], comments)
+		if err != nil {
+			return nil, err
+		}
 		w.write(" ")
 		w.write(kw)
 		w.write(" ")
-		comments = w.writeTermParens(true, operands[2], comments)
+		comments, err = w.writeTermParens(true, operands[2], comments)
+		if err != nil {
+			return nil, err
+		}
 	}
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeObject(obj ast.Object, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeObject(obj ast.Object, loc *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 	w.write("{")
 	defer w.write("}")
 
@@ -1039,7 +1393,7 @@ func (w *writer) writeObject(obj ast.Object, loc *ast.Location, comments []*ast.
 	return w.writeIterable(s, loc, closingLoc(0, 0, '{', '}', loc), comments, w.objectWriter())
 }
 
-func (w *writer) writeArray(arr *ast.Array, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeArray(arr *ast.Array, loc *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 	w.write("[")
 	defer w.write("]")
 
@@ -1047,14 +1401,24 @@ func (w *writer) writeArray(arr *ast.Array, loc *ast.Location, comments []*ast.C
 	arr.Foreach(func(t *ast.Term) {
 		s = append(s, t)
 	})
-	return w.writeIterable(s, loc, closingLoc(0, 0, '[', ']', loc), comments, w.listWriter())
+	var err error
+	comments, err = w.writeIterable(s, loc, closingLoc(0, 0, '[', ']', loc), comments, w.listWriter())
+	if err != nil {
+		return nil, err
+	}
+	return comments, nil
 }
 
-func (w *writer) writeSet(set ast.Set, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeSet(set ast.Set, loc *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 
 	if set.Len() == 0 {
 		w.write("set()")
-		return w.insertComments(comments, closingLoc(0, 0, '(', ')', loc))
+		var err error
+		comments, err = w.insertComments(comments, closingLoc(0, 0, '(', ')', loc))
+		if err != nil {
+			return nil, err
+		}
+		return comments, nil
 	}
 
 	w.write("{")
@@ -1064,24 +1428,29 @@ func (w *writer) writeSet(set ast.Set, loc *ast.Location, comments []*ast.Commen
 	set.Foreach(func(t *ast.Term) {
 		s = append(s, t)
 	})
-	return w.writeIterable(s, loc, closingLoc(0, 0, '{', '}', loc), comments, w.listWriter())
+	var err error
+	comments, err = w.writeIterable(s, loc, closingLoc(0, 0, '{', '}', loc), comments, w.listWriter())
+	if err != nil {
+		return nil, err
+	}
+	return comments, nil
 }
 
-func (w *writer) writeArrayComprehension(arr *ast.ArrayComprehension, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeArrayComprehension(arr *ast.ArrayComprehension, loc *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 	w.write("[")
 	defer w.write("]")
 
 	return w.writeComprehension('[', ']', arr.Term, arr.Body, loc, comments)
 }
 
-func (w *writer) writeSetComprehension(set *ast.SetComprehension, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeSetComprehension(set *ast.SetComprehension, loc *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 	w.write("{")
 	defer w.write("}")
 
 	return w.writeComprehension('{', '}', set.Term, set.Body, loc, comments)
 }
 
-func (w *writer) writeObjectComprehension(object *ast.ObjectComprehension, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeObjectComprehension(object *ast.ObjectComprehension, loc *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 	w.write("{")
 	defer w.write("}")
 
@@ -1091,12 +1460,16 @@ func (w *writer) writeObjectComprehension(object *ast.ObjectComprehension, loc *
 		w.startLine()
 	}
 
-	comments = w.writeTerm(object.Key, comments)
+	var err error
+	comments, err = w.writeTerm(object.Key, comments)
+	if err != nil {
+		return nil, err
+	}
 	w.write(": ")
 	return w.writeComprehension('{', '}', object.Value, object.Body, loc, comments)
 }
 
-func (w *writer) writeComprehension(openChar, closeChar byte, term *ast.Term, body ast.Body, loc *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeComprehension(openChar, closeChar byte, term *ast.Term, body ast.Body, loc *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 	if term.Location.Row-loc.Row >= 1 {
 		w.endLine()
 		w.startLine()
@@ -1107,51 +1480,82 @@ func (w *writer) writeComprehension(openChar, closeChar byte, term *ast.Term, bo
 	if ok {
 		parens = term.Location.Text[0] == 40 // Starts with "("
 	}
-	comments = w.writeTermParens(parens, term, comments)
+	var err error
+	comments, err = w.writeTermParens(parens, term, comments)
+	if err != nil {
+		return nil, err
+	}
 	w.write(" |")
 
 	return w.writeComprehensionBody(openChar, closeChar, body, term.Location, loc, comments)
 }
 
-func (w *writer) writeComprehensionBody(openChar, closeChar byte, body ast.Body, term, compr *ast.Location, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeComprehensionBody(openChar, closeChar byte, body ast.Body, term, compr *ast.Location, comments []*ast.Comment) ([]*ast.Comment, error) {
 	exprs := make([]interface{}, 0, len(body))
 	for _, expr := range body {
 		exprs = append(exprs, expr)
 	}
-	lines := groupIterable(exprs, term)
+	lines, err := w.groupIterable(exprs, term)
+	if err != nil {
+		return nil, err
+	}
 
 	if body.Loc().Row-term.Row > 0 || len(lines) > 1 {
 		w.endLine()
 		w.up()
 		defer w.startLine()
-		defer w.down()
+		defer func() {
+			if err := w.down(); err != nil {
+				w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+			}
+		}()
 
-		comments = w.writeBody(body, comments)
+		var err error
+		comments, err = w.writeBody(body, comments)
+		if err != nil {
+			return comments, err
+		}
 	} else {
 		w.write(" ")
 		i := 0
 		for ; i < len(body)-1; i++ {
-			comments = w.writeExpr(body[i], comments)
+			comments, err = w.writeExpr(body[i], comments)
+			if err != nil {
+				return comments, err
+			}
 			w.write("; ")
 		}
-		comments = w.writeExpr(body[i], comments)
+		comments, err = w.writeExpr(body[i], comments)
+		if err != nil {
+			return comments, err
+		}
 	}
-
-	return w.insertComments(comments, closingLoc(0, 0, openChar, closeChar, compr))
+	comments, err = w.insertComments(comments, closingLoc(0, 0, openChar, closeChar, compr))
+	if err != nil {
+		return nil, err
+	}
+	return comments, nil
 }
 
-func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) []*ast.Comment {
+func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) ([]*ast.Comment, error) {
 	m, comments := mapImportsToComments(imports, comments)
 
 	groups := groupImports(imports)
 	for _, group := range groups {
-		comments = w.insertComments(comments, group[0].Loc())
+		var err error
+		comments, err = w.insertComments(comments, group[0].Loc())
+		if err != nil {
+			return nil, err
+		}
 
 		// Sort imports within a newline grouping.
 		slices.SortFunc(group, (*ast.Import).Compare)
 		for _, i := range group {
 			w.startLine()
-			w.writeImport(i)
+			err = w.writeImport(i)
+			if err != nil {
+				return nil, err
+			}
 			if c, ok := m[i]; ok {
 				w.write(" " + c.String())
 			}
@@ -1160,10 +1564,10 @@ func (w *writer) writeImports(imports []*ast.Import, comments []*ast.Comment) []
 		w.blankLine()
 	}
 
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeImport(imp *ast.Import) {
+func (w *writer) writeImport(imp *ast.Import) error {
 	path := imp.Path.Value.(ast.Ref)
 
 	buf := []string{"import"}
@@ -1173,7 +1577,10 @@ func (w *writer) writeImport(imp *ast.Import) {
 		w2 := writer{
 			buf: bytes.Buffer{},
 		}
-		w2.writeRef(path, nil)
+		_, err := w2.writeRef(path, nil)
+		if err != nil {
+			return err
+		}
 		buf = append(buf, w2.buf.String())
 	} else {
 		buf = append(buf, path.String())
@@ -1183,12 +1590,17 @@ func (w *writer) writeImport(imp *ast.Import) {
 		buf = append(buf, "as "+imp.Alias.String())
 	}
 	w.write(strings.Join(buf, " "))
+
+	return nil
 }
 
-type entryWriter func(interface{}, []*ast.Comment) []*ast.Comment
+type entryWriter func(interface{}, []*ast.Comment) ([]*ast.Comment, error)
 
-func (w *writer) writeIterable(elements []interface{}, last *ast.Location, close *ast.Location, comments []*ast.Comment, fn entryWriter) []*ast.Comment {
-	lines := groupIterable(elements, last)
+func (w *writer) writeIterable(elements []interface{}, last *ast.Location, close *ast.Location, comments []*ast.Comment, fn entryWriter) ([]*ast.Comment, error) {
+	lines, err := w.groupIterable(elements, last)
+	if err != nil {
+		return nil, err
+	}
 	if len(lines) > 1 {
 		w.delayBeforeEnd()
 		w.startMultilineSeq()
@@ -1196,34 +1608,49 @@ func (w *writer) writeIterable(elements []interface{}, last *ast.Location, close
 
 	i := 0
 	for ; i < len(lines)-1; i++ {
-		comments = w.writeIterableLine(lines[i], comments, fn)
+		comments, err = w.writeIterableLine(lines[i], comments, fn)
+		if err != nil {
+			return nil, err
+		}
 		w.write(",")
 
 		w.endLine()
 		w.startLine()
 	}
 
-	comments = w.writeIterableLine(lines[i], comments, fn)
+	comments, err = w.writeIterableLine(lines[i], comments, fn)
+	if err != nil {
+		return nil, err
+	}
 
 	if len(lines) > 1 {
 		w.write(",")
 		w.endLine()
-		comments = w.insertComments(comments, close)
-		w.down()
+		comments, err = w.insertComments(comments, close)
+		if err != nil {
+			return nil, err
+		}
+		if err := w.down(); err != nil {
+			return nil, err
+		}
 		w.startLine()
 	}
 
-	return comments
+	return comments, nil
 }
 
-func (w *writer) writeIterableLine(elements []interface{}, comments []*ast.Comment, fn entryWriter) []*ast.Comment {
+func (w *writer) writeIterableLine(elements []interface{}, comments []*ast.Comment, fn entryWriter) ([]*ast.Comment, error) {
 	if len(elements) == 0 {
-		return comments
+		return comments, nil
 	}
 
 	i := 0
 	for ; i < len(elements)-1; i++ {
-		comments = fn(elements[i], comments)
+		var err error
+		comments, err = fn(elements[i], comments)
+		if err != nil {
+			return nil, err
+		}
 		w.write(", ")
 	}
 
@@ -1231,7 +1658,7 @@ func (w *writer) writeIterableLine(elements []interface{}, comments []*ast.Comme
 }
 
 func (w *writer) objectWriter() entryWriter {
-	return func(x interface{}, comments []*ast.Comment) []*ast.Comment {
+	return func(x interface{}, comments []*ast.Comment) ([]*ast.Comment, error) {
 		entry := x.([2]*ast.Term)
 
 		call, isCall := entry[0].Value.(ast.Call)
@@ -1242,7 +1669,11 @@ func (w *writer) objectWriter() entryWriter {
 			w.write("(")
 		}
 
-		comments = w.writeTerm(entry[0], comments)
+		var err error
+		comments, err = w.writeTerm(entry[0], comments)
+		if err != nil {
+			return nil, err
+		}
 		if paren {
 			w.write(")")
 		}
@@ -1260,7 +1691,7 @@ func (w *writer) objectWriter() entryWriter {
 }
 
 func (w *writer) listWriter() entryWriter {
-	return func(x interface{}, comments []*ast.Comment) []*ast.Comment {
+	return func(x interface{}, comments []*ast.Comment) ([]*ast.Comment, error) {
 		t, ok := x.(*ast.Term)
 		if ok {
 			call, isCall := t.Value.(ast.Call)
@@ -1276,7 +1707,7 @@ func (w *writer) listWriter() entryWriter {
 
 // groupIterable will group the `elements` slice into slices according to their
 // location: anything on the same line will be put into a slice.
-func groupIterable(elements []interface{}, last *ast.Location) [][]interface{} {
+func (w *writer) groupIterable(elements []interface{}, last *ast.Location) ([][]interface{}, error) {
 	// Generated vars occur in the AST when we're rendering the result of
 	// partial evaluation in a bundle build with optimization.
 	// Those variables, and wildcard variables have the "default location",
@@ -1303,17 +1734,26 @@ func groupIterable(elements []interface{}, last *ast.Location) [][]interface{} {
 			return false
 		})
 		if def { // return as-is
-			return [][]interface{}{elements}
+			return [][]interface{}{elements}, nil
 		}
 	}
 
-	slices.SortFunc(elements, locCmp)
+	slices.SortFunc(elements, func(i, j any) int {
+		l, err := locCmp(i, j)
+		if err != nil {
+			w.errs = append(w.errs, ast.NewError(ast.FormatErr, &ast.Location{}, err.Error()))
+		}
+		return l
+	})
 
 	var lines [][]interface{}
 	cur := make([]interface{}, 0, len(elements))
 	for i, t := range elements {
 		elem := t
-		loc := getLoc(elem)
+		loc, err := getLoc(elem)
+		if err != nil {
+			return nil, err
+		}
 		lineDiff := loc.Row - last.Row
 		if lineDiff > 0 && i > 0 {
 			lines = append(lines, cur)
@@ -1323,7 +1763,7 @@ func groupIterable(elements []interface{}, last *ast.Location) [][]interface{} {
 		last = loc
 		cur = append(cur, elem)
 	}
-	return append(lines, cur)
+	return append(lines, cur), nil
 }
 
 func mapImportsToComments(imports []*ast.Import, comments []*ast.Comment) (map[*ast.Import]*ast.Comment, []*ast.Comment) {
@@ -1409,7 +1849,7 @@ func partitionComments(comments []*ast.Comment, l *ast.Location) ([]*ast.Comment
 			before = append(before, c)
 		case cmp > 0:
 			after = append(after, c)
-		case cmp == 0:
+		default:
 			at = c
 		}
 	}
@@ -1445,39 +1885,46 @@ loop:
 	return rules, others[i:]
 }
 
-func locLess(a, b interface{}) bool {
-	return locCmp(a, b) < 0
+func locLess(a, b interface{}) (bool, error) {
+	c, err := locCmp(a, b)
+	return c < 0, err
 }
 
-func locCmp(a, b interface{}) int {
-	al := getLoc(a)
-	bl := getLoc(b)
+func locCmp(a, b interface{}) (int, error) {
+	al, err := getLoc(a)
+	if err != nil {
+		return 0, err
+	}
+	bl, err := getLoc(b)
+	if err != nil {
+		return 0, err
+	}
 	switch {
 	case al == nil && bl == nil:
-		return 0
+		return 0, nil
 	case al == nil:
-		return -1
+		return -1, nil
 	case bl == nil:
-		return 1
+		return 1, nil
 	}
 
 	if cmp := al.Row - bl.Row; cmp != 0 {
-		return cmp
+		return cmp, nil
 
 	}
-	return al.Col - bl.Col
+	return al.Col - bl.Col, nil
 }
 
-func getLoc(x interface{}) *ast.Location {
+func getLoc(x interface{}) (*ast.Location, error) {
 	switch x := x.(type) {
 	case ast.Node: // *ast.Head, *ast.Expr, *ast.With, *ast.Term
-		return x.Loc()
+		return x.Loc(), nil
 	case *ast.Location:
-		return x
+		return x, nil
 	case [2]*ast.Term: // Special case to allow for easy printing of objects.
-		return x[0].Location
+		return x[0].Location, nil
 	default:
-		panic("Not reached")
+		return nil, fmt.Errorf("unable to get location for type %v", x)
 	}
 }
 
@@ -1569,15 +2016,46 @@ func (w *writer) endLine() {
 	w.write("\n")
 }
 
+type unexpectedCommentError struct {
+	newComment         string
+	newCommentRow      int
+	existingComment    string
+	existingCommentRow int
+}
+
+func (u unexpectedCommentError) Error() string {
+	return fmt.Sprintf("unexpected new comment (%s) on line %d because there is already a comment (%s) registered for line %d",
+		u.newComment, u.newCommentRow, u.existingComment, u.existingCommentRow)
+}
+
 // beforeLineEnd registers a comment to be printed at the end of the current line.
-func (w *writer) beforeLineEnd(c *ast.Comment) {
+func (w *writer) beforeLineEnd(c *ast.Comment) error {
 	if w.beforeEnd != nil {
 		if c == nil {
-			return
+			return nil
+		}
+
+		existingComment := truncatedString(w.beforeEnd.String(), 100)
+		existingCommentRow := w.beforeEnd.Location.Row
+		newComment := truncatedString(c.String(), 100)
+		w.beforeEnd = nil
+
+		return unexpectedCommentError{
+			newComment:         newComment,
+			newCommentRow:      c.Location.Row,
+			existingComment:    existingComment,
+			existingCommentRow: existingCommentRow,
 		}
-		panic("overwriting non-nil beforeEnd")
 	}
 	w.beforeEnd = c
+	return nil
+}
+
+func truncatedString(s string, max int) string {
+	if len(s) > max {
+		return s[:max-2] + "..."
+	}
+	return s
 }
 
 func (w *writer) delayBeforeEnd() {
@@ -1619,11 +2097,12 @@ func (w *writer) up() {
 }
 
 // down decreases the indentation level
-func (w *writer) down() {
+func (w *writer) down() error {
 	if w.level == 0 {
-		panic("negative indentation level")
+		return errors.New("negative indentation level")
 	}
 	w.level--
+	return nil
 }
 
 func ensureFutureKeywordImport(imps []*ast.Import, kw string) []*ast.Import {
@@ -1675,16 +2154,16 @@ func ensureImport(imps []*ast.Import, path ast.Ref) []*ast.Import {
 	return append(imps, imp)
 }
 
-// ArgErrDetail but for `fmt` checks since compiler has not run yet.
+// ArityFormatErrDetail but for `fmt` checks since compiler has not run yet.
 type ArityFormatErrDetail struct {
 	Have []string `json:"have"`
 	Want []string `json:"want"`
 }
 
-// arityMismatchError but for `fmt` checks since the compiler has not run yet.
+// ArityFormatMismatchError but for `fmt` checks since the compiler has not run yet.
 func ArityFormatMismatchError(operands []*ast.Term, operator string, loc *ast.Location, f *types.Function) *ast.Error {
 	want := make([]string, f.Arity())
-	for i, arg := range f.Args() {
+	for i, arg := range f.FuncArgs().Args {
 		want[i] = types.Sprint(arg)
 	}
 
diff --git a/v1/format/testfiles/v0/test_issue_6330.rego b/v1/format/testfiles/v0/test_issue_6330.rego
new file mode 100644
index 0000000000..c10f476363
--- /dev/null
+++ b/v1/format/testfiles/v0/test_issue_6330.rego
@@ -0,0 +1,45 @@
+package a
+
+value := {"a":
+{"b":
+{"c":
+"d"}}}
+
+value := {"a":  # test 1
+"b"} # test 2
+
+value := {"a":  # test 1
+"b"}
+
+value := {"a":
+{"b":
+{"c":
+"d"}}}
+
+value := {"a": # this is
+{"b": # my ridiculous
+{"c": # way of
+"d"}}} # commenting code
+
+value := {"a":
+{"b":
+{"c":
+"d"}}}
+
+p[
+{"a": #
+"b"} #
+] := true
+
+p.foo.bar[
+{"a": #
+"b"} #
+] := true
+
+p[
+{"a": #
+"b"} #
+][
+{"c": #
+"d"} #
+] := true
\ No newline at end of file
diff --git a/v1/format/testfiles/v0/test_issue_6330.rego.formatted b/v1/format/testfiles/v0/test_issue_6330.rego.formatted
new file mode 100644
index 0000000000..0209af21c7
--- /dev/null
+++ b/v1/format/testfiles/v0/test_issue_6330.rego.formatted
@@ -0,0 +1,31 @@
+package a
+
+value := {"a": {"b": {"c": "d"}}}
+
+value := {"a":  # test 1
+"b"} # test 2
+
+value := {"a": "b"} # test 1
+
+value := {"a": {"b": {"c": "d"}}}
+
+value := {"a": # this is
+{"b": # my ridiculous
+{"c": # way of
+"d"}}} # commenting code
+
+value := {"a": {"b": {"c": "d"}}}
+
+p[{"a": #
+"b"} #
+] := true
+
+p.foo.bar[{"a": #
+"b"} #
+] := true
+
+p[{"a": #
+"b"} #
+][{"c": #
+"d"} #
+] := true
diff --git a/v1/format/testfiles/v0/test_issue_6330_1.rego b/v1/format/testfiles/v0/test_issue_6330_1.rego
new file mode 100644
index 0000000000..996cf400dc
--- /dev/null
+++ b/v1/format/testfiles/v0/test_issue_6330_1.rego
@@ -0,0 +1,6 @@
+package a
+
+p[
+{"a": #
+"b"} #
+] := true
\ No newline at end of file
diff --git a/v1/format/testfiles/v0/test_issue_6330_1.rego.formatted b/v1/format/testfiles/v0/test_issue_6330_1.rego.formatted
new file mode 100644
index 0000000000..e185363e6f
--- /dev/null
+++ b/v1/format/testfiles/v0/test_issue_6330_1.rego.formatted
@@ -0,0 +1,6 @@
+package a
+
+p[
+{"a": #
+"b"} #
+] := true
diff --git a/v1/format/testfiles/v1/test_issue_6330.rego b/v1/format/testfiles/v1/test_issue_6330.rego
new file mode 100644
index 0000000000..14b738959f
--- /dev/null
+++ b/v1/format/testfiles/v1/test_issue_6330.rego
@@ -0,0 +1,135 @@
+package a
+
+p[
+{"a": #
+"b"} #
+] := true
+
+value := {"a":
+{"b":
+{"c":
+"d"}}}
+
+value := {"a":  # test 1
+"b"} # test 2
+
+value := {"a":  # test 1
+"b"}
+
+value := {"a":
+{"b":
+{"c":
+"d"}}}
+
+value := {"a": # this is
+{"b": # my ridiculous
+{"c": # way of
+"d"}}} # commenting code
+
+p := {"a":  #
+"b"} if { #
+	str := "my \n string"
+}
+
+value := {"a":
+{"b":
+{"c":
+"d"}}}
+
+f(_) := value if {
+	value := {"a": # this is
+        {"b": # my ridiculous
+        {"c": # way of
+        "d"}}} # commenting code
+}
+
+p := value if {
+	value := {"a": # this is
+        {"b": # my ridiculous
+        {"c": # way of
+        "d"}}} # commenting code
+}
+
+p := x if {
+	x := [v | v := {"a": # this is
+        {"b": # my ridiculous
+        {"c": # way of
+        "d"}}} # commenting code
+    ]
+}
+
+p if {
+	every x in input.foo {
+    	x == {"a": # this is
+            {"b": # my ridiculous
+            {"c": # way of
+            "d"}}} # commenting code
+    }
+}
+
+value contains {"a":  #
+"b"} #
+
+p := {"a":  #
+str} if {
+	str := "my \n string"
+}
+
+p := {"a":
+str} if {
+    #
+	str := "my \n string"
+}
+
+authorize := "allow" if {
+	value := {"a": # this is
+    {"b": # my ridiculous
+    {"c": # way of
+    "d"}}} # commenting code
+    input.user == "superuser"           # allow 'superuser' to perform any operation.
+} else := "deny" if {
+    value := {"a": # this is
+    {"b": # my ridiculous
+    {"c": # way of
+    "d"}}} # commenting code
+    input.path[0] == value              # disallow 'admin' operations...
+    input.source_network == "external"  # from external networks.
+} # ... more rules
+
+p[
+{"a": #
+"b"} #
+] := true
+
+p.foo.bar[
+{"a": #
+"b"} #
+] := true
+
+p[
+{"a": #
+"b"} #
+][
+{"c": #
+"d"} #
+] := true
+
+p if {
+    x := {"a": # do
+          "b"} # re
+    1 + 2      ==       3
+    y := {"c": # mi
+          "d"} # fa
+    x != y
+}
+
+p[x].r := y if {
+	x := "q"
+	y := 1
+	    y := {"c": # hello
+"d"} # world
+
+
+
+    y := 1
+}
\ No newline at end of file
diff --git a/v1/format/testfiles/v1/test_issue_6330.rego.formatted b/v1/format/testfiles/v1/test_issue_6330.rego.formatted
new file mode 100644
index 0000000000..0f3a7ea948
--- /dev/null
+++ b/v1/format/testfiles/v1/test_issue_6330.rego.formatted
@@ -0,0 +1,116 @@
+package a
+
+p[{"a": #
+"b"} #
+] := true
+
+value := {"a": {"b": {"c": "d"}}}
+
+value := {"a":  # test 1
+"b"} # test 2
+
+value := {"a": "b"} # test 1
+
+value := {"a": {"b": {"c": "d"}}}
+
+value := {"a": # this is
+{"b": # my ridiculous
+{"c": # way of
+"d"}}} # commenting code
+
+p := {"a":  #
+"b"} if { #
+	str := "my \n string"
+}
+
+value := {"a": {"b": {"c": "d"}}}
+
+f(_) := value if {
+	value := {"a": # this is
+        {"b": # my ridiculous
+        {"c": # way of
+        "d"}}} # commenting code
+}
+
+p := value if {
+	value := {"a": # this is
+        {"b": # my ridiculous
+        {"c": # way of
+        "d"}}} # commenting code
+}
+
+p := x if {
+	x := [v | v := {"a": # this is
+        {"b": # my ridiculous
+        {"c": # way of
+        "d"}}} # commenting code
+    ]
+}
+
+p if {
+	every x in input.foo {
+		x == {"a": # this is
+            {"b": # my ridiculous
+            {"c": # way of
+            "d"}}} # commenting code
+	}
+}
+
+value contains {"a":  #
+"b"} #
+
+p := {"a": str} if { #
+	str := "my \n string"
+}
+
+p := {"a": str} if {
+	#
+	str := "my \n string"
+}
+
+authorize := "allow" if {
+	value := {"a": # this is
+    {"b": # my ridiculous
+    {"c": # way of
+    "d"}}} # commenting code
+	input.user == "superuser" # allow 'superuser' to perform any operation.
+} else := "deny" if {
+	value := {"a": # this is
+    {"b": # my ridiculous
+    {"c": # way of
+    "d"}}} # commenting code
+	input.path[0] == value # disallow 'admin' operations...
+	input.source_network == "external" # from external networks.
+} # ... more rules
+
+p[{"a": #
+"b"} #
+] := true
+
+p.foo.bar[{"a": #
+"b"} #
+] := true
+
+p[{"a": #
+"b"} #
+][{"c": #
+"d"} #
+] := true
+
+p if {
+	x := {"a": # do
+          "b"} # re
+	1 + 2 == 3
+	y := {"c": # mi
+          "d"} # fa
+	x != y
+}
+
+p[x].r := y if {
+	x := "q"
+	y := 1
+	y := {"c": # hello
+"d"} # world
+
+	y := 1
+}
diff --git a/v1/format/testfiles/v1/test_issue_6330_1.rego b/v1/format/testfiles/v1/test_issue_6330_1.rego
new file mode 100644
index 0000000000..414c3c1524
--- /dev/null
+++ b/v1/format/testfiles/v1/test_issue_6330_1.rego
@@ -0,0 +1,9 @@
+package a
+
+# A ref-head rule by itself is formatted differently when there are multiple ref-head rules in a file.
+# This file handles the case when there is a single ref-head rule, while `test_issue_6330.rego` handles multiple.
+
+p[
+{"a": #
+"b"} #
+] := true
\ No newline at end of file
diff --git a/v1/format/testfiles/v1/test_issue_6330_1.rego.formatted b/v1/format/testfiles/v1/test_issue_6330_1.rego.formatted
new file mode 100644
index 0000000000..72a0670d09
--- /dev/null
+++ b/v1/format/testfiles/v1/test_issue_6330_1.rego.formatted
@@ -0,0 +1,9 @@
+package a
+
+# A ref-head rule by itself is formatted differently when there are multiple ref-head rules in a file.
+# This file handles the case when there is a single ref-head rule, while `test_issue_6330.rego` handles multiple.
+
+p[
+{"a": #
+"b"} #
+] := true