improve odinfmt on if statements

2 files changed, 72 insertions, 54 deletions

diff --git a/src/odin/printer/document.odin b/src/odin/printer/document.odin
index 953a846..5ab8ea3 100644
--- a/src/odin/printer/document.odin
+++ b/src/odin/printer/document.odin
@@ -10,9 +10,8 @@ Document :: union {
 	Document_Break,
 	Document_Group,
 	Document_Cons,
-	Document_If_Break,
+	Document_If_Break_Or,
 	Document_Align,
-	Document_Nest_If_Break,
 	Document_Break_Parent,
 	Document_Line_Suffix,
 }
@@ -37,18 +36,15 @@ Document_Nest :: struct {
 	document:  ^Document,
 }
 
-Document_Nest_If_Break :: struct {
-	document: ^Document,
-	group_id: string,
-}
-
 Document_Break :: struct {
 	value:   string,
 	newline: bool,
 }
 
-Document_If_Break :: struct {
-	value: string,
+Document_If_Break_Or :: struct {
+	break_document: ^Document,
+	fit_document: ^Document,
+	group_id: string,
 }
 
 Document_Group :: struct {
@@ -127,12 +123,12 @@ nest_if_break :: proc(
 	group_id := "",
 	allocator := context.allocator,
 ) -> ^Document {
-	document := new(Document, allocator)
-	document^ = Document_Nest_If_Break {
-		document = nested_document,
-		group_id = group_id,
-	}
-	return document
+	return if_break_or_document(
+		nest(nested_document, allocator),
+		nested_document,
+		group_id,
+		allocator,
+	)
 }
 
 hang :: proc(
@@ -185,10 +181,38 @@ align :: proc(
 	return document
 }
 
-if_break :: proc(value: string, allocator := context.allocator) -> ^Document {
+if_break :: proc(
+	value: string,
+	allocator := context.allocator,
+) -> ^Document {
+	return if_break_or_document(text(value, allocator), nil, "", allocator)
+}
+
+if_break_or :: proc {
+	if_break_or_string,
+	if_break_or_document,
+}
+
+if_break_or_string :: proc(
+	break_value: string,
+	fit_value: string,
+	group_id := "",
+	allocator := context.allocator,
+) -> ^Document {
+	return if_break_or_document(text(break_value, allocator), text(fit_value, allocator), group_id, allocator)
+}
+
+if_break_or_document :: proc(
+	break_document: ^Document,
+	fit_document: ^Document,
+	group_id := "",
+	allocator := context.allocator,
+) -> ^Document {
 	document := new(Document, allocator)
-	document^ = Document_If_Break {
-		value = value,
+	document^ = Document_If_Break_Or {
+		break_document = break_document,
+		fit_document = fit_document,
+		group_id = group_id,
 	}
 	return document
 }
@@ -387,28 +411,24 @@ fits :: proc(width: int, list: ^[dynamic]Tuple) -> bool {
 			} else {
 				width -= len(v.value)
 			}
-		case Document_If_Break:
-			if data.mode == .Break {
-				width -= len(v.value)
-			}
-		case Document_Nest_If_Break:
+		case Document_If_Break_Or:
 			if data.mode == .Break {
 				append(
 					list,
 					Tuple{
-						indentation = data.indentation + 1,
+						indentation = data.indentation,
 						mode = data.mode,
-						document = v.document,
+						document = v.break_document,
 						alignment = data.alignment,
 					},
 				)
-			} else {
+			} else if v.fit_document != nil {
 				append(
 					list,
 					Tuple{
 						indentation = data.indentation,
 						mode = data.mode,
-						document = v.document,
+						document = v.fit_document,
 						alignment = data.alignment,
 					},
 				)
@@ -560,30 +580,25 @@ format :: proc(
 				strings.write_string(builder, v.value)
 				consumed += len(v.value)
 			}
-		case Document_If_Break:
-			if data.mode == .Break {
-				strings.write_string(builder, v.value)
-				consumed += len(v.value)
-			}
-		case Document_Nest_If_Break:
+		case Document_If_Break_Or:
 			mode := v.group_id != "" ? p.group_modes[v.group_id] : data.mode
 			if mode == .Break {
 				append(
 					list,
 					Tuple{
-						indentation = data.indentation + 1,
+						indentation = data.indentation,
 						mode = data.mode,
-						document = v.document,
+						document = v.break_document,
 						alignment = data.alignment,
 					},
 				)
-			} else {
+			} else if v.fit_document != nil {
 				append(
 					list,
 					Tuple{
 						indentation = data.indentation,
 						mode = data.mode,
-						document = v.document,
+						document = v.fit_document,
 						alignment = data.alignment,
 					},
 				)
@@ -630,7 +645,9 @@ format :: proc(
 						alignment = data.alignment,
 					},
 				)
-			} else if fits(width - consumed, &list_fits) && v.mode != .Break && v.mode != .Fit {
+			} else if fits(width - consumed, &list_fits) &&
+			   v.mode != .Break &&
+			   v.mode != .Fit {
 				append(
 					list,
 					Tuple{
diff --git a/src/odin/printer/visit.odin b/src/odin/printer/visit.odin
index 70cc8dd..b04fe1d 100644
--- a/src/odin/printer/visit.odin
+++ b/src/odin/printer/visit.odin
@@ -984,35 +984,36 @@ visit_stmt :: proc(
 			)
 		}
 
-		if_document := text("if")
+		begin_document := text("if")
+		end_document := empty()
 
 		if v.init != nil {
-			if_document = cons_with_nopl(
-				if_document,
-				cons(group(visit_stmt(p, v.init)), text(";")),
+			begin_document = cons_with_nopl(
+				begin_document,
+				cons(group(visit_stmt(p, v.init), Document_Group_Options {id = "init"}), text(";")),
 			)
 		}
 
 		if v.cond != nil && v.init != nil {
-			if_document = cons(
-				if_document,
+			end_document = cons(
 				group(cons(break_with_space(), group(visit_expr(p, v.cond)))),
 			)
 		} else if v.cond != nil {
-			if_document = cons_with_nopl(
-				if_document,
+			end_document = cons(
+				break_with_no_newline(),
 				group(visit_expr(p, v.cond)),
 			)
 		}
 
-		if v.init != nil && is_value_decl_statement_ending_with_call(v.init) {
-			document = cons(document, group(if_document))
-		} else if v.cond != nil &&
-		   v.init == nil &&
-		   is_value_expression_call(v.cond) {
-			document = cons(document, group(if_document))
-		} else {
-			document = cons(document, group(hang(3, if_document)))
+
+
+		if v.init != nil && is_value_decl_statement_ending_with_call(v.init) || 
+		   v.cond != nil &&
+		   v.init == nil && 
+		   is_value_expression_call(v.cond)  {
+			document = cons(document, group(cons(begin_document, if_break_or(end_document, hang(3, end_document), "init"))))
+		}  else {
+			document = cons(document, group(hang(3, cons(begin_document, end_document))))
 		}
 
 		set_source_position(p, v.body.pos)