Process the rest of the comment group when encountering an odinfmt enable

3 files changed, 30 insertions, 5 deletions

diff --git a/src/odin/printer/visit.odin b/src/odin/printer/visit.odin
index 9784052..e3a7e96 100644
--- a/src/odin/printer/visit.odin
+++ b/src/odin/printer/visit.odin
@@ -185,15 +185,26 @@ visit_disabled :: proc(p: ^Printer, node: ^ast.Node) -> ^Document {
 
 	move := cons(move_line(p, pos_one_line_before), escape_nest(move_line(p, node_pos)))
 
-	for comment_before_or_in_line(p, disabled_info.end_line + 1) {
-		next_comment_group(p)
-	}
-
 	p.disabled_until_line = disabled_info.end_line
 	p.source_position = node.end
 	p.source_position.line = disabled_info.end_line
 
-	return cons(move, text(disabled_info.text))
+	document := cons(move, text(disabled_info.text))
+
+	for comment_before_or_in_line(p, disabled_info.end_line + 1) {
+		// we need to handle the rest of the comment group
+		comment_group := p.comments[p.latest_comment_index]
+		for comment in comment_group.list {
+			if comment.pos.line <= disabled_info.end_line {
+				continue
+			}
+			newlined, tmp_document := visit_comment(p, comment)
+			document = cons(document, tmp_document)
+		}
+		next_comment_group(p)
+	}
+
+	return document
 }
 
 @(private)
diff --git a/tools/odinfmt/tests/.snapshots/comments.odin b/tools/odinfmt/tests/.snapshots/comments.odin
index 9550b08..82ac2dd 100644
--- a/tools/odinfmt/tests/.snapshots/comments.odin
+++ b/tools/odinfmt/tests/.snapshots/comments.odin
@@ -78,3 +78,10 @@ Instruction :: struct {
 	prefixes: LegacyPrefixes,
 	operands: []Operand,
 }
+
+// odinfmt: disable
+Bar :: struct{}
+// odinfmt: enable
+
+// Foo doc
+Foo :: struct {}
diff --git a/tools/odinfmt/tests/comments.odin b/tools/odinfmt/tests/comments.odin
index 658bc76..d9a5cf1 100644
--- a/tools/odinfmt/tests/comments.odin
+++ b/tools/odinfmt/tests/comments.odin
@@ -78,3 +78,10 @@ Instruction :: struct {
 	prefixes: LegacyPrefixes,
 	operands: []Operand,
 }
+
+// odinfmt: disable
+Bar :: struct{}
+// odinfmt: enable
+
+// Foo doc
+Foo :: struct   {}