Attempt to fix the odinfmt disable

2 files changed, 26 insertions, 13 deletions

diff --git a/src/odin/printer/printer.odin b/src/odin/printer/printer.odin
index 04aee66..2cf7d62 100644
--- a/src/odin/printer/printer.odin
+++ b/src/odin/printer/printer.odin
@@ -23,6 +23,7 @@ Printer :: struct {
 	newline:              string,
 	indentation_count:    int,
 	disabled_lines:       map[int]string,
+	last_disabled_line:   int,
 	src:                  string,
 }
 
@@ -101,10 +102,12 @@ build_disabled_lines_info :: proc(p: ^Printer) {
 				found_disable = true
 				disable_position = comment.pos
 			} else if strings.contains(comment.text[:], "//odinfmt: enable") && found_disable {
+				begin := disable_position.offset - (comment.pos.column - 1)
+				end := comment.pos.offset+len(comment.text)
 				for line := disable_position.line; line <= comment.pos.line; line += 1 {
-					p.disabled_lines[line] = ""
+					p.disabled_lines[line] = p.src[begin:end]
 				}
-				p.disabled_lines[disable_position.line+1] = p.src[disable_position.offset:comment.pos.offset+len(comment.text)]
+				
 				found_disable = false
 			}
 		}
diff --git a/src/odin/printer/visit.odin b/src/odin/printer/visit.odin
index 19353e8..33f085f 100644
--- a/src/odin/printer/visit.odin
+++ b/src/odin/printer/visit.odin
@@ -158,22 +158,28 @@ visit_comments :: proc(p: ^Printer, pos: tokenizer.Pos, end_newline := true) ->
 	return document, lines
 }
 
-visit_disabled_decl :: proc(p: ^Printer, decl: ^ast.Decl) -> ^Document {
-	disabled_text := p.disabled_lines[decl.pos.line]
+visit_disabled :: proc(p: ^Printer, node: ^ast.Node) -> ^Document {
+	disabled_text := p.disabled_lines[node.pos.line]
+
+	if p.last_disabled_line + 1 == node.pos.line {
+		return empty();
+	}
 
 	if disabled_text == "" {
 		return empty()
 	}
 
-	move := move_line(p, decl.pos)
+	move := move_line(p, node.pos)
 
-	for comment_before_or_in_line(p, decl.end.line + 1) {
+	for comment_before_or_in_line(p, node.end.line) {
 		next_comment_group(p)
 	}
+
+	p.last_disabled_line = node.pos.line
 	
-	p.source_position = decl.end
+	p.source_position = node.end
 
-	return cons(move, text(disabled_text))
+	return cons(nest(-p.indentation_count, move), text(disabled_text))
 }
 
 @(private)
@@ -185,7 +191,7 @@ visit_decl :: proc(p: ^Printer, decl: ^ast.Decl, called_in_stmt := false) -> ^Do
 	}
 
 	if decl.pos.line in p.disabled_lines {
-		return visit_disabled_decl(p, decl)
+		return visit_disabled(p, decl)
 	}
 	
 	switch v in &decl.derived {
@@ -440,6 +446,10 @@ visit_stmt :: proc(p: ^Printer, stmt: ^ast.Stmt, block_type: Block_Type = .Gener
 		return empty()
 	}
 
+	if stmt.pos.line in p.disabled_lines {
+		return visit_disabled(p, stmt)
+	}
+
 	switch v in stmt.derived {
 	case Import_Decl:
 		return visit_decl(p, cast(^Decl)stmt, true)	 
@@ -796,7 +806,7 @@ push_where_clauses :: proc(p: ^Printer, clauses: []^ast.Expr) -> ^Document {
 		return empty()
 	}
 	
-	return nest(p.indentation_count, cons(newline(1), cons_with_opl(text("where"), visit_exprs(p, clauses, {.Add_Comma, .Enforce_Newline}))))
+	return group(nest(p.indentation_count, cons_with_nopl(text("where"), visit_exprs(p, clauses, {.Add_Comma, .Enforce_Newline}))))
 }
 
 @(private)
@@ -926,7 +936,7 @@ visit_expr :: proc(p: ^Printer, expr: ^ast.Expr, options := List_Options{}) -> ^
 			document = cons_with_opl(document, text("#maybe"))
 		}
 
-		document = cons_with_opl(document, push_where_clauses(p, v.where_clauses))
+		document = cons_with_nopl(document, push_where_clauses(p, v.where_clauses))
 
 		if v.variants != nil && (len(v.variants) == 0 || v.pos.line == v.end.line) {
 			document = cons_with_nopl(document, text("{"))
@@ -979,7 +989,7 @@ visit_expr :: proc(p: ^Printer, expr: ^ast.Expr, options := List_Options{}) -> ^
 			document = cons_with_nopl(document, visit_expr(p, v.align))
 		}
 
-		document = cons_with_opl(document, push_where_clauses(p, v.where_clauses))
+		document = cons_with_nopl(document, push_where_clauses(p, v.where_clauses))
 
 		if v.fields != nil && (len(v.fields.list) == 0 || v.pos.line == v.end.line) {
 			document = cons_with_nopl(document, text("{"))
@@ -1008,7 +1018,7 @@ visit_expr :: proc(p: ^Printer, expr: ^ast.Expr, options := List_Options{}) -> ^
 		}
 
 		document = cons(document, visit_proc_type(p, v.type^))
-		document = cons_with_opl(document, push_where_clauses(p, v.where_clauses))
+		document = cons_with_nopl(document, push_where_clauses(p, v.where_clauses))
 
 		if v.body != nil {
 			set_source_position(p, v.body.pos)