odinfmt: add more tests and fix comp_lit

10 files changed, 132 insertions, 23 deletions

diff --git a/src/odin/printer/document.odin b/src/odin/printer/document.odin
index d77a490..2345140 100644
--- a/src/odin/printer/document.odin
+++ b/src/odin/printer/document.odin
@@ -233,7 +233,7 @@ fits :: proc(width: int, list: ^[dynamic]Tuple) -> bool {
 
 	if len(list) == 0 {
 		return true
-	} else if width < 0 {
+	} else if width <= 0 {
 		return false
 	}
 
@@ -276,7 +276,7 @@ fits :: proc(width: int, list: ^[dynamic]Tuple) -> bool {
 				append(list, Tuple {indentation = data.indentation, mode = data.mode, document = v.document, alignment = data.alignment})
 			}
 		case Document_Group:
-			append(list, Tuple {indentation = data.indentation, mode = data.mode, document = v.document, alignment = data.alignment})		
+			append(list, Tuple {indentation = data.indentation, mode = (v.mode == .Break ? .Break : data.mode), document = v.document, alignment = data.alignment})		
 		}
 	}
 
@@ -300,6 +300,7 @@ format :: proc(width: int, list: ^[dynamic]Tuple, builder: ^strings.Builder, p:
 	assert(builder != nil)
 
 	consumed := 0
+	recalculate := false;
 
 	for len(list) != 0 {
 		data: Tuple = pop(list)
@@ -318,7 +319,11 @@ format :: proc(width: int, list: ^[dynamic]Tuple, builder: ^strings.Builder, p:
 					strings.write_string(builder, " ")
 				}
 				consumed = data.indentation * p.indentation_width + data.alignment
-			}
+
+				if data.mode == .Flat {
+					recalculate = true
+				}
+			}	
 		case Document_Cons:
 			append(list, Tuple {indentation = data.indentation, mode = data.mode, document = v.rhs, alignment = data.alignment})
 			append(list, Tuple {indentation = data.indentation, mode = data.mode, document = v.lhs, alignment = data.alignment})
@@ -349,14 +354,21 @@ format :: proc(width: int, list: ^[dynamic]Tuple, builder: ^strings.Builder, p:
 				append(list, Tuple {indentation = data.indentation, mode = data.mode, document = v.document, alignment = data.alignment})
 			}
 		case Document_Group:
-			l := make([dynamic]Tuple, 0, len(list))
+			if data.mode == .Flat && !recalculate {
+				append(list, Tuple {indentation = data.indentation, mode = v.mode, document = v.document, alignment = data.alignment})
+				break
+			}
 
+			l := make([dynamic]Tuple, 0, len(list))
+	
 			for element in list {
 				append(&l, element)
-			}	
+			}
+
+			append(&l, Tuple {indentation = data.indentation, mode = .Fit, document = v.document, alignment = data.alignment})
+
+			recalculate = false
 
-			append(&l, Tuple {indentation = data.indentation, mode = .Flat, document = v.document, alignment = data.alignment})
-		
 			if data.mode == .Fit {
 				append(list, Tuple {indentation = data.indentation, mode = .Fit, document = v.document, alignment = data.alignment})
 			}
diff --git a/src/odin/printer/printer.odin b/src/odin/printer/printer.odin
index 6f39c7c..2f25724 100644
--- a/src/odin/printer/printer.odin
+++ b/src/odin/printer/printer.odin
@@ -66,6 +66,7 @@ Expr_Called_Type :: enum {
 	Generic,
 	Value_Decl,
 	Assignment_Stmt,
+	Call_Expr,
 }
 
 Newline_Style :: enum {
diff --git a/src/odin/printer/visit.odin b/src/odin/printer/visit.odin
index 14d5d4e..faec0dd 100644
--- a/src/odin/printer/visit.odin
+++ b/src/odin/printer/visit.odin
@@ -206,7 +206,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(p, decl)
 	}
-	
+
 	#partial switch v in decl.derived {
 	case ^Expr_Stmt:
 		document := move_line(p, decl.pos)
@@ -328,6 +328,19 @@ visit_decl :: proc(p: ^Printer, decl: ^ast.Decl, called_in_stmt := false) -> ^Do
 }
 
 @(private)
