Add ensure_fit function

2 files changed, 33 insertions, 6 deletions

diff --git a/src/odin/printer/document.odin b/src/odin/printer/document.odin
index 5c14704..0ae8e2c 100644
--- a/src/odin/printer/document.odin
+++ b/src/odin/printer/document.odin
@@ -45,6 +45,7 @@ Document_If_Break :: struct {
 Document_Group :: struct {
 	document: ^Document,
 	fill: bool,
+	fit: bool,
 }
 
 Document_Cons :: struct {
@@ -60,6 +61,7 @@ Document_Group_Mode :: enum {
 	Flat,
 	Break,
 	Fill,
+	Fit,
 }
 
 empty :: proc(allocator := context.allocator) -> ^Document {
@@ -102,6 +104,15 @@ hang :: proc(align: int, hanged_document: ^Document, allocator := context.alloca
 	return document
 }
 
+enforce_fit :: proc(fitted_document: ^Document, allocator := context.allocator) -> ^Document {
+	document := new(Document, allocator)
+	document^ = Document_Group { 
+		document = fitted_document,
+		fit = true,
+	}
+	return document
+}
+
 align :: proc(aligned_document: ^Document, allocator := context.allocator) -> ^Document {
 	document := new(Document, allocator)
 	document^ = Document_Align {
@@ -272,7 +283,7 @@ format :: proc(width: int, list: ^[dynamic]Tuple, builder: ^strings.Builder, p:
 		data: Tuple = pop(list)
 
 		switch v in data.document {
-		case Document_Nil: 
+		case Document_Nil:
 		case Document_Newline:
 			if v.amount > 0 {
 				for i := 0; i < v.amount; i += 1 {
@@ -324,11 +335,16 @@ format :: proc(width: int, list: ^[dynamic]Tuple, builder: ^strings.Builder, p:
 
 			fits_consumed := 0
 
-			if fits(width-consumed, &l, &fits_consumed) {
+			if data.mode == .Fit {
+				append(list, Tuple {indentation = data.indentation, mode = .Fit, document = v.document, alignment = data.alignment})
+			}
+			else if fits(width-consumed, &l, &fits_consumed) && !v.fit {
 				append(list, Tuple {indentation = data.indentation, mode = .Flat, document = v.document, alignment = data.alignment})
 			} else {
 				if v.fill {
 					append(list, Tuple {indentation = data.indentation, mode = .Fill, document = v.document, alignment = data.alignment})
+				} else if v.fit {
+					append(list, Tuple {indentation = data.indentation, mode = .Fit, document = v.document, alignment = data.alignment})
 				} else {
 					append(list, Tuple {indentation = data.indentation, mode = .Break, document = v.document, alignment = data.alignment})
 				}
diff --git a/src/odin/printer/visit.odin b/src/odin/printer/visit.odin
index 9a33933..da60e48 100644
--- a/src/odin/printer/visit.odin
+++ b/src/odin/printer/visit.odin
@@ -535,6 +535,17 @@ visit_state_flags :: proc(p: ^Printer, flags: ast.Node_State_Flags) -> ^Document
 }
 
 @(private)
+enforce_fit_if_do :: proc(stmt: ^ast.Stmt, document: ^Document) -> ^Document {
+	if v, ok := stmt.derived.(^ast.Block_Stmt); ok {
+		if v.uses_do {
+			return enforce_fit(document)
+		} 
+	}
+
+	return document
+}
+
+@(private)
 visit_stmt :: proc(p: ^Printer, stmt: ^ast.Stmt, block_type: Block_Type = .Generic, empty_block := false, block_stmt := false) -> ^Document {
 	using ast
 
@@ -637,7 +648,7 @@ visit_stmt :: proc(p: ^Printer, stmt: ^ast.Stmt, block_type: Block_Type = .Gener
 				document = cons_with_opl(document, cons_with_nopl(text("else"), visit_stmt(p, v.else_stmt)))
 			}
 		}
-		return document
+		return enforce_fit_if_do(v.body, document)
 	case ^Switch_Stmt:
 		document := move_line(p, v.pos)
 		document = cons(document, visit_state_flags(p, v.state_flags))
@@ -751,7 +762,7 @@ visit_stmt :: proc(p: ^Printer, stmt: ^ast.Stmt, block_type: Block_Type = .Gener
 		document = cons_with_nopl(document,  visit_stmt(p, v.body))
 		set_source_position(p, v.body.end)
 
-		return document
+		return enforce_fit_if_do(v.body, document)
 	case ^Inline_Range_Stmt:
 		document := move_line(p, v.pos)
 		document = cons(document, visit_state_flags(p, v.state_flags))
@@ -777,7 +788,7 @@ visit_stmt :: proc(p: ^Printer, stmt: ^ast.Stmt, block_type: Block_Type = .Gener
 		document = cons_with_nopl(document,  visit_stmt(p, v.body))
 		set_source_position(p, v.body.end)
 
-		return document
+		return enforce_fit_if_do(v.body, document)
 	case ^Range_Stmt:
 		document := move_line(p, v.pos)
 		document = cons(document, visit_state_flags(p, v.state_flags))
@@ -804,7 +815,7 @@ visit_stmt :: proc(p: ^Printer, stmt: ^ast.Stmt, block_type: Block_Type = .Gener
 		document = cons_with_nopl(document,  visit_stmt(p, v.body))
 		set_source_position(p, v.body.end)
 
-		return document
+		return enforce_fit_if_do(v.body, document)
 	case ^Return_Stmt:
 		document := move_line(p, v.pos)