+is_call_expr_nestable :: proc(list: []^ast.Expr) -> bool {
+
+	if len(list) == 1 {
+		if _, is_comp := list[0].derived.(^ast.Comp_Lit); is_comp {
+			return false
+		}
+	}
+
+	
+	return true
+}
+
+@(private)
 is_values_nestable_assign :: proc(list: []^ast.Expr) -> bool {
 	for expr in list {
 		#partial switch v in expr.derived {
@@ -495,7 +508,11 @@ visit_comp_lit_exprs :: proc(p: ^Printer, comp_lit: ast.Comp_Lit, options := Lis
 		if (.Enforce_Newline in options) {
 			alignment := get_possible_comp_lit_alignment(p, comp_lit.elems)
 			if value, ok := expr.derived.(^ast.Field_Value); ok && alignment > 0 {
-				document = cons(document, cons_with_nopl(visit_expr(p, value.field), cons_with_nopl(cons(repeat_space(alignment - get_node_length(value.field)), text_position(p, "=", value.sep)), visit_expr(p, value.value))))
+				align := empty()
+				if should_align_comp_lit(p, comp_lit) {
+					align = repeat_space(alignment - get_node_length(value.field))
+				}
+				document = cons(document, cons_with_nopl(visit_expr(p, value.field), cons_with_nopl(cons(align, text_position(p, "=", value.sep)), visit_expr(p, value.value))))
 			} else {
 				document = group(cons(document, visit_expr(p, expr, .Generic, options)))
 			}
@@ -860,6 +877,24 @@ visit_stmt :: proc(p: ^Printer, stmt: ^ast.Stmt, block_type: Block_Type = .Gener
 
 @(private)
 should_align_comp_lit :: proc(p: ^Printer, comp_lit: ast.Comp_Lit) -> bool {
+	if len(comp_lit.elems) == 0 {
+		return false
+	}
+
+	for expr in comp_lit.elems {
+		if field, ok := expr.derived.(^ast.Field_Value); ok {
+			#partial switch v in field.value.derived {
+				case ^ast.Proc_Type, ^ast.Proc_Lit:
+					return false
+			}
+		}
+	}
+	
+	return true
+}
+
+@(private)
+comp_lit_contains_fields :: proc(p: ^Printer, comp_lit: ast.Comp_Lit) -> bool {
 
 	if len(comp_lit.elems) == 0 {
 		return false
@@ -875,6 +910,24 @@ should_align_comp_lit :: proc(p: ^Printer, comp_lit: ast.Comp_Lit) -> bool {
 }
 
 @(private)
+comp_lit_contains_blocks :: proc(p: ^Printer, comp_lit: ast.Comp_Lit) -> bool {
+	if len(comp_lit.elems) == 0 {
+		return false
+	}
+
+	for expr in comp_lit.elems {
+		if field, ok := expr.derived.(^ast.Field_Value); ok {
+			#partial switch v in field.value.derived {
+				case ^ast.Proc_Type, ^ast.Proc_Lit:
+					return true
+			}
+		}
+	}
+	
+	return false
+}
+
+@(private)
 contains_comments_in_range :: proc(p: ^Printer, pos: tokenizer.Pos, end: tokenizer.Pos) -> bool {
 	for i := p.latest_comment_index; i < len(p.comments); i += 1 {
 		for c in p.comments[i].list {
@@ -1160,7 +1213,12 @@ visit_expr :: proc(p: ^Printer, expr: ^ast.Expr, called_from: Expr_Called_Type =
 			contains_do |= contains_do_in_expression(p, arg)
 		}
 
-		document = cons(document, nest(p.indentation_count, cons(break_with(""), visit_call_exprs(p, v))))
+		if is_call_expr_nestable(v.args) {
+			document = cons(document, nest(p.indentation_count, cons(break_with(""), visit_call_exprs(p, v))))
+		} else {
+			document = cons(document, cons(break_with(""), visit_call_exprs(p, v)))
+		}
+	
 		document = cons(document, cons(break_with(""), text(")")))
 
 		//We enforce a break if comments exists inside the call args
@@ -1219,8 +1277,12 @@ visit_expr :: proc(p: ^Printer, expr: ^ast.Expr, called_from: Expr_Called_Type =
 				break
 			}
 		}
+
+		should_newline := comp_lit_contains_fields(p, v^) || contains_comments_in_range(p, v.pos, v.end)
+		should_newline &= (called_from == .Value_Decl || called_from == .Assignment_Stmt || (called_from == .Call_Expr && comp_lit_contains_blocks(p, v^)))
+		should_newline &= len(v.elems) != 0
 	
-		if (should_align_comp_lit(p, v^) || contains_comments_in_range(p, v.pos, v.end)) && (called_from == .Value_Decl || called_from == .Assignment_Stmt) && len(v.elems) != 0 {
+		if should_newline {
 			document = cons_with_nopl(document, visit_begin_brace(p, v.pos, .Generic))
 			
 			set_source_position(p, v.open)
@@ -1565,12 +1627,9 @@ visit_binary_expr :: proc(p: ^Printer, binary: ast.Binary_Expr, first := false)
 		}
 	} 
 
-	if true {
-		if nest_first_expression {
-			document = nest(p.indentation_count, document)
-			document = group(document)
-		}
-		
+	if nest_first_expression {
+		document = nest(p.indentation_count, document)
+		document = group(document)
 	}
 
 	document = cons_with_nopl(document, text(binary.op.text))
@@ -1597,7 +1656,7 @@ visit_call_exprs :: proc(p: ^Printer, call_expr: ^ast.Call_Expr) -> ^Document {
 			document = cons(document, text(".."))
 		}
 
-		document = cons(document, group(visit_expr(p, expr)))
+		document = cons(document, group(visit_expr(p, expr, .Call_Expr)))
 				
 		if i != len(call_expr.args) - 1 {
 			document = cons(document, text(","))
diff --git a/tools/odinfmt/snapshot/snapshot.odin b/tools/odinfmt/snapshot/snapshot.odin
index c4f860d..81d48ad 100644
--- a/tools/odinfmt/snapshot/snapshot.odin
+++ b/tools/odinfmt/snapshot/snapshot.odin
@@ -60,7 +60,7 @@ snapshot_file :: proc(path: string) -> bool {
 	if os.exists(snapshot_path) {
 		if snapshot_data, ok := os.read_entire_file(snapshot_path, context.temp_allocator); ok {
 			if cast(string)snapshot_data != formatted {
-				fmt.eprintf("Formatted file was different from snapshot file: %v", snapshot_path)
+				fmt.eprintf("\nFormatted file was different from snapshot file: %v", snapshot_path)
 				os.write_entire_file(fmt.tprintf("%v_failed", snapshot_path), transmute([]u8)formatted)
 				return false
 			}
diff --git a/tools/odinfmt/tests/.snapshots/binary_expressions.odin b/tools/odinfmt/tests/.snapshots/binary_expressions.odin
index 7758a03..969f55c 100644
--- a/tools/odinfmt/tests/.snapshots/binary_expressions.odin
+++ b/tools/odinfmt/tests/.snapshots/binary_expressions.odin
@@ -1,4 +1,4 @@
-package odin_fmt_tests
+package odinfmt_test
 
 
 binary :: proc() {
diff --git a/tools/odinfmt/tests/.snapshots/calls.odin b/tools/odinfmt/tests/.snapshots/calls.odin
index 3d6a3d3..8cba026 100644
--- a/tools/odinfmt/tests/.snapshots/calls.odin
+++ b/tools/odinfmt/tests/.snapshots/calls.odin
@@ -1,4 +1,4 @@
-package odin_fmt_tests
+package odinfmt_test
 
 
 calls :: proc() {
diff --git a/tools/odinfmt/tests/.snapshots/comp_lit.odin b/tools/odinfmt/tests/.snapshots/comp_lit.odin
new file mode 100644
index 0000000..757b211
--- /dev/null
+++ b/tools/odinfmt/tests/.snapshots/comp_lit.odin
@@ -0,0 +1,18 @@
+package odinfmt_test
+
+
+main :: proc() {
+	sort.sort(sort.Interface {
+		len = proc(it: sort.Interface) -> int {
+			c := 2
+		},
+		less = proc(it: sort.Interface, i, j: int) -> bool {
+			b := 1
+		},
+		swap = proc(it: sort.Interface, i, j: int) {
+			a := 2
+			c := 3
+		},
+		collection = &polygon,
+	})
+}
diff --git a/tools/odinfmt/tests/binary_expressions.odin b/tools/odinfmt/tests/binary_expressions.odin
index 85dc63c..6dd0761 100644
--- a/tools/odinfmt/tests/binary_expressions.odin
+++ b/tools/odinfmt/tests/binary_expressions.odin
@@ -1,4 +1,4 @@
-package odin_fmt_tests
+package odinfmt_test
 
 
 binary :: proc() {
diff --git a/tools/odinfmt/tests/calls.odin b/tools/odinfmt/tests/calls.odin
index f450368..59002fb 100644
--- a/tools/odinfmt/tests/calls.odin
+++ b/tools/odinfmt/tests/calls.odin
@@ -1,4 +1,4 @@
-package odin_fmt_tests
+package odinfmt_test
 
 
 calls :: proc() {
diff --git a/tools/odinfmt/tests/comp_lit.odin b/tools/odinfmt/tests/comp_lit.odin
new file mode 100644
index 0000000..c8a4aff
--- /dev/null
+++ b/tools/odinfmt/tests/comp_lit.odin
@@ -0,0 +1,19 @@
+package odinfmt_test 
+
+
+main :: proc() {
+	sort.sort(sort.Interface{len = proc(it: sort.Interface) -> int {
+		c := 2
+		}, 
+		less = proc(it: sort.Interface, i, j: int) -> bool {
+			b := 1
+		}, 
+		swap = proc(it: sort.Interface, i, j: int) {
+			a := 2
+			c := 3
+		}, 
+		collection = &polygon,
+	})
+}
+
+